[Paper Review] Text and Code Embeddings by Contrastive Pre-Training

Arvind Neelakantan, Tao Xu+23 more|arXiv (Cornell University)|Jan 24, 2022

Topic ModelingComputer Science149 citations

TL;DR

The paper shows that large-scale contrastive pre-training on unlabeled text and code with large batches yields high-quality embeddings, achieving strong results in classification, retrieval, and code search, often rivaling or surpassing supervised fine-tuned models.

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.

Motivation & Objective

• Develop high-quality unsupervised text and code embeddings using contrastive pre-training on paired data.
• Demonstrate that large batch sizes and scale improve performance across classification, search, and code retrieval.
• Evaluate transfer learning and zero-shot capabilities on standard benchmarks.
• Analyze training dynamics and the effects of model size and data although performance varies by task.
• Discuss broader impacts and computational considerations of large-scale embedding models.

Proposed method

• Use Transformer encoders to map inputs to embeddings via the last hidden state of a special EOS token.
• Train with a contrastive objective using in-batch negatives and a trainable temperature parameter tau.
• Construct positive pairs from naturally occurring paired data (neighboring text pieces for text; (text, code) pairs for code).
• Initialize CPT-text from GPT-family models and CPT-code from Codex models to bootstrap learning.
• Leverage very large batch sizes to enable effective hard negative mining and improve representations.
• Evaluate embeddings across linear-probe classification, sentence similarity, and large-scale semantic search, plus code search.

Experimental results

Research questions

• RQ1 Can unsupervised contrastive pre-training at scale produce high-quality text and code embeddings?
• RQ2 How do these embeddings perform on linear-probe classification, zero-shot, k-NN, and semantic/search tasks?
• RQ3 What is the impact of model size and batch size on embedding quality across tasks?
• RQ4 How well do the learned embeddings transfer under fine-tuning and in BEIR-style zero-shot evaluations?

Key findings

• Largest unsupervised text model achieves a 4% relative improvement over the previous best unsupervised text embeddings and 1.8% over the supervised text embeddings on 7 linear-probe tasks.
• On large-scale text search benchmarks, CPT-text attains a relative improvement of 23.4% on MSMARCO, 14.7% on Natural Questions, and 10.6% on TriviaQA over prior unsupervised methods.
• CPT-code achieves a 20.8% relative improvement over the previous best code search results on CodeSearchNet.
• Zero-shot SST-2 results with 6B CPT-text reach 88.1% accuracy, 89.1% with prompting, 93.3% with 256-NN, 95.7% linear-probe, and approaching full SOTA at 97.5% in supervised settings.
• In BEIR, unsupervised CPT-text transfers yield a 5.2% relative improvement over prior embedding methods, and compete with heavier test-time compute baselines.
• Code-search results show CPT-code surpassing prior methods across languages, with CPT-code S/M achieving state-of-the-art in the CodeSearchNet benchmark.