논문 리뷰) Improving Text Embeddings with Large Language Model

본 글은 [Open-Up] 오픈소스 소프트웨어 통합지원센터로부터 지원받아 작성하였습니다.

 

23년에 공개된 ‘Improving Text Embeddings with Large Language Model’ 에서는 오픈소스 LLM인 Mistral 7B 모델과 ChatGPT4로 생성한 학습 데이터를 활용한 방식을 제안하였다.
LLM 모델을 생성 데이터로 Contrastive learning(InfoNCE loss)한다는 심플한 아이디어로 SOTA(24년 5월 기준,  e5-mistral-7b-instruct 모델로 아직까지 6위)를 달성했다. 

MTEB 리더보드

 
[Abstract]
제안자는 합성데이터와 1k 미만의 training step(가중치 업데이트 횟수)만으로 높은 품질의 텍스트 임베딩을 얻었다.

현존하는 방법은 주로 몇단계에 걸친 semi-supervised 사전학습과 fine-tuning을 위한 소량의 label 데이터를 요구해왔다.
그러나, 제안된 방법은 데이터셋을 수집하거나 복잡한 훈련 단계를 거칠필요가 없다.

closed LLM을 데이터 합성에 활용하였으며, 일반적인 Contrastive loss(infoNCE Loss)로 LLM을 파인튜닝하였다.

합성 데이터만으로 학습했을때도 쓸만한 모델이 되었으며, 파인튜닝 데이터까지 활용하니깐 SOTA가 되었다.


[1. Intro]
제안자는 closed LLM을 활용하여, 여러 태스크에 걸친 데이터셋을 생성하였다.
2단계로 데이터를 합성하였는데, 첫번째 단계는 단순히 LLM이 주어진 조건에 해당하는 Task가 어떤 것들이 있는지 brainstorm하게 헀으며,
두번째 단계에는 LLM이 brainstorm한 task들에 대해 구체적인 데이터셋을 뽑아내도록했다. 

다양한 데이터셋 생성을 위해, 여러 prompt template을 적용하였으며, 작은 BERT 모델보다는, 
파워풀한 오픈소스 LLM(Mistral 7B)을 모델로 선택하였다.
해당 LLM의 경우, web data로 사전학습된 점과 contrastive 러닝을 했던 것이 이점이 있을거라 판단했다.
합성 데이터만으로 학습했을때도 쓸만한 모델이 되었으며, 파인튜닝 데이터까지 활용하니깐 SOTA가 되었다.

또한, rotation position embedding으로 바꿈으로써 최대 32000 토큰까지 효과적으로 검색을 수행할수 있음을 확인하였다.
(그러나, 해당 모델의 max_input_sequence는 4096으로 확인됨)


[3.1. Synthetic Data Generation]
데이터 생성 방법이 이 논문의 주된 포커스이다. 데이터 생성 방법만 이해해도 논문의 반을 이해한 셈.
제안자는 데이터를 Asymmetric Tasks vs Symmetric Tasks의 2가지로 나누어 생성하였음.

Asymmetric Tasks의 경우, short-long match, long-short match, short-short match, and long-long match의 4가지 방법으로 분류하였고 이에 따라 Prompt를 달리 설계하였음.
Symmetric Task의 경우, STS and bitext retrieval로 분류하였고, 이에따라 Prompt를 달리 설계하였음


<Asymmetric Task>
비대칭 Task는 크게 아래와 같이 적용되었다.
- short-long : query와 검색 
- long-short : text classification 
- short-short : 문장 간 유사도 검색
- long-long : 문서/문단 간 유사도 검색

<Symmetric Task>
- STS : 두 문장 생성하고, 유사도 점수 매기기
- biText 검색 : 1문장(A)을 A언어로 만들고, 추가적으로 A언어의 문장을 B언어로 번역하여 2문장(B1,B2) 생성. 
                       그리고 유사도 점수를 각각 매기기 (A와 B1, A와 B2)

[short query & Long Document Matching의 Prompt 중에서 Task Brainstorming 지시하기]

 
 
 
 
전체 Full-Prompt를 참고하면, {query_length} 을 {less than 5 words, 5-10 words, at least 10 words}에서 선택되게 하고, {num_words}와 {difficulty}도 마찬가지 방식으로 랜덤하게 설정되게 함으로써, Prompt의 다양성을 키웠다.

 
 
 
 
아래 그림은 STS 데이터셋을 만드는 프롬프트 예시이다.

 
 
 
[3.2 Training]
query 문장은 task_instruction + query로 정의하였다.

[EOS] token을 query와 document 끝에 붙여주었고,

임베딩 값을 취할때 last hidden state에서 [EOS] token의 임베딩 값을 활용하였다.

(논문에서는 last layer [EOS] vector라는 워딩이 있어 헷갈렸지만, code 확인시 last hidden state가 맞음)

 # https://github.com/kamalkraj/e5-mistral-7b-instruct/blob/master/peft_lora_embedding_semantic_search.py
 
 class MistralForSequenceEmbedding(MistralPreTrainedModel):
    def __init__(self, config):
        	super().__init__(config)
        	self.model = MistralModel(config)
       		self.normalize = True
        
        .....
        
    def forward(...)
 		last_hidden_state = transformer_outputs.last_hidden_state
        	embeddings = self.last_token_pool(last_hidden_state, attention_mask)
        
        .....

잘 알려진 InfoNCE loss를 활용하였으며,
in-batch negative와 hard negative를 모두 적용하였다.

infoNCE Loss 식

 

InfoNCE Loss 식에서 빨간색 항은 positive sample간의 거리,

초록색 항은 negative sample간의 거리를 의미한다.

분자(빨간색)를 minimize하고, 분모(초록색)를 maximize할수록 InfoNCE Loss는 최소화된다.

즉, 의미적으로 가까운 문장의 거리는 가깝게,

의미적으로 먼 문장 간의 거리는 멀게하는 것이 핵심이다.

 

그리고 'in-batch negative'란, batch 내에 있는 다른 문장을 모두 negative sample로 취급하는 것이다.

또한 Hard-negative는 질의(Query)와는 관련이 멀거나 반대인 negatvie sample을 말한다. 

해당 논문에서는 아래 그림과 같이, LLM(ChatGPT4)를 통해 Hard-negative를 생성하였다.

Hard-negative 생성 Prompt

 

[4. Evaluation]
나름 유명한 모델이니, 간략히 서술하자면

MTEB, BEIR 벤치마크에서 SOTA 내지는 좋은 성능을 보였고, (사이즈가 클수록 좋은 성능, E5_large)
Multilingual Retrieval에서도 좋은 성능을 보임
 

[5. Analysis]
사전학습에서의 Contrastive learning의 효과

사전학습에서의 Contrastive learning의 효과

XLM-R-large의 경우, 사전학습때  contrastive learning으로 성능이 향상이 되었으나 
LLM(mistral 7B)의 경우, 그렇지 않았다.
이는 광범위한 auto-regressive pre-traing이 LLM을 이미 좋은 임베딩 형성에 충분히 기여했기때문이라고 볼수 있다. 
따라서, 최소한의 fine-tuning 정도만이 필요하다고 볼수 있겠다. 혹은 필요가 없거나...
(그렇다고 사전학습때 constrastive learning을 안해봐도 된다라고 판단하기에는 실험의 양과 종류가 부족하게 보인다. 그냥 auto-regressive pre-traing만으로 임베딩 형성이 잘된 케이스로 보면 좋을 것 같다. 물론 prediction 타입의 semi-supervised가 더 성능이 좋아보이긴 한다.)