Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

By Luyu Gao, Yunyi Zhang, Jiawei Han, Jamie Callan. 2021

167 citations recorded by the source corpus.

Contrastive learning has been applied successfully to learn vector representations of text. Previous research demonstrated that learning high-quality representations benefits from batch-wise contrastive loss with a large number of negatives. In practice, the technique of in-batch negative is used, where for each example in a batch, other batch examples’ positives will be taken as its negatives, avoiding encoding extra negatives. This, however, still conditions each example’s loss on all batch examples and requires fitting the entire large batch into GPU memory. This paper introduces a gradient caching technique that decouples backpropagation between contrastive loss and the encoder, removing encoder backward pass data dependency along the batch dimension. As a result, gradients can be computed for one subset of the batch at a time, leading to almost constant memory usage.

Subfields: Representation & Self-Supervised Learning, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Representation & Self-Supervised Learning

In Optimization & Training

Citations in the corpus

Florence: A New Foundation Model for Computer Vision (2021, methodology): “It consists of several key techniques such as ZeRO (Rajbhandari et al., 2019), activation checkpointing, mixed-precision training, gradient cache (Gao et al., 2021) to greatly reduce the memory consumption and thus improves the training throughput.”
C-Pack: Packaged Resources To Advance General Chinese Embedding (2023, methodology): “In our implementation, we use a compound strategy of gradient checkpointing and cross-device embedding sharing (Gao et al., 2021b), which results in a maximum batch size of 19,200.”
Dense Text Retrieval Based on Pretrained Language Models: A Survey (2022, methodology): “The idea of crossbatch negatives can be also extended to a single-GPU setting with the technique of gradient caching [109], where one can accumulate multiple mini-batches for increasing the number of negatives (while taking a larger computational cost).”
SGPT: GPT Sentence Embeddings for Semantic Search (2022, methodology): “We use GradCache [14] to train with large batch sizes in a limited memory setting.”
Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a Single Transformer (2022, methodology): “…Dense retrieval models can be categorized into two groups, passage-level retrievers (Karpukhin et al., 2020; Oguz et al., 2021; Xiong et al., 2021; Gao and Callan, 2022) and token/phrase-level retrievers (Khattab and Za-haria, 2020; Khattab et al., 2020; Gao et al., 2021a; Lee et al., 2021b).”
Nomic Embed: Training a Reproducible Long Context Text Embedder (2024, unknown): “Unfortunately, due to the nature of closed-source models, we have no indication regarding whether closed-source models trained on these datasets.”
C-Pack: Packed Resources For General Chinese Embeddings (2023, methodology): “In our implementation, we use a compound strategy of gradient checkpointing and cross-device embedding sharing [18], which results in a maximum batch size of 19,200.”
Unsupervised Corpus Aware Language Model Pre-training for Dense Passage Retrieval (2021, methodology): “To break the memory constraint and perform effective contrastive learning, we incorporate the gradient caching technique (Gao et al., 2021b).”
Condenser: a Pre-training Architecture for Dense Retrieval (2021, methodology): “An alternative is the recently proposed contextualized sparse retrieval model (Gao et al., 2021a).”
MIMO: Multilingual Information Retrieval via Monolingual Objectives (2026, unknown): “To enable the large-scale in-batch negatives essential for contrastive learning, we employ gradient caching (Gao et al., 2021) to achieve a global batch size of 2048.”
TurkEmbed: Turkish Embedding Model on Natural Language Inference & Sentence Text Similarity Tasks (2025, unknown): “This approach enables efficient learning from a large number of negative examples without requiring explicit negative sampling strategies [31], and its contrastive nature aligns naturally with NLI data structure where positive pairs (entailment relationships) must be distinguished from negative…”
Benchmarking Patent Embeddings: A Multi-Task Evaluation of 22 Models Across Retrieval, Classification, and Clustering (2026, unknown): “Fine-tuning uses the SentenceTransformerTrainer (Reimers and Gurevych, 2019) with CachedMultiple NegativesRankingLoss (Gao and Callan, 2021), which samples in-batch negatives from a cached pool to scale the effective negative set size beyond the mini-batch.”
ColBERT-Zero: To Pre-train Or Not To Pre-train ColBERT models (2026, unknown): “We use the PyLate [10] library and leverage the implementation of GradCache [11] to scale the per-GPU batch size arbitrarily without VRAM constraints (as standard gradient accumulation is not applicable for contrastive learning).”
Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report (2025, unknown): “This training is done with a large batch size, facilitated by GradCache [9], and resulting in a better embedding alignment.”
PyLate: Flexible Training and Retrieval for Late Interaction Models (2025, unknown): “To overcome the memory requirement of large batches, PyLate implements two key features to scale the effective batch size: • GradCache: Since standard gradient accumulation is not equivalent to a larger batch size in contrastive learning, we provide a GradCache [13] implementation.”
ChEmbed: Enhancing Chemical Literature Search Through Domain-Specific Text Embeddings (2025, unknown): “Models were trained on 4×NVIDIA A100 40GB GPUs using GradCache [59] and mixed precision [60], which enabled the use of a large total batch size of 16 , 384 for more effective contrastive learning.”
ReasonIR: Training Retrievers for Reasoning Tasks (2025, unknown): “We include popular public training data , including MS MARCO (Nguyen et al., 2016), Natural Questions (Kwiatkowski et al., 2019), DUReader (He et al., 2017), FEVER (Thorne et al., 2018), HotpotQA (Yang et al., 2018), MIRACL (Zhang et al., 2023a), Mr. Tydi (Zhang et al., 2021), QUORA (DataCanary et a”