NV-Embed: Técnicas Aprimoradas para Treinar LLMs como Modelos de Embedding Generalistas

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

May 27, 2024

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

cs.AI

Resumo

Modelos de embedding baseados em grandes modelos de linguagem (LLM) do tipo decoder-only estão começando a superar modelos de embedding baseados em BERT ou T5 em tarefas gerais de embedding de texto, incluindo recuperação baseada em vetores densos. Neste trabalho, apresentamos o modelo NV-Embed, que incorpora uma variedade de designs arquitetônicos e procedimentos de treinamento para aprimorar significativamente o desempenho de LLMs como modelos de embedding versáteis, mantendo sua simplicidade e reprodutibilidade. Para a arquitetura do modelo, propomos uma camada de atenção latente para obter embeddings agrupados, o que melhora consistentemente a precisão em tarefas de recuperação e tarefas subsequentes em comparação com o uso de pooling médio ou do embedding do último token <EOS> de LLMs. Para aprimorar o aprendizado de representações, removemos a máscara de atenção causal dos LLMs durante o treinamento contrastivo. Para o treinamento do modelo, introduzimos um método de ajuste fino por instrução contrastivo em duas etapas. Na primeira etapa, aplicamos treinamento contrastivo com instruções em conjuntos de dados de recuperação, utilizando negativos intra-lote e exemplos negativos difíceis selecionados. Na segunda etapa, integramos diversos conjuntos de dados não relacionados à recuperação no ajuste fino por instrução, o que não apenas melhora a precisão em tarefas não relacionadas à recuperação, mas também aprimora o desempenho em tarefas de recuperação. Combinando essas técnicas, nosso modelo NV-Embed, utilizando apenas dados publicamente disponíveis, alcançou uma pontuação recorde de 69,32, ocupando o primeiro lugar no Massive Text Embedding Benchmark (MTEB) (em 24 de maio de 2024), com 56 tarefas, abrangendo recuperação, reclassificação, classificação, agrupamento e tarefas de similaridade textual semântica. Notavelmente, nosso modelo também atingiu a pontuação mais alta de 59,36 em 15 tarefas de recuperação no benchmark MTEB (também conhecido como BEIR). Disponibilizaremos o modelo em código aberto em: 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.