Can Contrastive Learning Use Gradient Accumulation?

对比学习可以使用梯度累积吗？ 苏剑林 · 2021-06-17 · Machine Learning model optimization gradient

In an earlier post, Trading Time for Effect: Keras Gradient Accumulation Optimizer], we introduced "gradient accumulation," a trick for achieving the effect of a large batch_size under limited GPU memory. Generally speaking, gradient accumulation applies to scenarios where the loss is i.i.d., meaning each sample's loss is computed individually, and the total loss is the average or sum of all the individual losses. However, not every task satisfies this condition — for example, the currently popular contrastive learning, where each sample's loss also depends on the other samples.

So, in the contrastive learning setting, can we still use gradient accumulation to achieve the effect of a large batch_size? This post analyzes this question.

Introduction

In general, the loss for contrastive learning can be written as

\begin{equation}\mathcal{L}=-\sum_{i,j=1}^b t_{i,j}\log p_{i,j} = -\sum_{i,j=1}^b t_{i,j}\log \frac{e^{s_{i,j}}}{\sum\limits_j e^{s_{i,j}}}=-\sum_{i,j=1}^b t_{i,j}s_{i,j} + \sum_{i=1}^b \log\sum_{j=1}^b e^{s_{i,j}}\label{eq:loss}\end{equation}

Here $b$ is the batch_size; $t_{i,j}$ is a given label, satisfying $t_{i,j}=t_{j,i}$ — it is a one-hot matrix, with each column having exactly one 1 and the rest 0; and $s_{i,j}$ is the similarity between sample $i$ and sample $j$, satisfying $s_{i,j}=s_{j,i}$. Generally there's also a temperature parameter, but here we assume the temperature parameter has already been folded into $s_{i,j}$, to simplify notation. The model parameters appear in $s_{i,j}$, which we assume equals $\theta$. more

It can be verified that, in general:

\begin{equation}-\sum_{i,j=1}^{2b} t_{i,j}\log p_{i,j} \neq -\sum_{i,j=1}^{b} t_{i,j}\log p_{i,j}-\sum_{i,j=b+1}^{2b} t_{i,j}\log p_{i,j}\end{equation}

So directly accumulating the gradients of small-batch_size contrastive learning is not equivalent to contrastive learning with a large batch_size. A similar issue exists for models with BN (Batch Normalization).

Gradient

Note that what we just said applies to the conventional, simple form of gradient accumulation being non-equivalent — but there may exist a slightly more sophisticated accumulation scheme that works. To this end, let's analyze the gradient of equation $\eqref{eq:loss}$:

\begin{equation}\begin{aligned} \nabla_{\theta}\mathcal{L} =&\, -\sum_{i,j=1}^b t_{i,j}\nabla_{\theta}s_{i,j} + \sum_{i=1}^b \nabla_{\theta}\log\sum_{j=1}^b e^{s_{i,j}} \\ =&\, -\sum_{i,j=1}^b t_{i,j}\nabla_{\theta}s_{i,j} + \sum_{i,j=1}^b p_{i,j}\nabla_{\theta} s_{i,j} \\ =&\,\nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j} \end{aligned}\end{equation}

where $p_{i,j}^{(sg)}$ denotes that we don't need to take the gradient of $\theta$ with respect to $p_{i,j}$ — that is, the stop_gradient operator in deep learning frameworks. The above equation shows that, if we use a gradient-based optimizer, then using equation $\eqref{eq:loss}$ as the loss is completely equivalent to using $\sum\limits_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j}$ as the loss (since the resulting gradients are identical).

Inner Product

Next, consider the computation of $\nabla_{\theta}s_{i,j}$. Generally speaking, it takes the form of an inner product of vectors, i.e., $s_{i,j}=\langle h_i, h_j\rangle$, with the parameters $\theta$ contained in $h_i,h_j$. In that case:

\begin{equation}\nabla_{\theta}s_{i,j}=\langle \nabla_{\theta}h_i, h_j\rangle + \langle h_i, \nabla_{\theta}h_j\rangle = \nabla_{\theta}\left(\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle\right)\end{equation}

So the term $s_{i,j}$ in the loss can be replaced by $\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle$ without changing the effect:

\begin{equation}\begin{aligned} \nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)s_{i,j} =&\, \nabla_{\theta}\sum_{i,j=1}^b \left(p_{i,j}^{(sg)} - t_{i,j}\right)\left(\langle h_i, h_j^{(sg)}\rangle + \langle h_i^{(sg)}, h_j\rangle\right)\\ =&\, 2\nabla_{\theta}\sum_{i,j=1}^b \left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)\langle h_i, h_j^{(sg)}\rangle\\ =&\,\nabla_{\theta}\sum_{i=1}^b \left\langle h_i, 2\sum_{j=1}^b\left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)h_j^{(sg)}\right\rangle \end{aligned}\label{eq:g}\end{equation}

where $2\overline{p_{i,j}^{(sg)}}=p_{i,j}^{(sg)} + p_{j,i}^{(sg)}$, and the second equality follows from swapping the summation indices $i,j$ in the $\langle h_i^{(sg)}, h_j\rangle$ term, which doesn't change the result of the summation.

Procedure

Equation $\eqref{eq:g}$ has, in fact, already given us the final scheme, which can be split into two steps. The first step is computing the vector

\begin{equation}\tilde{h}_i = 2\sum_{j=1}^b\left(\overline{p_{i,j}^{(sg)}} - t_{i,j}\right)h_j^{(sg)}\label{eq:h}\end{equation}

This step doesn't require taking gradients — it's a pure forward pass — so the batch_size can be fairly large. The second step is to feed $\tilde{h}_i$ into the model as "labels," and optimize the model using $\langle h_i, \tilde{h}_i\rangle$ as the loss for each individual sample. This step does require taking gradients, but it has already been reduced to a sum of per-sample gradients, so at this point ordinary gradient accumulation can be used.

Suppose the maximum batch_size for backpropagation is $b$, and the maximum batch_size for the forward pass is $nb$. Then, to achieve the effect of contrastive learning with batch_size $nb$ via gradient accumulation, the formalized procedure is as follows:

1. Sample a batch of data $\{x_i\}_{i=1}^{nb}$, with corresponding label matrix $\{t_{i,j}\}_{i,j=1}^{nb}$; initialize the accumulated gradient as $g=0$;
2. Perform a forward pass of the model to obtain the encoded vectors $\{h_i\}_{i=1}^{nb}$ and the corresponding probability matrix $\{p_{i,j}\}_{i,j=1}^{nb}$;
3. Compute the label vector $\{\tilde{h}_i\}_{i=1}^{nb}$ according to equation $\eqref{eq:h}$;
4. For $k=1,2,\cdots,n$, execute:
$g \leftarrow g + \nabla_{\theta}\sum\limits_{i=(k-1)b+1}^{kb} \langle h_i, \tilde{h}_i\rangle$
5. Use $g$ as the final gradient to update the model, then go back to step 1.

Overall, this requires one extra forward pass compared to conventional gradient accumulation, in terms of computation. Of course, if even the maximum batch_size for the forward pass can't meet our needs, we can also do the forward pass in batches, since we only need to compute and store each $\{h_i\}_{i=1}^{nb}$, and $\{p_{i,j}\}_{i,j=1}^{nb}$ can then be computed based on $\{h_i\}_{i=1}^{nb}$.

One last reminder: the above procedure is only equivalent to the large-batch_size model in terms of optimization — that is, the gradient of $\langle h_i, \tilde{h}_i\rangle$ is equivalent to the gradient of the original loss, but its value is not equal to the value of the original loss. Therefore $\langle h_i, \tilde{h}_i\rangle$ cannot be used as a loss to evaluate the model: it need not be monotonic, need not be non-negative, and doesn't have a strict correlation with the original loss either.

A Caveat

The procedure above has the same issue as the "recomputation" technique introduced in The Memory-Saving Recomputation Trick Now Has a Keras Version]: it's not compatible with Dropout. This is because each update involves multiple forward passes, and each forward pass has a different Dropout mask. This means that the $h_i$ used when computing the label vector $\tilde{h}_i$ is not the same as the $h_i$ used when computing the gradient, so the resulting gradient is not the most sensible one.

There's no good solution to this. The simplest and most effective approach is just to remove Dropout from the model. This isn't much of a problem for CV, since CV models generally don't use Dropout anyway. For NLP, the first thing that comes to mind is that SimCSE] can't use gradient accumulation, because Dropout is the very foundation of SimCSE~

Summary

This post analyzed how to apply gradient accumulation to contrastive learning. The result shows that gradient accumulation can indeed be used for contrastive learning, at the cost of one extra forward pass, and it requires removing Dropout from the model. The same line of reasoning can also be used to analyze how to apply gradient accumulation to BN — interested readers are welcome to give it a try.