[논문 리뷰] Towards General Text Embeddings with Multi-stage Contrastive Learning

Zehan Li, Xin Zhang외 4명|arXiv (Cornell University)|2023. 08. 07.

Topic Modeling인용 수 57

한 줄 요약

GTE는 다양한 오픈 데이터 혼합에서 다단계 대조 학습을 통해 일반 목적 텍스트 임베딩 모델을 학습하고, 텍스트, 코드, 검색 작업에서 비교적 작은 모델 크기로 제로샷 및 감독 학습 성능 모두에서 강력한 성능을 달성합니다.

ABSTRACT

We present GTE, a general-purpose text embedding model trained with multi-stage contrastive learning. In line with recent advancements in unifying various NLP tasks into a single format, we train a unified text embedding model by employing contrastive learning over a diverse mixture of datasets from multiple sources. By significantly increasing the number of training data during both unsupervised pre-training and supervised fine-tuning stages, we achieve substantial performance gains over existing embedding models. Notably, even with a relatively modest parameter count of 110M, GTE$_ ext{base}$ outperforms the black-box embedding API provided by OpenAI and even surpasses 10x larger text embedding models on the massive text embedding benchmark. Furthermore, without additional fine-tuning on each programming language individually, our model outperforms previous best code retrievers of similar size by treating code as text. In summary, our model achieves impressive results by effectively harnessing multi-stage contrastive learning, offering a powerful and efficient text embedding model with broad applicability across various NLP and code-related tasks.

연구 동기 및 목표

• 다양한 데이터 소스에서 다단계 대조 학습으로 학습된 단일 텍스트 임베딩 모델 개발.
• 일반화 향상을 위해 대규모 비지도 사전 학습에 이어 감독 미세 조정을 활용.
• 텍스트 및 코드 작업에서 임베딩 모델의 대형 또는 특정 작업 모델에 비해 경쟁력이나 우수성을 보여주기.

제안 방법

• 평균 풀링으로 텍스트 임베딩을 얻는 이중 인코더 Transformer 백본 사용.
• 다양한 공개 소스에서 ~800M 텍스트 쌍에 대해 비지도 대조 학습으로 사전 학습.
• 다중 작업의 ~3M 주석 텍스트 트리플을 이용한 감독 대조 학습으로 미세 조정, 하드 네거티브 포함.
• 배치 내 쿼리/문서와 양방향 용어로 네거티브를 확장하여 대조 손실 개선.
• 소스를 균형 있게 샘플링하기 위해 다중 소스에서 몬티노미얼 분포로 데이터 배치(α = 0.5).
• 제로샷 검색, MTEB, 코드 검색 벤치마크에서 작업 특화 프롬프트 없이 평가.]

실험 결과

연구 질문

• RQ1 다양한 오픈 데이터로 학습된 다단계 대조 학습을 이용한 일반 텍스트 임베딩 모델이 표준 벤치마크에서 작업 특이적이거나 더 큰 임베딩 모델보다 우수할 수 있는가?
• RQ2 데이터 다양성, 배치 크기, 모델 규모가 비지도 사전 학습과 감독 미세 조정 성능에 어떤 영향을 미치는가?
• RQ3 언어 특화 미세 조정 없이 텍스트 임베딩으로 학습될 때 코드 검색이 실제로 향상되는가?

주요 결과

• GTEbase(110M)은 OpenAI의 임베딩 API를 능가하고 대형 모델의 10배에 달하는 규모의 모델들보다 일부를 능가할 수 있다.
• 오픈 데이터에 대한 비지도 사전 학습과 감독 미세 조정은 MTEB 작업, BEIR, 코드 검색 전반에서 강력한 성능을 낳는다.
• GTEbase는 작업 특화 프롬프트 없이 제로샷 텍스트 검색 및 텍스트 분류 작업에서 더 큰 모델의 성능과 같거나 우수하다.
• GTEbase는 CodeSearchNet의 유사 규모의 최첨단 코드 리트리버보다도 언어 특화 미세 조정 없이도 크게 우수하다.
• 데이터 다양성과 모델 규모를 확장하면 선형 이득이 발생하며, 배치 크기 네거티브 수가 약 10,000에서 포화되고 다단계 학습은 단일 단계 학습을 능가한다.