NV-Embed: Tecniche Avanzate per l'Addestramento di LLM come Modelli di Embedding Generalisti

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

May 27, 2024

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

cs.AI

Abstract

I modelli di embedding basati su grandi modelli linguistici (LLM) di tipo decoder-only stanno iniziando a superare i modelli di embedding basati su BERT o T5 nelle attività generali di embedding del testo, inclusa la ricerca basata su vettori densi. In questo lavoro, introduciamo il modello NV-Embed con una varietà di design architetturali e procedure di addestramento per migliorare significativamente le prestazioni degli LLM come modelli di embedding versatili, mantenendo al contempo la loro semplicità e riproducibilità. Per l'architettura del modello, proponiamo un livello di attenzione latente per ottenere embedding aggregati, che migliora costantemente l'accuratezza nella ricerca e nelle attività downstream rispetto al pooling medio o all'uso dell'embedding del token <EOS> finale degli LLM. Per migliorare l'apprendimento delle rappresentazioni, rimuoviamo la maschera di attenzione causale degli LLM durante l'addestramento contrastivo. Per l'addestramento del modello, introduciamo un metodo di instruction-tuning contrastivo in due fasi. Nella prima fase, si applica l'addestramento contrastivo con istruzioni su dataset di ricerca, utilizzando negativi in-batch ed esempi negativi difficili selezionati. Nella seconda fase, si integrano vari dataset non di ricerca nell'instruction tuning, il che non solo migliora l'accuratezza nelle attività non di ricerca, ma anche le prestazioni nella ricerca. Combinando queste tecniche, il nostro modello NV-Embed, utilizzando solo dati pubblicamente disponibili, ha raggiunto un punteggio record di 69.32, classificandosi al primo posto nel Massive Text Embedding Benchmark (MTEB) (al 24 maggio 2024), con 56 attività che includono ricerca, riordinamento, classificazione, clustering e attività di similarità semantica testuale. In particolare, il nostro modello ottiene anche il punteggio più alto di 59.36 su 15 attività di ricerca nel benchmark MTEB (noto anche come BEIR). Renderemo il modello open-source all'indirizzo: 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.