NV-Embed: Verbesserte Techniken zur Schulung von LLMs als Embedding-Modellen für Allgemeinanwendungen

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

May 27, 2024

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

cs.AI

papers.abstract

Decoder-only Large Language Model (LLM)-basierte Einbettungsmodelle beginnen, BERT- oder T5-basierte Einbettungsmodelle in allgemeinen Texteinbettungsaufgaben, einschließlich dichter Vektor-basierter Rückgewinnung, zu übertreffen. In dieser Arbeit stellen wir das NV-Embed-Modell mit einer Vielzahl von architektonischen Designs und Trainingsverfahren vor, um die Leistung des LLM als vielseitiges Einbettungsmodell signifikant zu verbessern, während wir seine Einfachheit und Reproduzierbarkeit beibehalten. Für die Modellarchitektur schlagen wir eine latente Aufmerksamkeitsschicht vor, um gepoolte Einbettungen zu erhalten, die im Vergleich zu Mittelwert-Pooling oder der Verwendung der letzten <EOS>-Token-Einbettung aus LLMs die Rückgewinnungs- und Genauigkeit der nachgelagerten Aufgaben kontinuierlich verbessert. Um die Repräsentationslernen zu verbessern, entfernen wir die kausale Aufmerksamkeitsmaske von LLMs während des kontrastiven Trainings. Für das Modelltraining führen wir eine zweistufige kontrastive Anweisungsabstimmungsmethode ein. Zunächst wird kontrastives Training mit Anweisungen auf Rückgewinnungsdatensätzen durchgeführt, wobei In-Batch-Negative und kuratierte harte negative Beispiele verwendet werden. In Stufe 2 werden verschiedene Nicht-Rückgewinnungsdatensätze in die Anweisungsabstimmung integriert, was nicht nur die Genauigkeit von Nicht-Rückgewinnungsaufgaben verbessert, sondern auch die Rückgewinnungsleistung steigert. Durch die Kombination dieser Techniken hat unser NV-Embed-Modell unter Verwendung nur öffentlich verfügbarer Daten einen Rekordwert von 69,32 erzielt und belegt damit den ersten Platz im Massive Text Embedding Benchmark (MTEB) (Stand 24. Mai 2024) mit 56 Aufgaben, einschließlich Rückgewinnung, Neusortierung, Klassifizierung, Clustering und semantischer Textähnlichkeitsaufgaben. Beachtenswert ist, dass unser Modell auch den höchsten Wert von 59,36 bei 15 Rückgewinnungsaufgaben im MTEB-Benchmark (auch bekannt als BEIR) erreicht. Wir werden das Modell unter folgendem Link als Open Source veröffentlichen: 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.