Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

TL;DR

This paper proposes a gradient caching method for contrastive learning that enables training with large batch sizes under limited memory by decoupling backpropagation, thus reducing memory requirements significantly.

Contribution

It introduces a novel gradient caching technique that decouples backpropagation from the encoder, allowing large batch contrastive learning with limited GPU memory.

Findings

Memory usage becomes almost constant regardless of batch size

Enables training with larger batches than traditional methods

Maintains high-quality representations despite memory constraints

Abstract

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.