원제: M3-Embedding: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation

1. Introduction

Pre-training Language Model이 등장하며 텍스트 임베딩(텍스트 데이터 -> latent space로 인코딩)의 품질이 향상 되었고, 이런 임베딩 중 특히 Dense retrieval로 Query와 Document간 유사도를 검색하는 IR에 필수적으로 사용된다. 이 외에도 Multi-vector retrieval이나 Sparse retreival에도 사용된다.
하지만 이런 임베딩 모델들은 영어에만 특화되어있으며, 하나의 검색 기능만을 목표로 학습된다. 또한 학습 비용이 크기 때문에 짧은 입력만을 지원한다.
M3-Embedding은 다양한 언어를 지원하고, 서로 다른 언어 간 교차 언어 검색도 가능하다. 동시에 Dense + Sparse + multi-vector 까지 지원하는 범용 임베딩이 가능하다. 또한 문장의 길이에도 비교적 자유롭다.
이를 위해 사용한 기법으로
1) self-knowledge distillation

2. Related Work

Pre-training Language Model의 대중화로 데이터를 효과적으로 인코딩 하는 것에 더 유리해졌고, Contrastive Learning의 Negative Sampling의 개선과 Knowledge Distillation의 활용이 두드러진 성과를 보였다.
이를 주로 응용하는 분야는 Neural Retrieval로, 임베딩을 통해 의미적 관계를 측정해서 임베딩 유사도를 기반으로 문서를 검색할 수 있다.
1) Dense reetrieval: text encoder의 출력값을 [CLS] token 또는 average pooling과 같은 방식으로 집계해서 임베딩 유사도를 계산한다.
2) Multi-vector retrieval: text encoder의 출력 임베딩간의 상호작용으로 유사도를 계산한다.
3) Sparse retrieval: 각 단어의 가중치로 변화하여 단어를 기준으로 유사도를 계산한다.
이 모든 기능을 하나로 통합한 모델은 존재하지 않고, 대부분 영어에만 초점이 맞추어져있다.

3. M3-Embedding

다양한 언어를 지원하고, 서로 다른 세분성의 데이터를 처리하며, 일반적 검색 기능을 통합하기 위해서 만들어졌디.
$D^y: d^y \leftarrow fn^{*}(q^x,D^y)$, 이때 $fn^*(\cdot)$의 경우에는 위의 dense, sprase, multi-vector 중 어느 하나이다.

Data Curation

라벨이 없는 corpus에서 추출한 비지도 데이터, 라벨이 부여된 corpus에서 수집한 fine-tuning 데이터, 합성을 통해 생성된 fine-tuning data 3개를 사용한다.

1) 비지도 데이터: 다언어 corpus에 내재된 의미 구조를 추출함으로써 구축한다.. MTP에서 제공되는 정제된 데이터와 교차언어를 위한 NLLB의 병렬 문장도 데이터로 활용했다.
2) 고품질 fine-tuning 데이터: 영어는 8개, 중국어는 7개의 데이터셋을 통합했고, 그 외의 기타 언어는 MIRACL의 학습 데이터를 활용했다고 한다.
3) 합성 fine-tuning 데이터: 긴 문서를 샘플링한 후 임의의 단락을 선택하고, GPT-3.5를 사용하여 해당 단락을 기반으로 질문을 생성했다고 한다.

Hybrid Retrieval

Dense, Sparse, multi-vector retrieval을 한개로 통합한다.

Dense Retrieval

Query q는 text encoder을 통해서 hidden state $H_q$로 변환된다. 특수토큰 [CLS]의 hidden state를 Query 표현으로 사용한다.
$e_q = norm(H_q[0])$. 여기서 norm은 보통 L2 정규화라고 보면 된다.
Document passage p에 대해서도 동일하게 한다.
$e_p = norm(H_p[0])$
관련도 점수는 내적으로 계산한다.
$s_{dense} \leftarrow <e_p, e_q>$

Lexical Retrieval(Sparse)

Query 내 각 토큰 t에 대해 토큰의 가중치는
$w_{qt} \leftarrow ReLU(W^{\top}_{lex}H_q[i])$
이때, $W_{qt}\in R^{d \times 1}$은 d차원의 hidden state를 한개의 실수값으로 매핑해주는 행렬이다. 만약 t가 여러 번 등장하는 경우에는 최대 가중치만을 유지한다.
이를 바탕으로 Query와 Passage에 공통으로 등장하는 단어 집합 $q \cap p$의 중요도는
$s_{lex} \leftarrow \displaystyle\sum_{t \in q \cap p}(w_{qt} \cdot w_{pt})$

Multi-Vector Retrieval

Dense Retrieval의 확장으로 Query, Passage를 표현하기 위해 전체 출력 임베딩을 활용한다. Dense와 달리 벡터 하나로 요약하는 것이 아닌, 토큰마다 벡터를 가진채로 Query 토큰과 가장 잘 맞는 문서 토큰을 대응해 점수를 매긴다.
$E_q = norm(W^{\top}_{mul}H_q)$, $E_p = norm(W^{\top}_{mul}H_p)$
$W_{mul} \in R^{d \times d}$는 학습 가능한 projection matrix이다.
여기서 이전에 배운 ColBERT의 Late Interaction을 적용한다.

$s_{mul} \leftarrow {{1}\over{N}}\displaystyle\sum^N_{i=1} \displaystyle \max^M_{j=1} E_q[i] \cdot E^{\top}_p[j]$, (N, M은 Query와 Passage의 길이)

통합

각 방법을 사용하여 개별적으로 후보 결과를 검색하고, 해당 후보들을 합집합한 후 최종 검색 결과는 통합 관련도 점수를 기반으로 reranking한다. 이 과정에 multi-vector은 너무 무겁기에 후보 결과 검색땐 제외하고, reranking시에만 사용한다.
$s_{rank} \leftarrow w_1 \cdot s_{dense} + w_2 \cdot s_{lex} + w_3 \cdot s_{mul}$
이때 $w_1, w_2, w_3$는 hyperparameter인데 task에 따라서 다르게 배분했다고 한다.

Self-Knowledge Distillation

Knowledge Distillation(지식 증류)

우선 Knowledge Distillation부터 알아보자.
간단하게 잘 아는 모델(Teacher)가 잘 모르는 모델(Student)에게 단순한 답이 아닌 발상까지 넘겨주는 방법이다.

우선 기존의 학습 방식은 정답이 A라고 주어지면 모델은 A일 확률이 0.80이라고 예측하는 방식이었다.
지식 증류에선 Teacher가 어떤 데이터에 대해 A(0.80), B(0.15), C(0.05)라고 예측했다면, 이 분포를 Student에게 그대로 넘겨준다. 즉, 단순한 답 뿐만 아니라 Decision Boundary를 전달해주는 방식이다.
해당 방식을 통해 일반화(Generalization) 성능이 증가하고, 작은 모델을 더 효율적으로 학습시킬 수 있다.

이를 수식적으로 나타내보면 교사 분포 P, 학생 분포 Q에 대해 loss는
$\mathcal L_{KD} = CrossEntropy(P, Q) = - \displaystyle \sum_i P_i log Q_i$

CrossEntropy: 두 분포가 가장 비슷해지도록(확률 분포가 가까워 지도록) 설정하기 위해 KL Divergence의 최소화를 정리한 식을 Cross-Entropy라고 한다.

Embedding Loss

임베딩 모델에선 Positive sample과 Negative sample을 구분하도록 학습된다. 각 검색에 대해서 Positive가 더 높은 점수를 받도록 부여되도록 학습한다. 수식으로 보면
$\mathcal L_s(\cdot) = -\log{{\exp(s(q,p^{*}))/\tau}\over{\sum_{p\in\{p^{*},P'\}}exp(s(q,p))/\tau}}$
이때 p*와 P'는 Query q에 대한 positive sample과 negative sample이고, $s(\cdot)$은 {$s_{dense}(\cdot), s_{lex}(\cdot), s_{mul}(\cdot)$} 중 하나이다.

Self-Knowledge

지금 해당 모델이 전체적으로 dense, lexical, multi-vector 3가지를 섞어쓰는만큼 학습 목표가 충돌할 가능성도 존재한다. 이를 해결하기 위해서 Self-Knowledge distillation을 기반으로 학습 과정을 통합한다.
Priciple of ensemble learning에 따라서 여러 검색 방식의 예측 결과를 결합하면 더 정확한 관련성 점수를 얻을 수 있다.

Principle of ensemble learning: 작은 모델 여러 개를 묶으면 개별 모델보다 안정적이다.
Bagging: 같은 알고리즘, 랜덤 샘플링으로 과적합 방지
Boosting: 이전 모델이 틀린 문제에 가중치를 둬서 다음 모델이 집중 학습
Voting: 서로 다른 알고리즘을 섞어서 투표

Distillation을 적용시키지 않은 모델(단순 weight sum)을 보면,
$s_{inter} \leftarrow w_1 \cdot s_{dense} + w_2 \cdot s_{lex} + w_3 \cdot s_{mul}$
$\mathcal L \leftarrow (\lambda_1 \cdot \mathcal L_{dense} + \lambda_2 \cdot \mathcal L_{lex} + \lambda_3 \cdot \mathcal L_{mul} + \mathcal L_{inter}) / 4$

여기서 다른 ranking모델에게 얻은 soft label을 활용해서 distillation을 할 수 있다. 해당 모델에선 $s_{inter}$을 teacher로 사용하며, 이때 loss는 (p($\cdot$)은 softmax)
$\mathcal L'_{*} \leftarrow -p(s_{inter}) * \log p(s_*)$
식을 살펴보면, Query가 주어졌을 때 후보 문서들 중 어느 문서가 정답인지를 맞추는 문제이기에 이를 softmax()==p($\cdot$)으로 바꾼다.
이를 위에서 본 식에 대입하면
$P_i = p(s_{inter, i})$, $Q_i = p(s_{*, i})$
$\mathcal L'_* = -\displaystyle\sum_i p(s_{inter, i})\log p(s_{*, i})$
해당 논문에선 한 Query에 한 번만 사용하니 위의 식이 나오게 된다.

이제 위의 식을 바탕으로 loss function을 재정의하면
$\mathcal L' \leftarrow (\lambda_1 \cdot \mathcal L'_{dense} + \lambda_2 \cdot \mathcal L'_{lex} + \lambda_3 \cdot \mathcal L'_{mul}) / 3$

최종적으로 Self-Knowledge Distillation의 최종 loss function은 L, L'의 선형결합인
$\mathcal L_{final} \leftarrow (\mathcal L + \mathcal L')/2$
로 한다.

Learning

1. text encoder은 비지도 데이터로 pre-training되며 dense retrieval만 학습한다.
2. 다른 지도학습 데이터까지 전부 활용하여 Self-Knowledge Distillatio을 통해 임베딩 모델이 세 가지 검색 기능을 모두 학습하도록 fine-tuning된다.

• 이때 $W_{lex}$는 무작위 초기화로 인해 학습 초기는 $s_{lex}$ 정확도가 낮고, $L_{lex}$가 크다.
  이를 방지하기 위해 학습에선 $w_1 = 1, w_2 = 0.3, w_3 = 1$로,
  $\lambda_1 = 1, \lambda_2 = 0.1, \lambda_3 = 1$로 설정한다.
• 또한 각 Query에 대해 ANCE에 따라 Hard-Negative Sample을 도입한다.
• ANCE는 BM25와 달리 현재 모델을 기준으로 정답은 아니지만 점수가 높은 문서를 negative를 사용한다.

Efficient Batching

다양한 언어에 대한 일반적인 의미를 얻기 위해서 대규모의 다양한 다국어 데이터를 학습해야 한다. 추가로 텍스트 임베딩의 판별력을 보장하기 위해서 가능한 큰 batch size를 유지해야 한다.
하지만 GPU 메모리의 제약이 있기에 batch size를 유지하기 위해선 입력 데이터를 짧은 시퀀스로 잘라 학습해야 한다.
M3-Embedding의 경우는 서로 다른 입력 단위의 데이터를 효과적으로 처리하기 위해서 짧은 시퀀스에서 긴 시퀀스까지 전부 다 학습해야하기 때문에 batching 전략을 최적화해야 한다.

우선 입력 데이터를 시퀀스 길이에 따라서 그룹화한다. mini-batch를 적용할 경우 동일한 길이의 그룹에서만 추출한다. 이로인해 padding으로 인한 낭비가 줄어들며, GPU 자원을 효율적으로 활용할 수 있다.
또한 서로 다른 GPU에서 데이터를 샘플링할 때 동일한 시드를 사용하여 부하 균형을 유지한다. 즉, 비슷한 난이도의 샘플을 뽑게 만드는 과정을 통해 가장 느린 GPU 한개로 전부가 멈추는 현상을 방지한다.

또한 긴 시퀀스의 경우 mini-batch를 더 작은 sub-batch로 분할한다. gradient checkpoint를 활용해 sub-batch를 순차적으로 인코딩하고, 모든 임베딩을 수집한다. 이후 서로 다른 GPU에서 생성된 임베딩을 broadcasting한다. 이를 통해 in-batch negative sample의 규모가 확장된다.

gradient checkpoint: 중간 activation을 저장하지 않고 backward때 다시 계산하는 방식. 연산을 조금 포기하고 메모리를 줄임.

추가로 논문에선 현실적인 방안으로 MCLS도 제안하고 있다.

MCLS: 긴 문서를 일정 간격으로 쪼개 [CLS]를 추가로 삽입하여 [CLS] embedding들을 평균내는 방식

4. Experiment

1) Multi-Lingual Retrieval: 18개 언어에서 각 언어 내에서 QA를 검색했을 때 Dense, Sparse, Multi-vector 각각의 성능도 뛰어나고 세개를 합쳤을 때 최고 성능이 나왔다.
2) Cross-Lingual Retrieval: Query는 비영어, 문서는 영어일 때 언어 전반에서 일관성을 보이고 특히 저자원 언어에서 M3가 훨씬 안정적이었다.
3) Multilingual Long-Document Retrieval: Sparse가 훨씬 유리했고 Dense+ Sparse, All조합이 최고성능이었다.
4) Ablation Study
a) Distillation 제거시 Sparse가 폭락
b) 대규모 비지도 사전학습이 다국어 안정성의 기반

5. Conclusion

이전의 SPLADE도 Sparse 방식과 Dense방식을 합치려고 했다. 하지만 M3의 경우 단순히 dense를 보조로 사용한 것이 아닌 각각 다른 학습 방식을 같이 학습시킨다.
뭔가 난이도가 엄청 높은 느낌은 아니였지만 내용이 좀 많은 느낌이었다. 천천히 정리해보고 이젠 조금 더 실용적인? 쪽의 논문과 함께 해보고싶던걸 해볼 생각이다.