Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup (2021)

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.

## introduction contrastive learning learns to encode data into an embedding space such that related data points have closer representations and unrelated ones have further apart ones. recent works in nlp adopt deep neural nets as encoders and use unsupervised contrastive learning on sentence representation @xcite ), text retrieval @xcite , and language model pre-training tasks @xcite . supervised contrastive learning @xcite has also been shown effective in training dense retrievers @xcite . these works typically use batch-wise contrastive loss, sharing target texts as in-batch negatives. with such a technique, previous works have empirically shown that larger batches help learn better representations. however, computing loss and updating model parameters with respect 1 our code is at github.com/luyug/gradcache. to a big batch require encoding all batch data and storing all activation, so batch size is limited by total available gpu memory. this limits application and research of contrastive learning methods under memory limited setup, e.g. academia. for example, @xcite pre-train a bert @xcite passage encoder with a batch size of 4096 while a high-end commercial gpu rtx 2080ti can only fit a batch of 8. the gradient accumulation technique, splitting a large batch into chunks and summing gradients across several backwards, cannot emulate a large batch as each smaller chunk has fewer in-batch negatives. in this paper, we present a simple technique that thresholds peak memory usage for contrastive learning to almost constant regardless of the batch size. for deep contrastive learning, the memory bottlenecks are at the deep neural network based encoder. we observe that we can separate the backpropagation process of contrastive loss into two parts, from loss to representation, and from representation to model parameter, with the latter being independent across batch examples given the former, detailed in subsection 3.2. we then show in subsection 3.3 that by separately pre-computing the representations' gradient and store them in a cache, we can break the update of the encoder into multiple sub-updates that can fit into the gpu memory. this pre-computation of gradients allows our method to produce the exact same gradient update as training with large batch. experiments show that with about 20% increase in runtime, our technique enables a single consumer-grade gpu to reproduce the state-of-the-art large batch trained models that used to require multiple professional gpus. deep network memory reduction many existing techniques deal with large and deep models. the gradient checkpoint method attempts to emulate training deep networks by training shallower layers and connecting them with gradient checkpoints and re-computation @xcite . some methods also use reversible activation functions, allowing internal activation in the network to be recovered throughout back propagation @xcite . however, their effectiveness as part of contrastive encoders has not been confirmed. recent work also attempts to remove the redundancy in optimizer tracked parameters on each gpu @xcite . compared with the aforementioned methods, our method is designed for scaling over the batch size dimension for contrastive learning. ## methodologies in this section, we formally introduce the notations for contrastive loss and analyze the difficulties of using it on limited hardware. we then show how we can use a gradient cache technique to factor the loss so that large batch gradient update can be broken into several sub-updates. ## experiments to examine the reliability and computation cost of our method, we implement our method into dense passage retriever (dpr; karpukhin et al. ( 2020 )) 2 . we use gradient cache to compute dpr's supervised contrastive loss on a single gpu. following dpr paper, we measure top hit accuracy on the natural question dataset @xcite for different methods. we then examine the training speed of various batch sizes. ## extend to deep distance function previous discussion assumes a simple parameterless dot product similarity. in general it can also be deep distance function φ richly parameterized by ω, formally, this can still scale by introducing an extra distance gradient cache. in the first forward we collect all representations as well as all distances. we compute loss with d ij s and back propagate to get w @xmath0, and store them in distance gradient cache, [w 00 , w 01 , .., w 10 , ..]. we can then update ω in a sub-batch manner, with which we can build up the representation gradient cache. when all representations' gradients are computed and stored, encoder gradient can be computed with step3 described in subsection 3.3. in philosophy this method links up two caches. note this covers early interaction f (s) = s, g(t) = t as a special case. ## conclusion in this paper, we introduce a gradient cache technique that breaks gpu memory limitations for large batch contrastive learning. we propose to construct a representation gradient cache that removes in-batch data dependency in encoder optimization. our method produces the exact same gradient update as training with a large batch. we show the 3 we used the gradient checkpoint implemented in huggingface transformers package method is efficient and capable of preserving accuracy on resource-limited hardware. we believe a critical contribution of our work is providing a large population in the nlp community with access to batch-wise contrastive learning. while many previous works come from people with industry-grade hardware, researchers with limited hardware can now use our technique to reproduce state-of-the-art models and further advance the research without being constrained by available gpu memory.