Contrastive training

A teacher has a stack of question cards and a stack of answer cards. She lays out a few questions, deals out all the answers face up, and asks the student to pair each question with its answer. Every wrong pairing, she corrects. Over many rounds the student learns what makes an answer fit a question.

Then she gets sneaky. Alongside each right answer she slips in a look-alike wrong card: same subject, wrong answer. "How do I reset my password?" now faces both the reset steps and the password policy. A student who only ever saw easy wrong cards (answers about printers or holidays) would happily pick the policy card because it says "password". The look-alikes force the student to learn the difference that actually matters.

That's contrastive training, and it's how nearly every modern embedding model is taught. The student is the model; "pairing" means making a question's vector point the same way as its answer's vector (see primer.ml.embeddings.similarity); the wrong cards are negatives; the look-alikes are hard negatives.

A question's vector is q = (1, 0). The right answer is p₊ = (1, 0) and a wrong one is p₋ = (0, 1). All three have length 1, so a dot product is a cosine.

1. Score each card with the dot product: q·p₊ = 1, q·p₋ = 0.
2. Divide by the temperature τ (tau), a sharpness dial explained below. With τ = 1 the scores stay (1, 0).
3. Softmax (from primer.ml.attention): exponentiate and share out. e¹ = 2.718 and e⁰ = 1, so the right card gets 2.718 / 3.718 = 0.731.
4. Loss = −ln 0.731 = 0.313. It would be 0 if the right card got 100%.

At τ = 0.1 the scores become (10, 0), the right card gets 99.995%, and the loss is 0.0000454. Same vectors, far more confident: that's what the temperature does.

flowchart LR
  Q["Question: reset my password"] --> E1[Encoder]
  P["Right answer: reset steps"] --> E2[Encoder]
  N["Look-alike: password policy"] --> E3[Encoder]
  E1 --> PULL[Pull these two<br/>vectors together]
  E2 --> PULL
  E1 --> PUSH[Push these two<br/>vectors apart]
  E3 --> PUSH

In practice the teacher deals a whole batch. B questions sit in rows, all the answer cards in the batch sit in columns, and every card that isn't a question's own answer is a free negative for it: in-batch negatives.

flowchart LR
  subgraph Batch["Similarity table for a batch of 3"]
    direction TB
    r1["q₁:  ✔ ✗ ✗ | ✗"]
    r2["q₂:  ✗ ✔ ✗ | ✗"]
    r3["q₃:  ✗ ✗ ✔ | ✗"]
  end
  Batch --> SM["softmax along<br/>each row"] --> L["loss: −log of the<br/>✔ cell's share"]

On the example: B = 1, M = 2, τ = 1: L = −log(e¹ / (e¹ + e⁰)) = −log 0.731 = 0.313.

In Python:

import math
def dot(a, b):
    return sum(a_i * b_i for a_i, b_i in zip(a, b))
def info_nce(q, p, tau):
    B = len(q)
    total = 0
    for i in range(B):
        # exp(q_i · p_j / τ) for every card j
        exps = [math.exp(dot(q[i], p_j) / tau) for p_j in p]
        # log of the ✔ card's share
        total += math.log(exps[i] / sum(exps))
    # -(1/B) Σ_i
    return -total / B
# B = 1 question
q = [(1, 0)]
# M = 2 cards; card 0 is the question's own answer
p = [(1, 0), (0, 1)]
round(info_nce(q, p, tau=1), 3)  # → 0.313
print(f"{info_nce(q, p, tau=0.1):.7f}")  # → 0.0000454

The name InfoNCE comes from "noise-contrastive estimation": telling the true pair apart from noise. It's the softmax cross-entropy loss from primer.ml.losses, where the "classes" are the answer cards in the batch.

The model trained here is a deliberately tiny bi-encoder (the same encoder applied separately to questions and answers): count the words in a text (a bag of words), multiply by one learned matrix W, and scale the result to length 1. _contrastive_grads and _normalize_backward derive the gradient (the direction in which each number in W should move to lower the loss) by hand, and encoder_gradient_check confirms it against a brute-force numerical estimate.

In code: info_nce_loss computes L for a batch, treating any rows of P beyond B as extra hard negatives. BagOfWords turns text into word counts, BiEncoder holds the learned matrix W, and BiEncoder.encode maps texts to unit vectors.

Why it matters the embedding models behind search and RAG (Sentence-BERT, DPR, E5 and their successors) are trained with exactly this loss on millions of (question, answer) pairs. When a retrieval system confuses "reset my password" with "password policy", this is the training signal that was missing.

Everyday picture grading on a curve. With a gentle curve (high τ), a slightly better answer gets slightly more credit. With a steep curve (low τ), the best answer takes nearly all the credit, so small differences in score become big differences in outcome.

Tiny example the right card has cosine 0.8 and three wrong cards have 0.6 each.

On the example: at τ = 1, 2.2255 / (2.2255 + 3 × 1.8221) = 0.289; at τ = 0.05 the scores become 16 and 12, and 1 / (1 + 3e⁻⁴) = 0.948.

In Python:

import math
def share_plus(tau):
    # e^(0.8/τ)
    right = math.exp(0.8 / tau)
    # three wrong cards, e^(0.6/τ) each
    wrong = 3 * math.exp(0.6 / tau)
    return right / (right + wrong)
round(share_plus(1), 3), round(share_plus(0.05), 3)  # → (0.289, 0.948)

In code: positive_probability computes share₊ for given cosines and τ; this curve is that function swept across τ.

Why it matters too high a temperature and the model can't become confident; too low and a few hard pairs dominate training and it becomes unstable. It's one of the handful of settings that really matters when fine-tuning an embedding model.

Everyday picture a driving test on an empty road teaches you to steer. It teaches nothing about merging, because merging never came up. Negatives only teach the distinctions they contain.

The setup (train_bi_encoder): questions come in two intents for each topic: "how do I fix it?" (answered by a how-to card: "restart, reinstall and follow the setup guide") and "what are the rules?" (answered by a policy card: "approval, compliance and usage limits"). The question and answer cards share almost no words, so matching intent must be learned ("fix" goes with "restart"), not read off the overlap. Each batch holds one intent across different topics, so its in-batch negatives differ from the right card only by topic. We train two identical models; the only difference is that one also gets each question's same-topic, other-intent card as a hard negative. Then we test on four topics neither model has seen, including "password".

flowchart TD
  D[Training pairs] --> A[Batches: one intent,<br/>8 different topics]
  A --> M1[Model 1: in-batch negatives only<br/>wrong cards differ by topic]
  A --> H[Add each question's<br/>look-alike card]
  H --> M2[Model 2: plus hard negatives<br/>wrong cards also differ by intent]
  M1 --> T[Test on unseen topics:<br/>right card vs. look-alike]
  M2 --> T

In code: make_pairs builds the (question, right card, look-alike) triples; look_alike_accuracy runs the exam on unseen topics, and training_curve records that score at checkpoints during training.

Why it matters hard negatives are the biggest single driver of retrieval quality. In practice you mine them: search your corpus with BM25 or the current model, and take high-ranking results that are not the right answer.

flowchart LR
  L[Logs: real questions +<br/>the document that resolved each] --> P[Positive pairs]
  P --> MINE[Mine hard negatives:<br/>top results that are wrong]
  MINE --> T[Fine-tune with InfoNCE,<br/>in-batch + hard negatives]
  T --> E["Evaluate recall@k<br/>on a held-out set"]
  E -->|better than base model| D[Re-embed corpus, deploy]
  E -->|not better| P

A few thousand good pairs is usually enough to adapt a general model to legal, medical or internal jargon. Libraries such as sentence-transformers implement this loss as MultipleNegativesRankingLoss.

Everyday picture two people describing the same holiday, one through photos and one through postcards. CLIP teaches a photo reader and a text reader to put matching photos and captions at the same spot on one shared map, so a photo of a dog and the words "a photo of a dog" land together.

Tiny example two images and two captions, each a unit vector. Correctly paired, image 1 scores (1, 0) against the captions and caption 1 scores (1, 0) against the images: each direction's loss is ln(1 + 1/e) = 0.313. Swap the captions and every right pair scores 0 against a wrong pair's 1: the loss rises to ln(1 + e) = 1.313.

On the example: ½(0.313 + 0.313) = 0.313 correctly paired; ½(1.313 + 1.313) = 1.313 with the captions swapped.

In Python:

import math
# InfoNCE with τ = 1: row i must pick column i
def rows_pick_diagonal(S):
    return -sum(math.log(math.exp(row[i]) / sum(math.exp(s) for s in row))
                for i, row in enumerate(S)) / len(S)
def clip_loss(S):
    # captions picking images
    columns = [list(col) for col in zip(*S)]
    # ½(L_image→text + L_text→image)
    return (rows_pick_diagonal(S) + rows_pick_diagonal(columns)) / 2
# correctly paired
round(clip_loss([[1, 0], [0, 1]]), 3)  # → 0.313
# captions swapped
round(clip_loss([[0, 1], [1, 0]]), 3)  # → 1.313

In code: clip_loss computes L_CLIP by averaging info_nce_loss over the rows (images pick captions) and the columns (captions pick images).

flowchart LR
  I[Images] --> IE[Image encoder] --> IV[Image vectors]
  T[Captions] --> TE[Text encoder] --> TV[Text vectors]
  IV --> S["Similarity table<br/>(images × captions)"]
  TV --> S
  S --> L[Symmetric InfoNCE:<br/>rows pick captions,<br/>columns pick images]
  Z["'a photo of a {label}'"] --> TE

train_clip_toy trains two small linear encoders on toy "images" and "captions" that describe the same hidden meaning through different, unrelated feature spaces. Before training, zero-shot labeling is near chance (20% for five classes); after training it's above 90%.

Why it matters the shared space enables search across modalities (find images with text, or text with images), zero-shot classification, and mixing scanned document images with text in one retrieval index. Multimodal models use the same recipe.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.