Source-linked AI summary
Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup
Luyu Gao, Yunyi Zhang, Jiawei Han, Jamie Callan
TL;DR
Large-batch contrastive learning is difficult on memory-limited hardware because in-batch negatives require the whole batch and gradient accumulation changes the negative set. The paper introduces gradient caching to separate representation-gradient computation from encoder backpropagation, enabling exact large-batch updates through sub-batches with nearly constant memory; it reports preserved retrieval accuracy and roughly 20% additional runtime.
Problem
Large-batch contrastive learning requires fitting all batch activations into GPU memory, while gradient accumulation cannot preserve the same in-batch negatives.
Method
Gradient caching separates contrastive-loss backpropagation from encoder backpropagation, caches representation gradients, and accumulates encoder gradients over memory-fitting sub-batches.
Results
The method produces the exact same gradient update as large-batch training, preserves retrieval accuracy, and uses about 20% more runtime.
Takeaways & Limitations
Gradient caching enables resource-limited hardware to reproduce state-of-the-art large-batch models and broadens access to batch-wise contrastive learning.
Takeaways & Limitations
The discussion assumes temperature τ = 1, noting that temperature generally adds only a constant multiplier to the gradient.
Abstract
from arXiv · showhide
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.
1 Introduction
Large-batch contrastive learning improves representation training but is constrained by GPU memory because in-batch negatives require encoding and storing the entire batch. The paper introduces gradient caching to preserve large-batch updates with nearly constant memory usage.
- Motivation: In-batch negatives make each example’s loss depend on all batch examples, requiring the entire batch’s activations to fit in GPU memory.A BERT passage encoder used batch size 4096, while an RTX 2080ti could fit only 8 examples.
- Motivation: Gradient accumulation cannot reproduce a large batch because each smaller chunk contains fewer in-batch negatives.
- Approach: Gradient caching separates loss-to-representation backpropagation from representation-to-parameter backpropagation, removing encoder backward-pass dependency across batch examples.
- Approach: The method precomputes and caches representation gradients, enabling sub-batch encoder updates that produce the exact same gradient update as a large batch.
- Outcome: With about 20% more runtime, one consumer-grade GPU can reproduce state-of-the-art models previously requiring multiple professional GPUs.
2 Related Work
Prior work established contrastive learning as a method for learning representations across language modeling, retrieval, sentence encoding, and pretraining. Related memory-reduction methods target deep networks, but their effectiveness for contrastive encoders had not been confirmed.
- Contrastive Learning: Noise Contrastive Estimation and Word2Vec are earlier examples of contrastive methods for learning word embeddings.
- Contrastive Learning: Contrastive learning has been used for unsupervised pretraining, dense retrieval, sentence representation, and language-model pretraining.
- Deep Network Memory Reduction: Gradient checkpointing and reversible activation functions reduce memory for large or deep networks through recomputation or activation recovery.
- Deep Network Memory Reduction: The effectiveness of these deep-network memory techniques as components of contrastive encoders had not been confirmed.
3 Methodologies
The method factors contrastive-loss backpropagation into representation-gradient computation and encoder-gradient accumulation. It caches full-batch representations and gradients while processing encoder updates in memory-fitting sub-batches, including across multiple GPUs.
- Preliminaries: The contrastive objective depends on the full target set, so direct large-batch training requires fitting all target representations into memory.
- Gradient Analysis: Backpropagation is divided into loss-to-representation and representation-to-encoder stages, allowing encoder parameter updates to avoid batch-wide data dependency.
- Gradient Analysis: Representation gradients require encoded representations but not encoder parameters, enabling their computation before encoder backpropagation.
- Gradient Cache Technique: The full batch is split into memory-fitting sub-batches, whose encoder gradients are accumulated before a single optimizer step.
- Gradient Cache Technique: The cache stores representations and representation gradients, requiring only (|S|d + |T|d) floating points beyond the computation graph.
- Multi-GPU Training: The method combines full-batch loss computation with local representation-gradient storage on each GPU, followed by standard cross-GPU gradient reduction.
4 Experiments
The experiments evaluate gradient cache for dense passage retrieval on one GPU, comparing accuracy, batch-size scaling, and training speed against standard baselines.
- Experimental setup: Cache training uses DPR’s batch size of 128 with smaller sub-batches, while Sequential fits batch size 8 and Accumulation combines sixteen size-8 batches.All experiments use a single RTX 2080ti; Cache also tests batch size 512.
- Retrieval accuracy: Cache matches large-batch training performance and improves over standard methods, while Sequential and Accumulation perform worse with fewer negatives or smaller batches.Accumulation matches the number of examples but has fewer negatives, and Sequential uses a smaller batch with higher variance.
- Training speed: Cache steadily scales to larger batch updates, whereas gradient accumulation is compared across update sizes from 64 through 4096 examples.The comparison measures training speed against the number of examples per update.
- Training speed: 20% more time for representation pre-computation enables gradient cache to create much larger batch updates.The added pre-computation cost is presented as supporting the larger updates associated with strong performance.
- Training speed: Gradient cache trains a dense retriever in 31 hours on one RTX2080ti, compared with the original DPR report of roughly one day on 8 V100 GPUs.The reported single-GPU result uses improved data loading.
- Training speed: Gradient checkpointing runs only up to batch 64 and takes twice as long as accumulation.This comparison is reported for the tested hardware and implementation.
5 Extend to Deep Distance Function
The method extends gradient caching to parameterized deep distance functions by caching distance gradients and then computing representation gradients in sub-batches.
- Deep distance function: A deep distance function Φ can be richly parameterized by Ω rather than being only a parameterless dot-product similarity.The extension addresses distance functions with trainable parameters.
- Distance Gradient Cache: The method stores w_ij = ∂L/∂d_ij in a Distance Gradient Cache after computing all representations and distances.The cached distance gradients are obtained by backpropagating the loss through the distances.
- Representation Gradient Cache: It simultaneously computes distance derivatives with respect to f(s_i) and g(t_j), accumulating them across sub-batches to build the Representation Gradient Cache.Encoder gradients are computed after all representation gradients have been collected.
- Scope: The linked-cache procedure also covers the early-interaction case f(s)=s and g(t)=t.This case is identified as a special case of the general method.
6 Conclusion
The paper concludes that gradient caching removes large-batch GPU memory limitations while preserving the exact large-batch gradient update and accuracy on resource-limited hardware.
- Conclusion: Gradient caching removes in-batch data dependency in encoder optimization and produces the exact same gradient update as large-batch training.The technique is introduced specifically for large-batch contrastive learning under GPU memory limits.
- Conclusion: The method is efficient and capable of preserving accuracy on resource-limited hardware.This conclusion is tied to making batch-wise contrastive learning more accessible beyond industry-grade hardware.
- Implications: The authors argue that researchers with limited hardware can use the technique to reproduce state-of-the-art models and continue research without available GPU memory being the constraint.This stated consequence motivates broader access to batch-wise contrastive learning in the NLP community.