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

General-purpose embeddings need decoder-only LLMs that represent retrieval and non-retrieval tasks effectively while remaining reproducible. NV-Embed combines latent-attention pooling, mask removal, two-stage contrastive instruction tuning, and curated public and synthetic data; its models reached No. 1 on MTEB across 56 tasks and achieved strong AIR-Benchmark results.

  • Problem

    NV-Embed addresses how to improve decoder-only LLMs as versatile embedding models for retrieval and other downstream tasks while maintaining simplicity and reproducibility.

  • Method

    The method combines latent-attention pooling, removal of causal masking during contrastive training, two-stage instruction tuning, and curated public, hard-negative, synthetic, and non-retrieval data.

  • Results

    NV-Embed-v1 and NV-Embed-v2 ranked No. 1 on MTEB across 56 tasks, with NV-Embed-v2 reaching 72.31; the models also led AIR-Benchmark Long Doc and ranked second in QA.

  • Takeaways & Limitations

    The results support NV-Embed as a sustained generalist embedding approach across retrieval, clustering, classification, long-document, and out-of-domain QA evaluations.

  • Takeaways & Limitations

    Some training datasets, including MSMARCO, are training splits of MTEB, limiting the independence of those evaluations from the training data.

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.

1 INTRODUCTION

NV-Embed enhances decoder-only LLMs for general-purpose embedding through architectural changes, two-stage instruction tuning, and reproducible data curation. Its models reached the top of MTEB and performed strongly across retrieval, clustering, classification, long-document, and QA evaluations.

  • Contributions: NV-Embed introduces a latent attention pooling layer and removes the causal attention mask during contrastive training.The authors report consistent improvements over mean pooling or the last <EOS> embedding, with simpler training than related approaches.
  • Contributions: A two-stage contrastive instruction-tuning method first targets retrieval with in-batch and curated hard negatives, then blends non-retrieval tasks without in-batch negatives.The second stage improves classification, clustering, and semantic textual similarity while also enhancing retrieval.
  • Contributions: NV-Embed emphasizes reproducible training-data curation using public datasets, positive-aware hard-negative mining, synthetic generation, and example-based multi-class labeling.The paper presents this recipe as intended to help the community reproduce or surpass the model.
  • Evaluation: 72.31 was NV-Embed-v2’s MTEB score, reclaiming No. 1 on August 30, 2024, after NV-Embed-v1 ranked No. 1 across 56 tasks on May 24, 2024.NV-Embed also achieved the highest scores in 15 retrieval, 11 clustering, and 12 classification tasks, plus strong AIR-Benchmark section results.
  • Evaluation: The paper additionally studies pruning, quantization, and knowledge distillation for compressing LLM-based embedding models.Compared with smaller models built directly on Llama3.2-3B, Qwen2.5-3B, and Minitron-4B, the approach shows superior accuracy and quantization robustness.

2 RELATED WORK

Related work traces a shift from dominant BERT- and T5-based embedding models toward decoder-only LLMs. NV-Embed is positioned as a public-data, directly trained alternative with architectural and reproducibility distinctions from recent systems.

  • Bidirectional embedding models: BERT- and T5-based embedding models traditionally used contrastive learning after initialization from pretrained bidirectional encoders.Representative systems include Sentence-BERT and SimCSE.
  • Decoder-only embedding models: Decoder-only LLMs were historically viewed as weaker for embeddings because causal attention limits representation learning and large models produce very high-dimensional embeddings.Recent work nevertheless showed decoder-only models could outperform frontier bidirectional models.
  • Recent systems: E5-Mistral and later systems established decoder-only LLMs as competitive embedding backbones using contrastive training and mixtures of retrieval and non-retrieval data.The cited progression includes MTEB scores of 66.63, 67.56, and 70.31 for E5-Mistral, SFR-Embedding-Mistral, and SFR-Embedding-2R.
  • NV-Embed’s position: NV-Embed differs by training directly from Mistral 7B with public data rather than depending on an existing embedding model or proprietary synthetic data.The paper links this distinction to its new architecture and training choices.
  • Reproducibility: Many recent top-performing MTEB models have not disclosed key technical details, especially their training-data mixtures, motivating attention to reproducibility.The paper contrasts this with recently disclosed work.

3 METHODS

NV-Embed combines a latent-attention pooling architecture with two-stage contrastive instruction tuning. The method removes causal masking during contrastive training and adapts negative-sampling choices to retrieval versus non-retrieval tasks.

  • Representation learning: Causal masking is removed during contrastive training because unidirectional attention can limit representation power.The causal mask was originally designed for autoregressive next-token prediction.
  • Latent attention layer: Mean pooling can dilute key-phrase information, while the last <EOS> embedding can suffer from recency bias.The latent attention layer is introduced to address these pooling limitations.
  • Latent attention layer: The decoder-only LLM’s final hidden states serve as queries that attend to trainable latent keys and values, producing a richer sequence representation.The latent array acts as a trainable dictionary, with sequence length l, hidden dimension d, and r latent vectors.
  • Latent attention layer: After latent attention, an MLP with two linear transformations and GELU is followed by mean pooling to obtain the sequence embedding.The described configuration uses 512 latents and 8 attention heads.
  • Two-stage instruction tuning: Stage one uses in-batch and curated hard negatives for retrieval, whereas stage two mixes retrieval and non-retrieval data without in-batch negatives.In-batch negatives can be misleading when examples from the same class or cluster are not true negatives.

4 TRAINING DATA

The training recipe combines public retrieval and non-retrieval datasets with mined hard negatives and synthetic instruction data. Non-retrieval examples are converted into contrastive query-positive-negative formats for joint embedding training.

  • Data sources: Training uses public retrieval and non-retrieval datasets plus synthetic samples spanning retrieval, classification, clustering, and semantic textual similarity.These sources are formatted for contrastive embedding training.
  • Retrieval data: The retrieval collection includes datasets such as MSMARCO, HotpotQA, Natural Question, PAQ, Stack Exchange, SQuAD, BioASQ, FiQA, FEVER, SciFact, NFCorpus, and MIR.The paper notes that some datasets, including MSMARCO, overlap with MTEB training splits.
  • Retrieval data: Public retrieval datasets generally provide positive query-passage pairs but require mined hard negatives for contrastive learning.The paper applies positive-aware hard-negative mining to obtain these examples.
  • Non-retrieval data: Classification and clustering data are converted into query-positive-negative examples using labels or examples sampled from the same and different classes or clusters.Binary classification uses label texts as documents; multi-class and clustering tasks use sampled examples.
  • Non-retrieval data: For semantic similarity, text pairs with scores of at least 4 become bidirectional positive examples, while hard negatives are mined from other texts.Task instructions are appended to documents because the task is symmetric with the query.
  • Synthetic data: Mixtral-8x22B-Instruct-v0.1 generates 120,000 synthetic examples across 60,000 synthetic tasks to broaden task-wise generalization.The synthetic-data process follows a two-step prompting approach adapted from E5-mistral-7b-instruct.

5 EXPERIMENTS

NV-Embed is evaluated across 56 MTEB tasks, where staged training, bidirectional attention, latent-attention pooling, and curated data are examined through benchmark results and ablations. The models and design choices improve retrieval and broader embedding performance, with NV-Embed-v2 reaching a 72.31 MTEB score.

  • MTEB results: NV-Embed-v1 scored 69.32 and ranked No.1 on MTEB across 56 tasks, while NV-Embed-v2 reached 72.31 and reclaimed No.1.The reported rankings were dated May 24 and August 30, 2024, respectively.
  • Two-stage training: Two-stage training achieved 72.31 MTEB and 62.65 BEIR, compared with 71.94 and 61.37 for single-stage training without in-batch negatives.The proposed ordering enables in-batch negatives for retrieval in stage one while disabling them for non-retrieval tasks in stage two.
  • Attention masks: Bidirectional attention consistently outperformed causal masking across pooling types on average MTEB scores over 56 tasks.The comparison evaluates self-attention mask choices for decoder-only LLM embedding applications.
  • Pooling methods: Latent-attention pooling improved 15-task retrieval nDCG@10 from 61.82 with mean pooling to 62.65.The authors attribute the improvement to a more expressive latent-array representation that mitigates information dilution from averaging.
  • Multi-class labels: Example-based positives and negatives significantly improved both multi-class classification and clustering over label-based construction.The example-based approach was used for all 11 clustering datasets and 5 multi-class classification datasets.
  • Training-data curation: Hard-negative mining increased BEIR from 59.22 to 61.52 during stepwise training-data curation.Further curation added public retrieval datasets and synthetically generated data.

6 CONCLUSION

The paper presents NV-Embed as a decoder-only LLM embedding model combining architectural changes, two-stage training, and curated data. It reports sustained No.1 MTEB rankings and examines compression methods for deployment-related challenges.

  • Contributions: NV-Embed combines latent-attention pooling, removal of causal masking, two-stage contrastive instruction tuning, and curated training data.The data recipe includes hard-negative mining, synthetic generation, and example-based multi-class labeling.
  • Results: NV-Embed-v1 and NV-Embed-v2 achieved and maintained No.1 MTEB rankings while also demonstrating superior accuracy on out-of-domain AIR Benchmark tasks.The conclusion summarizes performance across diverse embedding tasks and beyond-MTEB evaluation.
  • Model compression: Compression addresses the computational and memory demands that limit the scalability and accessibility of LLM-based embedding models.The analysis focuses on post-training pruning and quantization for generalist embedding models.
  • Model compression: The compression pipeline prunes NV-Embed-v2 from 8 billion to 3.5 billion parameters, applies 8-bit quantization, and uses LoRA retraining to restore accuracy.Evaluation is conducted on the MTEB benchmark.

A.1 PRUNING

The pruning study compares magnitude-based, WANDA, and SparseGPT methods across semi-structured and unstructured sparsity. SparseGPT generally performs best, while 2:4 semi-structured pruning has the lowest scores but the greatest recovery after retraining.

  • Sparsity structures: Semi-structured 4:8 pruning removes four of every eight tensor elements, reducing matrix size and computational cost while preserving hardware-friendly regularity.The paper contrasts this with unstructured pruning, which removes individual network elements.
  • Pruning methods: SparseGPT generally delivers the best MTEB results among the evaluated pruning techniques.Magnitude-based and WANDA produce comparable performance during pruning and after retraining.
  • Pruning methods: 2:4 semi-structured pruning yields the lowest MTEB scores among the evaluated sparsity settings.The study evaluates semi-structured 2:4 and 4:8 sparsity alongside unstructured pruning.
  • Retraining effects: 2:4 semi-structured pruning shows the greatest accuracy recovery following retraining on MTEB benchmarks.This recovery pattern motivates focusing on SparseGPT pruning for subsequent ablation studies.

A.2 KNOWLEDGE DISTILLATION

Knowledge distillation transfers behavior from the uncompressed embedding model to a compressed student during accuracy recovery. The study combines state-level distillation with contrastive loss and finds consistent gains over ground-truth-only retraining.

  • Distillation objective: Mean-squared-error losses match both output and intermediate student states to the corresponding teacher states.The losses cover the output state and intermediate states across the model layers.
  • Distillation setup: The uncompressed embedding model serves as teacher, while the compressed model serves as student during knowledge distillation.The latent attention block is removed from the compressed model, and distillation compensates for the resulting accuracy degradation.
  • Distillation objective: The total training loss combines contrastive loss with a weighted knowledge-distillation loss.The weighting term is denoted by α.
  • Results: Knowledge distillation consistently outperforms ground-truth-only retraining across evaluated MTEB compression approaches.The strongest reported gain is 0.76 on MTEB for 2:4 semi-structured pruning.

A.3 QUANTIZATION

The quantization study evaluates GPTQ-based low-precision compressed embedding models against smaller LLM-based embedding models. NV-Embed retains MTEB performance more robustly under INT8 quantization, while FP8 formats remain comparable to INT8.

  • Method: GPTQ is used as a post-training weight-quantization method that reduces weight precision using approximate Hessian information.The compressed models are evaluated on the MTEB benchmark.
  • Quantization results: NV-Embed quantized to INT8 shows MTEB drops of 0.0% for 2:4 semi-structured, 0.01% for 4:8 semi-structured, and 0.01% for unstructured pruning.These drops are reported for the compressed NV-Embed models after quantization.
  • Quantization results: Smaller models show larger INT8 MTEB drops: Llama-3B -0.47%, Qwen-3B -0.14%, and Minitron-4B -0.84%.The comparison includes Llama3.2-3B, Qwen2.5-3B, and Minitron-4B embedding models.
  • Quantization formats: E4M3 and E5M2 8-bit floating-point formats achieve MTEB scores comparable to INT8.E4M3 uses a 4-bit exponent and 3-bit mantissa, whereas E5M2 uses a 5-bit exponent and 2-bit mantissa.

B AIR BENCHMARK

AIR-Bench evaluates embedding-model generalization on diverse information-retrieval domains beyond most MTEB samples. NV-Embed-v2 ranks second in QA and first in Long-Doc, with a Long-Doc score of 74.78.

  • Benchmark scope: AIR-Bench 24.04 contains QA and Long-Doc tasks spanning healthcare, law, news, books, arXiv, finance, and synthetic samples.The evaluation covers 8 English QA datasets and 15 English Long-Doc datasets.
  • Results: NV-Embed-v2 achieves the second-highest score in the AIR-Bench QA section.The QA evaluation uses nDCG@10 scores.
  • Results: 74.78 is NV-Embed-v2’s highest reported Long-Doc score on AIR-Bench 24.04.It surpasses Bge-en-icl, which requires adding in-context examples to queries during training.
  • Cross-benchmark comparison: NV-Embed-v2 improves over NV-Embed-v1 on both AIR-Bench QA and Long-Doc tasks.The passage contrasts this with cases where higher MTEB performance does not correspond to better AIR-Bench performance.

C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION

The section specifies NV-Embed’s architecture, optimization setup, evaluation constraints, instruction resources, benchmark tables, prompt templates, and latent-attention visualization.

  • C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION: NV-Embed uses Mistral 7B with LoRA, bidirectional attention, and a latent attention layer containing 512 latents, 4096 hidden dimensions, and 8 heads.The causal attention mask is replaced with a bidirectional mask during finetuning.
  • C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION: The model is trained end-to-end with contrastive loss using LoRA rank 16, alpha 32, dropout 0.1, Adam, linear decay, and a batch size of 128.Learning rates are 2e-5 for stage one and 1.5e-5 for stage two, with optimizer restart and 50 warm-up steps in stage two.
  • C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION: Evaluation uses a maximum sequence length of 512 tokens to ensure fair comparisons with prior work.Evaluation instruction templates are provided in Table 13.
  • C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION: Table 11 records the parameters used in the experiments.The table is identified as the parameter reference for the experimental setup.
  • C EXPERIMENTAL DETAILS AND INSTRUCTION TEMPLATES FOR TRAINING AND EVALUATION: Table 12 lists the instructions and sample counts for each training dataset.It provides the training-data instruction reference alongside the experimental setup.
  • D LATENT-ATTENTION VISUALIZATION: Figure 2 visualizes attention over 4096 latents across 8 heads for positive and negative AmazonReviewsClassification examples.The visualization indicates that latents specialize in query features; one highlighted latent attends strongly to positive reviews and weakly to negative reviews.
  • D LATENT-ATTENTION VISUALIZATION: Table 13 provides MTEB evaluation instructions and defines “STS*” as shared instructions across all STS tasks.The table is specifically used for evaluation on the MTEB benchmark.
  • D LATENT-ATTENTION VISUALIZATION: Tables 14–16 provide the full BEIR and MTEB benchmark, short-long and long-short matching prompt templates.The benchmark table and the two subgroup prompt-template tables document evaluation resources.
Loading 2405.17428v3…