Contrastive Learning for Retrieval: InfoNCE & DPR Guide › Code Implementation
Language AI HandbookAll chapters Contrastive Learning for Retrieval: InfoNCE & DPR Guide › Code ImplementationOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a simplified version of DPR training to make these concepts concrete. We'll use the transformers library for the encoder and build the contrastive training loop step by step.
Setting Up the Encoders
First, let's set up the dual encoder architecture using a pretrained model:
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
class DualEncoder(nn.Module):
"""Dual encoder with separate query and passage BERT models."""
def __init__(self, model_name):
super().__init__()
self.query_encoder = AutoModel.from_pretrained(model_name)
self.passage_encoder = AutoModel.from_pretrained(model_name)
def encode_queries(self, input_ids, attention_mask):
outputs = self.query_encoder(
input_ids=input_ids, attention_mask=attention_mask
)
return outputs.last_hidden_state[:, 0, :] # [CLS] token
def encode_passages(self, input_ids, attention_mask):
outputs = self.passage_encoder(
input_ids=input_ids, attention_mask=attention_mask
)
return outputs.last_hidden_state[:, 0, :] # [CLS] tokenEach encoder extracts the [CLS] token representation from BERT's final layer, producing a 768-dimensional vector. As we discussed in Part XXIV, the [CLS] token is designed to aggregate sequence-level information during pretraining. At initialization, both encoders start from the same pretrained BERT weights, but because they are separate nn.Module instances, they will diverge during fine-tuning as the contrastive loss pushes each one to specialize for its respective input type: short queries versus longer passages.
Building the Contrastive Loss
Now let's implement the InfoNCE loss with in-batch negatives and an optional hard negative. The implementation follows the mathematical formulation closely, constructing the score matrix and then delegating the softmax and negative log-likelihood computation to PyTorch's built-in cross_entropy function:
def contrastive_loss_with_hard_negatives(
query_embeds, pos_passage_embeds, hard_neg_embeds=None, temperature=1.0
):
"""
Compute InfoNCE loss with in-batch negatives and optional hard negatives.
Args:
query_embeds: (B, d) query embeddings
pos_passage_embeds: (B, d) positive passage embeddings
hard_neg_embeds: (B, d) hard negative embeddings, or None
temperature: softmax temperature
Returns:
scalar loss
"""
B = query_embeds.size(0)
# In-batch similarity: (B, B) matrix
# Each query scored against all positive passages
in_batch_scores = query_embeds @ pos_passage_embeds.T / temperature
if hard_neg_embeds is not None:
# Hard negative scores: (B, 1) - each query vs its own hard negative
hard_neg_scores = (query_embeds * hard_neg_embeds).sum(
dim=-1, keepdim=True
) / temperature
# Concatenate: (B, B+1) - in-batch scores + hard negative score
all_scores = torch.cat([in_batch_scores, hard_neg_scores], dim=1)
else:
all_scores = in_batch_scores
# Labels: the positive is always at index i for query i (the diagonal)
labels = torch.arange(B, device=query_embeds.device)
return F.cross_entropy(all_scores, labels)The key insight in this implementation is how the labels work and how the score matrix is structured. For query , the correct positive passage is at index in the score vector, because it corresponds to the diagonal of the in-batch similarity matrix. The hard negative, if present, is appended as an additional column at index , so the score tensor for each query has shape : the first entries are the in-batch similarities (one of which is the positive), and the last entry is the hard negative. The cross_entropy function computes the softmax over this entire vector and applies the negative log-likelihood to the position indicated by the label, which is always the diagonal index . This single function call handles the entire InfoNCE computation, including the numerator, denominator, exponentiation, normalization, and log, in a numerically stable way.
Demonstrating the Training Step
Let's create a small example to see the loss in action. We define four query-passage-negative triplets that illustrate the kind of data a real DPR training pipeline would process:
# Create sample training data
queries = [
"What is the capital of France?",
"Who wrote Romeo and Juliet?",
"What is photosynthesis?",
"When was the moon landing?",
]
positive_passages = [
"Paris is the capital and most populous city of France.",
"Romeo and Juliet is a tragedy written by William Shakespeare.",
"Photosynthesis is the process by which plants convert light energy into chemical energy.",
"Apollo 11 landed on the Moon on July 20, 1969.",
]
hard_negatives = [
"France is a country in Western Europe known for its cuisine and culture.",
"Shakespeare's plays include Hamlet, Macbeth, and A Midsummer Night's Dream.",
"Plants require water, sunlight, and carbon dioxide to grow properly.",
"NASA's Apollo program conducted several missions to explore outer space.",
]
def tokenize_batch(texts, max_length=64):
return tokenizer(
texts,
padding=True,
truncation=True,
max_length=max_length,
return_tensors="pt",
)
query_tokens = tokenize_batch(queries, max_length=32)
pos_tokens = tokenize_batch(positive_passages, max_length=128)
hard_neg_tokens = tokenize_batch(hard_negatives, max_length=128)Notice how each hard negative is topically related to its query but does not answer the question. The first hard negative mentions France but says nothing about its capital. The second mentions Shakespeare but does not attribute Romeo and Juliet to him. These are exactly the kind of BM25-style hard negatives that DPR uses: high lexical overlap, low semantic relevance to the specific information need.
# Initialize model and run a forward pass
dual_encoder = DualEncoder(model_name)
dual_encoder.eval()
with torch.no_grad():
q_embeds = dual_encoder.encode_queries(
query_tokens["input_ids"], query_tokens["attention_mask"]
)
p_embeds = dual_encoder.encode_passages(
pos_tokens["input_ids"], pos_tokens["attention_mask"]
)
hn_embeds = dual_encoder.encode_passages(
hard_neg_tokens["input_ids"], hard_neg_tokens["attention_mask"]
)Query embeddings shape: torch.Size([4, 768]) Passage embeddings shape: torch.Size([4, 768]) Hard negative embeddings shape: torch.Size([4, 768])
The output confirms that the encoder produces 768-dimensional vectors for all inputs, matching the hidden size of the BERT-base model. Each query, positive passage, and hard negative has been compressed into a single point in this 768-dimensional space. Before training, these points are placed essentially at random with respect to the relevance relationships we care about.
Now let's compute the similarity matrix and loss to see how the untrained model performs:
# Compute similarity matrix (before any contrastive training)
sim_matrix = q_embeds @ p_embeds.TSimilarity matrix (queries x positive passages):
Pass 0 Pass 1 Pass 2 Pass 3
Query 0: 179.456 *113.254 139.480 158.916
Query 1: 180.976 121.717 *152.937 180.686
Query 2: 178.976 100.527 159.343 *166.903
Query 3: 176.282 108.444 143.652 169.107 *
* = positive pair (diagonal)Before contrastive training, the similarity scores are essentially random. The diagonal (positive pairs) doesn't stand out from the off-diagonal (negatives). This is the expected baseline: the pretrained BERT encoders produce embeddings that capture general linguistic features, but they have received no signal about which query-passage pairs should be considered relevant. The contrastive training process will reshape this embedding space so that the diagonal entries (true positives) consistently dominate their respective rows. Let's see the loss:
import math
# Number of queries in the batch
B = len(queries)
# Loss without hard negatives (in-batch only)
loss_ib = contrastive_loss_with_hard_negatives(
q_embeds, p_embeds, temperature=1.0
)
baseline_ib = -math.log(1 / B)
# Loss with hard negatives
loss_ib_hn = contrastive_loss_with_hard_negatives(
q_embeds, p_embeds, hn_embeds, temperature=1.0
)
baseline_ib_hn = -math.log(1 / (B + 1))Loss (in-batch negatives only): 21.6566 Random baseline: -log(1/4) = 1.3863 Loss (in-batch + hard negatives): 22.3673 Random baseline: -log(1/5) = 1.6094
The random baseline is for a uniform distribution over candidates: for four candidates and for five candidates with the hard negative. Our initial loss is much higher because DPR uses unnormalized dot products. Even modest differences in the magnitudes and directions of the pretrained [CLS] vectors can therefore produce very confident but incorrect softmax rankings before retrieval training. The relevant diagnostic is not proximity to the uniform baseline, but whether the diagonal eventually becomes dominant and the loss trends downward under the contrastive objective.
Training Loop
Let's run a few training steps to show the loss decreasing:
# Set up optimizer and training
dual_encoder.train()
optimizer = torch.optim.AdamW(
dual_encoder.parameters(), lr=2e-5, weight_decay=0.01
)
losses = []
num_steps = 20
for step in range(num_steps):
optimizer.zero_grad()
# Forward pass
q_embeds = dual_encoder.encode_queries(
query_tokens["input_ids"], query_tokens["attention_mask"]
)
p_embeds = dual_encoder.encode_passages(
pos_tokens["input_ids"], pos_tokens["attention_mask"]
)
hn_embeds = dual_encoder.encode_passages(
hard_neg_tokens["input_ids"], hard_neg_tokens["attention_mask"]
)
# Compute loss
loss = contrastive_loss_with_hard_negatives(q_embeds, p_embeds, hn_embeds)
# Backward pass
loss.backward()
optimizer.step()
losses.append(loss.item())
Let's examine the similarity matrix after training to see if the model has learned to differentiate positives from negatives:
# Evaluate after training
dual_encoder.eval()
with torch.no_grad():
q_embeds = dual_encoder.encode_queries(
query_tokens["input_ids"], query_tokens["attention_mask"]
)
p_embeds = dual_encoder.encode_passages(
pos_tokens["input_ids"], pos_tokens["attention_mask"]
)
sim_matrix_after = q_embeds @ p_embeds.T
# Get predictions (index of highest similarity passage for each query)
predicted_indices = sim_matrix_after.argmax(dim=1)Similarity matrix AFTER training:
Pass 0 Pass 1 Pass 2 Pass 3
Query 0: 50.641 * 27.969 35.681 28.449
Query 1: -0.697 11.861 *-12.083 3.324
Query 2: 32.829 18.353 54.858 * 21.022
Query 3: 9.244 9.108 4.739 25.987 *
* = positive pair (diagonal)
Query 0: best match = passage 0 ✓
Query 1: best match = passage 1 ✓
Query 2: best match = passage 2 ✓
Query 3: best match = passage 3 ✓Even with just 20 steps on 4 examples (far less than real training), you can see the model beginning to push diagonal scores higher than off-diagonal ones. In a production DPR training run with 60,000 examples over 40 epochs and batch sizes of 128, the model would fully converge, producing a similarity matrix where the diagonal entries clearly dominate every row.
Simulating Hard Negative Mining
Let's demonstrate the BM25 hard negative mining process that DPR uses. We'll simulate it with a simple TF-IDF retriever, which shares BM25's core mechanism of favoring passages with high lexical overlap with the query:
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
# Corpus of passages (including positives and many distractors)
corpus = (
positive_passages
+ hard_negatives
+ [
"The Louvre Museum in Paris houses thousands of artworks.",
"Shakespeare was born in Stratford-upon-Avon in 1564.",
"Chlorophyll gives plants their green color.",
"Neil Armstrong was the first person to walk on the Moon.",
"Berlin is the capital of Germany.",
"Charles Dickens wrote Oliver Twist and Great Expectations.",
"Cellular respiration converts glucose into ATP energy.",
"The Space Shuttle program ran from 1981 to 2011.",
]
)
# Fit TF-IDF vectorizer
vectorizer = TfidfVectorizer()
corpus_vectors = vectorizer.fit_transform(corpus)
def mine_bm25_negatives(query, positive_idx, top_k=3):
"""Find hard negatives: high TF-IDF similarity but not the positive."""
query_vec = vectorizer.transform([query])
scores = cosine_similarity(query_vec, corpus_vectors).flatten()
# Rank passages by similarity, excluding the known positive
ranked_indices = np.argsort(-scores)
negatives = []
for idx in ranked_indices:
if idx != positive_idx and len(negatives) < top_k:
negatives.append((idx, corpus[idx], scores[idx]))
return negativesmining_results = []
for i, (query, pos) in enumerate(zip(queries, positive_passages)):
hard_negs = mine_bm25_negatives(query, i, top_k=2)
mining_results.append({"query": query, "pos": pos, "hard_negs": hard_negs})Query: "What is the capital of France?" Positive: "Paris is the capital and most populous city of France...." Hard neg 1 (sim=0.613): "Berlin is the capital of Germany...." Hard neg 2 (sim=0.219): "France is a country in Western Europe known for its cuisine and cultur..." Query: "Who wrote Romeo and Juliet?" Positive: "Romeo and Juliet is a tragedy written by William Shakespeare...." Hard neg 1 (sim=0.274): "Charles Dickens wrote Oliver Twist and Great Expectations...." Hard neg 2 (sim=0.075): "Paris is the capital and most populous city of France...." Query: "What is photosynthesis?" Positive: "Photosynthesis is the process by which plants convert light energy int..." Hard neg 1 (sim=0.175): "Berlin is the capital of Germany...." Hard neg 2 (sim=0.134): "Paris is the capital and most populous city of France...." Query: "When was the moon landing?" Positive: "Apollo 11 landed on the Moon on July 20, 1969...." Hard neg 1 (sim=0.527): "Neil Armstrong was the first person to walk on the Moon...." Hard neg 2 (sim=0.189): "Shakespeare was born in Stratford-upon-Avon in 1564...."
Notice how the mined hard negatives are topically relevant to each query but don't answer the question. For the query about France's capital, the retriever returns passages that mention France or Paris but in a different context. For the Shakespeare query, it finds passages about Shakespeare that discuss other works rather than Romeo and Juliet. This is exactly the property that makes them valuable for training: they force the model to go beyond keyword matching and learn to distinguish between passages that share surface features with the true answer and passages that contain the information being sought.