Gradient Accumulation Hidden Inside Momentum: Fewer Update Steps, Better Results?

隐藏在动量中的梯度累积：少更新几步，效果反而更好？ 苏剑林 · 2021-08-24 · Machine Learning model analysis gradient optimizer

As we know, gradient accumulation is a common trick for achieving large-batch-size training under limited GPU memory. In an earlier post, Trading Time for Effectiveness: A Keras Gradient Accumulation Optimizer, I gave a brief introduction to implementing gradient accumulation, the basic idea being to add an extra set of parameters to cache the gradients, and finally use the cached gradients to update the model. The downside is that this extra set of parameters brings additional GPU memory overhead.

A few days ago, while thinking about optimizers, I suddenly realized: gradient accumulation can actually be built directly into optimizers with momentum! With this idea in mind, I did some derivation and experiments on the optimizer, and eventually arrived at an interesting but somewhat counterintuitive conclusion: updating the parameters a few times less often can actually make the final model better!

Note: the results presented below appeared, almost verbatim and without citation, in Google's paper Combined Scaling for Zero-shot Transfer Learning.

SGDM

Before getting into the discussion proper, let's define the function

\begin{equation}\chi_{t/k} = \left\{ \begin{aligned}&1,\quad t \equiv 0\,(\text{mod}\, k) \\ &0,\quad t \not\equiv 0\,(\text{mod}\, k) \end{aligned}\right.\end{equation}

That is, $t$ is an integer, and when it is a multiple of $k$, $\chi_{t/k}=1$, otherwise $\chi_{t/k}=0$ — this is essentially the indicator function for whether $t$ is divisible by $k$. We will make repeated use of this function in what follows. more

Now, let's discuss SGD with momentum (SGDM), whose general form is

\begin{equation}\left\{\begin{aligned} m_t =&\, \beta m_{t-1} + \left(1 - \beta\right) g_t\\ \theta_t =&\, \theta_{t-1} - \alpha_t m_t \end{aligned}\right.\end{equation}

Here $g_t$ is the gradient at step $t$, $\theta$ is the parameter, and $\beta$ is the moving-average coefficient. Suppose we accumulate gradients over $k$ steps — that is, we actually update the parameters only once every $k$ steps, and each update uses the average of $k$ steps' worth of gradients, i.e.

\begin{equation}\left\{\begin{aligned} m_{kt} =&\, \beta m_{k(t-1)} + \left(1 - \beta\right) \frac{1}{k}\sum_{i=1}^k g_{k(t-1) + i}\\ \theta_{kt} =&\, \theta_{k(t-1)} - \alpha_{kt} m_{kt} \end{aligned}\right.\end{equation}

We can split apart the update of the momentum $m$:

\begin{equation}\begin{aligned} m_{k(t-1)+1} &\,= \beta m_{k(t-1)} + \frac{1}{k}\left(1 - \beta\right)g_{k(t-1) + 1}\\ m_{k(t-1)+2} &\,= m_{k(t-1) + 1} + \frac{1}{k}\left(1 - \beta\right)g_{k(t-1) + 2} \\ m_{k(t-1)+3} &\,= m_{k(t-1) + 2} + \frac{1}{k}\left(1 - \beta\right)g_{k(t-1) + 3} \\ &\,\,\,\vdots \\ m_{kt} &\,= m_{kt-1} + \frac{1}{k}\left(1 - \beta\right)g_{kt} \\ \end{aligned}\label{eq:m-part}\end{equation}

or write it as a general formula:

\begin{equation}m_t = \big[(\beta - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta\right) g_t\end{equation}

Likewise, the update of the parameter $\theta$ can also be split apart:

\begin{equation}\begin{aligned} \theta_{k(t-1)+1} &\,= \theta_{k(t-1)}\\ \theta_{k(t-1)+2} &\,= \theta_{k(t-1) + 1}\\ &\,\,\,\vdots \\ \theta_{kt-1} &\,= \theta_{kt-2}\\ \theta_{kt} &\,= \theta_{kt-1} - \alpha_{kt} m_{kt} \\ \end{aligned}\end{equation}

or written as a general formula:

\begin{equation}\theta_t = \theta_{t-1} - \chi_{t/k} \alpha_t m_t \end{equation}

So, for gradient descent with momentum, if we want to accumulate gradients over $k$ steps, we only need to update as follows:

\begin{equation}\left\{\begin{aligned} m_t =&\, \big[(\beta - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta\right) g_t\\ \theta_t =&\, \theta_{t-1} - \chi_{t/k} \alpha_t m_t \end{aligned}\right.\end{equation}

without needing to introduce any new set of parameters.

Adam

For Adam, the update formulas are as follows:

\begin{equation}\left\{\begin{aligned} m_t =&\, \beta_1 m_{t-1} + \left(1 - \beta_1\right) g_t\\ v_t =&\, \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\ \hat{m}_t =&\, m_t\left/\left(1 - \beta_1^t\right)\right.\\ \hat{v}_t =&\, v_t\left/\left(1 - \beta_2^t\right)\right.\\ \theta_t =&\, \theta_{t-1} - \alpha_t \hat{m}_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right.\end{equation}

The momentum $m$ is handled the same way as in SGDM above, so the key issue is the second moment $v$. By definition, when accumulating gradients over $k$ steps, the update formula for $v$ is

\begin{equation} v_{kt} = \beta_2 v_{k(t-1)} + \left(1 - \beta_2\right) \left(\frac{1}{k}\sum_{i=1}^k g_{k(t-1) + i}\right)^2\end{equation}

Unfortunately, because of the squaring involved, $v$ cannot be split apart in the same way as equation $\eqref{eq:m-part}$. So strictly speaking, it is not possible to implement Adam's gradient accumulation without adding an extra set of cache parameters.

However, if we assume that "the square of the average" can be estimated by "the average of the squares" — that is, if we assume

\begin{equation}\left(\frac{1}{k}\sum_{i=1}^k g_{k(t-1) + i}\right)^2\sim \frac{1}{k}\sum_{i=1}^k g_{k(t-1) + i}^2\end{equation}

(here $\sim$ denotes a fairly tight linear relationship, not necessarily approximate equality — they may differ by a constant factor) — then we can still modify the formula for $v$ in the same way as $m$:

\begin{equation}v_t = \big[(\beta_2 - 1)\chi_{(t - 1)/k} + 1\big] v_{t-1} + \frac{1}{k}\left(1 - \beta_2\right) g_t^2\end{equation}

Putting it all together, we get the modified Adam formula:

\begin{equation}\left\{\begin{aligned} m_t =&\, \big[(\beta_1 - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta_1\right) g_t\\ v_t =&\, \big[(\beta_2 - 1)\chi_{(t - 1)/k} + 1\big] v_{t-1} + \frac{1}{k}\left(1 - \beta_2\right) g_t^2\\ \hat{m}_t =&\, m_t\left/\left(1 - \beta_1^{t/k}\right)\right.\\ \hat{v}_t =&\, v_t\left/\left(1 - \beta_2^{t/k}\right)\right.\\ \theta_t =&\, \theta_{t-1} - \chi_{t/k}\alpha_t \hat{m}_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right.\end{equation}

After experimenting, I found that this modified version of Adam does indeed achieve an effect similar to gradient accumulation.

Reflecting on the Conclusion

Overall, the two modified optimizers above involve two main changes:

1. Modify the update formulas for $m$ and $v$;
2. Update the parameters only once every $k$ steps (or equivalently, change the learning rate from $\alpha_t$ to $\chi_{t/k}\alpha_t$).

Here the update formula for $m,v$ is changed to (without loss of generality, taking $m$ as an example):

\begin{equation}m_t = \big[(\beta - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta\right) g_t\end{equation}

To match this against the original format of $m_t = \beta m_{t-1} + (1 - \beta) g_t$, we might guess that the above iteration is equivalent to replacing $\beta$ with $\tilde{\beta}=1 - \frac{1}{k}(1-\beta)$.

Indeed this turns out to be the case: we can show that the moving-average momentum obtained by replacing $\beta$ with $\tilde{\beta}=1 - \frac{1}{k}(1-\beta)$ is approximately equal to the moving-average momentum obtained from the original $\beta$ when accumulating gradients over $k$ steps:

\begin{equation}\begin{aligned} m_{kt} =&\, \tilde{\beta}m_{kt-1} + (1 - \tilde{\beta})g_{kt} \\ =&\, \tilde{\beta}^2 m_{kt-2} + \tilde{\beta}(1 - \tilde{\beta})g_{kt - 1} + (1 - \tilde{\beta})g_{kt}\\ =&\,\cdots\\ =&\, \tilde{\beta}^k m_{k(t-1)} + (1 - \tilde{\beta})\sum_{i=1}^k \tilde{\beta}^{i-1} g_{kt-i+1} \end{aligned}\end{equation}

Since $\tilde{\beta}$ is usually quite close to 1, we can approximate $\tilde{\beta}^{i-1}\approx 1, \forall i=1,2,\cdots,k$, and

\begin{equation}\tilde{\beta}^k = (1 - (1 - \tilde{\beta}))^k \approx 1 - k(1 - \tilde{\beta}) = \beta\end{equation}

giving us the approximation

\begin{equation}\begin{aligned} m_{kt} \approx &\, \beta m_{k(t-1)} + (1 - \tilde{\beta})\sum_{i=1}^k g_{kt-i+1} \\ = &\, \beta m_{k(t-1)} + (1 - \beta) \frac{1}{k}\sum_{i=1}^k g_{kt-i+1} \end{aligned}\end{equation}

which is exactly the form of gradient accumulation.

So, if you're facing a performance problem caused by a small batch size, but don't want to implement gradient accumulation yourself or modify your optimizer, you can try the following:

1. Change all $\beta$ parameters to $1-\frac{1}{k}(1-\beta)$;
2. Multiply the learning rate by $\chi_{t/k}$, so that the parameters are updated only once every $k$ steps.

Interestingly, I found that point 2 matters more than point 1: even if we don't modify $\beta$ at all, and simply change the parameter update frequency to once every $k$ steps, the results improve to some extent (provided that a small batch size is indeed the current bottleneck for the model). This is the "counterintuitive" phenomenon mentioned in the title:

Without changing anything else, updating a bit less often actually gives better results.

Summary

This post described a discovery I made while debugging a model — that gradient accumulation is hidden inside the momentum term of an optimizer. I then carried out some further simple analysis and experiments on this, finding that adjusting the moving-average coefficient and the learning rate can also achieve an effect approximating gradient accumulation, which in turn led to the counterintuitive phenomenon that "updating less often can actually give better results."