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

Decoder-only 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 NV-Embed, incorporating architectural designs, training procedures, and curated datasets 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 lasttoken embedding from LLMs. To enhance representation learning, we remove the causal attention mask of LLMs during contrastive training. For training algorithm, 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 into instruction tuning, which not only enhances non-retrieval task accuracy but also improves retrieval performance. For training data, we utilize the hard-negative mining, synthetic data generation and existing public available datasets to boost the performance of embedding model. By combining these techniques, our NV-Embed-v1 and NV-Embed-v2 models obtained the No.1 position on the MTEB leaderboard (as of May 24 and August 30, 2024, respectively) across 56 tasks, demonstrating the sustained effectiveness of the proposed methods over time. It also achieved the highest scores in the Long Doc section and the second-highest scores in the QA section of the AIR Benchmark, which covers a range of out-of-domain information retrieval topics beyond those in MTEB. We further provide the analysis of model compression techniques for generalist embedding models.

한 줄 요약

NV-Embed는 디코더-온리 LLM을 활용한 다목적 임베딩 모델로, MTEB 벤치마크에서 69.32 점을 달성하며 1위를 기록했다.

핵심 기여도

**latent attention layer**를 도입하여 mean pooling 대비 retrieval 및 downstream task 정확도를 향상.
디코더-온리 LLM의 **causal attention mask 제거**로 대조 학습 시 표현 학습을 강화.
**2단계 대조 인스트럭션 튜닝**(retrieval 데이터 + non-retrieval 데이터)을 통해 retrieval 및 non-retrieval 성능을 동시에 향상.
**공개 데이터만 사용**하여 GPT-4 기반 합성 데이터 없이도 E5-mistral-7b-instruct 대비 +2.69% 개선.

핵심 아이디어

NV-Embed는 디코더-온리 LLM을 기반으로, 기존 임베딩 모델에서 사용되는 양방향 모델과 달리 단일 방향 구조를 유지하면서도 임베딩 성능을 향상시키는 새로운 접근법을 제시한다. 특히, **latent attention layer**를 통해 토큰 시퀀스의 표현을 효과적으로 풀링하여, 기존의 mean pooling이나 <EOS> 토큰 기반 임베딩보다 더 정확한 벡터를 생성한다. 또한, 대조 학습 중 **causal attention mask를 제거**함으로써, 모델이 텍스트 전체의 의미를 더 잘 포착하도록 유도한다.

이를 기반으로 **2단계 대조 인스트럭션 튜닝**을 도입한다. 1단계에서는 retrieval 데이터셋에 대해 in-batch negatives와 hand-crafted hard negatives를 활용한 대조 학습을 수행하고, 2단계에서는 non-retrieval 데이터를 추가하여 분류, 클러스터링, 유사도 등 다양한 임베딩 작업을 동시에 향상시킨다. 이는 기존 연구에서 볼 수 없었던 **단계별 학습 전략**으로, retrieval 성능을 유지하면서 non-retrieval 작업에서도 향상된 결과를 보인다.

기술적 접근법

**모델 아키텍처**: latent attention layer를 사용하여 풀링 임베딩 생성.
**학습 알고리즘**: 2단계 대조 인스트럭션 튜닝 (Stage 1: retrieval 데이터 + in-batch/hard negatives, Stage 2: non-retrieval 데이터 추가).
**데이터셋**: 공개 데이터만 사용 (GPT-4 기반 합성 데이터 미사용).
**하이퍼파라미터**: 미세 조정 없이 Mistral-7B 기반으로 학습.

주요 결과

**MTEB 벤치마크**에서 56개 작업에서 69.32 점 (2024년 5월 24일 기준 1위).
**BEIR (MTEB 내 15개 retrieval 작업)**에서 59.35 점 (1위).
E5-mistral-7b-instruct 대비 +2.69% 개선.
SFR-Embedding 대비 +1.76% 개선.
Voyage-large-2-instruct 대비 +1.04% 개선.

의의 및 한계

NV-Embed는 디코더-온리 LLM을 기반으로, 기존 양방향 모델을 대체할 수 있는 다목적 임베딩 모델로 제시하며, 특히 **공개 데이터만으로도 뛰어난 성능**을 보이는 점에서 실용적 가치가 높다. 또한, **2단계 대조 인스트럭션 튜닝**은 retrieval과 non-retrieval 작업을 모두 고려한 새로운 학습 전략으로, 임베딩 모델 개발에 새로운 방향성을 제시한다.

하지만, 학습 데이터가 **전적으로 공개 데이터에 의존**하므로, 특정 도메인에 최적화된 합성 데이터가 필요한 상황에서는 한계가 있을 수 있다. 또한, **모델 압축 기법**에 대한 분석은 제공되었으나, 실제 압축 모델의 성능은 명시되지 않았다.

실용적 활용

NV-Embed는 검색, 재랭킹, 분류, 클러스터링, 유사도 비교 등 다양한 자연어 처리 작업에서 활용 가능하다. 특히, **RAG**(Retrieval-Augmented Generation) 시스템에서 외부 지식을 효과적으로 검색하는 데 유용하며, **사용자 지정 학습 없이도 높은 성능**을 유지하므로, 데이터 접근성이 제한된 산업 현장에서도 실용적이다.