본문 바로가기
NLP/논문리뷰

[논문 Review] 12. LLM2Vec : Large language models are secretly powerful text encoders

by ㅣlㅣl 2024. 4. 18.
논문의 핵심 내용 한줄 요약
LLM을 text encoder로 이용해보자

Abstract

이 논문에서는 decoder-only LLM을 강력한 성능의 text encoder로 변환하는 간단한 비지도 접근 방식 LLM2Vec을 소개한다.

 

 

1.3B ~ 7B 크기의 인기 있는 LLM 3개 (LLaMA 1.3B, LLaMA 7B, Mistral 7B) 에 공개된 데이터만 사용하여 해당 학습 방식을 적용해보았고,

그 결과 대규모 텍스트 임베딩 벤치마크 (MTEB) 에서 비지도 학습 방식 중 최고 성능을 달성했다. 

또한 지도 대조학습과 결합한 경우, 공개된 데이터로만 학습한 모델 중에서는 최고 성능을 달성했다.

 

2024.04.18 기준 확인 결과

더보기

 

 

 

 

별도의 리더보드를 제공하는 topic 중에서

 

Pair Classification, Retrieval task에서의 성능이 매우 높게 나타났다.

 

이를 통해 비용이 드는 adaptation이나 추가적으로 생성한 데이터를 사용하지 않고도 LLM을 효율적으로 범용 텍스트 인코더로 변환하는 방식을 입증했다.

 

 

1. Introduction

Text embedding model은 자연어 텍스트의 의미를 벡터 표현으로 인코딩해주는 모델로, 임베딩된 벡터는 STS, IR, Clustering과 같은 다양한 NLP 작업에 사용된다.

그동안은 pretrained bidirectional encoder (BERT), 혹은 encoder-decoder (T5) 구조를 가지는 모델이 이러한 임베딩 모델로 사용되었다.

 

[논문 Review] 08. T5 : Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

논문의 핵심 내용 한줄 요약 모든 NLP task를 text to text 형식으로 바꿔서 해결해보자! Abstract 풍부한 데이터를 통해 모델을 Pre training하고, 다운스트림 작업에서 Fine tuning을 진행하는 전이학습은 NLP

ll2ll.tistory.com

 

최근에서야 decoder only LLM을 text embedding model로 쓰려는 시도가 보이고 있다.

특히 아래 2개의 논문은 본 논문에서도 많이 인용되었다.

 

이러한 시도가 최근에서야 보인 이유는 decoder only model에서 쓰이는 causal attention mechanism으로 인해,
전체 input sequence를 참조하지 못해서 풍부한 컨텍스트 표현에 제약이 있기 때문이라고 지적하고 있다.

 

  • causal attention (=masked multi head attention)
    미래의 토큰을 볼 수 없도록 마스킹 한 후 진행되는 attention 연산

이러한 제약은 생성 작업을 위해서 필요하지만, 텍스트 임베딩에서는 전체 입력 시퀀스에 대한 정보를 필요로 하기 때문에 적합하지 않다.

 

 

그렇다면 그냥 Encoder only model 쓰면 되지 않아?

굳이 Decoder의 구조까지 변경해가면서 써야 하는 이유는 뭘까? 그냥 Encoder-only 모델을 쓰면 되는 거 아닐까?

본 논문에서는 decoder only model를 사용하는 것애 몇 가지 장점이 있다고 주장한다.

  • 사전 학습 과정에서 Decoder-only LLM은 일부 토큰이 아닌 모든 입력 토큰에서 학습하기 때문에 동일한 양의 학습 데이터가 있을 때 더 효율적으로 학습 가능
    -> Encoder only 모델은 일반적으로 pretraining을 진행할 때 input sequence의 15%를 마스킹하기 때문 (BERT)
    -> 효율적인 학습에 대한 실험은 챕터 5에 기재되어 있다

  • 최근 연구에서 뛰어난 성능의 instruction decoder model이 많이 등장했기 때문
    -> Instruction을 사용해 다양한 작업에 일반화가 가능해졌다 

 

 

2. LLM2Vec

사전 훈련된 Decoder only LLM을 범용적인 text encoder로 변환하기 위한 비지도 접근방식이다.

LLM2Vec 과정은 다음과 같이 3단계로 구성된다.

1.Enabling Bidirectional Attention

2. masked next token Prediction

3. unsupervised contrastive learning

 

 

 

Enabling bidirectional attention (Bi)

첫 번째로 Decoder only LLM의 causal attention mask를 all-ones matrix로 대체하는 것이다.

이렇게 하면 각 토큰이 시퀀스의 다른 모든 토큰에 액세스하는 것이 가능해져, Bidirectional LLM으로의 변환이 가능하다.

하지만 사전학습 시 Decoder-only LLM이 미래 토큰을 고려하도록 훈련되지 않았기 때문에, 이러한 나이브한 접근 방식은 더 나쁜 representation을 생성할 수도 있다.

실제로 대부분의 모델에서는 bidirectional attention만 enable 했을 때, embedding performance가 감소하는 모습을 보였다.

 

따라서 아래의 추가적인 전략들을 사용해, Decoder 모델이 미래 토큰을 고려하도록 학습을 진행한다.

 

 

Masked next token prediction (MNTP)

모델이 bidirectional attention을 인식하도록 하는 전략이다.

MNTP는 next token prediction + Masked language modeling을 결합한 새로운 학습 목표이다.

임의의 시퀀스 x = (x1, ...x_n) 이 주어진다고 가정할 때, 먼저 입력 토큰의 일부를 마스킹한 다음 마스킹된 토큰을 예측하도록 모델을 훈련한다. 

마스킹 과정은 huggingface의 MLM script를 가져왔다고 한다.

masked token percentage를 탐색하기 위해 하이퍼파라미터 서치를 수행했고, LLaMA에서는 BERT의 masking 방식을 사용해 20%를, Mistral에서는 RoBERTa의 masking 방식을 사용해 80%로 훈련

BERT : 15% 마스킹 이후, 마스킹 된 토큰의 10%는 무작위 토큰으로 교체
RoBERTa : 15% 마스킹 이후 후처리 X

또한 위치 i에서 masked token을 예측할 때에는 이전 위치인 i-1의 token representation에서 얻은 logit을 기반으로 loss를 계산한다.

 

  • next token prediction
    다음 토큰은 이전 토큰 출력 벡터를 linear layer에 입력으로 전달해서 예측
  • Masked Language modeling (MLM)
    임의의 input token을 마스킹하고, 시퀀스 내의 다른 토큰들을 사용해서 마스킹 토큰을 학습

 

세부 작동 방식 예시

subtoken representation을 평균내서 한 단어의 임베딩을 계산한다.

예를 들어 "\(w_1 w_2 w_3\)"을 토큰화 한 문장이 "\(BOS t_{11}t_{12} t_{21}t_{22}t_{23} t_{31}\)" 일 경우, 평균은 다음과 같이 계산될 것이다.

e는 각 w, t의 최종 표현

 

하지만 MNTP 학습 방식에서는 다음과 같이 계산된다.

 

 

 

Unsupervised contrastive learning (SimCSE)

앞선 2단계를 거치면 word-level Encoder로 변환이 가능하지만, sequence level representation을 위해서는 부족할 수 있다. 

사전학습 과정에서 Object에 Next sentence prediction을 포함하는 bidirectional encoder와 달리, decoder-only LLM은 전체 문맥을 파악하도록 명시적으로 학습되지 않기 대문이다.

 

이 간극을 메우기 위해서 SimCSE 방식을 통해 Unsupervised contrastive learning을 적용한다.

입력 시퀀스가 주어지면, 독립적으로 sampling된 dropout mask을 적용한 모델에 동일한 input sequence를 2번 통과시켜 2개의 다른 표현을 생성한다.

일반적으로 사용되는 낮은 dropout 확률 (0.1) 보다 더 높은 dropout)0.3)을 적용한 성능이 더 좋았다

 

모델은 생성된 2개 표현의 유사성을 최대화하는 동시에, 배치 내 다른 시퀀스 표현과의 유사성을 최소화하도록 훈련된다.

최종적으로 산출된 token representation에 pooling operation을 적용하여 sequence representation을 얻었다.

pooling operation 방법에 따른 성능 비교도 하단 실험에 나와있다.

 

 

상세 실험과정

  • 모델 : 1.3B ~ 7B 파라미터에 이르는 3가지 Decoder-only LLM으로 실험한다
    • Sheared-LLaMA-1.3B
    • Llama-2-7B-chat
    • Mistral-7B-Instruct-v0.2
  • 학습 데이터 : 영어 위키피디아 데이터를 사용해 MNTP, SimCSE를 수행했다
    • 위키피디아 데이터는 실험하는 모든 모델의 사전 학습 데이터에 포함되어 있기에 선택
    • MNTP, SimCSE는 새로운 지식을 가르치는 단계가 아니기 때문
    • decoder model이 다음 토큰을 이용하는 방법과 sequence representation 방법만 가르치도록 함
  • Masked next token prediction
    • BERT에서 사용된 방법처럼 토큰의 일부를 무작위로 마스킹
    • 실험 모델에서 마스킹을 위한 special token이 없기 때문에, '_'를 special token으로 사용했음
    • LoRA를 사용해 모델을 Fine tuning해서 training objective를 이전 토큰 표현을 사용해 마스킹된 토큰을 예측
      -> decoder-only LLM의 사전 훈련 세팅과 최대한 일치하도록 함
    • 단일 80GB A100 GPU / batch size : 32 / 1000 step 동안 훈련
    • 7B 모델의 경우 훈련에 1.5시간 소요
  • Unsupervised contrastive learning
    • SimCSE approach 사용
    • positive example은 dropout mask를 사용해 동일한 input sequence를 모델에 2번 통과시켜 2개의 다른 표현을 생성
    • 배치 내의 다른 시퀀스가 negative sample로 사용됨
    • 모델이 이전 단계에서 학습한 지식을 유지할 수 있도록 simCSE 훈련 이전 MNTP LoRA 가중치를 기본 모델에 병합하고 새로운 LoRA 파라미터 초기화
    • 단일 80GB A100 GPU / batch size : 128
    • 7B 모델의 경우 훈련에 2.5시간 소요

 

 

3. LLM2Vec-transformed models are strong unsupervised text embedders

Evaluation on word-level tasks

실험 세팅

  • CoNLL-2003 벤치마크 : named entity recognition data, 영어 & 독일어 (https://paperswithcode.com/dataset/conll-2003)
  • 각 입력 문장을 임베딩한 후 frozen하고,representation 상단에 task-specific linear classifer를 훈련해서 부착함
  • 최신 encoder only model인 DeBERTa-v3-large와 비교 진행

 

실험 결과

검정색 선 : Causal attention (Uni)
빨간 선 : DeBERTa-v3-large

 

  • Causal attention (Uni) 를 사용하는 것만으로 이미 Encoder-only baseline을 뛰어 넘었다.
    -> 이는 모델 크기의 차이로 인해 어떻게 보면 당연한 결과이다. (Deberta-large-v3 : 304M / Decoder model : 1.3 ~ 7B)
  • Naive하게 bidirectional attention을 적용할 경우(Bi) 성능이 크게 저하되었다
    • 특이하게도 Mistral의 경우에는 성능 저하가 훨씬 덜했고, NER 태스크에서는 오히려 성능이 오르기도 했다
  • MNTP만 사용한 경우 (w/o SimCSE) 모든 모델과 작업에서 성능이 향상되엇다
  • SimCSE를 사용한 경우, 오히려 성능 하락이 일어났다
    -> SimCSE는 sequence level에서 representation을 조정하기 때문에, word-level에서는 성능 저하가 보이는 것으로 판단

 

 

Evaluation on sequence-level tasks

실험 세팅

  • MTEB 벤치마크 : 7개의 다양한 임베딩 작업의 모음이며, 총 56개의 데이터셋을 포괄하는 대규모 text embedding 벤치마크
    (https://huggingface.co/blog/mteb)
  • 선행 연구들에 따라 query에만 task-specific instruction을 사용해 평가
  • 사용한 instruction은 선행 논문 Improving text embeddings with large language models[각주:1] 와 같음

사용된 Instructions

  • mean pooling, weighted mean pooling을 적용할 때에는 instruction token을 제외함
  • 베이스라인으로는 아래 논문에서 사용된 unsupervised BERT 모델 활용
    • SimCSE: Simple contrastive learning of sentence embeddings[각주:2]
    • Repetition improves language model embeddings[각주:3]
  • 가장 성능이 우수한 풀링 방법 선택을 위해 / 평가 속도를 높이기 위해 각 MTEB 카테고리의 대표 작업으로 구성된 15개 작업 하위 집합에 대해서 실험 수행
    -> 편향되지 않기 위해 전체 MTEB와 거의 동일한 비율로 하위 작업 추출
    MTEB subset



  • Pooling 방법은 선행 연구[각주:4]에서 사용된 방법 활용
    • 여기서도 SBERT의 방식을 따왔다
      - [EOS] : 가장 마지막 토큰만 사용
      - [mean] : 평균
      - [weighted mean] : position에 따른 가중치

 

 

15개 하위 집합에 대한 실험 결과

EOS
mean : mean pooling
w.mean : weighted mean pooling

3가지 모델 모두에 대한 다양한 pooling 방법이 MTEB 작업에 미치는 영향을 보여준다.

  • Text Embedding을 구성할 때 Causal attention이 최적의 방법이 아님을 분명하게 알수 있음
  • 단어 수준에서의 결과와 마찬가지로, 학습 없이 bidirectional을 적용할 경우 성능이 저하됨
  • MNTP, SimCSE 적용 시 성능 향상

 

전체 MTEB에 대한 실험 결과

  • bidirectional attention + MNTP만 적용 (=w/o simCSE) 해도 상당한 성능 향상 관찰됨
    • Echo embedding과 성능이 비슷하나, echo embedding은 입력을 반복해 시퀀스 길이가 두 배로 늘어나 inference 속도가 상당히 느려지므로 LLM2Vec이 훨씬 효율적
      Echo embedding과의 Evaluation time 비교
  • SimCSE 추가하면 세 모델 모두 큰 폭으로 향상됨
  • 특히 Mistral-7B 모델은 모든 비지도 모델 중 최고의 성능을 발휘함

 

 

 

 

4.  LLM2Vec은 모델에 어떤 영향을 미치는가?

모델이 future token에 대한 정보를 습득하는 데에 도움

LLM2Vec으로 변환을 거친 모델이 미래 토큰의 정보를 얼마나 잘 통합하는 지 분석하기 위해 테스트를 진행했다.

 

실험 세팅

선행 연구[각주:5]에서 수집된 합성 데이터셋을 통해 실험을 진행했다.

데이터 구성은 35개의 문장 triple set \(\{(q_i, s_i^+, s_i^-)\}^{35}_{i=1}\)으로 구성되어 있다.

데이터셋 예시

  • \(q_i\) = \((A_i, B_i)\)
  • \(s_i^+\) = \((A_i, C_i)\)
  • \(s_i^-\) = \((A_i, D_i)\)
\(A_i, B_i, C_i\) 는 모두 비슷한 의미를 지님
\(B_i, D_i\) 는 서로 다른 의미

 

문장의 첫 부분 \(A_i\) 에만 pooling을 진행해서 각 문장에 대한 시퀀스 표현을 생성한 다음, 결과 임베딩 간의 코사인 유사도를 계산

이 때 prefix \(A_i\) 의 표현에 대해 미래 토큰인 \(B_i, C_i, D_i\)의 정보를 통합하는 모델은 positive example에 더 높은 similarity를 할당할 것이다.

 

실험 결과

 

  • S-LLaMA-1.3B의 경우 bidirectional attention과 MNTP 훈련만으로도 positive / negative example에 대한 명확한 구분 가능
  • Mistral-7B의 경우 모든 설정에서 query - positive example 유사성이 query - negative eample 유사성보다 큼

 

 

 

training (MNTP, SimCSE) 없이도 Mistral model에서 bidriectional attention이 잘 먹히는 이유는?

실험 세팅

  • 각 모델의 단일 입력 시퀀스 (위키피디아의 무작위 문장)을 넣고 모든 레이어 l에서 causal \(H_l^c\), bidirectional attention \(H_l^{bi}\) 로 hidden representation을 계산

 

  • 모든 레이어에 대해 causal attention, bidirectional attention으로 구성된 표현 사이의 코사인 유사도, 즉 \(sim( H_l^c , H_l^{bi}\)을 계산
    -> 대부분의 레이어에서 이 유사도는 낮을 것으로 예상되는데, 학습 없이 bidirectional attention을 활성화하면 상당히 다른 표현이 생성되기 때문

 

실험 결과

LLaMa 기반 모델에서는 유사도가 매우 다르게 나타나는 반면, Mistral 에서는 높은 코사인 유사도를 보임

-> Mistral을 훈련할 때 어떤 형태의 bidirectional attention이 학습된 것으로 보이나, 자세한 조사는 향후 연구로 남겨둠

 

 

5. Combining LLM2Vec with supervised contrastive learning

현재까지는 unsupervised contrastive learning을 사용한 방식만 평가했다.

그렇다면 supervised contrastive learning을 사용하면 어떻게 될까?

 

실험 세팅

선행 연구[각주:6]에서 큐레이션한 E5 데이터셋의 일부를 그대로 지도학습에 사용한다.

  • E5 dataset
    • 여러 데이터셋 (ELI5, SQuAD, NLI...) 이 합쳐진 약 150만개의 sample로 구성
    • hard negative, in-batch negative를 이용한 contrastive learning
    • 마찬가지로 LoRA fine-tuning을 활용
      MNTP 가중치는 기본 모델에 병합, 훈련 가능한 LoRA 가중치는 SimCSE 가중치로 초기화됨
    • batch size 512 / training step 1000

 

 

실험 결과

세 모델 모두 LLM2Vec을 사용했을 때 baseline (Causal attention + Weighted mean) 보다 더 좋은 성능을 보였다

SimCSE를 수행하는 것은 unsupervised contrasitive learning을 수행할 때보다, 심지어 MNTP만 수행하는 것보다 성능이 떨어진 경우도 있었다.

하지만 SimCSE를 사용하는 LLM2Vec은 훨씬 sample-efficient learning을 할 수 있기 때문에, 컴퓨팅 또는 데이터 제약이 있는 환경에서 매우 유용하다.

 

 

LLM2Vec leads to more sample-efficient training

실험 세팅

LLM2Vec 방식의 샘플 효율성을 입증하기 위해, 25 step마다 체크포인트를 저장하고 15개의 하위집합 MTEB에서 평가를 진행했다.

 

실험 결과

세 모델 모두 훈련 초기에 LLM2Vec을 적용한 버전에서 더 나은 성능을 달성하는 것으로 그 효율성이 입증되었다.

 


참고 문헌

https://www.wolfram.com/language/12/neural-network-framework/use-transformer-neural-nets.html?product=mathematica

https://magazine.sebastianraschka.com/p/understanding-and-coding-self-attention

https://huggingface.co/blog/alonsosilva/nexttokenprediction