[논문 리뷰] Text and Code Embeddings by Contrastive Pre-Training

Arvind Neelakantan, Tao Xu외 23명|arXiv (Cornell University)|2022. 01. 24.

Topic Modeling인용 수 149

한 줄 요약

이 논문은 대규모 배치로 비표기 텍스트와 코드에 대한 대규모 대비(contrastive) 사전학습이 고품질 임베딩을 생성하고 분류, 검색, 코드 검색에서 강력한 성과를 보이며 감독형 파인튜닝 모델에 거의 대등하거나 우수한 경우가 많음을 보여준다.

ABSTRACT

Text embeddings are useful features in many applications such as semantic search and computing text similarity. Previous work typically trains models customized for different use cases, varying in dataset choice, training objective and model architecture. In this work, we show that contrastive pre-training on unsupervised data at scale leads to high quality vector representations of text and code. The same unsupervised text embeddings that achieve new state-of-the-art results in linear-probe classification also display impressive semantic search capabilities and sometimes even perform competitively with fine-tuned models. On linear-probe classification accuracy averaging over 7 tasks, our best unsupervised model achieves a relative improvement of 4% and 1.8% over previous best unsupervised and supervised text embedding models respectively. The same text embeddings when evaluated on large-scale semantic search attains a relative improvement of 23.4%, 14.7%, and 10.6% over previous best unsupervised methods on MSMARCO, Natural Questions and TriviaQA benchmarks, respectively. Similarly to text embeddings, we train code embedding models on (text, code) pairs, obtaining a 20.8% relative improvement over prior best work on code search.

연구 동기 및 목표

• 쌍으로 구성된 데이터에서 대비 학습을 사용하여 고품질의 비지도 텍스트 및 코드 임베딩을 개발한다.
• 대규모 배치 크기와 규모가 분류, 검색, 코드 검색 전반의 성능을 향상시킨다는 것을 입증한다.
• 표준 벤치마크에서 전이 학습 및 제로샷 능력을 평가한다.
• 작업에 따라 성능이 달라지더라도 학습 역학과 모델 크기 및 데이터의 영향 분석을 수행한다.
• 대규모 임베딩 모델의 광범위한 영향력과 계산적 고려사항에 대해 논의한다.

제안 방법

• 입력 데이터를 임베딩으로 매핑하기 위해 Transformer 인코더를 사용하고, 특수 EOS 토큰의 마지막 히든 상태를 통해 임베딩을 생성한다.
• 배치 내 부정 샘플과 학습 가능한 온도 매개변수 tau를 사용하는 대비 목표로 학습한다.
• 자연스럽게 발생하는 쌍 데이터로부터 양의 쌍을 구성한다(텍스트의 이웃하는 텍스트 조각; 코드의 경우 (텍스트, 코드) 쌍).
• GPT-계열 모델에서 CPT-text를, Codex 모델에서 CPT-code를 초기화하여 학습을 부트스트랩한다.
• 매우 큰 배치 크기를 활용하여 효과적인 하드 네거티브 마이닝을 가능하게 하고 표현을 개선한다.
• 선형 탐침 분류, 문장 유사도, 대규모 의미적 검색 및 코드 검색을 포함한 임베딩을 평가한다.

실험 결과

연구 질문

• RQ1 대규모에서의 비지도 대비 사전 학습이 고품질 텍스트 및 코드 임베딩을 생성할 수 있는가?
• RQ2 이 임베딩이 선형 탐침 분류, 제로샷, k-NN 및 의미적/검색 작업에서 어떻게 수행하는가?
• RQ3 모델 크기와 배치 크기가 작업 전반의 임베딩 품질에 어떤 영향을 미치는가?
• RQ4 학습을 미세 조정하고 BEIR 스타일의 제로샷 평가에서 학습된 임베딩이 얼마나 잘 전이되는가?

주요 결과

• 가장 큰 비지도 텍스트 모델이 이전의 최상의 비지도 텍스트 임베딩 대비 상대적 4%의 향상 및 감독형 텍스트 임베딩 대비 1.8%의 향상을 7개의 선형 탐침 작업에서 달성한다.
• 대규모 텍스트 검색 벤치마크에서 CPT-text는 MSMARCO에서 상대적 향상 23.4%, Natural Questions에서 14.7%, TriviaQA에서 10.6%를 달성하며 이전의 비지도 방법보다 우수하다.
• CPT-code는 CodeSearchNet에서 이전 최상의 코드 검색 결과 대비 상대적 향상 20.8%를 달성한다.
• 제로샷 SST-2에서 6B CPT-text는 88.1% 정확도에 도달하고, 프롬프트를 사용할 때 89.1%, 256-NN으로는 93.3%, 선형 탐침으로는 95.7%에 이르며, 감독 설정에서 전체 SOTA에 근접하여 97.5%에 도달한다.
• BEIR에서 비지도 CPT-text의 전이는 이전 임베딩 방법 대비 상대적 5.2%의 향상을 보여주고, 더 무거운 테스트 시점 계산 베이스라인과 경쟁한다.
• 코드 검색 결과는 CPT-code가 언어 간 우수성을 보여주며 CPT-code S/M이 CodeSearchNet 벤치마크에서 최첨단 성능을 달성한다.