<aside>
📄
Paper Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup (GradCache)
Authors Luyu Gao, Yunyi Zhang, Jiawei Han, Jamie Callan
Venue RepL4NLP 2021 · arXiv 2021.01
Links arXiv
Keywords Contrastive Learning Memory Efficiency Large Batch
</aside>
TL;DR
Larger batches give better contrastive representations (more in-batch negatives) but don't fit in GPU memory. Gradient caching splits the backward pass so the batch can be processed in small chunks, keeping peak memory almost constant regardless of batch size while computing exactly the same gradients.
1. Problem Definition
- Contrastive learning of text representations benefits from many negatives, and in-batch negatives are the standard way to get them cheaply.
- But each example's loss depends on every other example in the batch, so the entire large batch must fit in GPU memory at once — limiting batch size on typical hardware.
2. Method: Gradient Cache
- Representation pass (no graph): encode the full batch in chunks without storing activations, and collect all representations.
- Loss and representation gradients: compute the contrastive loss over the full batch and the gradient of the loss with respect to each representation; cache these gradients.
- Chunked encoder backward: re-encode each chunk with the computation graph and back-propagate the cached representation gradients through the encoder, accumulating parameter gradients.
This decouples back-propagation between the contrastive loss and the encoder, removing the encoder's data dependency along the batch dimension.
3. Findings
- Peak memory becomes nearly constant in batch size, so large-batch contrastive training (e.g., DPR-style) can be reproduced on a single GPU.
- The gradients are exact, so results match true large-batch training, at the cost of roughly one extra forward pass.
4. Takeaways
<aside>
💭
- My note at the time: not closely related to the few-shot retriever idea itself, but a practical tool whenever in-batch negatives are used on limited hardware.
- Worth remembering as an engineering trick rather than a modeling contribution; it is now built into common retrieval toolkits (e.g., Tevatron).
</aside>
5. Related Papers
- Dense Passage Retrieval for Open-Domain Question Answering (DPR)