[2025-1] 정유림 - NV-Embed : Improved Techniques for Training Decoder Models for General Text Embeddings

논문 개요

• 논문 제목: NV-Embed: Improved Techniques for Training Decoder Models for General Text Embeddings
• 게재 연도: 2024년
• 인용 횟수: 2025.01.10 기준 54회 인용
• 주요 성과:
  • 디코더 모델의 representation 성능 한계 극복(Decoder 기반 text embedding model)
  • MTEB leader board에서 SOTA 달성(2024.05.22 기준 MTEB 결과 SOTA)
• DOI: https://doi.org/10.48550/arXiv.2405.17428

---

2. 연구 배경

• Decoder model의 한계:
  • 기존 디코더 기반 모델은 단방향(Unidirectional) Attention을 사용하여 Representation 학습 과정에서 정보 손실 발생.
  • 고차원 임베딩의 정보 희석 가능성 증가.
  • 텍스트 임베딩 작업에서는 인코더 기반 모델(BERT, T5 등) 대비 성능이 부족하다는 지적.
• 최근 동향:
  • 디코더 기반 모델 개선을 위한 새로운 방법론(E5, Mistral, Llama2 등)이 등장하는 추세.

---

3. NV-Embed의 주요 특징 요약

1. Bidirectional Attention:
   • 기존 디코더 모델이 사용하는 causual attention mask를 제거하여 Bidirectional attention으로 학습을 진행.
   • 모든 토큰이 서로 영향을 주고받아 더 나은 representation 가능.
2. Latent Attention Layer 추가:
   • 기존의 pooling 방식(mean pooling, EOS token 사용)이 가진 정보 희석 문제를 극복.
   • 기존 디코더 출력값을 Latent Query로 사용
   • 학습가능한 Key, Value를 통해 더 풍부한 representation 생성.
   • 최종 출력값 : softmax를 이용한 output→ MLP(Multi-Layer Perceptron)→ mean pooling 을 통한 최종 embedding 추출.
3. Contrastive Learning (Two-Stage Contrastive Instruction-Tuning):
   • 동일한 시퀀스는 유사성을 최대화하고, 다른 시퀀스 간 유사성은 최소화.
   • 텍스트 임베딩 품질을 높이기 위해 Two stage 학습 전략 도입:
     • (Stage 1) Retrieval 데이터셋 기반 학습. Hard negative sampling 및 In-batch negative trick사용.
     • (Stage 2) Non-retrieval 데이터셋(분류,클러스터링,의미적 유사도(STS)) 기반 학습. In-batch negative 비활성화.
4. 공개 데이터셋 활용:
   • 공개 데이터셋(MTEB 등)만으로 학습.

(참고) PEFT method중에서 LoRA(rank 16)를 사용해서 Finetuning. base decoder only LLM으로 는 Mistal 7B를 사용함.

---

4. 실험 및 결과(2024.05.22 기준 MTEB 결과 SOTA)

1. M1 리더보드 결과:

• 평균 점수: 69.3 (최고 성능).
• 15개 검색 작업 (BEIR) 성능: 59.30 (최고 성능 달성).

---

 

 

 

 

 

 

2. Ablation study :

• Two-stage Training: 단일 스테이지보다 성능 개선 확인.

• Bidirectional Attention vs. Causal Attention : 모든 task에서 Bidirectional Attention 이 우수.

• Latent Attention Layer vs. 기존 방식  : Latent Attention Layer가 더 나은 임베딩 제공.

5. Reference

• NV-Embed: Improved Techniques for Training Decoder Models for General Text Embeddings (2024)
• MTEB leaderboard