NV-Embed

Abstract

• 24년 5월 30일 현재 MTEB 기준 1위
• Latent attention layer를 통해 pooled embedding을 얻어 retrieval 및 downstream task 성능 향상
• Contrastive 학습 시 causal attention mask를 제거하여 representation learning 개선
• Retrieval dataset에 대한 instruction을 활용한 2단계 contrastive instruction-tuning 방식

Introduction

최근에는 쌍방향 언어 모델(bidirectional language models)에 기반한 임베딩 모델이 주를 이루었으나, 최신 연구에 따르면 디코더 전용 대형 언어 모델(decoder-only LLMs)이 검색 및 범용 임베딩 작업에서 더 나은 성능을 보이는 것으로 나타났습니다.

하지만 이전의 선도적인 연구들은 GPT-4와 같은 대규모 독점 합성 데이터를 사용하여 LLM을 미세 조정(fine-tuning)하는 데 의존했기 때문에 일반적으로 활용하기 어려웠습니다.

이에 따라 저자들은 NV-Embed라는 새로운 범용 임베딩 모델을 제안합니다. NV-Embed는 다음과 같은 주요 특징을 가지고 있습니다.

1. 모델 구조 측면에서, 토큰 시퀀스에 대한 pooled embedding을 얻기 위해 새로운 latent attention layer를 제안했습니다. 이는 기존의 평균 풀링이나 마지막 토큰 임베딩 방식보다 검색 및 하위 작업의 정확도를 일관되게 향상시켰습니다. 또한, 디코더 전용 LLM의 대조 학습(contrastive training) 중 인과 어텐션 마스크(causal attention mask)를 제거하여 표현 학습을 강화했습니다.
2. 모델 학습 측면에서는 2단계 대조 지시 튜닝(two-stage contrastive instruction-tuning) 방법을 도입했습니다. 1단계에서는 검색 데이터셋에 대한 지시문(instruction)과 함께 대조 학습을 적용하고, 2단계에서는 비검색 데이터셋을 1단계 학습 데이터와 혼합합니다. 이는 분류, 클러스터링, 의미적 유사성 작업의 정확도를 높일 뿐만 아니라 검색 성능도 향상시키는 것으로 나타났습니다.
3. 이러한 기술들을 모두 결합한 NV-Embed 모델은 56개의 임베딩 작업으로 구성된 Massive Text Embedding Benchmark (MTEB)에서 새로운 최고 점수인 69.32점을 기록하며 1위를 차지했습니다(2024년 5월 22일 기준). 특히 15개의 검색 작업에서 59.35점으로 가장 높은 점수를 얻었습니다.

Bidirectional Embedding Models

• BERT나 T5 기반의 임베딩 모델이 오랫동안 범용 임베딩 작업에서 지배적인 접근법이었습니다.
• 초기 모델로는 SentenceBERT와 SimCSE 등이 있으며, 이들은 사전 학습된 BERT를 NLI 데이터셋에 대해 미세 조정합니다.
• 일반적으로 이러한 임베딩 모델은 사전 학습된 BERT나 T5 인코더로 초기화된 후, curated unsupervised 또는 weakly-supervised 텍스트 쌍에 대한 대조 학습을 통해 추가로 사전 학습됩니다.
• 마지막으로 MS MARCO 등 다양한 supervised 데이터에 대해 미세 조정되어 검색 및 기타 하위 작업에 사용됩니다.
• 최신 모델로는 mxbai-embed-large-v1, UAE-Large-V1, voyage-large-2-instruct 등이 있습니다.

Decoder-only LLM-based Embedding Models

• Decoder-only LLMs은 일반적으로 bidirectional 모델보다 범용 임베딩 작업에서 성능이 낮을 것으로 여겨졌습니다. 그 이유는 단방향 어텐션이 표현 학습 능력을 제한하고, LLM의 스케일링으로 인해 매우 고차원의 임베딩이 생성되어 차원의 저주 문제가 발생할 수 있기 때문입니다.
• Neelakantan et al. (2022)은 사전 학습된 GPT-3 모델로 임베딩 모델을 초기화하고 지속적인 대조 학습을 적용했습니다. 시퀀스 끝에 있는 특수 토큰에 해당하는 마지막 레이어의 hidden state가 입력 시퀀스의 임베딩으로 사용됩니다.
• 최근 E5-Mistral은 Mistral 7B에 대해 task-specific instruction을 사용하여 대조 학습을 적용하였고, 독점 GPT-4 모델에서 생성된 대량의 합성 데이터를 활용하여 종합적인 임베딩 벤치마크에서 최신 bidirectional 모델을 능가하기 시작했습니다.
• LLM2Vec은 공개적으로 사용 가능한 데이터만을 사용하여 LLM에서 임베딩 모델을 구축하려고 했지만, 여전히 E5-Mistral보다 성능이 낮았습니다.
• SFR-Embedding-Mistral은 E5-Mistral을 비검색 및 검색 데이터셋의 혼합물에 대해 추가로 미세 조정하여 두 작업 모두에서 정확도를 향상시켰습니다. 이는 NV-Embed와 밀접한 관련이 있습니다.
• NV-Embed는 공개 데이터를 사용하여 Mistral 7B LLM에서 처음부터 학습되며, 다른 임베딩 모델이나 독점 합성 데이터에 의존하지 않습니다.
• SFR-Embedding-Mistral은 task-homogeneous batching을 사용하는 반면, NV-Embed는 모든 작업의 샘플로 구성된 잘 혼합된 배치를 사용하여 “zigzag” 경사 업데이트를 방지합니다.
• SFR-Embedding-Mistral은 task-homogeneous batching을 사용합니다. 이는 각 배치가 단일 작업(task)에서 나온 샘플들로만 구성된다는 것을 의미합니다. 예를 들어, 배치 1은 검색 작업의 샘플들로만 구성되고, 배치 2는 분류 작업의 샘플들로만 구성되는 식입니다. 이러한 배치 구성 방식은 각 작업에 특화된 모델 업데이트를 가능하게 합니다.
• 반면에 NV-Embed는 모든 작업의 샘플들이 골고루 섞인 배치를 사용합니다. 즉, 각 배치 내에 검색, 분류, 클러스터링 등 다양한 작업의 샘플들이 함께 포함됩니다. 이는 모델이 여러 작업을 동시에 학습할 수 있도록 합니다.
• NV-Embed 논문에서는 이러한 배치 구성 방식이 “zigzag” 경사 업데이트를 방지한다고 언급하고 있습니다.
• “Zigzag” 경사 업데이트란, 모델이 특정 작업에 과도하게 최적화되어 다른 작업에 대한 성능이 저하되는 현상을 말합니다.
• 예를 들어, 검색 작업에 최적화된 업데이트 이후에 분류 작업에 최적화된 업데이트가 이루어지면, 모델이 검색 작업에 대한 성능을 일부 잃을 수 있습니다.
• 이러한 현상이 반복되면 모델의 전반적인 성능 향상이 제한될 수 있습니다.
• NV-Embed에서 사용한 잘 혼합된 배치는 이러한 문제를 완화할 수 있습니다. 다양한 작업의 샘플들이 함께 포함되어 있으므로, 모델은 각 배치에서 여러 작업을 동시에 학습할 수 있습니다.
• 이는 특정 작업에 편향되지 않은 더 일반화된 임베딩을 학습할 수 있도록 도와주며, 결과적으로 전체 MTEB 벤치마크와 검색 작업 모두에서 새로운 최고 성능을 달성할 수 있었다고 합니다.
• 다른 최근 연구로는 Gecko와 GritLM 등이 있습니다.

Bidirectional Attention

• Decoder-only LLM에서 사용되는 causal attention mask는 next-token prediction 작업을 위해 도입되었습니다. 이는 auto-regressive 텍스트 생성 중에 이전 위치에만 attend할 수 있도록 하여 정보 누출을 방지합니다.
• 그러나 단방향 어텐션이 모델의 표현력을 제한한다는 것이 관찰되었습니다. 이는 GPT 모델이 유사한 크기의 BERT나 T5 모델에 비해 자연어 이해 벤치마크에서 성능이 좋지 않다는 것으로 증명됩니다.
• 최근 LLM2Vec과 GRIT은 양방향 어텐션을 활용하기 위한 추가 학습 단계나 하이브리드 목적 함수를 도입했습니다. 반면 NV-Embed는 대조 학습 중에 causal attention mask를 제거하는 간단한 방법을 사용하며, 이것이 매우 잘 작동한다는 것을 결과로 보여줍니다.

Latent Attention Layer

• 토큰 시퀀스의 임베딩을 얻기 위한 두 가지 일반적인 방법은 i) mean pooling과 ii) 마지막 <EOS> 토큰 임베딩입니다.
• 이전의 양방향 임베딩 모델은 일반적으로 mean pooling을 사용하는 반면, decoder-only LLM 기반 임베딩 모델은 마지막 <EOS> 토큰 임베딩을 더 많이 사용합니다.
• 그러나 이 두 방법 모두 특정 한계가 있습니다. Mean pooling은 토큰 임베딩의 평균을 취하므로 중요한 정보가 희석될 수 있고, 마지막 토큰 임베딩은 마지막 토큰의 출력 임베딩에 크게 의존하는 recency bias 문제가 있을 수 있습니다.
• 이 논문에서는 범용 임베딩 작업을 위해 시퀀스를 더 잘 pooling할 수 있는 latent attention layer를 제안합니다. 이는 trainable “dictionary”인 latent array를 사용하여 더 나은 표현을 얻는 cross-attention의 한 형태로 작동합니다.
• Latent attention layer의 출력은 MLP를 거쳐 최종적으로 전체 시퀀스의 임베딩을 얻습니다. 이는 더 나은 표현을 얻기 위한 dictionary learning의 정신을 따르고 있습니다.

일반적으로 토큰 시퀀스의 임베딩을 얻는 방법은 2가지가 있다.

시퀀스를 mean pooling 하거나 마지막 <EOS> 토큰의 임베딩을 얻는 것이다. 하지만 평균 풀링은 중요한 정보를 희석할 수 있고, <EOS> 임베딩은 후반 토큰에 의존하는 recency bias가 생길 수 있다.

대안으로 보다 표현적인 풀링을 위해 latent attention layer를 제안한다.

LLM의 출력을 query로 간주하고 trainable latent array와 attention을 수행한 뒤 MLP 이후 평균 풀링을 적용해 임베딩을 얻는다.

Decoder-only LLM은 입력 토큰 임베딩(Input Token Emb.)을 받아 각 토큰에 대한 출력 임베딩을 생성, 이 때 단방향(unidirectional) 어텐션을 사용하여 이전 토큰들만 참조

이렇게 생성된 Decoder Model Output은 Latent Attention Block으로 전달됩니다. Latent Attention Block은 크게 세 부분으로 구성되어 있습니다.

1. Latent Array: 학습 가능한 “dictionary” 역할을 하는 latent 벡터들의 집합입니다. 이 latent array는 Key (K)와 Value (V) 역할을 동시에 수행합니다.
2. Cross-Attention: Decoder Model Output을 Query (Q)로, Latent Array를 Key (K)와 Value (V)로 사용하여 어텐션 연산을 수행합니다. 이를 통해 Decoder Model Output은 Latent Array가 제공하는 새로운 표현 공간으로 변환됩니다.
3. MLP (Multi-Layer Perceptron): Cross-Attention 결과를 입력으로 받아 최종 출력 임베딩을 생성합니다.

이렇게 생성된 최종 출력 임베딩(Output Emb.)은 mean pooling을 거쳐 전체 입력 시퀀스를 나타내는 단일 벡터로 요약됩니다.

이러한 구조를 통해 Latent Attention Block은 Decoder-only LLM의 출력을 더 표현력 있는 공간으로 변환하고, 단순 평균으로 인한 정보 손실을 완화할 수 있습니다. 이는 검색, 분류, 유사도 측정 등 다양한 자연어 처리 작업에서 모델의 성능을 향상시키는 데 도움을 줄 수 있습니다.

Two-stage Instruction-Tuning

검색 및 비검색 작업 (e.g. 분류, 클러스터링)을 적절하게 수행하려면 다양한 작업의 특성을 고려해야 한다.

예를 들어 배치 내의 다른 데이터를 부정 예시로 취급하는 in-batch negatives trick의 경우 비검색 작업에 악영향을 끼칠 수 있다.

Instruction-tuning은 LLM이 지시를 따르도록 학습시키고 검색 기반 생성을 수행하는 데 널리 사용되어 왔습니다. 최근에는 검색기와 범용 임베딩 모델을 학습시키는 데에도 적용되고 있습니다.

검색 및 비검색 작업에서 적절한 성능을 발휘하는 일반화된 임베딩 모델을 얻기 위해서는 서로 다른 작업의 특성을 고려해야 합니다.

예를 들어, in-batch negatives의 사용은 dense-embedding 기반 검색기를 학습시키는 데 매우 효율적인 것으로 알려져 있습니다. 그러나 분류나 클러스터링 작업의 경우 in-batch negatives를 적용하면 오히려 임베딩 모델을 오도할 수 있습니다.

이러한 점을 고려하여, NV-Embed는 2단계 지시 튜닝 방법을 도입합니다. 1단계에서는 다양한 검색 데이터셋에 대해 in-batch negatives와 curated hard-negative 예제를 활용하여 대조 학습을 수행합니다. 2단계에서는 in-batch negatives 없이 검색 및 비검색 데이터셋의 조합에 대해 대조 지시 튜닝을 수행합니다.

검색 작업이 다른 작업에 비해 더 큰 어려움을 나타내기 때문에, NV-Embed의 학습 전략은 초기에 검색을 위해 모델을 미세 조정하는 데 중점을 둡니다. 2단계에서는 나머지 임베딩 작업들을 지시 튜닝에 혼합합니다.

• 1-stage: 검색 데이터셋, in-batch negatives trick, 대조 훈련
• 2-stage: 엄선된 hard negatives 예제를 활용하여 각 작업에 대해 contrastive instruction-tuning을 수행

Causal Attention vs. Bidirectional Attention

• Decoder-only LLM 모델에서 임베딩 적용을 위한 self-attention mask의 영향을 조사하기 위해, bidirectional mask와 causal mask를 비교하는 실험을 진행했습니다.
• 실험 결과, 모든 pooling 방법에 대해 bidirectional mask가 causal mask보다 56개 작업의 평균 MTEB 점수 기준으로 일관되게 더 좋은 성능을 보였습니다. 이는 causal attention mask로 생성된 임베딩이 bidirectional attention mask로 생성된 임베딩보다 효과가 현저히 떨어진다는 것을 나타냅니다.

Pooling Methods

• 임베딩 모델에 대한 다양한 pooling 방법의 영향을 조사하기 위해, -last, mean, latent-attention, self-attention pooling 방법을 비교하는 실험을 진행했습니다.
• 실험 결과, 56개 작업의 평균 MTEB 점수 기준으로 mean pooling이 -last token embedding보다 일관되게 더 좋은 성능을 보였습니다. 이는 마지막 토큰 임베딩이 최종 토큰의 출력에 지나치게 의존하는 recency bias의 영향을 받기 때문일 수 있습니다.
• Mean pooling의 성능을 더 향상시키기 위해, 중요한 정보가 희석되는 문제를 해결하고자 mean pooling 전에 제안된 latent-attention이나 self-attention layer를 추가하는 실험을 진행했습니다.
• 실험 결과, self-attention은 decoder-only LLM의 임베딩 성능에 추가적인 정확도 향상을 제공하지 않았습니다. 반면, latent-attention layer는 검색, 분류, STS 하위 작업에 도움이 되는 것으로 나타났습니다. 특히, 더 어려운 15개 검색 작업의 nDCG@10 정확도가 개선되었습니다.
• 이는 latent array에 의해 제공되는 “dictionary learning”이 더 표현력 있는 representation을 제공하기 때문일 것으로 추측됩니다. Latent-attention layer는 decoder-only LLM에서 출력 임베딩 representation을 효과적으로 학습하여 출력 임베딩의 평균으로 인한 정보 희석을 완화합니다.

Experiments

Mistral-7B에 LoRA, latent attention layer를 통합하여 fine-tuning 한다.

!https://blog.kakaocdn.net/dn/T5j6d/btsHGuK66E4/2Rv1yVhVl9Bm8qUcp1FR0K/img.png