NV-Embed: Verbeterde Technieken voor het Trainen van LLMs als Algemene Embeddingmodellen

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

Samenvatting

Decoder-only grote taalmodellen (LLM)-gebaseerde embeddingmodellen beginnen BERT- of T5-gebaseerde embeddingmodellen te overtreffen in algemene tekstembeddingstaken, inclusief dense vector-gebaseerde retrieval. In dit werk introduceren we het NV-Embed-model met een verscheidenheid aan architectonische ontwerpen en trainingsprocedures om de prestaties van LLM als een veelzijdig embeddingmodel aanzienlijk te verbeteren, terwijl de eenvoud en reproduceerbaarheid behouden blijven. Voor de modelarchitectuur stellen we een latente aandachtlaag voor om gepoolde embeddings te verkrijgen, wat consistent de retrieval- en downstreamtaaknauwkeurigheid verbetert in vergelijking met mean pooling of het gebruik van de laatste <EOS>-tokenembedding van LLM's. Om de representatielearning te verbeteren, verwijderen we het causale aandachtmasker van LLM's tijdens contrastieve training. Voor de modeltraining introduceren we een tweefasen contrastieve instructieafstemmingsmethode. Deze past eerst contrastieve training toe met instructies op retrievaldatasets, waarbij gebruik wordt gemaakt van in-batch negatieven en geselecteerde harde negatieve voorbeelden. In fase 2 worden verschillende niet-retrievaldatasets gemengd in de instructieafstemming, wat niet alleen de nauwkeurigheid van niet-retrievaltaken verbetert, maar ook de retrievalprestaties. Door deze technieken te combineren, heeft ons NV-Embed-model, met alleen openbaar beschikbare data, een recordscore van 69,32 behaald, wat het op de eerste plaats brengt op de Massive Text Embedding Benchmark (MTEB) (per 24 mei 2024), met 56 taken, waaronder retrieval, reranking, classificatie, clustering en semantische tekstuele gelijkenistaken. Opmerkelijk is dat ons model ook de hoogste score van 59,36 behaalt op 15 retrievaltaken in de MTEB-benchmark (ook bekend als BEIR). We zullen het model opensourcen op: 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.