Towards General Text Embeddings with Multi-stage Contrastive Learning

TL;DR

그림 1. 2단계의 Multi Stage Contrastive Learning에 대한 설명. Hard negative의 라벨 가용성 여부로 비지도/지도가 나뉨

• 문제 의식 : 기존의 "일반적인" 목적의 텍스트 임베딩 모델은 사실 일반적이지 않다! (특정 태스크만을 고려한 훈련과정 등)
• 2단계의 빅데이터를 활용한 "양방향의" Constrastive Learning을 통해 훈련한 모델은 "작은 모델 사이즈"에도 불구하고 "좋다"
• 어떤 측면에서 "좋다"? 모든 NLU 태스크에 대해서?
  • 다양한 태스크 별로 (인간 기준 make-sense 한) positive/negative labeling scheme 정의 및 적용
  • 다양한 태스크 벤치마크 데이터셋 MTEB에 평가 결과, GTE_large는 새로운 SOTA로 자리매김하여 다양한 태스크에도 사용 가능한 임베딩 모델을 제시함.

---

Multi Stage Contrastive Learning

Contrastive Learning : commutable comparison

일반적인 텍스트 특징벡터를 대조학습을 통해 얻을 때  query -> doc의 단 방향 손실함수를 사용함.

본 논문은 양방향, 즉, query -> doc // doc -> query의 양 방향 참조를 함.

Unsupervised Contrastive Pre-training

Unsupervised라고 하는 건 일단 전문가의 손길을 거쳐 라벨링 된 것은 아니기 때문임.
그럼에도 불구하고 positivity 자체는 데이터셋 특징을 고려했을 때 쉽게 쌍을 지을 수 있음. 표 1 참고.

표 1. positivity는 어떤 식으로 짝 지어지는지 보는 거임.

또한, 도합 33개의 이질적인 데이터셋을 합쳐 훈련에 활용하기에 배치 샘플링 전략 역시 중요하다. 이를 위해서 각 데이터셋의 종류를 다항분포를 따르는 다항 분포로 모델링하여 데이터셋 크기에 비례한 확률로 추출한다.

Supervised Contrastive Fine-tuning

Experimental Result

 

비지도 대조 학습 -> 지도 대조 학습 순서로 진행해야 가장 최적의 결과를 얻음

표 2. multi-stage 관련 ablation study. 비지도 대조학습 -> 지도 대조학습 순서로 하는 게 제일 좋은 결과를 보임 (MTEB)

MTBE 내의 다양한 태스크들에 대해서 좋은 성능을 보임을 확인함.  

• 진정한 의미의 General Text Embedding이 드디어 구현되었다는 의의가 있음. 
• 특히 GTE_large는 해당 벤치마크 데이터셋에 대해서 SOTA를 찍어 성능 최적화 면에서도 의의가 있음.

표 3. MTEB 평가결과.

양방향의 참조를 통한 대조학습 손실함수는 성능 향상 (MTEB 평균값)에 기여함

• 사전학습/추가학습 가릴 것 없이 개선된 손실함수를 사용하는 것이 가장 좋았음.

---

Reference

[1] Li, Z., Zhang, X., Zhang, Y., Long, D., Xie, P., & Zhang, M. (2023). Towards general text embeddings with multi-stage contrastive learning. arXiv preprint arXiv:2308.03281.