NV-Embed: Técnicas Mejoradas para Entrenar Modelos de Lenguaje como Modelos de Incrustación 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

Resumen

Los modelos de embeddings basados en grandes modelos de lenguaje (LLM) de solo decodificador están comenzando a superar a los modelos de embeddings basados en BERT o T5 en tareas generales de generación de embeddings de texto, incluyendo la recuperación basada en vectores densos. En este trabajo, presentamos el modelo NV-Embed, que incorpora una variedad de diseños arquitectónicos y procedimientos de entrenamiento para mejorar significativamente el rendimiento de los LLM como modelos de embeddings versátiles, manteniendo su simplicidad y reproducibilidad. Para la arquitectura del modelo, proponemos una capa de atención latente para obtener embeddings agrupados, lo que mejora consistentemente la precisión en tareas de recuperación y otras tareas relacionadas en comparación con el agrupamiento por media o el uso del embedding del último token <EOS> de los LLM. Para mejorar el aprendizaje de representaciones, eliminamos la máscara de atención causal de los LLM durante el entrenamiento contrastivo. En cuanto al entrenamiento del modelo, introducimos un método de ajuste por instrucciones contrastivo en dos etapas. En la primera etapa, se aplica entrenamiento contrastivo con instrucciones en conjuntos de datos de recuperación, utilizando negativos dentro del lote y ejemplos negativos difíciles seleccionados. En la segunda etapa, se integran diversos conjuntos de datos no relacionados con la recuperación en el ajuste por instrucciones, lo que no solo mejora la precisión en tareas no relacionadas con la recuperación, sino que también refuerza el rendimiento en tareas de recuperación. Combinando estas técnicas, nuestro modelo NV-Embed, utilizando únicamente datos disponibles públicamente, ha alcanzado una puntuación récord de 69.32, posicionándose en el primer lugar en el Massive Text Embedding Benchmark (MTEB) (al 24 de mayo de 2024), con 56 tareas que abarcan recuperación, reordenamiento, clasificación, agrupamiento y similitud semántica de textos. Cabe destacar que nuestro modelo también obtiene la puntuación más alta de 59.36 en 15 tareas de recuperación dentro del benchmark MTEB (también conocido como BEIR). El modelo estará disponible en código abierto en: 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.