Text and Code Embeddings by Contrastive Pre-Training

In plain words

Many tools need to turn a piece of text, or a snippet of code, into a list of numbers so that similar things get similar lists. Search engines, clustering and simple classifiers all work on such lists. This paper shows that one recipe can produce good lists for both text and code, without hand-written labels.

The recipe uses pairs of things that naturally belong together: a passage and the text next to it on a web page, or a code function and its documentation string. A neural network reads each side separately and produces a list of numbers for it. Training pushes the two lists of a true pair together and pushes them away from the lists of the other pairs in the same batch. The network starts from a model that was already trained to write text or code, and the batches are very large, because more "wrong" examples per batch make the task harder and the result better.

The authors report that one such model is a strong general-purpose tool. It does well on classification with a simple classifier on top, on finding passages that answer a question across millions of documents, and on finding the right function for a description.

TL;DR

• Contrastive pre-training on unlabeled paired data: web text neighbours for text models §1 §3, docstring and function pairs for code models §1.
• Text and code encoders are initialised from GPT-family and Codex models respectively §2.2; the embedding is the last-layer state at the end-of-sequence token §2.1.
• Author-reported: best unsupervised model reaches 91.8 average linear-probe accuracy over 7 SentEval tasks, against 88.3 for DECLUTR tab. 2; single-embedding search beats unsupervised BM25 and Contriever on MSMARCO, NQ and TriviaQA tab. 5.
• Weak spot: sentence-similarity benchmarks, where it trails SimCSE tab. 4. Code embeddings do not improve with size from 300M to 1.2B §3.3.
• No official code was found, so the method can only be re-implemented from this card and the paper.

Problem

Dense text embeddings are needed for semantic search, clustering and feature extraction, and each embedding would otherwise be built for one use case §1. Earlier work split the field into sentence embedding (evaluated on classification and similarity benchmarks) and neural information retrieval (evaluated on search benchmarks), and rarely reported both §1 §3.1 §3.2.

Previous search methods usually needed fine-tuning on a search dataset, or a multi-step pipeline with an expensive cross-attention re-ranker at query time §3.2. The paper asks whether one unsupervised embedding model can serve both kinds of task, and whether the same recipe works for code search §1.

Assumptions: large amounts of naturally paired data are available, and the compute to train very large batches is affordable. The authors state that the training regime is costly in energy and compute and do not quantify it §5.

Method

Data. Training pairs are naturally occurring. For text, two neighbouring pieces of text on the Internet form a positive pair §1 §3. For code, the top-level docstring of a function and its implementation form a (text, code) pair, extracted from open-source code §1 §3. The paper gives no data size, filtering rule or deduplication step for either source §3.

Encoder. A Transformer encoder E embeds each side of a pair independently §2.1 fig. 3. Each input gets a start and an end delimiter. The embedding is the last-layer hidden state at the end delimiter §2.1 fig. 2. The left side x uses [ and ]; the right side y uses { and }. The authors say different delimiters for the two sides made training more stable §2.1.

Similarity. The similarity of a pair is the cosine of the two embeddings §2.1 (unnumbered display).

Objective. In a batch of M pairs, every other pair serves as a negative (in-batch negatives) §2.2. The M by M logit matrix is the cosine similarity times exp(tau), where tau is a trainable temperature §2.2 (unnumbered display). Only the diagonal is positive. The loss is the cross-entropy over rows and over columns, combined as in the listing in §2.2.

# Contrastive training step with in-batch negatives  [§2.1] [§2.2]
# Input:  M text pairs (xs, ys); scalar temperature tau
# Output: scalar loss

def training_step(xs, ys, tau):
    # Encode both sides with the same encoder  [§2.1]
    v_x = encode_x(xs)                          # (M, d)
    v_y = encode_y(ys)                          # (M, d)
    # Cosine similarity, scaled by exp(tau), for all M x M pairs  [§2.2]
    logits = cosine(v_x, v_y) * exp(tau)        # (M, M)
    # Diagonal pairs are positive; other pairs in the batch are negatives  [§2.2]
    diagonal_targets = range(M)
    # One cross-entropy per direction: over rows, and over columns  [§2.2]
    row_loss = row_cross_entropy(logits, diagonal_targets)
    col_loss = column_cross_entropy(logits, diagonal_targets)
    # Listing averages the two directions; the prose says "sum"  [§2.2]
    return (row_loss + col_loss) / 2

Initialisation. cpt-text starts from GPT-family models; cpt-code starts from Codex models §2.2. The paper says the contrastive signal alone is not enough to learn useful representations, which is why the pre-trained start is used §1.

Fine-tuning (optional, reported separately).

• NLI transfer: fine-tune on SNLI and MNLI, with entailment pairs as positives and contradiction pairs as explicit hard negatives, used alongside the in-batch negatives §3.1.1 §2.2.
• MSMARCO transfer: fine-tune the unsupervised model on MSMARCO before BEIR evaluation §1 §3.2.2.

Inference.

• Classification: a linear probe on the embedding, or zero-shot label matching, or k-nearest-neighbour voting with k = 256 §3.1.1 §3.1.2.
• Search: build an approximate nearest-neighbour index with FAISS §3.2.1. Each document is embedded once offline; the query is embedded at query time §3.2.2.

# Zero-shot classification with label descriptions  [§3.1.2]
# Input:  text x; label descriptions for the classes
# Output: predicted label

def classify_zero_shot(x, descriptions):
    # Embed the input and each label sentence with the same model  [§3.1.2]
    v_x = embed(x)
    # The label sentence is a template, e.g. "this is an example of a
    # positive movie review." versus the negative one  [§3.1.2]
    label_vectors = [embed(d) for d in descriptions]
    # Pick the label whose embedding is closest to the input  [§3.1.2]
    return argmax(cosine(v_x, v) for v in label_vectors)

Hyperparameters

Numbers are from the paper's text and tables. Everything not listed is not stated.

Results

Author-reported. Setup: SentEval (7 linear-probe tasks) §3.1.1; SST-2 with a 6B NLI-tuned model §3.1.2; MSMARCO (4M passages), NQ and TriviaQA (21M Wikipedia documents) §3.2.1; BEIR (11 zero-shot tasks) §3.2.2; CodeSearchNet (6 languages, 1K and 10K candidates) §3.3.

Relative gains quoted in the text: 20.8% on code search over GraphCodeBERT §1 tab. 7; 23.4% on MSMARCO §1 §3.2.1; 10.6% on TriviaQA §3.2.1; 5.2% on BEIR transfer over Contriever §3.2.2.

Figure 1 summarises average performance over 22 tasks by model size fig. 1. Figure 4 shows that search and classification keep improving with training steps while sentence similarity degrades fig. 4 §3.4.2.

Paper vs code

No official code. The paper links no repository and no project page §1–§6. The source does not mention a code or weights release. The searches were: the paper and its metadata for code links; GitHub repository search for cpt-text, cpt-code, and text and code embeddings with OpenAI as owner. None returned an official repository.

Pitfalls

• Loss direction is inconsistent. The prose says the loss is the sum of the row and column cross-entropies; the listing averages them §2.2. Pick one and state it in your implementation.
• Softmax axis is ambiguous. The listing calls cross_entropy(logits, labels, axis=0) for rows and axis=1 for columns, but the function is not defined, and the text does not say which axis is the softmax direction §2.2. Check that the diagonal is the target in both directions.
• Temperature is unspecified. tau is trainable, and its initial value and any clamping are not stated §2.2. The logit scale exp(tau) can drift, so log it.
• Batch size is the main lever, and the setup is not reproducible from the paper. Batch size sets the number of in-batch negatives, and the ablation shows a large gain from 1536 to 12288 tab. 9. The table lists batch sizes that shrink as the model grows tab. 1, and the reason is not stated. Whether the batch is global or per device is not stated.
• Table 8 average for cpt-code M does not match its row values. The six language values in that row average to about 86.6, but the printed average is 85.0 tab. 8. The S row is consistent. Check against the PDF before quoting it.
• NQ improvement does not reproduce from Table 5. The text gives 14.7% for NQ over the previous unsupervised methods §1 §3.2.1. From the table's Contriever row, cpt-text XL is 17.3% higher on Recall@20 and 6.8% higher on Recall@100. No combination of the printed values reproduces 14.7%. The MSMARCO (23.4%) and TriviaQA (10.6%) figures do reproduce.
• "1.8% over supervised" is a cross-setting comparison. It is the unsupervised XL score (91.8) against the supervised SimCSE score (90.2) tab. 2 §1. Within the NLI-transfer setting, the gain is 2.2% §1 tab. 2.
• Sentence similarity is weak. cpt-text trails SimCSE on STS tab. 4. The authors argue search and similarity may define relevance differently §3.1.3 §3.4.2. Do not use these models for similarity tasks without testing.
• Checkpoints were picked for search and classification. The paper says it chose checkpoints by search and classification scores, not similarity §3.4.2. Results on other checkpoints may differ.
• Code models do not scale with size here. cpt-code S and M are within about 0.1 points of each other on the 1K task tab. 7 and are not larger than that §3.3. The 10K task is lower for M than for S tab. 8.
• Text models are competitive on Python search at 10K. cpt-text M reaches 91.4 on Python at 10K candidates, close to cpt-code tab. 8.
• Search metric is not stated. Training uses cosine §2.1, but the FAISS search metric and whether vectors are normalised are not stated §3.2.1.
• Pooling and tokenisation details are thin. Only the end-delimiter state is used. Maximum input length, tokeniser, and handling of long documents are not stated.
• Training data is undocumented. The paper names the sources but not their size, licensing, filtering or deduplication §3. Replicating the data is the largest gap.
• Costs are hidden. Very large batches need many accelerators; the paper calls the regime environmentally and computationally costly and does not quantify it §5.
• Social bias. The authors warn that embedding models encode bias from their training data and that representational harms follow §5. Restrict and monitor downstream uses.
• Figure and table numbering. This card numbers figures and tables in source order (figs. 1 to 4, tabs. 1 to 9). Equations in the source are unnumbered, so they are cited by section. Verify against the rendered PDF before quoting numbers.

Lineage

• Builds on: in-batch contrastive learning §2.2 §4; GPT-family models for cpt-text and Codex for cpt-code §2.2.
• Compared against: BM25, ICT and masked salient spans tab. 5; Contriever, DECLUTR, SimCSE and SBERT tab. 2 tab. 5 tab. 6; ColBERT v2, Splade v2 and BM25+CE tab. 6; docT5query tab. 6; CodeBERT and GraphCodeBERT tab. 7.
• Superseded by: not stated in this wiki.