[논문리뷰]Gecko: Versatile Text Embeddings Distilled from Large Language Models(DeepMinde, 2024)

Jinhyuk Lee, Zhuyun Dai, Xiaoqi Ren, Blair Chen, Daniel Cer, Jeremy R. Cole, Kai Hui, Michael Boratko, Rajvi Kapadia, Wen Ding, Yi Luan, Sai Meher Karthik Duddu, Gustavo Hernandez Abrego, Weiqiang Shi, Nithi Gupta, Aditya Kusupati,Prateek Jain, Siddhartha Reddy Jonnalagadda, Ming-Wei Chang, and Iftekhar Naim.2024. Gecko: Versatile Text Embeddings Distilled from Large Language Models.arXiv:2403.20327 [cs.CL] https://arxiv.org/abs/2403.20327

Problem Statement

[다목적 텍스트 임베딩 모델의 한계] 기존의 임베딩 모델들은 특정 작업에 맞춰 개별적으로 구축되어 왔으며, 다양한 작업에 걸쳐 포괄적인 지식을 다루기 위해선 방대한 양의 훈련 데이터가 필요하다는 한계점을 지니고 있다.

[범용성 확보의 한계] 기존의 synthetic data 구축 방식은 주로 도메인 적응(domain adaptation)에 초점을 맞췄고, 범용성을 고려하여 데이터셋을 구축하지 않았다. 또한 이전 연구들은 주로 사람이 레이블링한 데이터를 보강하거나 특정 성능을 개선하는데 초점을 맞추었다.

Methodology

Overview

Gecko는 대규모 언어 모델(LLM)의 방대한 지식을 활용하여 텍스트 임베딩 모델의 성능을 향상시키는 두 단계 증류 과정을 통해 개발되었다. 첫 번째 단계인 pre-finetuning에서는 모델을 대규모 커뮤니티 QA 데이터셋과 웹 텍스트 쌍으로 훈련시켜 다양한 텍스트에 대한 광범위한 지식을 습득하게 한다. 두 번째 단계, fine-tuning에서는 LLM을 사용하여 합성 데이터셋인 FRet을 생성하고 이를 통해 모델을 미세조정한다. FRet (Few-shot Prompted Retreival dataset) 데이터셋은 LLM이 다양한 쿼리를 생성하고, 기존 임베딩 모델을 사용하여 후보 문서를 검색한 후, 다시 LLM이 이 문서들의 관련성을 평가하여 긍정 문서와 하드 네거티브 문서를 선별함으로써 구축된다. 이 두 단계를 거쳐 Gecko는 MTEB 벤치마크에서 뛰어난 제로샷 성능을 달성하는 다목적 임베딩 모델로 완성된다.

Gecko의 학습 과정은 Pre-finetuning → FRet(LLM Distillation) → Unified Fine-tuning의 세 단계로 구성된다.

Pre-finetuning

• Input: question-answering 쌍 or title-body 쌍 (\(q_i, p_i\))
• Output: Mean Pooling된 토큰 임베딩 벡터 (\(\mathbf{q}_i, \mathbf{p}_i\))

사전 학습을 위해서 사용되는 데이터셋은 Community QA와 Web에서 크롤링한 title-body 쌍을 학습에 이용한다. 먼저 사전학습된 언어 모델 \(\mathcal{M}\)에 주어진 각 쿼리-패세지 쌍의 집합 \(\mathcal{D}_{pre} = \{(q_i, p_i) \}_{i=1}^N\)을 입력한다. 이 때, 쿼리의 경우 각 쿼리별로 어떤 task를 위한 쿼리인지 명시한 \(t_i\)를 같이 입력시킨다. 예를 들어,

• 질문 (\(q_i\)): What is the color of lemon?
• 테스크 (\(t_i\)): Question-answering task
• 입력: Question-answering task. What is the color of lemon?

와 같은 형태로 모델에 입력된다. 이후, 쿼리와 패세지의 각 토큰 임베딩을 mean pooling하여 각각 쿼리 임베딩 (\(\mathbf{q}_i\))과 패세지 임베딩(\(\mathbf{p}_i\))으로 정의한다.

$$\mathcal{L}_{pre} = \frac{1}{B}\displaystyle\sum_{i=1}^B \Big[ -\log\frac{\exp(\text{sim}(\mathbf{q}_i, \mathbf{p}_i)/\tau)}{\sum_{j=1}^B \exp(\text{sim}(\mathbf{q}_i, \mathbf{p}_j)/\tau)} \Big]$$

Pre-finetuning을 위해 InfoNCE loss를 사용한다. 두 mean pooling 임베딩 간의 코사인 유사도를 logit으로 하고, mini-batch(\(B\))내의 다른 문서들이 한 쿼리의 negative로 학습에 사용된다.

FRet: Two-Step LLM Distillation

Step 1. LLM-based Diverse Query Generation

• 입력: 웹 코퍼스에서 무작위로 선택한 passage \(p_{\text{seed}}\)
• 출력: 태스크 설명 \(t\), 쿼리 \(q\)

$$\text{LLM}(\mathbb{P}, p_{\text{seed}}) \rightarrow (t, q)$$

LLM에 few-shot 예시와 instruction으로 구성된 고정 프롬프트 \(\mathbb{P}\)를 웹 코퍼스에서 무작위로 선택한 passage와 함께 같이 입력시킨다. 그러면 LLM은 이 passage와 관련된 태스크와 쿼리를 생성하게 된다. 이를 통해 뉴스, 블로그, 위키 등 다양한 도메인 데이터를 얻을 수 있다.

Step 2. LLM-based Positive and Negative Mining

기존 방식은 LLM이 생성한 쿼리 \(q\)와 입력으로 들어간 패세지 \(p_{\text{seed}}\)를 positive로 사용하였다. 하지만, 이 방식은 항상 \(q\)에 대한 최적의 패세지가 \(p_{\text{seed}}\)가 아닐 수 있음을 간과한다.

따라서, Passage Retrieval과 LLM Reranking 과정을 통해 정답 문서를 선택하게 된다. 사전학습된 인코더를 사용해 쿼리와 후보 패세지 집합을 임베딩하고, 유사도를 계산하여 상위 \(N\)개의 패세지 \(P = {p^{(1)}, \dots, p^{(N)}}\)를 검색한다. 이후 LLM을 이용하여 리랭킹을 하게 된다. 이 때, 리랭킹을 위해서는 두 가지 score를 도입한다.

$$QL(q, p) = \text{LLM}(q \mid p, \mathbb{P}_{QL})$$

Query Likelihood (QL)는 패세지가 쿼리를 “얼마나 잘 설명해줄 수 있는지”를 평가한다. 예를 들어 쿼리가 “Who made the atomic bomb?”이고, 패세지가 “Oppenheimer developed the atomic bomb”일 경우, LLM은 쿼리가 패세지로부터 자연스럽게 생성될 가능성이 높다고 평가해 높은 점수를 부여할 것이다.

$$RC(q, p) = \text{LLM}(\text{label} \mid q,p,\mathbb{P}_{RC})$$

Relevance Classification(RC)는 쿼리-패세지 쌍이 주어졌을 때, LLM이 이 쌍을 relevant / irrelevant로 분류하는 방식이다.

$$R(q,p) = \frac{1}{r_{QL}(q, p)} + \frac{1}{r_{RC}(q, p)}$$

\(r_{QL}(q,p)\)와 \(r_{RC}(q,p)\)는 각 스코어를 기반으로 산출된 순위이다. 이 순위를 앙상블하여 리랭킹을 위한 최종 스코어, Reciprocal Rank Fusion (RRF)를 계산한다. 이를 통해 각 쿼리마다 하나의 정답 문서를 최종적으로 결정하게 된다.

Unified Fine-tuning Mixture

FRet 방식을 통해 만들어진 데이터셋은 물론 여러 도메인과 여러 테스크를 포함하는 범 도메인 데이터셋이지만, 그럼에도 불구하고 여전히 데이터 편향이 나타날 수 있다. 이를 위해, Gecko는 FRet데이터를 중심으로 기존의 여러 벤치마크 데이터셋 (e.g., HotpotQA, NQ, FEVER, MedMCQA, SNLI 등)을 합쳐 하나의 학습을 위한 Mixture로 구성한다.

$$(t_i, q_i, p^+_i, p^-_i)$$

학습을 위해서는 한 쿼리에 대해 위와 같이 태스크, 정답 문서, 하드 네거티브 문서 (= distractor)로 만든다. 이 통일된 포맷 덕분에 retrieval, classification, STS 등 이질적인 태스크를 하나의 contrastive 학습으로 학습 가능하다.

다만, classification 태스크의 경우 원래 패세지가 없는 구조이므로 변환이 필요하다.

• label from RC (relevant / irrelevant)
• 같은 label \(y\)를 가진 다른 입력 \(x^+\)를 positive
• 다른 label을 가진 입력 \(x^-\)를 negative로 설정

추가적으로 (\(x\), label \(y\))에 대해서 false negative 문제 방지를 위해 \((x, x^+, x^-)\)마다 unique ID를 부여한다. 이 ID를 query/positive/negative 모두에 추가해서, batch 내 다른 label 샘플이 헷갈리지 않도록 했다.

Training Objective

$$L_{main} = \frac{1}{B}\sum_{i=1}^{B}\left[-\log\frac{e^{sim(q_i,p_i^+)/\tau}}{\sum_{j=1}^{B}\left(e^{sim(q_i,p_j^+)/\tau} + \mathbb{I}_{[j \ne i]}e^{sim(q_i,q_j)/\tau}\right) + e^{sim(q_i,p_i^-)/\tau}}\right]$$

Fine-tuning을 위해서 InfoNCE 기반의 loss를 사용한다. 주어진 mini-batch 크기 \(B\)에서, 각 쿼리 \(q_i\)는 positive \(p_i^+\)를 정답으로 구분해야 하고, hard negative \(p_i^-\)뿐 아니라 같은 배치 내 다른 positive와 다른 query들도 부정 예시로 취급한다 (in-batch negative). 위 식에서 \(\mathbb{I}_{[j \neq i]}\)는 쿼리-쿼리 쌍에 대해서도 비교를 진행한다. 이를 통해 두 query가 비슷한 문장을 가지면, 모델이 혼동하는 경우를 방지한다. 이로써 retrieval + similarity + classification을 동시에 학습할 수 있다.

Experiments

Main Result

Table 1은 MTEB 벤치마크에서 Gecko와 다른 기준 모델들의 성능을 요약한 결과다. Gecko는 256 차원 임베딩을 사용하는 gecko-1b-256 모델과 768 차원 임베딩을 사용하는 gecko-1b-768 모델의 두 가지 버전이 있다.

gecko-1b-256은 256 차원 임베딩 차원의 기존 모델들을 모두 능가하며, 768 차원 임베딩을 사용하는 모델들보다도 뛰어난 성능을 보인다. gecko-1b-768은 MTEB의 7가지 작업에서 평균 점수 66.31을 달성하며, 7배 더 큰 모델이나 5배 더 높은 차원의 임베딩을 사용하는 모델들과 경쟁한다. 특히, Retrieval, STS, Classification, Summary에서 새로운 SOTA(State-of-the-Art) 성능을 기록하였다.

Ablation Study

[LLM을 labler로 사용] LLM이 긍정 문서와 하드 네거티브 문서를 선택하는 전략을 실험한 결과다. 실험 결과, LLM이 선정한 가장 관련성 높은 문서 \(p_1\)를 긍정 문서로 사용했을 때, 쿼리를 생성하는 데 사용된 원본 문서 \(p_{\text{seed}}\)를 사용했을 때보다 항상 더 좋은 성능을 보였다. 이는 원본 문서가 쿼리에 대한 최적의 긍정 문서가 아닐 수 있음을 의미한다.

개별 태스크별 학습(QA, Fact-checking 등)보다 다양한 태스크를 균등하게 샘플링한 FRet-all-tasks가 더 높은 성능을 보였다. Unified format이 중요하며, naive concatenation은 성능 저하 유발하며, Human data와 결합 시 Classification과 STS가 크게 개선된다.

Conclusion

Contributions

• LLM을 활용한 2단계 distillation(FRet)을 통해 Positive/Negative relabeling을 자동화하고 데이터 품질을 향상시켰다.
• 다양한 태스크와 데이터셋을 Unified Format으로 통합하여 범용 임베딩 모델을 학습할 수 있음을 보였다.
• Gecko는 1.2B 파라미터, 256/768 차원 임베딩만으로도 대규모 모델(7B+, 3-4k 차원)을 능가하는 성능을 보여, 모델 효율성과 범용성을 동시에 달성했다.

Limitations

• FRet는 현재 영어 기반으로만 구축되어 다국어 지원이 제한적이다.
• LLM을 labeler로 사용하는 과정에서 비용 및 효율성 문제가 존재한다.
• 완벽한 gold standard가 아니므로, LLM judgment에 따른 bias 가능성이 남아 있다.

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