NV-Embed: Улучшенные техники обучения LLM в качестве обобщенных моделей встраивания

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

May 27, 2024

Авторы: Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, Wei Ping

cs.AI

Аннотация

Модели встраивания на основе больших языковых моделей (LLM) только для декодера начинают превосходить модели встраивания на основе BERT или T5 в общих задачах встраивания текста, включая поиск на основе плотных векторов. В данной работе мы представляем модель NV-Embed с различными архитектурными конструкциями и процедурами обучения для значительного улучшения производительности LLM в качестве универсальной модели встраивания, сохраняя ее простоту и воспроизводимость. Для архитектуры модели мы предлагаем слой латентного внимания для получения объединенных встраиваний, что последовательно улучшает поиск и точность задач на уровне последующих этапов по сравнению с усреднением или использованием встраивания последнего токена <EOS> из LLM. Для улучшения обучения представлений мы удаляем маску причинного внимания LLM во время контрастного обучения. Для обучения модели мы представляем двухэтапный метод настройки инструкций контраста. Сначала он применяет контрастное обучение с инструкциями на наборах данных поиска, используя отрицательные примеры внутри пакета и отобранные сложные негативные примеры. На втором этапе он объединяет различные нерелевантные наборы данных в настройку инструкций, что не только улучшает точность задач без поиска, но также повышает производительность поиска. Комбинируя эти техники, наша модель NV-Embed, используя только общедоступные данные, достигла рекордного результата 69.32, занимая первое место в Massive Text Embedding Benchmark (MTEB) (на момент 24 мая 2024 года), с 56 задачами, включая поиск, переранжирование, классификацию, кластеризацию и задачи семантической текстовой схожести. Заметим, что наша модель также достигла наивысшего результата 59.36 на 15 задачах поиска в бенчмарке MTEB (также известном как BEIR). Мы планируем предоставить модель с открытым исходным кодом по ссылке: 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.