Source-linked AI summary
Text and Code Embeddings by Contrastive Pre-Training
Arvind Neelakantan, Tao Xu, Raul Puri, Alec Radford, Jesse Michael Han, Jerry Tworek, Qiming Yuan, Nikolas Tezak, Jong Wook Kim, Chris Hallacy, Johannes Heidecke, Pranav Shyam, Boris Power, Tyna Eloundou Nekoul, Girish Sastry, Gretchen Krueger, David Schnurr, Felipe Petroski Such, Kenny Hsu, Madeleine Thompson, Tabarak Khan, Toki Sherbakov, Joanne Jang, Peter Welinder, Lilian Weng
TL;DR
Existing text embeddings are often specialized by use case, while this paper asks whether unsupervised contrastive pre-training can produce strong shared representations for text and code. It trains on naturally paired unlabeled data and reports state-of-the-art classification, text-search, and code-search results, while sentence similarity is weaker and training is costly.
Problem
Previous text-embedding work typically customized data, objectives, and architectures for different use cases such as classification, search, and similarity.
Method
The models use Transformer encoders with contrastive learning and in-batch negatives on naturally occurring unlabeled text pairs and text–code pairs.
Results
The models achieve state-of-the-art results in linear-probe classification, text search, and code search, while underperforming on sentence similarity tasks.
Takeaways & Limitations
A single unsupervised text embedding model can perform strongly across classification and large-scale search, and the approach also transfers to code search.
Takeaways & Limitations
Training uses large batch sizes and substantial computation, and performance decreases on sentence similarity as training continues while search and classification improve.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper develops unsupervised contrastive text and code embeddings intended to unify representation learning across classification, search, and similarity applications. These models achieve strong results across text and code benchmarks, though sentence similarity remains a weaker and behaviorally distinct setting.
- Approach: Contrastive pre-training uses unlabeled naturally paired data, with neighboring Internet text for text models and docstring–implementation pairs for code models.The input is encoded with a Transformer encoder and trained with in-batch negatives.
- Text Embeddings: 4% and 1.8% relative improvements over previous best unsupervised and supervised text embedding models, respectively, were achieved on average linear-probe classification accuracy across 7 SentEval tasks.The cpt-text model series ranges from 300M to 175B parameters, with performance consistently improving as model size increases.
- Motivation: Previous text-embedding work studied sentence embedding and information retrieval separately, despite targeting high-quality text representations.Prior approaches varied in data, training objective, and architecture across domains.
- Text Search: 23.4%, 14.7%, and 10.6% relative improvements over previous best unsupervised methods were achieved on MSMARCO, Natural Questions, and TriviaQA, respectively.The same text embeddings were evaluated for large-scale search over 4M passages and 21M Wikipedia documents.
- Code Search: 20.8% relative improvement over the previous best result was achieved on CodeSearchNet code search using text–code pairs.Increasing cpt-code from 300M to 1.2B parameters did not improve code-search performance.
- Transfer Learning: Fine-tuning further improved transfer performance, including a 5.2% relative improvement over the previous best embedding method on BEIR search.The models remained competitive with methods requiring substantially more computation at test time.
2. Approach
The approach independently encodes paired inputs into dense vectors and trains them with contrastive learning using efficient in-batch negatives. Text and code models use naturally occurring paired data, while similarity is computed from the resulting embeddings.
- Model: The Transformer encoder independently maps each input in a positive pair to a dense embedding, using the final [EOS] hidden state as the sequence representation.[SOS] and [EOS] delimit each input sequence.
- Model: Cosine similarity between the two embedding vectors quantifies the similarity of the paired inputs.
- Training Objective: In-batch negatives contrast each positive pair against the other M −1 examples in a mini-batch, reusing computation for efficient training.The batch logits form an M × M matrix.
- Training Objective: The logits scale cosine similarity by exp(τ), where τ is a trainable temperature parameter.
- Training Objective: Only diagonal logits are positive, and the loss averages cross-entropy terms computed across both matrix directions.
- Training Data: Text models use neighboring Internet text as positive pairs, while code models pair function docstrings with implementations and initialize from pretrained generative models.
3. Results
The text embeddings are evaluated across linear-probe classification, sentence similarity, and semantic search using one unsupervised model across these task types. This unified evaluation contrasts with prior work’s domain-specific benchmarking.
- Evaluation: The evaluation covers linear-probe classification, sentence similarity, and semantic search.
- Evaluation: The same unsupervised model is evaluated across embedding and search tasks, unlike prior methods that report results within separate benchmark domains.
- Comparison: Prior sentence-embedding methods focus on embedding benchmarks, whereas neural information-retrieval methods focus on search benchmarks.
3.1. Text Embedding
The text embeddings achieve state-of-the-art results on linear-probe classification and support zero-shot and k-nearest-neighbor classification. However, they perform worse than prior state-of-the-art methods on SentEval sentence similarity.
- Evaluation: SentEval evaluates sentence embeddings through linear-probe classification and sentence-similarity tasks.
- Linear-Probe Classification: Larger models produce better features for linear-probe classification, and both unsupervised and transfer-learning settings achieve state-of-the-art results.
- Zero-Shot and k-NN Classification: Zero-shot classification assigns labels by embedding proximity, while prompting uses descriptive label sentences instead of single-word labels.
- Zero-Shot and k-NN Classification: A k-nearest-neighbor classifier using 256 nearby training examples and no task-specific parameter tuning achieves results comparable to a linear-probe classifier.
- Sentence Similarity: The models perform worse than previous state-of-the-art methods on SentEval sentence-similarity tasks.The authors suggest these benchmarks may use a definition of similarity for which the models were not optimized.
3.2. Text Search
cpt-text uses a single embedding model for large-scale semantic search and achieves strong results across multiple benchmarks, including competitive transfer performance.
- 3.2. Text Search: cpt-text evaluates semantic search across MSMARCO, Natural Questions, TriviaQA, and 11 zero-shot BEIR tasks.The benchmarks use vector indices for approximate k-nearest-neighbor search, with collections ranging from 4M to 21M documents.
- 3.2. Text Search: cpt-text achieves the best results among unsupervised methods on MSMARCO, Natural Questions, and TriviaQA.It surpasses keyword search on MSMARCO and embedding-based methods on Natural Questions and TriviaQA.
- 3.2. Text Search: cpt-text achieves the best transfer-learning results in BEIR, while BM25 achieves the best results in the unsupervised setting.The comparison spans 11 zero-shot search tasks and includes methods using supervised MSMARCO data.
- 3.2. Text Search: 5.2% relative improvement over the previous best embedding method is achieved in the BEIR transfer setting.cpt-text also outperforms docT5query and remains competitive with methods using substantially more computation at query time.
3.3. Code Search
The code embedding models achieve state-of-the-art CodeSearchNet performance, but unlike text embeddings, larger models do not improve code-search results.
- 3.3. Code Search: 20.8% relative improvement over the previous best result is achieved by cpt-code on CodeSearchNet.The benchmark retrieves relevant code from natural-language queries across six programming languages.
- 3.3. Code Search: cpt-code achieves state-of-the-art results when retrieving relevant code blocks among 1K candidates.The models are evaluated on six programming languages.
- 3.3. Code Search: Increasing cpt-code model size does not improve code-search performance, including in the harder 10K-candidate setting.Performance drops with more distractors, and larger models still do not provide a search boost.
3.4. Analysis
Training behavior differs across tasks: larger batches improve search, whereas longer training improves search and classification but reduces sentence-similarity performance.
- 3.4. Analysis: Larger batches significantly improve NQ search performance by increasing the chance of hard in-batch negatives.The ablation uses the 300M cpt-text model trained with different batch sizes.
- 3.4. Analysis: Longer training increases search and classification performance while decreasing sentence-similarity performance.The authors prioritize search and classification checkpoints because those tasks are associated with clearly defined real-world applications.
4. Related Work
Related work spans contrastive representation learning, sentence embeddings, neural information retrieval, and semantic code search. Prior approaches vary in data, objectives, architectures, and task-specific training, while the paper compares cpt models across these settings.
- Contrastive representation learning: Contrastive learning frames representation learning as distinguishing positive examples from negatives, with positives formed through contexts, augmentations, or multiple views.The related work describes several strategies for constructing positive samples in self-supervised settings.
- Benchmark comparisons: The paper evaluates cpt-text on BEIR search and cpt-code on multilingual code search, comparing them with established retrieval and embedding baselines.The cited table descriptions report comparisons across 11 BEIR tasks and six programming languages, including large-candidate retrieval settings.
- Sentence embeddings: Sentence embedding research commonly uses supervised NLI pairs, siamese architectures, cosine similarity, and post-processing methods such as normalization or whitening.Unsupervised alternatives include contextual pairing and text augmentation.
- Training behavior: Training duration produces divergent task behavior: search and classification improve with longer training, whereas sentence similarity degrades.This result is shown for the 1.2B cpt-text model across training steps.
- Neural information retrieval: Neural information retrieval encodes documents for nearest-neighbor search, often combining dense retrieval with BM25 candidate narrowing and later reranking.Prior unsupervised retrievers also use hard negatives, the Inverse Cloze Task, or masked salient spans.
- Semantic code search: Semantic code search retrieves code relevant to a natural-language query, with prior methods learning from code structure, functional transformations, or text-code contrastive pairs.CodeSearchNet evaluates this task across programming languages, while CodeBERT and GraphCodeBERT provide prominent baselines.
5. Broader Impacts
The broader-impact discussion highlights risks from bias encoded in training data and costs from large-scale training. It calls for stronger bias evaluation, mitigation, monitoring, and more efficient or accessible infrastructure.
- Bias and representational harm: Embedding models can encode discriminatory biases from training data and propagate them into downstream applications.The discussion specifically identifies potential harms affecting groups such as Black people or women.
- Mitigation and governance: The paper encourages robust evaluation across bias classes and improved mitigation methods, including fine-tuning, while recommending restricted and monitored downstream use until evaluation improves.These recommendations address both pretrained models and their training data.
- Bias and representational harm: Representational harms include reinforcing societal subordination through stereotyping or denigration, especially when biased representations affect resources or opportunities.The paper defines representational harm and identifies higher-risk downstream uses.
- Compute costs: Large-batch training requires substantial computation and is environmentally and computationally costly.Suggested paths forward include safe public access to pretrained models and more efficient architectures and training schemes.
6. Conclusion
The paper concludes that contrastive pre-training on unsupervised data with sufficiently large batches yields high-quality text and code representations. The models achieve strong classification and search results, but sentence similarity remains weaker and shows unexpected training behavior.
- Conclusion: Contrastive pre-training on unsupervised data with a sufficiently large batch size can produce high-quality vector representations of text and code.This is the paper’s central conclusion about its training approach.
- Conclusion: The models achieve new state-of-the-art results in linear-probe classification, text search, and code search.The conclusion summarizes these outcomes jointly rather than attributing them to a single benchmark.
- Conclusion: The models underperform on sentence similarity tasks and exhibit unexpected training behavior for those tasks.This contrasts with their reported classification and search performance.