LLM2Vec: Large Language Models are Secretly Powerful Text Encoders 리뷰

Background

Text embedding model을 구축하기 위해서 보통 pre-trained bidirectional encoder나 encoder-decoder 구조 모델을 사용한다. 반면, decoder only 구조 모델로 text embedding model을 구성하는 경우는 매우 드물다. Decoder only model은 causal attention을 하기 때문에 전체 context를 보지 못한 상태로 embedding을 생성해서 성능이 안 좋기 때문이다.

LLM2Vec

하지만 본 논문은 decoder only LLM의 이런 구조적 한계에도 불구하고 encoder기반 LLM보다 text embedding에 있어서 이점을 갖고 있다고 주장한다.

• Decoder LLM은 모든 token에 대해서 학습을 하지만 encoder-only model은 15%의 masking을 pretraining에 적용하기 때문에 sample efficient하지 않다.
• Decoder only LLM 자체의 연구가 더 활발하기 때문에 더 많은 정보가 알려져있다. (다양한 pre-training 방법 등.. )
• Instruction fine-tuning, human preference을 학습할 때 모두 decoder-only LLM을 사용하기 대문에 instrctuctions을 통해 universal embedding model들을 구축할 수 있을 것이라고 본다

이런 가설을 통해 본 논문은 간단한 unsupervised approach을 통해 pre-trained된 아무 decoder-only LLM을 text embedding 모델로 바꾸는 방법론 LLM2Vec을 제시한다. LLM2Vec는 labeled data가 하나도 필요로 하지 않고 data / parameter efficient하다.

Three simple ingredients

LLM2Vec을 decoder-only model에 적용하기 위해서는 세 단계를 거쳐야 한다.

** Enabling Bidirectional Attention**

Decoder only LLM은 unidirectional attention을 적용하기 위해서 future token들을 보지 못하도록 causal attention mask M을 사용한다.

\[O = \text{softmax} ( \frac {M_{j \le i} Q K^T} {\sqrt d}) \cdot V\]

O는 self attention layer의 output이다.

[1,1,1,-inf,-inf,-inf … ] 과 같은 masking이라고 생각하면 된다.

하지만 Causal attention mask을 적용하게 된다면 future context을 보지 못하기 때문에 제대로 된 representation을 얻지 못한다. 따라서, causal attention mask을 all ones matrix으로 바꾸어 bidirectional attention이 가능해지도록, 즉 미래의 token들까지 참고하여 token representation을 학습할 수 있도록 하였다.

\[O = \text{softmax} ( \frac {Q K^T} {\sqrt d}) \cdot V\]

Masked Next Token Prediction (MNTP)

이때, 본 논문은 decoder-only model은 엄청 많은 데이터로 unidirection으로만 pretrain되었기 때문에 바로 bidirectional attention으로 predict하게 하면 embedding 성능이 좋지 않을 것이라고 가정하고 bidirectional attention에 적응할 수 있는 추가적인 pre-train이 필요할 것이라고 예측하였다.

따라서, Masked Next Token Prediction task으로 추가적인 pre-train을 하여 bidirectional attention 적용에 익숙해지도록 하였다.

Text sequence $x = (x_1, x_2, … x_N)$이 있을 대, input tokens 중 일부를 masking한 뒤, masked token들을 과거 + 미래 context를 예측하도록 학습하는 과정을 거치도록 한다. 이때, causal language modeling처럼 위치 i의 masked token을 예측한다고 할 때, loss 계산을 이전 위치 i-1의 token representation을 기반으로 계산한다. 위치 i-1의 token representation에는 causal attention mask을 제거하였기 때문에 past & future context 모두 반영되어 있다.

Unsupervised Contrastive Learning

MNTP로 adaptation을 하면, 단어 level의 representation을 잘 생성할 수 있게 되지만 (past, future context를 모두 고려) 그렇다고 해서 전체적인 sequence representation 성능과는 직결되지 않는다. 예를 들어, BERT는 Next Sentence Prediction task을 통해 [SEP]으로 분리된 두 text 간의 관계를 학습하지만, decoder only는 전체 context를 capture하도록 pre-train되지 않았기 때문이다.

이를 해결하기 위해 SimSCE에서 제시한 unsupervisal contrastive learning을 적용한다. 이 과정은 다음과 같다. Input text가 주어졌을 때, 독립적인 dropout mask을 적용하면서 model을 두 번 통과하게 하여 서로 다른 representation (rep1, rep2)을 얻는다. Rep1과 rep2는 같은 text로부터 생성된 것이기 때문에 둘의 관계는 연관성이 높다(positive)고 볼 수 있다. 이런 방법으로 label없이 text pair dataset을 구축하여 query rep1가 주어졌을 때, batch 내의 example들 중에서 negative example들과 rep1의 연관성 정도는 최대한 줄이고 positive example rep2와의 연관성 정도를 최대한 높이는 contrastive learning을 적용한다.

\[L = \frac{e^{\lambda s(q, d^+)}}{e^{\lambda s(q, d^+)} + \sum_{d^- \in N} e^{\lambda s(q, d^-)}}\]

이때, sequence의 representation은 sequence을 구성하는 token들에 pooling operation을 적용하여 계산하였다.

Training Details

Models
사용한 decoder-only LLM은 Sheared-LlaMa-1.3B, Llama-2-7B-chat, 그리고 Mistral-7B-Instruct-v0.2이다.

Training data
English wikipedia를 사용하여 MNTP와 SimCSE step을 진행하였다. English wikipedia를 학습 데이터로 사용한 이유는 대부분의 모델들의 pre-training 단계에서 wikipedia 데이터를 사용하기 때문이다. 이미 pre-train 단계에서 사용한 Wikipedia 데이터로 다시 pre-train을 한다면 새로운 world knowledge을 학습하는 것에 집중하지 않고 bidirectional context에 적응하는 것과 sequence representation 학습에 집중할 것이라고 가정하였기 때문이다.

Masked next token prediction
Masking을 하기 위해서는 underscore ‘_‘을 mask token으로 사용하였다.

LoRA로 training을 하였다.

Unsupervised contrastive learning
Positive example들은 SimCSE 방법을 따르고 negative example들은 in-batch negative을 사용한다.

이때, MNTP의 LoRA 가중치를 base model에 반영시킨 뒤, 새로운 LoRA 가중치들을 초기화하여 SimSCE training을 적용하였다. 이 방법을 통해 이전 training에서 얻은 지식이 model에 남아있도록 하였다.

LLM2Vec transformed models are strong unsupervised text embedders

Evaluation on word-level tasks

LLM2Vec의 word-level task에 대해서 성능을 측정하여 bidirectional attention으로 각 단어를 더 rich한 context representation으로 mapping이 가능한지 보았다.

사용한 word level task는 chunking, name-entity recognition, 그리고 part-of-speech tagging이다. 사용한 benchmark는 CoNLL-2003이다.

• 가정한 것과 동일하게 bidirectional attention을 추가적인 adaptation 없이 바로 적용하였을 때, 성능이 매우 저하되었다.
• LLM2Vec을 적용한 모델들을 보면 MNTP로 adaptation을 한 결과, 성능이 큰 폭으로 상승하였다.
• 하지만 SimCSE까지 적용했을 때, 성능이 저하되는 것을 관찰할 수 있는데 이는 SimCSE가 sequence-level task에 집중하는 training objective를 사용하기 때문이다.

Evaluation on sequence-level tasks

Massive Text Embedding Benchmark (MTEB)으로 sequence-level task의 성능을 측정하였다.

Decoder-only LLM은 instruction tuning / preference 학습이 잘 되어있기 때문에 task-specific instructions을 leverage한다. STS(Semantic Textual Similarity)처럼 symmetric한 task에는 query와 passage 모두에 instruction을 추가하였고 QA처럼 asymmetric task에는 query에만 instruction을 추가하였다.

Instruction 예시는 다음과 같다.

이때, 사용된 baseline은 unsupervised BERT models와 Echo embedding이다.

Echo embedding은 본 논문처럼 decoder-only model로 text embedding을 하기 위해 고안된 방법이다. Causal 정보의 흐름은 단방향이기 때문에, input을 한번 model에 넣어 embedding 값을 얻은 뒤 (이 embedding에는 이제 모든 context를 거쳤기 때문에 전체 context의 정보를 지니고 있다), 다시 model에 넣어서 나온 output embedding을 pooling하여 최종 sequence representation으로 사용하는 방법이다.

Results on our 15 task subset of MTEB
실험 전에 최적의 pooling method을 알아내기 위해 MTEB task 중 15개의 subset만 사용하여 다양한 pooling method들을 비교하였다.

• Causal attention으로 text representation을 얻을 때 보통 맨 마지막 token인 [EOS]을 사용하는데, [EOS] token을 사용하는 방법이 가장 성능이 안 좋다는 것을 통해 causal attention은 text embedding에 최적화된 방법이 아니라는 것을 알 수 있다.

• MNTP와 SimCSE을 적용할 때, 성능이 점점 향상된다는 것을 관찰할 수 있다.

• Mean pooling이 LLM2Vec에 최적인 pooling 방법이라는 것을 알 수 있다.

Results on full MTEB

전체 MTEB dataset으로 각 embedding model들을 평가한 결과는 다음과 같다.

• Bidirectional attention과 MNTP를 적용하는 것만으로도 성능 향상을 야기하였다.
• Mistral 7B에서 몇 task에서 LLM2Vec와 Echo embedding의 성능이 유사한 것을 관찰할 수 있는데 LLM2Vec는 Echo embedding보다 더 효율적이기 때문에 LLM2Vec가 더 우수하다고 할 수 있다. Echo는 input을 다시 처리하기 때문에 input sequence의 두 배의 길이를 매번 처리해야 한다. 이로 인해 Echo embedding을 적용하면 inference 속도가 매우 느려진다.
• SimCSE를 적용하면 모든 모델에서 embedding의 성능이 큰 폭으로 향상된다.

결과적으로 LLM2Vec는 decoder-only LLM을 encoder 기반 text embedding model보다 더 좋은 embedding model로 바꿀 수 있음을 보여준다.

How does LLM2Vec affect a model?

LLM2Vec helps models to capture information from future tokens

LLM2Vec가 진짜 future tokens의 정보를 가져오는지 알아내기 위해 동일한 prefix을 가졌지만 뒤 내용은 다른 문장들의 연관성 계산을 정확하게 수행하는지 평가하였다.

Evaluation dataset은 35개의 문장 triples ${ (q_i, s_i^{+}, s_i^{-})}^{35}_{\text{i=1}}$으로 되어 있으며 prefix를 $A_i$라고 할 때, $q_i = (A_i, B_i), s^+_i = (A_i, C_i), s^+_i = (A_i, D_i)$이다. 이때, $B_i$와 $C_i$는 비슷한 의미를 갖고 있지만 $B_i$와 $D_i$는 그렇지 않다.

각 문장의 representation 계산은 전체 token embedding의 mean pooling을 하는 것이 아니라 $A_i$ 부분만 pooling을 적용한다. 즉, LLM2Vec가 정말로 future context까지 고려한다면 $A_i$의 representation에는 future 내용 ($B_i, C_i, D_i$) 이 반영되어 있어야 하고 $s^+_i$와 $s^-_i$을 구분할 수 있어야 한다.

• 모든 모델에서 Bidirectional attention만 적용하거나 MNTP objective으로 학습하기만 하면 positive와 negative sample들을 뚜렷하게 잘 구분한다는 것을 볼 수 있다.

$\rightarrow$ Bidirectional context를 실제로 고려한다는 것을 알 수 있다.

Why does bidirectional attention without training work for Mistral models?

실험 결과들을 보면 Mistral-7B는 다른 모델들과 다르게 bidirectional attention만 할 수 있게 해주면 큰 성능 저하 없이 잘 작동한다는 것을 볼 수 있다. 이런 이유를 알아내기 위해 추가적인 실험을 하였다.

하나의 input sequence를 각 모델에 서로 다른 attention method (causal, bidirectional)을 적용하여 모든 layer마다 각 token의 hidden representation $H_{l}^c, H_{l}^{\text{bi}}$을 계산하였다. 그리고 모든 layer마다 두 representation의 similarity sim($H_{l}^c, H_{l}^{\text{bi}}$)를 계산하였다.

본 논문은 추가적인 adaptation 없이 바로 bidirectional attention을 적용한다면 아예 다른 성격의 representation을 생성할 것이기 때문에 대부분의 layer들에서 similarity score가 매우 낮을 것이라고 가정하였다.

• S-LLaMA-1.3B and LLaMA-2-7B에서는 가정과 들어맞는 결과를 보여주었다. 거의 모든 layer에서 cosine similarity 값이 낮았다.

• 반면 Mistral-7B에서는 cosine similarity 값이 꾸준히 높은 것을 관찰할 수 있다.이를 통해 Mistral-7B는 causal language modeling만 적용하여 pre-train하지 않았고 bidirectional attention으로 pre-train을 하였음을 추론할 수 있다.