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
2024-05-28
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 with a variety of architectural designs and training procedures 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 model training, 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 datasets into instruction tuning, which not only enhances non-retrieval task accuracy but also improves retrieval performance. Combining these techniques, our NV-Embed model, using only publicly available data, has achieved a record-high score of 69.32, ranking No. 1 on the Massive Text Embedding Benchmark (MTEB) (as of May 24, 2024), with 56 tasks, encompassing retrieval, reranking, classification, clustering, and semantic textual similarity tasks. Notably, our model also attains the highest score of 59.36 on 15 retrieval tasks in the MTEB benchmark (also known as BEIR). We will open-source the model at: <a class="link-external link-https" href="https://huggingface.co/nvidia/NV-Embed-v1" rel="external noopener nofollow">this https URL</a>.
Computation and Language,Artificial Intelligence,Information Retrieval,Machine Learning
What problem does this paper attempt to address?
The main problem that this paper attempts to solve is to improve the performance of decoder - only large language models (LLMs) in general - purpose text embedding tasks, especially in dense vector retrieval. Specifically, the authors propose a new model named NV - Embed. Through a series of improvements in architecture design and training methods, decoder - only LLMs can significantly outperform traditional bidirectional language models (such as BERT or T5) in general - purpose text embedding tasks. ### Main Contributions 1. **Model Architecture Improvement**: - A new latent attention layer is proposed to obtain pooled embeddings from the tokens of a sequence. Compared with the commonly used average pooling and using the last <EOS> token embedding, this method performs better in retrieval and other downstream tasks. - The causal attention mask in decoder - only LLMs during contrastive training is removed to enhance the representation learning ability. 2. **Training Method Improvement**: - A two - stage contrastive instruction - tuning method is introduced. In the first stage, contrastive training is carried out on the retrieval dataset, using in - batch negative samples and carefully selected hard negative samples. In the second stage, non - retrieval datasets are mixed into the training data without using in - batch negative samples to improve the accuracy of classification, clustering, and semantic text similarity tasks, while unexpectedly enhancing the retrieval performance. 3. **Performance Improvement**: - Combining the above techniques, the NV - Embed model has achieved the highest score of 69.32 in 56 tasks of the Massive Text Embedding Benchmark (MTEB), ranking No. 1 (as of May 24, 2024). In particular, in 15 retrieval tasks, NV - Embed has reached the highest score of 59.36. ### Background and Motivation - **Existing Problems**: Traditional bidirectional embedding models (such as BERT or T5) perform well in general - purpose text embedding tasks, but rely on a large amount of privately synthesized data for fine - tuning, which limits their wide application in the community. - **Solution**: The NV - Embed model is trained entirely with publicly available data, without relying on any privately synthesized data, thus providing a more transparent and reproducible solution. ### Experimental Results - **MTEB Benchmark**: NV - Embed performs well in multiple subtasks, including retrieval, re - ranking, classification, clustering, and semantic text similarity tasks. - **Ablation Study**: The experimental results show that the use of bidirectional attention masks and latent attention layers significantly improves the model performance. ### Conclusion The NV - Embed model significantly improves the performance of decoder - only LLMs in general - purpose text embedding tasks, especially in dense vector retrieval, through a series of innovative architecture designs and training methods. These improvements not only enhance the accuracy of the model but also ensure the simplicity and reproducibility of the model.