Meaningful

Meaningful

Greeting :)

[논문리뷰]LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders(COLM, 2024)

Date:     Updated:

카테고리:

Parishad BehnamGhader, Vaibhav Adlakha, Marius Mosbach, Dzmitry Bahdanau, Nicolas Chapados, and Siva Reddy. 2024. *LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders**. arXiv:2404.05961 [cs.CL] https://arxiv.org/abs/2404.05961

Problem Statment

[Causal Attention 제약] Decoder-only LLM은 causal attention을 사용하여 각 토큰이 이전 토큰만 볼 수 있어, 시퀀스 전체의 문맥 정보를 반영하는 데 제한이 있었다. 이로 인해 풍부한 컨텍스트 임베딩 생성에 적합하지 않았다.

[복잡하고 비용 높은 적응(adaptation)] 기존의 텍스트 임베딩 학습은 대규모 weakly supervised & fully-supervised learning 파이프라인을 거쳐야 했으며, 종종 GPT-4와 같은 외부 생성 데이터를 필요로 했다. 이는 데이터 및 연산 자원이 많이 소모되는 방식이었다.

[Causal Mask 제거의 제한적 연구] 일부 연구는 supervised fine-tuning 중 특정 레이어에서 causal mask를 제거하거나 입력을 복제하여 future token을 참조하도록 했으나, 이는 계산 비용이 크거나 제한된 태스크에만 효과적이었다.



Methodology

Overview

1

Figure 1은 LLM2Vec의 세 가지 핵심 단계를 보여준다. 먼저, ‘Enabling Bidirectional Attention(Bi)’ 단계는 기존 디코더 전용 언어 모델의 인과적 어텐션 제약을 극복하기 위해 Bidirectional Attention을 활성화하는 과정이다. 다음으로, ‘Masked Next Token Prediction(MNTP)’ 단계는 마스킹된 다음 토큰 예측 학습을 통해 모델이 Bidirectional Attention을 활용하도록 적응시킨다. 마지막으로, ‘Unsupervised Contrastive Learning’ 단계에서는 비지도 대조 학습(SimCSE)과 평균 풀링(mean pooling)을 적용하여 더 나은 시퀀스 표현을 학습하게 된다.

Enabling Bidirectional Attention

1

첫 번째 단계는 디코더 전용 LLM의 본질적인 한계인 인과적 어텐션 마스크를 극복하는 데 중점을 둔다. 인과적 어텐션 마스크(causal attention mask)는 토큰이 이전 토큰에만 주의를 기울이도록 제한하여, 전체 시퀀스에 걸친 문맥적 정보를 완전히 파악하기 어렵게 만든다. LLM2Vec은 이 어텐션 마스크를 모든 토큰이 시퀀스 내의 다른 모든 토큰에 접근할 수 있도록 하는 전체 1 행렬(all-ones matrix)로 대체한다. 이로써 모델은 양방향 LLM으로 변환된다. 그러나 이 단순한 접근 방식은 모델이 미래 토큰에 주의를 기울이도록 학습되지 않았기 때문에, 단독으로 사용하면 성능이 저하될 수 있다.

Masked Next Token Prediction

1

두 번째 단계는 Bidirectional Attention을 효과적으로 활용하도록 모델을 적응시키는 학습 과정이다. MNTP는 다음 토큰 예측과 마스크된 언어 모델링(Masked Language Modeling, MLM)을 결합한 학습 목표를 사용한다. 입력 시퀀스 \(x=(x_{1}, x_{2},...,x_{N})\)에서 무작위로 토큰의 일부를 마스킹한 다음, 모델은 마스킹된 토큰을 과거 및 미래 문맥을 기반으로 예측하도록 학습된다. 특히, MNTP 학습 시, 위치 i의 마스킹된 토큰을 예측하기 위한 손실은 마스킹된 위치 자체의 로짓이 아닌 이전 위치 i−1의 토큰 표현에서 얻은 로짓을 기반으로 계산된다. 이 과정은 모델이 양방향 문맥 정보를 활용하도록 학습시킨다.

Unsupervised Contrastive Learning

1

세 번째 단계는 문장 표현의 품질을 향상시키는 데 초점을 맞춘다. 디코더 전용 LLM은 전체 시퀀스의 문맥을 포착하도록 명시적으로 학습되지 않았기 때문에, 이 단계는 이 격차를 메운다. LLM2Vec은 SimCSE라는 비지도 대조 학습 접근 방식을 사용하여, 동일한 문장에 대해 두 개의 다른 표현을 생성한다. 이는 모델에 독립적으로 샘플링된 dropout 마스크를 두 번 통과시켜 데이터 증강을 수행하고, 이를 통해 positive pair를 만들고, in-batch negative를 활용해 contrastive 학습을 수행한다. 여기서 문장 임베딩은 mean pooling 또는 weighted mean pooling을 통해 계산된다.

학습 데이터는 Wikipedia(Wikitext-103, Gao et al. 2021 subset)로만 수행되어, 새로운 지식을 주입하지 않고 모델의 표현 방식을 적응시킨다. MNTP와 SimCSE 학습 모두 LoRA를 활용하여 파라미터 효율성을 유지한다



Experiments

Main Results

1

Table 1은 MTEB(Massive Text Embeddings Benchmark)에 대한 LLM2Vec의 비지도 학습 성능을 보여준다. 이 표에 따르면, LLM2Vec의 첫 두 단계인

Bidirectional Attention와 MNTP(Masked Next Token Prediction)를 적용하는 것만으로도 모든 LLM 모델의 성능이 크게 향상된다. 예를 들어, Mistral-7B의 경우, 이 두 단계만으로 16.4%의 성능 향상을 보였다. 이는 단순히 인과적 어텐션을 사용하는 것보다 더 우수한 성능이다. 또한, LLM2Vec은 경쟁 모델인 Echo Embeddings와 비교했을 때, S-LLAMA-1.3B, LLAMA-2-7B, 그리고 Meta-LLAMA-3-8B 모델에서 더 나은 성능을 보였다. Echo Embeddings는 입력 시퀀스를 복제하여 추론 비용을 두 배로 증가시키는 반면, LLM2Vec은 효율적이다.

LLM2Vec의 마지막 단계인 비지도 SimCSE(Unsupervised SimCSE)를 추가하면 모든 모델의 성능이 대폭 향상된다. 특히, SimCSE가 적용된 LLM2Vec Mistral-7B 모델은 56.80점의 점수로, MTEB 벤치마크에서 모든 비지도 모델 중 새로운 최첨단(SOTA) 성능을 달성했다. 이는 LLM2Vec이 단순하고 효율적인 방법만으로도 디코더 전용 LLM을 강력한 텍스트 임베딩 모델로 효과적으로 변환할 수 있음을 보여준다.

Ablation Study

1

Table 5는 MTEB 하위 데이터셋에 대한 모델별, 풀링 방식별 성능을 보여준다.

  • Bi: Bidirectional Attention을 단순히 활성화하면 대부분의 모델에서 성능이 하락했지만, Mistral-7B는 오히려 성능이 향상되었다.
  • Bi + MNTP: MNTP 학습을 추가하면 Bidirectional Attention을 활성화했을 때 감소했던 성능이 크게 회복되고, 인과적 어텐션을 사용한 모델보다 높은 성능을 기록했다.
  • Bi + MNTP + SimCSE: SimCSE를 추가하는 것은 시퀀스 수준 작업에서 성능을 추가로 향상시켰다.

실험 결과 Bidirectional Attention을 하지 않을 경우 성능 감소가 매우 크다는 것을 알 수 있다.

Analysis of Future Token Information Capture

1

Figure 4는 LLM2Vec이 모델의 미래 토큰 정보 포착 능력을 어떻게 향상시키는지 보여준다. 이 그림은 query와 positive 및 negative 샘플 간의 코사인 유사도 분포를 히스토그램으로 시각화한다.

S-LLaMA-1.3B와 Mistral-7B 모델에서, Bidirectional Attention을 활성화하고 MNTP 학습을 적용하면 positive와 negative 사이의 유사도 분포가 명확하게 분리되는 것을 관찰할 수 있다. 특히, 학습 없이 단순히 Bidirectional Attention만 활성화한 경우보다 MNTP 학습을 거친 모델이 positive에 더 높은 유사도를 부여하며, 이는 LLM2Vec이 모델이 문장 끝의 정보를 접두사 표현에 통합하도록 돕는다는 것을 의미한다

Analysis of Mistral-7B’s Intriguing Property

1

Figure 5는 학습 없이 Bidirectional Attention을 활성화했을 때 모델의 내부 표현이 어떻게 변화하는지 보여준다. 이는 Causal Attention으로 생성된 표현과 Bidirectional Attention으로 생성된 표현 간의 코사인 유사도를 계층별 및 토큰 위치별로 시각화한 것이다.

S-LLaMA-1.3B와 LLAMA-2-7B의 경우, Bidirectional Attention을 활성화하면 대부분의 계층과 토큰 위치에서 표현의 유사도가 크게 떨어지는 것을 볼 수 있다. 이는 Bidirectional Attention이 모델의 표현을 근본적으로 변화시켰음을 의미한다. 반면, Mistral-7B는 학습 없이 Bidirectional Attention을 활성화하더라도 인과적 어텐션과 매우 높은 코사인 유사도를 유지한다. 저자들은 이 현상이 Mistral-7B가 사전 학습 과정에서 이미 접두사 언어 모델링(prefix language modeling)과 같은 Bidirectional Attention의 한 형태를 포함했을 가능성을 시사한다고 추측한다.



Conclusion

Contributions

  • 간단하면서도 파라미터 효율적인 LLM2Vec 프레임워크 제안: Bidirectional Attention + MNTP + SimCSE.
  • 효율성 입증: 1000 스텝, LoRA 기반 파라미터 효율적 학습으로 SOTA 달성

Limitations

  • 본 논문은 주로 Wikipedia 데이터만으로 실험을 수행했기 때문에, 도메인 다양성에 따른 일반화 성능은 아직 검증되지 않았다.
  • SimCSE 적용은 word-level 태스크에 불리하게 작용하여, 모든 유형의 태스크에 보편적으로 좋은 효과를 보장하지 않는다.
  • Mistral-7B의 특이한 동작 원인(prefix LM 활용 여부 등)은 아직 불명확하며, 향후 추가 분석이 필요하다.

NR 카테고리 내 다른 글 보러가기

댓글 남기기

최근 글 10 개 :)