NV-Embed : Techniques améliorées pour entraîner des LLM en tant que modèles d'embedding généralistes

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

May 27, 2024

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

cs.AI

Résumé

Les modèles d'embedding basés sur des grands modèles de langage (LLM) de type decoder-only commencent à surpasser les modèles d'embedding basés sur BERT ou T5 dans les tâches générales d'embedding de texte, y compris la recherche basée sur des vecteurs denses. Dans ce travail, nous présentons le modèle NV-Embed, qui intègre diverses conceptions architecturales et procédures d'entraînement pour améliorer significativement les performances des LLM en tant que modèles d'embedding polyvalents, tout en conservant leur simplicité et reproductibilité. Pour l'architecture du modèle, nous proposons une couche d'attention latente pour obtenir des embeddings regroupés, ce qui améliore systématiquement la précision de la recherche et des tâches en aval par rapport au regroupement moyen ou à l'utilisation de l'embedding du dernier token <EOS> des LLM. Pour améliorer l'apprentissage des représentations, nous supprimons le masque d'attention causal des LLM pendant l'entraînement contrastif. Pour l'entraînement du modèle, nous introduisons une méthode d'ajustement par instruction contrastive en deux étapes. La première étape applique un entraînement contrastif avec des instructions sur des ensembles de données de recherche, en utilisant des négatifs intra-lot et des exemples négatifs difficiles sélectionnés. À la deuxième étape, elle intègre divers ensembles de données non liés à la recherche dans l'ajustement par instruction, ce qui améliore non seulement la précision des tâches non liées à la recherche, mais aussi les performances de recherche. En combinant ces techniques, notre modèle NV-Embed, utilisant uniquement des données publiquement disponibles, a atteint un score record de 69,32, se classant premier au Massive Text Embedding Benchmark (MTEB) (au 24 mai 2024), avec 56 tâches couvrant la recherche, le reranking, la classification, le clustering et les tâches de similarité textuelle sémantique. Notamment, notre modèle obtient également le score le plus élevé de 59,36 sur 15 tâches de recherche dans le benchmark MTEB (également connu sous le nom de BEIR). Nous rendrons le modèle open-source à l'adresse suivante : https://huggingface.co/nvidia/NV-Embed-v1.

English

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: https://huggingface.co/nvidia/NV-Embed-v1.