Log In Sign Up

Scaling Deep Contrastive Learning Batch Size with Almost Constant Peak Memory Usage

by   Luyu Gao, et al.

Contrastive learning has been applied successfully to learn numerical vector representations of various forms of data, such as texts and images. Learned encoders exhibit versatile transfer capabilities to many downstream tasks. Representation based search is highly efficient with state-of-the-art performance. Previous researches demonstrated that learning high-quality representations requires a large number of negatives in contrastive loss. 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 re-computation technique that decouples back propagation 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 an almost constant peak GPU memory usage for batches of different sizes.


page 1

page 2

page 3

page 4


Watermarking Pre-trained Encoders in Contrastive Learning

Contrastive learning has become a popular technique to pre-train image e...

Provable Stochastic Optimization for Global Contrastive Learning: Small Batch Does Not Harm Performance

In this paper, we study contrastive learning from an optimization perspe...

MetAug: Contrastive Learning via Meta Feature Augmentation

What matters for contrastive learning? We argue that contrastive learnin...

Rethinking Prototypical Contrastive Learning through Alignment, Uniformity and Correlation

Contrastive self-supervised learning (CSL) with a prototypical regulariz...

Investigating the Role of Negatives in Contrastive Representation Learning

Noise contrastive learning is a popular technique for unsupervised repre...

Decoupled Contrastive Learning

Contrastive learning (CL) is one of the most successful paradigms for se...

Batch-iFDD for Representation Expansion in Large MDPs

Matching pursuit (MP) methods are a promising class of feature construct...

Code Repositories