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

Authors: Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, Wei Ping
Primary area: foundation or frontier models, including LLMs
Keywords: llm, embedding model, retriever

Abstract:
Decoder-only large language model (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 the NV-Embed model, 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 model secured the No.1 position on the Massive Text Embedding Benchmark (MTEB) (as of May 24, 2024), across 56 embedding tasks. NV-Embed-v2 has reclaimed and maintained the top spot on MTEB since August 30, 2024, demonstrating the sustained effectiveness of the proposed methods over time.  Also, it 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.  We open-source the model at: https://huggingface.co/nvidia/NV-Embed-v2 .

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730589747468

thread_type:
official_review

<thread_object>
creator: reviewer_wm7z

<content>
summary:
This paper presents the NV-Embed-v1 and v2 systems for general purpose text embeddings. The paper presents the learning algorithm along with architecture variations (latent attention), ywhich describes how the system achieves #1 on the MTEB leaderboard at the time of submission. This paper documents the high performing system, which is likely to have gathered much attention given its position in the leaderboard.

strengths:
This paper severs a very clear and important purpose -- documenting the current top performing system on the extremely competitive text embedding benchmark, MTEB. The paper presents an approach, which like other approaches on the MTEB leaderboard, takes a pretrained transformer model and trains using a mixture of datasets/tasks, including a hard negative mining pipeline. The authors also introduce a simple architectural change, latent attention layer.

The key strength of the paper, is its empirical gains and empirical analysis. The authors provide benchmark performance on a wide range of settings with highly performing models. This shows the capabilities of models on a wide variety of tasks, including important ablations regarding attention type, stages of training, etc.

weaknesses:
As this paper is primarily the documentation of a very highly performing empirical system, main weakness I would point out is about innovation. On one level the paper is ground breaking because of its empirical gains, on another, the core methodological techniques are well known, only that they are more effectively performed and analyzed here.

The latent attention mechanism, while a nice architectural change, only has a small effect on the average performance (Table 2). The hard negative mining pipeline, a more established technique, on the other hand makes much more of a difference (Table 4).

questions:
For each stage of training, do you use the same optimizer? Restarting it for each stage?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710112106

thread_type:
official_review

<thread_object>
creator: reviewer_9j2v

<content>
summary:
This paper proposes a new embedding model: NV-embedder. The model uses a latent attention layer to obtain pooled embeddings instead of the common EOS token. The authors also propose a two-stage training method and employ hard example mining and data synthesis.

strengths:
- The authors tested the model on a leaderboard outside of MTEB. Since the training data is related to MTEB, out-of-domain testing is necessary. 
- The authors conducted extensive ablation studies to verify the effectiveness of each module.

weaknesses:
The biggest drawback is the complexity of the entire training process, which includes multi-stage training, hard example mining, and data synthesis. This involves too many operations, making it difficult to reproduce.

questions:
no
</content>
</thread_object>

</thread>
<thread>
creation date:
1730734693553

thread_type:
official_review

<thread_object>
creator: reviewer_rurm

<content>
summary:
This paper presents techniques for leveraging pretrained decoder-only large transformers in retrieval tasks, achieving state-of-the-art results on the standard retrieval benchmark (MTEB). The core methods include:
* A latent attention layer for creating pooled embeddings, which surpasses traditional mean pooling and last-token embedding approaches.
* A two-stage training process: the first stage focuses on retrieval datasets, while the second integrates non-retrieval tasks for broader versatility.
* The use of curated datasets (e.g. hard-negative mining) to further refine embedding quality.

strengths:
* Provides very strong empirical results. This work achieves state-of-the-art results, securing the #1 position on the MTEB benchmark. This demonstrates the effectiveness of the proposed approach.
* Provides nice ablation studies (e.g., Table 4) to analyze the effects of various design choices, including single-stage vs. two-stage, hard negative mining (HN), public retrieval set (AD), and synthetic data generation (SD).
* Provides detailed descriptions of the data curation process, including specific techniques like hard negative mining and synthetic data generation. This transparency allows for better understanding and potential reproducibility of the work.

weaknesses:
* Using decoder-only pretrained models for retrieval is not entirely novel (e.g., as seen in GritLM [1] and discussed in Section 2.2). While the authors note differences in their specific training approach, this limits the novelty of this aspect.
* Given this, the main contributions center on the latent attention layer and two-stage instruction tuning. However, the performance improvement from latent attention over mean pooling appears modest (see Table 2, bidirectional columns), raising questions about the added complexity for minimal gains. 
* Although the authors provide empirical support for the benefits of two-stage training, there is limited explanation or intuition behind why this approach works effectively.

[1] Muennighoff, Niklas, et al. "Generative representational instruction tuning." arXiv preprint arXiv:2402.09906 (2024).

questions:
* What would happen if the order of the two-stage training were reversed?
* Could you expand on the points raised in the weaknesses section, offering more discussion on these limitations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731084772544

thread_type:
official_review

<thread_object>
creator: reviewer_mueh

<content>
summary:
This paper gives a summary of the NV-Embed model that achieved the top performance in the MTEB benchmark. 
The techniques used are
1. latent attention layer that achieves better pooling/combination of the last layer embeddings. Causal attention mask is removed during contrastive learning. 
2. a two-stage contrastive instruction tuning method. First step tuning with in-batch negative and hard negative on retrieval datasets, and the second step tuning on non-retrieval datasets.
3. large amount of efforts on training data curation.

strengths:
1. This work gives a good summary of the STOA NV-Embed model that leads the MTEB benchmark. I think the community will very much appreciate this paper.
2. Impressive experimental results with good ablations to justify the design choices.
3. Clear presentation.

weaknesses:
The NV-Embed model is a result of a combination of methods/tricks/datasets etc. 
There does not seem to be single innovative algorithm piece. This by itself is not a weakness of the work. However, as a result, the technical depth of the work is limited.

questions:
The v2 model largely outperformed v1: 
"We then further improve the model through the curation
of training dataset, including adding more retrieval datasets, applying positive-aware hard-negative
mining technique, using synthetic data generation process and constructing example-based multi-class
labels." 
It would be nice to have more discussions and ablations specifically regarding the v2 vs. v1.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a suite of techniques to improve the general purpose embeddings produced by large decoder-only LLMs. This includes a latent attention layer to create the summary embedding, a contrastive training objective for instruction tuning, and careful creation of the training data. These techniques are proven to be empirically effective, with the resulting method securing the first place in the MTEB benchmark. The paper also presents a series of ablations to identify the importance of the individual components of their method.

Building versatile and powerful text embeddings is of wide interest and import. Reviewers were unanimously favourable in their view of the paper, citing its strong empirical results and ablations as being of considerable interest to the community. The one lacuna is that the technical novelty of the individual ingredients may be somewhat limited. However, as the solid empirical results (with technical details) are expected to be useful for most future work in this space, this is not seen as diminishing the value of the work.