Contrastive Learning for Retrieval: InfoNCE & DPR Guide › In-Batch Negatives
Language AI HandbookAll chapters Contrastive Learning for Retrieval: InfoNCE & DPR Guide › In-Batch NegativesOne section · Sections in this chapter4 / 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
Training with InfoNCE naturally raises a question: where do the negatives come from? The naive approach is to sample random passages from the corpus for each query. But this requires encoding extra passages per training step, which is computationally expensive, since each passage must flow through a full BERT encoder to produce its embedding. If you want 127 negatives per query and your batch has 128 queries, the naive approach requires encoding additional passages per step. In-batch negatives provide an elegant solution that eliminates this overhead entirely.
The Core Idea
The insight behind in-batch negatives comes from recognizing a simple but powerful fact about mini-batch training. Consider a mini-batch of query-positive pairs: . For query , its positive passage is . But the positive passages of all other queries in the batch, , are almost certainly irrelevant to . A passage about Shakespeare is unlikely to be relevant to a query about photosynthesis, even though both happen to be in the same mini-batch. So we can reuse them as negatives, extracting additional training signal from data we have already encoded.
This is the in-batch negatives trick: each query in a batch of size automatically gets negatives without encoding a single additional passage. The computational cost is essentially zero beyond what you would already pay for a standard forward pass, because all the passage embeddings are computed as part of the normal training step. The similarity computation for the entire batch can be expressed as a single matrix multiplication:
where:
- : the similarity matrix where entry is the similarity between query and passage
- : the matrix of query embeddings ( is batch size, is embedding dimension)
- : the matrix of passage embeddings
- : the transpose operation
This single matrix multiplication produces all pairwise similarity scores at once, exploiting the massive parallelism of modern GPUs. The structure of the resulting matrix is elegant and informative. The diagonal entries are the positive scores, capturing the similarity between each query and its own matched passage . The off-diagonal entries where are the negative scores, representing the similarity between query and the passage that belongs to some other query . This means the labels for the cross-entropy loss are simply the indices : query 0 should match passage 0, query 1 should match passage 1, and so on. The entire loss computation reduces to a standard cross-entropy call with these diagonal labels.
import torch
import torch.nn.functional as F
# Simulated batch of B=4 queries and passages, each with d=8 dimensional embeddings
B, d = 4, 8
query_embeds = F.normalize(torch.randn(B, d), dim=-1)
passage_embeds = F.normalize(torch.randn(B, d), dim=-1)
# Compute full B x B similarity matrix
similarity_matrix = query_embeds @ passage_embeds.T
# Labels: query i should match passage i (the diagonal)
labels = torch.arange(B)Similarity matrix (B x B): [[-0.568 -0.264 -0.524 -0.72 ] [-0.504 0.449 -0.247 0.145] [-0.406 -0.108 -0.237 0.167] [-0.243 -0.199 -0.291 0.092]] Labels (positive passage index for each query): [0, 1, 2, 3] Positive scores (diagonal): [-0.5680000185966492, 0.4490000009536743, -0.2370000034570694, 0.09200000017881393]

The similarity matrix makes the structure clear: each row corresponds to a query's similarity against all passages in the batch. The diagonal contains the positive scores (query with passage ), and every off-diagonal element is a negative. With random initialization, all scores are small and roughly similar. This reflects the fact that the encoders have not yet learned any meaningful structure. The InfoNCE loss for the entire batch is simply cross-entropy with the diagonal as the target:
temperature = 0.05
scaled_similarities = similarity_matrix / temperature
# Cross-entropy loss with labels pointing to the diagonal
loss = F.cross_entropy(scaled_similarities, labels)Temperature: 0.05 InfoNCE loss: 3.5402 Softmax probabilities for query 0: P(passage 0 | query 0) = 0.0023 ← positive P(passage 1 | query 0) = 0.9921 P(passage 2 | query 0) = 0.0055 P(passage 3 | query 0) = 0.0001
With random initialization, the similarity scores are small and roughly equal, resulting in a uniform probability distribution (approximately 0.25 for each passage). The model has not yet learned to assign higher probability to the positive passage. This is exactly what we expect: before contrastive training, the embeddings carry no information about query-passage relevance, so every candidate looks equally plausible. Training will break this symmetry, pushing the probability mass toward the correct positive.
Why Batch Size Matters Enormously
With in-batch negatives, the number of negatives equals . This creates a direct and consequential link between batch size and training quality. A batch size of 32 gives 31 negatives per query, while a batch size of 512 gives 511 negatives. The difference is not merely quantitative; it is qualitative. With 31 negatives, the -way classification task is relatively easy, and the model can achieve a low loss without learning particularly fine-grained representations. With 511 negatives, the model must correctly identify the positive among hundreds of alternatives, many of which may be somewhat topically related by chance. Research on DPR and similar models consistently shows that larger batch sizes produce substantially better retrieval models, with gains that are especially pronounced when moving from small batches (16 or 32) to medium batches (128 or 256).
However, there's a practical tension. Larger batches require more GPU memory because you must simultaneously hold all query and passage embeddings and compute the similarity matrix. For with embeddings in 32-bit floating point, the embeddings alone require about 0.75 MB, which is modest. But the BERT forward and backward passes for 128 queries and 128 passages are far more memory-intensive, often requiring multiple gigabytes of GPU RAM. This is one reason why contrastive retrieval training often uses multi-GPU setups with techniques like gradient accumulation or cross-GPU negative sharing, where negatives are gathered from all GPUs to create effectively larger batches without requiring a single GPU to hold everything. For example, with 8 GPUs each processing 16 examples, cross-GPU sharing yields in-batch negatives per query, matching the effective batch size of 128 without any single GPU needing to encode more than 16 queries and 16 passages.
The Random Negative Problem
In-batch negatives are efficient but have a flaw: they're essentially random negatives. Because the queries in a mini-batch are drawn from the dataset without regard to topical similarity, the passages associated with other queries are, in most cases, completely unrelated to any given query. In a large corpus, a randomly selected passage is overwhelmingly likely to be completely irrelevant to a query. This makes the contrastive task too easy. The model quickly learns to distinguish obviously irrelevant passages (e.g., a passage about cooking when the query is about quantum physics) but struggles with subtle distinctions (e.g., two passages about quantum entanglement where only one answers the question).
To understand why easy negatives are wasteful, consider what happens to the loss and its gradients. When a negative passage is obviously irrelevant, its similarity score with the query is already much lower than the positive's score. After exponentiation in the softmax, such a negative contributes almost nothing to the denominator, and therefore almost nothing to the gradient. The model receives virtually no learning signal from this example. It is analogous to giving a student a multiple-choice exam where the wrong answers are absurdly implausible: they will get every question right without learning anything about the subject.
This is analogous to a problem we saw in Part IV with negative sampling for Word2Vec: randomly sampled negatives are useful for initial training but don't push the model to learn fine-grained distinctions. The solution there was the same as the solution here: we need harder negatives.