Elamur Research⌕
☰
home / papers / 2021 / Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

Paper · 2021 · Luyu Gao, Yunyi Zhang, Jiawei Han et al.

Summary

Batch-wise contrastive training of text encoders, such as dense retrievers, works better with large batches because every example gains more in-batch negatives. The whole batch must still fit in GPU memory, though, which shuts out researchers with modest hardware. Plain gradient accumulation does not solve this, since each chunk only sees its own negatives.

The authors propose Gradient Cache. They observe that backpropagation splits into two stages: loss to representations, and representations to encoder parameters. The second stage is independent per example once the representation gradients are known. The procedure is:

  • encode the full batch without storing activations;
  • compute the loss and the gradients with respect to the embeddings only, and cache them;
  • re-encode small sub-batches and backpropagate the cached gradients through the encoder, accumulating parameter gradients.

The result is mathematically identical to a full-batch update, while encoder memory stays nearly constant.

Applied to DPR on Natural Questions with a single RTX 2080ti, the method matches or slightly beats the original 8-GPU results at about 20% extra time per update. This makes large-batch contrastive learning practical for academic labs.

Key points

  • Top-20/top-100 hit accuracy on Natural Questions, single RTX 2080ti:
    • Cache (batch 128): 79.3/86.0
    • Cache (batch 512): 79.9/86.6
    • DPR reference (8 GPUs): 78.4/85.4
    • Accumulation (16 × batch 8): 77.2/84.9
    • Sequential (batch 8): 71.9/80.9
  • Cache runs used sub-batches of 16 questions and 8 passages. Each question was paired with one positive and one BM25 hard negative.
  • Overhead is roughly 20% more time per update, which pays for the extra graph-free forward pass. Updates scale to 4096 examples per batch.
  • A full DPR model trained in about 31 hours on one RTX 2080ti; the original DPR reports roughly a day on 8 V100s.
  • Gradient checkpointing (Huggingface transformers) only reached batch 64 and was about twice as slow as accumulation.
  • The only extra persistent storage is one vector per anchor and target plus its gradient, i.e. (|S|+|T|)·d floats.
  • Code is released as GradCache and GC-DPR.

Interesting ideas

  • I1. Cut the computation graph at the embedding layer. Any loss whose coupling across examples exists only through cheap representations can be split this way. Training memory then depends on sub-batch size rather than batch size, with no approximation.
  • I2. Make the trade-off explicit. The method spends one extra forward pass, done without storing activations, to buy exact large-batch gradients. This compute-for-memory exchange is useful wherever GPUs are small but training time is flexible.
  • I3. Rely on autodiff for the loss step. Only the embedding-level loss is differentiated in the caching step, so variants of the contrastive objective work without hand-derived gradients.
  • I4. Scale across GPUs cheaply. Gathering representations once across devices gives every example negatives from all GPUs. This adds only one communication step on top of standard data-parallel training.
  • I5. Chain caches for learned similarity functions. A parameterized similarity function can be handled with a second cache holding gradients of the pairwise scores. The two caches are linked, which could extend the trick to re-ranker-style or interaction models.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Carnegie Mellon UniversityPittsburgh · United States365 works · 158,048 citations🎓University of Illinois Urbana-ChampaignUrbana · United States108 works · 67,432 citations

📚 Shelves 4 Computer Science · 9,086 works Artificial Intelligence · 8,889 works Arts and Humanities · 9,296 works Language and Linguistics · 9,206 works

Cite

Cite as Luyu Gao, Yunyi Zhang, Jiawei Han and Jamie Callan. 2021. Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup. Proceedings of the 6th Workshop on Representation Learning for NLP (RepL4NLP-2021).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2021.repl4nlp-1.31

@inproceedings{gao2021scaling,
  title = {{Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup}},
  author = {Luyu Gao, Yunyi Zhang, Jiawei Han and Jamie Callan},
  booktitle = {Proceedings of the 6th Workshop on Representation Learning for NLP (RepL4NLP-2021)},
  year = {2021},
  doi = {10.18653/v1/2021.repl4nlp-1.31},
  url = {https://doi.org/10.18653/v1/2021.repl4nlp-1.31}
}

References

15 references, 7 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Wei-Cheng Chang, Felix X. Yu, Yin-Wen Chang, Yiming Yang, and Sanjiv Kumar. 2020. Pre-training tasks for embedding-based large-scale retrieval. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net. 57th Annual Meeting of the… find
  2. Matthew MacKay, Paul Vicol, Jimmy Ba, and Roger B. Grosse. 2018. Reversible recurrent neural networks. In NeurIPS. find
  3. T. Chen, B. Xu, C. Zhang, and Carlos Guestrin. 2016. Training deep nets with sublinear memory cost. ArXiv, abs/1604.06174. find
  4. Tomas Mikolov, Ilya Sutskever, Kai Chen, G. Corrado, and J. Dean. 2013. Distributed representations of words and phrases and their compositionality. In NIPS. our card
  5. Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. 2020. A simple framework for contrastive learning of visual representations. ArXiv, abs/2002.05709. our card
  6. A. Mnih and Y. Teh. 2012. A fast and simple algorithm for training neural probabilistic language models. In ICML. find
  7. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,… our card
  8. John Michael Giorgi, Osvald Nitski, Gary D Bader, and Bo Wang. 2020. Declutr: Deep contrastive learning for unsupervised textual representations. ArXiv, abs/2006.03659. our card
  9. Aidan N. Gomez, Mengye Ren, R. Urtasun, and Roger B. Grosse. 2017. The reversible residual network: Backpropagation without storing activations. In NIPS. find
  10. Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6769– 6781, Onl… our card
  11. Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. 2020. Supervised contrastive learning. arXiv preprint arXiv:2004.11362. — [arXiv:2004.11362](https://arxiv.org/abs/2004.11362) find
  12. T. Kwiatkowski, J. Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, C. Alberti, D. Epstein, Illia Polosukhin, J. Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Q. Le, and Slav Petrov. 2019. Natural questions: A benchma… our card
  13. Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. 2019. Latent retrieval for weakly supervised open domain question answering. In Proceedings of the Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. 2020. Rocketqa: An optimized training appro… our card
  14. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. find
  15. Z. Wu, Sinong Wang, Jiatao Gu, Madian Khabsa, Fei Sun, and Hao Ma. 2020. Clear: Contrastive learning for sentence representation. ArXiv, abs/2012.15466. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/gao-2021-scaling-deep-contrastive-learning-batch/