GTE 논문 리뷰

Towards General Text Embeddings with Multi-stage Contrastive Learning

---

Architecture

Dual-encoder Transformer

본 논문은 BERT 기반 dual-encoder 구조를 채택하며, mean pooling을 적용한 representation을 만든다. 우선 mean pooling은, 토큰  $x_1$, $x_2$, ..., $x_n$의 hidden representation $h_i$ 에 대한 평균을 적용하여 전체 문장에 대한 representation을 얻는 방법이다. 즉 다음과 같다.

여기서 dual encoder을 사용한다는 것은, query에 대한 임베딩 모델과 document에 대한 임베딩 모델을 따로 두어 mean pooling을 적용한다는 것이다. 이렇게 하여 positive pair, negative pair에 대해 InfoNCE loss를 사용하는데, 본 논문은 새로운 improved contrastive loss를 사용한다.

Improved Contrastive Loss

Positive text pair sample로 이루어진 배치 $B = \{ (q_1, d_1), ..., (q_n, d_n) \}$에 대하여 본 논문은 다음과 같은 loss function을 제안한다:

유사도 함수 $s(\cdot, \cdot)$는 cosine similarity이다. 여기서 InfoNCE가 단순히 query-document pair의 유사도에 관련된 term만을 분모로 둔다면 본 논문은 다음과 같은 네 가지 구성요소를 추가한다:

• query → document: 기존 InfoNCE와 동일한 term
• query → query: query끼리도 서로 다른 의미임을 학습
• document → query: document가 올바른 query을 맞히도록 하는 inverse contrast
• document → document: document끼리도 서로 다른 의미임을 학습

이렇게 양방향 유사도 및 같은 modality끼리의 대조학습을 통해서 embedding space의 표현력을 높일 수 있다.

Pre-training → Fine-tuning

Pretraining의 경우 웹/공개 데이터에서 weak supervised correlation text pairs를 바탕으로 진행된다. Weak supervised correlation text pairs는 positive pair으로 위키피디아 문서 안에서 문장-문장, Q&A 사이트에서의 질문-답변 등과 같은 데이터를 가져오며, negative pair은 명시적으로 두지 않았으며, batch 안의 나머지 샘플들을 negative sample로 사용한다.

Fine-Tuning의 경우 human annotation이 들어간 작은 양의 데이터셋으로 진행된다.

---

Experiments

BM25나 SimCSE, Contriever 등 다른 방법론에 비해서 성능이 우수함을 알 수 있다.

패러미터 수가 압도적으로 적음에도 불구하고 MTEB task에서 견줄 만하거나, 혹은 더 우수한 성능을 보이는 것을 확인할 수 있다.

단순히 Pretraining만 적용하는 경우 negative pair에 대한 명시적인 학습이 되지 않고, finetuning만 적용하는 경우 문장에 대한 표현력을 충분히 얻지 못한다는 문제점이 있다. 본 논문은 실험 결과를 통해 둘 다 적용하는 경우 좋은 성능을 달성할 수 있음을 보인다.

---

Result

본 논문의 의의는

• 적은 패러미터 수로도 학습 방법 설계, loss function 설계만으로도 우수한 성능을 낼 수 있다는 것과
• Contrastive learning을 위한 데이터셋 확보 과정에서 explicit negative pair 데이터가 적은 것에 대해 pretrain단계와 finetuning단계로 나누어 문제를 해결한 것

에 있다.