PaperlayerView exact revision on arXiv

Source-linked AI summary

NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models

Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, Wei Ping

arXiv:2405.17428v3cs.CLcs.AIcs.IRcs.LG

TL;DR

Generalist embedding models must accommodate the differing training characteristics of retrieval, classification, and clustering. NV-Embed combines latent attention pooling, non-causal attention, curated data, and two-stage contrastive instruction tuning. It leads MTEB across these task categories, with NV-Embed-v2 ranking first overall at 72.31 as of August 30, 2024.

  • Problem

    Training generalist embeddings for retrieval, classification, and clustering requires accommodating the differing characteristics of these tasks.

    Source: §3.3 · p. 5

  • Method

    NV-Embed combines latent attention pooling, non-causal attention, two-stage contrastive instruction tuning, and curated training data for retrieval and non-retrieval tasks.

    Source: §1 · p. 2

  • Results

    NV-Embed leads MTEB's retrieval, clustering, and classification categories, with NV-Embed-v2 ranking first overall at 72.31 as of August 30, 2024.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The data-curation recipe is intended to help the community reproduce and potentially surpass NV-Embed.

    Source: §1 · p. 2

    NV-Embed uses some MTEB training splits, including MSMARCO, as training data.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Decoder-only LLM-based embedding models are beginning to outperform BERT or T5-based embedding models in general-purpose text embedding tasks, including dense vector-based retrieval. In this work, we introduce NV-Embed, incorporating architectural designs, training procedures, and curated datasets to significantly enhance the performance of LLM as a versatile embedding model, while maintaining its simplicity and reproducibility. For model architecture, we propose a latent attention layer to obtain pooled embeddings, which consistently improves retrieval and downstream task accuracy compared to mean pooling or using the last <EOS> token embedding from LLMs. To enhance representation learning, we remove the causal attention mask of LLMs during contrastive training. For training algorithm, we introduce a two-stage contrastive instruction-tuning method. It first applies contrastive training with instructions on retrieval datasets, utilizing in-batch negatives and curated hard negative examples. At stage-2, it blends various non-retrieval into instruction tuning, which not only enhances non-retrieval task accuracy but also improves retrieval performance. For training data, we utilize the hard-negative mining, synthetic data generation and existing public available datasets to boost the performance of embedding model. By combining these techniques, our NV-Embed-v1 and NV-Embed-v2 models obtained the No.1 position on the MTEB leaderboard (as of May 24 and August 30, 2024, respectively) across 56 tasks, demonstrating the sustained effectiveness of the proposed methods over time. It also achieved the highest scores in the Long Doc section and the second-highest scores in the QA section of the AIR Benchmark, which covers a range of out-of-domain information retrieval topics beyond those in MTEB. We further provide the analysis of model compression techniques for generalist embedding models.

NVIDIA · 1 INTRODUCTION

NV-Embed improves decoder-only embedding models through architectural changes, two-stage contrastive instruction tuning, and curated training data. Its models reached first place on MTEB, while its compression approach improved accuracy and quantization robustness over directly built smaller embedding models.

1 INTRODUCTION

  • 72.31 on MTEB (Muennighoff et al., 2022) earned NV-Embed-v2 first place on August 30, 2024, following improved training-data curation.NV-Embed-v1 ranked first across 56 tasks on May 24, 2024; the model also achieved the highest scores across 15 retrieval, 11 clustering, and 12 classification tasks. On AIR-Benchmark, covering out-of-domain retrieval beyond MTEB, it ranked first for Long Doc and second for QA.Source: Table 1 · p. 2, §1 · p. 2
  • Latent attention pooling consistently improves retrieval and downstream-task accuracy over average pooling and last-<EOS>-token embeddings.The baselines are average pooling (Wang et al., 2022) and last-<EOS>-token embeddings (Neelakantan et al., 2022; Wang et al., 2023b).Source: §1 · p. 2
  • Removing the causal attention mask during contrastive training improves representation learning without adding a masked-token-prediction phase or mixed training objective.The authors describe their design as simpler and more effective than related approaches (BehnamGhader et al., 2024; Muennighoff et al., 2024).Source: §1 · p. 2
  • Two-stage contrastive instruction tuning first trains retrieval, then adds non-retrieval data, improving classification, clustering, semantic textual similarity, and retrieval.Training starts from pretrained Mistral-7B (Jiang et al., 2023), not an existing embedding model, using instructions, in-batch negatives, and curated hard negatives. Stage two disables in-batch negatives because they can mislead non-retrieval training.Source: §1 · p. 2
  • The training-data recipe details dataset curation, positive-aware hard-negative mining, synthetic data generation, and example-based multi-class labeling to support reproduction.Positive-aware hard-negative mining is intended to enhance contrastive training.Source: §1 · p. 2
  • Model compression achieves superior accuracy and quantization robustness compared with smaller embedding models built directly on Llama3.2-3B, Qwen2.5-3B, and Minitron-4B.The study examines pruning, quantization, and knowledge distillation for LLM-based embedding models.Source: §1 · p. 2

2 RELATED WORK

Related work spans supervised bidirectional encoders, decoder-only embeddings, and data-curation methods. NV-Embed distinguishes itself through public-data training, architectural changes, mixed-task batches, and positive-aware hard-negative mining.

2.1 BIDIRECTIONAL EMBEDDING MODELS

  • Bidirectional embedding models typically combine pretrained encoders, contrastive pretraining on unsupervised or weakly supervised pairs, and supervised fine-tuning for retrieval and other downstream tasks.Early examples include Sentence-BERT (Reimers & Gurevych, 2019) and SimCSE (Gao et al., 2021), which fine-tune BERT on natural language inference datasets.Source: §2.1 · p. 3

2.2 DECODER-ONLY LLM-BASED EMBEDDING MODELS

  • 66.63 MTEB is reported for E5-Mistral (Wang et al., 2023b), which combines task-specific contrastive instruction training with extensive proprietary GPT-4 synthetic data.Early decoder embedding work (Neelakantan et al., 2022) used contrastive training and final-layer <EOS> pooling; historical concerns included unidirectional attention and high-dimensional embeddings. LLM2Vec (BehnamGhader et al., 2024) uses only public data but achieves 65.01 MTEB, below E5-Mistral.Source: §2.2 · p. 3
  • NV-Embed trains directly from Mistral using public data, removes causal masking, adds latent-attention pooling, and mixes tasks within batches rather than relying on another embedding model.SFR-Embedding-Mistral and SFR-Embedding-2R instead fine-tune E5-Mistral on retrieval and non-retrieval datasets. Unlike SFR-Embedding-Mistral's task-homogeneous batches, NV-Embed's blended batches avoid potential zigzag gradient updates and achieve higher full-MTEB and retrieval scores.Source: §2.2 · p. 3
  • NV-Embed evaluates bge-en-icl (Li et al., 2024) in zero-shot conditions for fair comparison, rather than relying on its query-side few-shot examples.That model trains with both zero-shot and few-shot samples, with query-side examples introducing overhead. The survey also notes missing reproducibility details among many leading models and identifies GritLM (Muennighoff et al., 2024) as unifying embedding and generation.Source: §2.2 · pp. 3–4
  • NV-Embed adopts positive-aware hard-negative mining from NV-Retriever (Moreira et al., 2024), using positive relevance scores to eliminate false negatives and improve contrastive training.Gecko (Lee et al., 2024a) distills a smaller bidirectional model using synthetic pairs and LLM relabeling of retrieved positives and hard negatives. Linq-embed-mistral also uses LLMs for data generation, filtering, and negative mining.Source: §2.2 · p. 4

3 METHODS

NV-Embed combines bidirectional attention and latent-attention pooling with two-stage contrastive instruction-tuning. Training first focuses on retrieval, then blends retrieval and non-retrieval tasks while disabling in-batch negatives.

3.1 BIDIRECTIONAL ATTENTION

  • NV-Embed removes the causal attention mask during contrastive learning, enabling bidirectional attention in decoder-only LLMs.The method addresses the representation limitations of unidirectional attention through this simple architectural change.Source: §3.1 · p. 4

3.2 LATENT ATTENTION LAYER

  • Mean pooling may dilute important key-phrase information, while using the last <EOS> token embedding may introduce recency bias, motivating more expressive pooling.The latter approach relies heavily on the output embedding of the last token.Source: §3.2 · p. 5
  • Latent attention uses decoder hidden states as queries over a trainable dictionary, followed by an MLP and mean pooling to produce sequence embeddings.The layer uses 512 latents and eight attention heads, with two linear transformations separated by GELU in the MLP. Ablations report consistent improvements over normal self-attention.Source: §3.2 · p. 5

3.3 TWO-STAGE INSTRUCTION-TUNING

  • Two-stage contrastive instruction-tuning first trains retrieval with in-batch and curated hard negatives, then mixes retrieval and non-retrieval datasets without in-batch negatives.Retrieval receives initial emphasis because it is harder than the other tasks. In-batch negatives can mislead classification or clustering training because examples may share a class rather than be true negatives.Source: §3.3 · p. 5

4 TRAINING DATA

NV-Embed combines public retrieval and non-retrieval datasets with synthetic examples for contrastive training. Its data preparation includes positive-aware hard-negative mining, task-specific conversion of non-retrieval datasets, and synthetic task generation to broaden task coverage.

4 TRAINING DATA

  • Training combines public retrieval datasets, classification, clustering, semantic textual similarity, and synthetically generated examples.Instruction tokens are excluded from pooled embeddings during training and evaluation but still influence representations through self-attention; document corpora receive no instruction prefix.Source: §4 · pp. 5–6

4.1 PUBLIC RETRIEVAL DATASETS

  • Public retrieval training spans diverse datasets, including some MTEB training splits rather than an entirely benchmark-independent collection.The paper follows established generalist-embedding training practices and separately examines zero-shot generalization on AIR-bench.Source: §4.1 · p. 6

4.1.1 HARDNEGATIVE MINING TECHNIQUE

  • 95% of each positive’s score sets the maximum negative-score threshold in positive-aware hard-negative mining (Moreira et al., 2024).E5-mistral-7b-instruct (Wang et al., 2023b) serves as the teacher retrieval model; considering positive relevance scores helps remove false negatives.Source: §4.1.1 · p. 6

4.2 PUBLIC NON-RETRIEVAL DATASETS

  • Classification and clustering become contrastive examples, using label texts for binary tasks and same-class positives versus other-class negatives for multiclass classification and clustering.Queries use example text, and multiclass or clustering documents are randomly sampled from the appropriate classes or clusters. Emotion and AmazonCounterfactual (O’Neill et al., 2021) use BM25 filtering against evaluation data, while clustering removes evaluation matches.Source: §4.2 · pp. 6–7
  • Scores ≥4 yield both query–positive directions for semantic textual similarity training, with negatives mined from other texts and task instructions applied symmetrically.Training uses the training splits of STS12, STS22, and STS-Benchmark (Cer et al., 2017).Source: §4.2 · pp. 6–7

4.3 SYNTHETIC TASKS DATASET

  • 120,000 synthetic examples across 60,000 tasks are generated with Mixtral-8x22B-Instruct-v0.1 (MistralAI) to broaden task coverage.Generation adapts the two-step prompting approach of E5-mistral-7b-instruct (Wang et al., 2023b) for English, producing 40,000 examples each of short-long, long-short, and short-short pairs. Public STS data and the exclusion of bitext retrieval evaluation motivate this selection.Source: §4.3 · p. 7

5 EXPERIMENTS

NV-Embed-v2 scored 72.31 across 56 MTEB tasks, reclaiming first place as of August 30, 2024. Ablations support retrieval-first two-stage training, bidirectional attention, latent-attention pooling, example-based labels, and training-data curation.

5.1 MTEB RESULTS

  • 72.31 on MTEB across 56 tasks placed NV-Embed-v2 first on the leaderboard as of August 30, 2024.NV-Embed-v1 previously ranked first with 69.32 as of May 24, 2024. NV-Embed starts from Mistral 7B (Jiang et al., 2023), removing causal masking and adding latent-attention pooling rather than fine-tuning e5-mistral-7b-instruct (Wang et al., 2023b).Source: §5.1 · p. 7, Table 3 · p. 8, Table 2 · p. 8

5.2.3 POOLING METHODS

  • 62.65 versus 61.82 nDCG@10 on 15 retrieval tasks favored latent-attention over mean pooling.The added latent-attention layer is followed by an MLP before mean pooling and benefits most embedding tasks. The authors hypothesize that its latent array provides dictionary learning, yielding more expressive representations while mitigating information dilution from averaging.Source: Table 3 · p. 8, §5.2.3 · pp. 9–10

5.2.1 TWO-STAGE TRAINING

  • 72.31 versus 71.94 MTEB scores favored two-stage training over single-stage training without in-batch negatives.BEIR scores were 62.65 versus 61.37; two-stage training enables useful retrieval in-batch negatives, then disables them for non-retrieval tasks to avoid same-class false negatives. Unlike single-stage e5-mistral-7b-instruct (Wang et al., 2023b) and google-gecko (Lee et al., 2024a), NV-Embed prioritizes retrieval before blending other tasks.Source: §5.1 · p. 7, §5.2.1 · p. 8, Table 4 · p. 9
  • 71.85 overall MTEB for reversed two-stage training trailed retrieval-first training’s 72.31 and single-stage training without in-batch negatives’ 71.94.Retrieval scores rose from 61.37 to 61.91 but remained below retrieval-first training’s 62.65. Clustering and STS scores declined relative to single-stage training without in-batch negatives.Source: §5.2.1 · p. 9, Table 4 · p. 9

5.2.2 CAUSAL ATTENTION VS. BIDIRECTIONAL ATTENTION

  • Bidirectional attention consistently outperformed causal attention on average MTEB scores across 56 tasks for every pooling method tested.These comparisons support removing the causal attention mask when adapting decoder-only LLMs for embeddings.Source: Table 2 · p. 8, Table 3 · p. 8, §5.2.2 · p. 9

5.2.3 POOLING METHODS

  • 68.97 versus 69.10 MTEB scores for mean and self-attention pooling showed little benefit from adding self-attention.On 15 retrieval tasks, mean pooling scored 58.71 versus self-attention’s 58.64; mean pooling also consistently outperformed <EOS>-last across 56 MTEB tasks. The authors suggest recency bias harms <EOS>-last, while another self-attention layer adds little to an LLM already containing many.Source: Table 2 · p. 8, §5.2.3 · pp. 9–10

5.2.4 MULTI-CLASS CLASSIFICATION AND CLUSTERING LABELS

  • 11 clustering and 5 multiclass classification datasets used example-based documents, significantly improving both task categories over label-based training.Positives are sampled examples from the query’s class or cluster, while negatives come from other classes or clusters, rather than using label text. Random sampling provides broad coverage across labels and examples.Source: Table 4 · p. 9, §5.2.4 · p. 10, Table 5 · p. 10

5.2.5 HARDNEGATIVE MINING AND SYNTHETICALLY GENERATED DATASET

  • 59.22 to 61.52 BEIR improvement followed hard-negative mining, with additional public retrieval data and synthetic data providing further gains.Adding public retrieval datasets subsequently increased the retrieval score by 0.7 points. Synthetic data then increased the overall MTEB score by 0.24 points.Source: Table 4 · p. 9, §5.2.5 · p. 10

6 CONCLUSION · A COMPREHENSIVE STUDY OF MODEL COMPRESSION TECHNIQUES FOR NV-EMBED

NV-Embed combines architectural changes, two-stage contrastive instruction-tuning, and curated data to achieve No.1 MTEB rankings and superior out-of-domain AIR accuracy. The compression study reports improved accuracy and robustness through pruning and quantization, with performance exceeding smaller-model alternatives.

6 CONCLUSION

  • No.1 rankings on the MTEB leaderboard were achieved and maintained by the NV-Embed series, which also demonstrated superior accuracy on out-of-domain AIR tasks.The conclusion attributes these results to architectural changes, training procedures, and carefully curated data.Source: §6 · p. 10
  • NV-Embed introduces latent-attention pooling for expressive embeddings and removes the causal attention mask from its decoder-only LLM.The model is designed for general-purpose text embedding tasks.Source: §6 · p. 10
  • Two-stage contrastive instruction-tuning improves embedding tasks alongside curated datasets, hard-negative mining, synthetic data generation, and example-based multi-class labeling.The two training stages sequentially improve embedding tasks.Source: §6 · p. 10

A COMPREHENSIVE STUDY OF MODEL COMPRESSION TECHNIQUES FOR NV-EMBED

  • Pruning and quantization enhance embedding accuracy and robustness, outperforming smaller embedding models based on Llama3.2-3B, Qwen2.5-3B, and Minitron-4B.The study addresses computational and memory demands that limit deployment, scalability, and accessibility.Source: A COMPREHENSIVE STUDY OF MODEL… · p. 16
  • 3.5 billion parameters remain after pruning NV-Embed-v2 from 8 billion by pruning decoder-only blocks and removing the latent attention block.Pruning precedes quantization and continual retraining in the compression pipeline.Source: A COMPREHENSIVE STUDY OF MODEL… · p. 16
  • 8-bit weight quantization uses integer and floating-point E4M3/E5M2 formats, followed by LoRA continual retraining to restore accuracy.The compressed model is evaluated on MTEB (Muennighoff et al., 2022).Source: A COMPREHENSIVE STUDY OF MODEL… · p. 16

A.1 PRUNING

SparseGPT generally performs best among the evaluated pruning methods on MTEB and is selected for subsequent ablations. Semi-structured 2:4 pruning produces the lowest scores but shows the greatest accuracy recovery after retraining.

A.1 PRUNING

  • Magnitude-based pruning retains larger absolute weights, WANDA (Sun et al., 2023) considers weights and activations, and SparseGPT (Frantar & Alistarh, 2023) uses approximate-Hessian optimization.Source: §A.1 · p. 16
  • The study evaluates 2:4 and 4:8 semi-structured pruning alongside unstructured pruning, which removes individual weights.Semi-structured N:M sparsity retains N nonzero weights per group of M, reducing matrix size and computational cost while preserving regularity for efficient hardware utilization.Source: §A.1 · p. 16
  • 2:4 semi-structured pruning yields the lowest MTEB scores but achieves the greatest accuracy recovery following retraining.Source: §A.1 · p. 16
  • SparseGPT generally achieves the best MTEB results, while magnitude-based pruning and WANDA perform comparably both before and after retraining.These findings motivate using SparseGPT for subsequent pruning ablations.Source: Table 6 · p. 16, §A.1 · p. 16

A.2 KNOWLEDGE DISTILLATION

Knowledge distillation uses the uncompressed embedding model to teach its compressed counterpart, combining state-matching losses with contrastive training. Adding distillation consistently outperforms retraining with ground truth labels alone on MTEB across the evaluated compression approaches.

A.2 KNOWLEDGE DISTILLATION

  • Distillation matches student and teacher output and intermediate states using mean-squared error, with total loss Ltotal = Lcontrastive + α × Lkd.The weight α controls the knowledge distillation loss contribution.Source: §A.2 · pp. 16–17
  • The uncompressed embedding model serves as teacher for the compressed student, with knowledge distillation used to compensate for accuracy degradation from removing the latent attention block.Source: §A.2 · p. 17
  • 0.76 improvement on MTEB is achieved by 2:4 semi-structured pruning with knowledge distillation, the largest gain among the evaluated compression approaches.Table 7 reports that ground-truth-plus-distillation training (GT+KD) consistently outperforms ground-truth-only training (GT). Despite benefiting most from distillation, 2:4 semi-structured pruning yields the worst results among the approaches.Source: Table 7 · p. 17, §A.2 · p. 17

A.3 QUANTIZATION

Compressed NV-Embed models show greater robustness to INT8 quantization than smaller comparison models on MTEB. FP8 formats achieve MTEB scores comparable to INT8.

A.3 QUANTIZATION

  • Weight quantization uses GPTQ (Frantar et al., 2022), a post-training method that exploits approximate Hessian information to reduce weight precision.Evaluation compares compressed embedding models against Llama3.2-3B, Qwen2.5-3B, and Minitron-4B, which have different parameter counts.Source: §A.3 · p. 17
  • 8-bit floating-point formats E4M3 and E5M2 achieve MTEB scores comparable to INT8, with robustness observed across different 8-bit precision cases.E4M3 uses four exponent bits and three mantissa bits; E5M2 uses five and two, respectively, each with one sign bit. Floating-point formats provide a larger numerical range and non-uniform precision compared with integers.Source: §A.3 · p. 17
  • 0.0%, 0.01%, and 0.01% MTEB differences for INT8 NV-Embed’s 2:4 semi-structured, 4:8 semi-structured, and unstructured variants show greater robustness than smaller models.The corresponding smaller-model performance drops are Llama-3B (-0.47%), Qwen-3B (-0.14%), and Minitron-4B (-0.84%).Source: Table 8 · p. 17, §A.3 · p. 17

B AIR BENCHMARK

NV-Embed-v2 ranks first on Long-Doc and second on QA in AIR-Bench 24.04, improving over NV-Embed-v1 on both tasks. The benchmark assesses generalization across domains largely absent from MTEB.

B AIR BENCHMARK

  • AIR-Bench keeps ground truth confidential as a closed-book benchmark, with scores obtainable only through HuggingFace Hub.Source: B AIR BENCHMARK · p. 18
  • 70.31 versus 67.56 MTEB scores favor SFR-Embedding-2R over SFR-Embedding-Mistral, but both AIR-Bench tasks favor SFR-Embedding-Mistral.QA scores are 49.47 for SFR-Embedding-2R versus 51.58 for SFR-Embedding-Mistral. Long-Doc scores are 67.45 for SFR-Embedding-2R versus 69.0 for SFR-Embedding-Mistral, showing that higher MTEB performance need not transfer.Source: B AIR BENCHMARK · p. 18
  • 8 QA and 15 Long-Doc English datasets were evaluated in AIR-Bench 24.04, whose diverse domains are largely absent from MTEB.Domains include healthcare, law, news, books, arXiv, and finance, alongside samples synthetically generated using diverse LLMs.Source: B AIR BENCHMARK · p. 18
  • 74.78 Recall@10 gives NV-Embed-v2 the highest Long-Doc score; it also ranks second on QA nDCG@10 in AIR-Bench 24.04.NV-Embed-v2 improves over NV-Embed-v1 on both AIR-Bench tasks as well as MTEB. Its Long-Doc result surpasses Bge-en-icl, which adds in-context examples to queries during training.Source: Table 10 · p. 18, Table 9 · p. 18, B AIR BENCHMARK · p. 18

C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION · D LATENT-ATTENTION VISUALIZATION

The experimental setup specifies bidirectional attention, latent-attention pooling, two-stage contrastive training, and 512-token evaluation. Attention visualizations show a latent specializing in review positivity.

C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION

D LATENT-ATTENTION VISUALIZATION

  • Latent-attention visualization identifies a latent with high attention to positive reviews and low attention to negative reviews.Figure 2 visualizes eight heads for 10 positive and 10 negative AmazonReviewsClassification reviews, with attention weights mean-pooled across tokens.Source: Figure 2 · p. 20

Related papers

Loading 2405.17428v3…