E5 Text Embedding 시리즈 논문 및 구현 설명

범용 Text Embeeding 을 평가하기 위한 MTEB 리더보드 ( https://huggingface.co/spaces/mteb/leaderboard ) 상에서 최근 최상위권을 기록한 Microsoft 의 E5 Text Embeeding 시리즌 논문 리뷰입니다. 크게 다음 3가지 종류의 모델이고,

Model Name	Backbone	Pre-Trianing Dataset	Fine-tuning Dataset	Training Method
e5	BERT-style Encoder ( < 1B)	자체 구축	기존 Text Pair Labeled Dataset 들 활용	2-Stage Training
e5-instruct			+ LLM (GPT3.5-Turbo&4) 을 사용해 자체 생성한 Text Pair 데이터셋
e5-mistral-7b	LLM (Mistral-7B)	-		Contrastive Loss + LoRA

 

아래 페이퍼들을 복합적으로 리뷰한 내용입니다.

[1] Text Embeddings by Weakly-Supervised Contrastive Pre-training (E5) Dec. 2022

  |_ [1-1] SimLM: Pre-training with Representation Bottleneck for Dense Passage Retrieval (SimLM) Jul. 2022

[2] Improving Text Embeddings with Large Language Models (E5-Mistral-7B) Jan. 2024

[3] Multilingual E5 Text Embeddings: A Technical Report Feb. 2024

 

---

1. E5

Text Embeddings by Weakly-Supervised Contrastive Pre-training (E5) Dec. 2022

 

1-1. 개요

• 이 Paper 들은 Retrieval, Clustering, Classification, Text Similarity 등 다양한 NLP Task 에서 Zero-Shot 또는 Few-Shot으로 범용적으로 강력한 성능을 보이는 Text Embedding 생성을 목적으로 합니다. E5 논문은 이 당시 (2022년 말) 적은 파라미터의 모델 ( < 1B) 로 MTEB(종합 Text Embedding 평가 리더보드) 에서 SOTA를 달성 하였었습니다.
• E5는 BERT를 기반으로, 아래의 2-Stage 로 Text Embedding 모델을 학습 시켰습니다.
  • 1st-stage : CCPairs 라는 Unsupervised Large-Scale Text Dataset 을 자체 구축하고 이 데이터셋으로 Pre-trainig
  • 2nd-stage (Fine-Tuning) : MS-Marco(Retrieval), NLI (Similarity), NQ 데이터 셋으로 Fine-tuning. Mined Hard Negatives + Distillation From Cross Encoder 기법 사용

1-2. 1st-stage : Pre-Training

이 논문의 Pre-Training 단계에서 주요한 부분은 자체 구축한 CCPairs 데이터셋입니다. 이를 구축한 과정은 다음과 같습니다. 

 

Step1. Web 에서 Large Scale 로 q-p Text Pair 데이터셋을 수집하였습니다. Reddit 에서 <post - comment> pair,  Wikipedia 에서 <tittle-passage> pair 와 같은 데이터들이 포함됩니다. 이를 Noisy Large Dataset 으로 부르겠습니다.

 

Step2. 이 Noisy Large Dataset 으로 모델을 학습 시킵니다.

 

Step3. Step2에서 학습된 모델로, Noisy Large Dataset 데이터를 필터링 합니다. 필터링 방식은 다음과 같습니다.

* 각 q에 대해서 <정답 passage + 100만 개의 random passage>를 후보군으로 하여 Step2 의 모델로 후보군을 랭킹합니다.

* top-2 안에 원래 정답이 있으면 좋은 quality라고 가정하고 이 pair 는 선택합니다.

 

이 방식이 작동하는 근거는 논문에서 다음과 같이 설명합니다.

"The intuition for this technique comes from the memorization behaviors of neural networks:
when trained on noisy datasets, neural networks tend to memorize the clean labels first and then
gradually overfit the noisy labels. Similar techniques [42, 15, 23] have been widely used for removing
dataset noises"

> 이에, 아마 STEP2 에서 모델을 적은 steps 동안만 학습한 것으로 보입니다. 

 

이렇게 필터링 된 데이터를 논문에서 CCPairs 라고 부릅니다. 이 CCPair 로 학습은 InfoNCE loss를 써서, in-batch negatives contrastive learning 으로 학습하였다고 합니다. 이 때, batch size 는 충분이 큰 32,768를 사용하였습니다. score 함수로는 cosine similarity(x,y) / 0.02 를 사용했다고 합니다.

 

1-3. 2nd-stage : Fine-Tuning

논문에는 MS-Marco (Retrieval), NLI, NQ 데이터셋을 이용하여 "mined hard negatives and knowledge distillation from a cross-encoder (CE) teacher model" 로 학습하였다고, 간략하게 인용만 하고 넘어갑니다. 인용된 논문은 "SimLM: Pre-training with Representation Bottleneck for Dense Passage Retrieval (SimLM) Jul. 2022" 입니다. 이 인용된 논문에 상세히 소개된 Fine-tuning 기법은 다음과 같습니다.

먼저 위 그림과 같이 순차적으로 학습됩니다. Retriever 1 이 학습이 끝나고 이를 활용해 Retirevel2 를 학습하고 하는 식입니다. 최종적으로는 Retirever_distill 이 사용되며, 각 단계의 학습 방식은 다음과 같습니다.

Retriever 1 :  먼저 주어진 query와 BM25 유사도는 높지만 negative 인 passages 들을 가지고 옵니다. 이를 mined hard negatives 라 부릅니다. mined hard negative 와 in-batch negative 를 활용해서 위와 같이 InfoNCE loss 를 통해 contrastive learning 을 합니다. 구체적인 구현 방식은, 저자들이 공개한 코드(https://github.com/microsoft/unilm/blob/master/simlm/src/utils.py#L71) 를 보면서, 아래 다이어그램으로 정리하였습니다. 

Retriever 2 : Retriever 1 과 똑같이 학습 되는데, mined hard negatives 가 Retriever 1 모델에서의 점수는 높지만 negatives 인 것 들을 활용합니다.

 

Re-ranker (Cross-Encoder) : Retriever 2 에서의 점수는 높지만 negatives 인 것들을 활용합니다. Cross Encoder 이며 각 q-p pair 에 대해서 binary classification 을 하는 것이 아닌 주어진 q- {정답p, mined hard negatives} 에 대해 Contrastive Loss 로 학습됩니다.

 

Retriever_distll : Retriever 2 에서의 점수는 높지만 negatives 인 것들을 mined hard negatives 로 활용하여 Retrieval 1&2 와 똑같은 contrastive loss term도 활용합니다(아래 식에서 L_cont). 추가적으로 이전 스텝의 Re-ranker 가 내어준 확률 분포와의 KL Divergence 도 loss 로 활용합니다.

1-4. 결과

먼저 다양한 Text-Embedding Task 에서, 1st-stage 만 했을 때, 2nd-stage 까지 했을 때 둘다 SOTA 를 기록했다는 내용이 있습니다. 

또한, 1st-stage pre-training 을 할 때 구축했던 Consistency-based filter 가 좋은 효과를 있음을 보여줍니다.

 

1-5. 주요 내용 정리 (개인적)

• "when trained on noisy datasets, neural networks tend to memorize the clean labels first and then
  gradually overfit the noisy labels." --> 이 원리를 이용하여 Large Noisy Data 에 모델을 살짝 학습시켜 데이터를 필터링을 하는 것이 효과적이라는 것이 인상깊었습니다.
• 인용된 SimLM (2022) 라는 Paper와 Code에 상세히 나왔던, fine-tuning 시 hard-negative 를 잘 선택하고 이를 in-batch negative 와 조합해서 contrastive learning 을 하는 구체적인 테크닉들, Contrastive Loss 스타일로 Cross Encoder 를 통해 distillation 받으면서 Text-Embeding 을 잘 학습 시킬 수 있다는 점이 인상 깊었습니다. 
• 학습시 사용되는 text 유사도 함수로 계산 할때 temperature-scaled cosine similarity ( cos(x,y)/temperature ) 가 사용된 점

 

 

2. E5-Mistral-7B

Improving Text Embeddings with Large Language Models (E5-Mistral-7B) Jan. 2024

 

2-1. 개요 

E5를 포함하여 기존의 (non-LLM) SOTA Text Embedding 방식들은 크게  2Stage 방식을 도입하였었습니다.

• 1st : Billons 단위의 Weakly-Supervised Text Pairs 로 Pre-training
• 2nd : Labeled Dataset 을 사용하여 여러가지 방식으로 Fine-tuning

이 방법들의 단점은, 

- Engineering Effort가 많이 든다 (Large-Data Pretraining, 복잡한 Fine-Tuning)

- Manually Collected Dataset 에 의존해서 task 나 language 의 diversity가 부족하다.

 

이에 논문에서 제안하는 방식의 핵심 아이디어는

->  Text-Pair 데이터셋을 LLM 을 통해 다양하게 많이 만들고, 텍스트 임베딩도 LLM이 생성하도록 이 데이터로 LLM을 심플하게 학습시키자 입니다.

 

2-2. LLM을 활용한 Text-Pair 데이터셋 제작 방식 

Asymmetric Tasks ( e.g. query-document retrieval )

 * 2-Step 으로, Task Description 도 LLM 에게 생성하게 하고, 다음 Step으로 해당 Task 에 속하는 query-document pair 데이터 셋을 만들도록 하였습니다.

Step1) LLMs 에게 Task 목록을 생성하도록 하기.

Step2) 주어진 {task} 에 맞는 query - document pair 생성하기. {num_words}, {language} 등 다양한 placeholder 를 주어서 diversity를 보장함.

 

Symmetric Tasks (e.g. STS)

이는, Task 가 직관적이므로 LLM 이 Task Definition 을 생각하는 과정은 건너 뛰고,  "Retrieve semantically similar text." 등의  task definition 을 미리 정해놓고 바로 데이터 생성을 하였다고 합니다.

 

생성할 때 사용한 LLM 은 GPT 3.5-Turbo 와 GPT4이며, 결과적으로 150K Unique Instructions / 500K examples / 180M Tokens  규모의 데이터가 생성 되었다고합니다. 

 

2-3. 학습 방식

• 먼저 backbone 으로는 mistral-7b를 사용했고, query 와 passage 문자열을 각각 LLM 에 입력으로 넣어서, [EOS] 토큰 위치의 Last Hidden States 를 Text Embedding 으로 사용합니다.
• passage 는 텍스트 그대로 넣고 query 부분은 아래와 같이 task_definition 과 함께 넣어 주었습니다.

• in-batch negative 와 LLM 이 생성한 hard negative를 활용하여 contrastive loss 로 학습하였습니다.

• Fine-tuning 방식은 LoRA 를 사용하였습니다.
• 학습 Step 은 1K 도 채 안해도 좋은 성능을 보였다고 합니다.

2-4. 결과

synthetic data 란 위에서 설명한 LLM 을 통해 만든 text pair 데이터를 의미합니다. LLM 이 만든 데이터셋만 활용해도 SOTA 급 성능을 보이며, E5 학습할때 사용했던 기 구축된 데이터(MS-Marco, NLI 등) 까지 활용하면 SOTA를 달성한다는 결과입니다.

LLM 을 통해서 다양한 언어의 데이터를 만들었기에, 여러 언어에서 좋은 성능을 보인다는 내용입니다. 단, Low-resource Languaes 에서는 성능이 낮은데, 이는 mistral-7b가 predominantly pre-trained on English data 이기 때문으로 추정된다고 합니다. 

추가적으로, LLM을 활용했기 때문에 E5 와 같은 BERT-style encoder 를 기반으로 하는 text embedding 모델과 다르게 large-scale pretraining 을 굳이 안해도 된다는 결과를 보여주며, LLM 이기 때문에 long text 도 embedding 가능하다는 내용도 설명하는 듯 합니다.

3. E5-Instruct

Multilingual E5 Text Embeddings: A Technical Report Feb. 2024

 

마지막으로, mteb 리더보드와 hugginfcae hub 상에 "multilingual-e5-large-instruct" 되어 있는 모델입니다. 위 E5 와 E5-Mistral-7B 를 읽으면서

 

• LLM 으로 Text-Pair 데이터 증강 생성하는게 좋다는데는 이해가 감. 실제 결과도 좋음
• 근데 이 데이터를 LLM 에 학습시켜서 LLM 으로 Text Embedding 하면, 실사용할 때 너무 비용이 큰데, BERT-style encoder 에 학습시키면 어떨까? 

하는 생각이 들게 됩니다. 이 생각을 Micorsoft E5 팀에서 구현한게 E5-Instruct 이고 위 paper 에 간략하게 설명 되어 있습니다.

 

요약하면, 다음과 같습니다.

• Backbone : BERT-Style Encoder
• 학습 : e5 와 같은 학습 방식 (2-stage; pre-training -> fine-tuning)
• Fine-tuning 데이터셋 : e5-msitral-7b 학습 할 때와 같은 데이터 (LLM으로 만든 데이터 + 기존 labeled corpus 데이터)

E5 보다도 성능이 제법 뛰어나며, multi lingual 에 강점을 가지고, LLM 을 백본으로 하는 e5-mistral-7b 에는 못 미쳐도 거의 준하는 성능을 보여줍니다!