Contrastive Learning for Retrieval: InfoNCE & DPR Guide › Training Dynamics and Considerations

Beyond the basic setup, several nuances in contrastive training significantly impact the final model quality. These considerations arise from the interaction between the loss function's mathematical properties and the practical realities of training large models on limited hardware.

The loss we've described is asymmetric: for each query, we find its positive among all passages. In other words, we compute a row-wise softmax over the similarity matrix, asking each query to select its correct passage from the columns. You can also add a symmetric term where, for each passage, you find its matching query among all queries. This corresponds to a column-wise softmax, asking each passage to select the query it belongs to from the rows:

$$\mathcal{L}_{\text{symmetric}} = \frac{1}{2}\bigl(\mathcal{L}(Q \to P) + \mathcal{L}(P \to Q)\bigr)$$

where:

• $\mathcal{L}_{\text{symmetric}}$: the bidirectional contrastive loss
• $\mathcal{L}(Q \to P)$: the standard InfoNCE loss computed row-wise (queries selecting passages)
• $\mathcal{L}(P \to Q)$: the InfoNCE loss computed column-wise (passages selecting queries)

The symmetric formulation provides additional training signal by exploiting the same similarity matrix from both directions. In the asymmetric case, the passage encoder only receives gradients based on how passages are scored by queries. In the symmetric case, the passage encoder also receives gradients from the task of selecting the correct query, which can help it learn representations that generalize better. Symmetric loss is used by models like CLIP and some sentence embedding models. DPR uses only the asymmetric query-to-passage direction, which is a natural choice for retrieval since the operational task is always "given a query, find the passage" rather than the reverse.

As we noted, larger batches produce better models because they provide more in-batch negatives. But large batches hit GPU memory limits quickly since each BERT encoder's forward and backward passes consume substantial memory, often several gigabytes for a batch of 128 sequences through a BERT-base model. Gradient caching (sometimes called GradCache) addresses this by decoupling the batch size used for contrastive comparison from the batch size used for the encoder forward pass. The procedure works in four stages:

1. Compute all embeddings in small chunks (with torch.no_grad() to save memory). For example, process 16 queries at a time rather than all 128, caching the resulting embeddings.
2. Compute the full similarity matrix and loss using the cached embeddings. Since the embeddings are just vectors, the $B \times B$ matrix multiplication is cheap even for large $B$.
3. Backpropagate through the loss to get gradients on embeddings. This tells us how each embedding should change to reduce the loss.
4. Re-run the encoder forward passes in small chunks, using the embedding gradients to compute parameter gradients via the chain rule.

This technique decouples the batch size for contrastive learning from the batch size that fits in GPU memory during the encoder forward pass. You can use a contrastive batch size of 4096 while only ever processing 32 sequences through the encoder at a time, achieving the representational benefits of massive in-batch negative sets without exceeding your GPU's memory capacity.

In iterative hard negative mining, you periodically re-encode the entire corpus to find new hard negatives using the updated model. This is expensive: for a corpus of 21 million passages (as in DPR's Wikipedia corpus), re-encoding takes hours even on multiple GPUs, because every passage must flow through the passage encoder to produce a fresh embedding that reflects the model's current parameters. A practical approach is asynchronous refresh, where you mine hard negatives using a slightly stale model checkpoint while training continues with the current model. The mined negatives are based on the model as it was several training steps ago, which means they may not perfectly reflect the current model's failure modes. The staleness introduces some noise, but the efficiency gain is worth it, and empirical results show that even moderately stale negatives provide substantial benefit over using no hard negative mining at all. More sophisticated systems use a rolling refresh schedule, re-encoding a fraction of the corpus each step so that the index is always partially fresh without ever requiring a full re-encoding pause.