As we know, gradient accumulation is a common technique used to achieve large batch size training under limited VRAM. In a previous article, "Trading Time for Results: Keras Gradient Accumulation Optimizer", we briefly introduced the implementation of gradient accumulation. The general idea is to add a new set of parameters to cache gradients and finally use the cached gradients to update the model. The downside is that adding a new set of parameters incurs additional VRAM usage.
While thinking about optimizers recently, I suddenly realized: gradient accumulation can actually be built directly into momentum-based optimizers! With this idea in mind, I performed some derivations and experiments on the optimizer. Ultimately, I reached an interesting and somewhat counter-intuitive conclusion: updating parameters for fewer steps might actually lead to better final model performance!
Note: The results presented below in this article appeared almost identically and without citation in Google’s paper "Combined Scaling for Zero-shot Transfer Learning". I will not comment further on this; readers may judge for themselves.
SGDM
Before the formal discussion, let us define the function: \chi_{t/k} = \left\{ \begin{aligned} &1,\quad t \equiv 0 \pmod{k} \\ &0,\quad t \not\equiv 0 \pmod{k} \end{aligned}\right. That is, t is an integer; when it is a multiple of k, \chi_{t/k}=1, otherwise \chi_{t/k}=0. This is essentially an indicator function for whether t is divisible by k. We will use this function repeatedly in the following discussion.
Now, let’s discuss Stochastic Gradient Descent with Momentum (SGDM). Its general form is: \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. Here g_t is the gradient at step t, \theta represents the parameters, and \beta is the smoothing coefficient. Suppose we accumulate gradients for k steps; this effectively means we only update once every k steps, and each update uses the average of the k steps of gradients, i.e.: \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. We can decompose the update of the momentum m: \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} Or written as a general formula: m_t = \big[(\beta - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta\right) g_t Similarly, the update of parameter \theta can also be decomposed: \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} Or written as a general formula: \theta_t = \theta_{t-1} - \chi_{t/k} \alpha_t m_t Therefore, for SGDM, if we want to accumulate k steps of gradients, we only need to update in the following way: \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. This does not require introducing a new set of parameters.
Adam
For Adam, the update formulas are as follows: \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. The treatment of momentum m is consistent with the aforementioned SGDM, so the key lies in the second moment v. According to the definition, in the case of k-step gradient accumulation, the update formula for v is: 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 Unfortunately, due to the presence of the square, v cannot be decomposed like in Equation [eq:m-part]. Therefore, strictly speaking, it is impossible to implement Adam’s gradient accumulation without adding a set of cache parameters.
However, if we assume that the "square of the average" can be estimated by the "average of the squares," i.e., assuming: \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 Here \sim represents a close linear correlation, but not necessarily approximate equality (for example, they could differ by a constant factor). Then we can still modify the formula for v just as we did for m: 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 Combining these, we get the modified Adam formulas: \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. After experimenting, I found that this modified Adam indeed achieves an effect similar to gradient accumulation.
Conclusion and Reflection
In summary, the two modified optimizers above mainly include two changes:
Modifying the update formulas for m and v;
Changing the parameter update to occur every k steps (or changing the learning rate from \alpha_t to \chi_{t/k}\alpha_t).
The update formula for m and v was changed to (taking m as an example without loss of generality): m_t = \big[(\beta - 1)\chi_{(t - 1)/k} + 1\big] m_{t-1} + \frac{1}{k}\left(1 - \beta\right) g_t If we want to map this back to the original m_t = \beta m_{t-1} + (1 - \beta) g_t format, we might guess that the above iteration is equivalent to replacing \beta with \tilde{\beta}=1 - \frac{1}{k}(1-\beta).
In fact, this is indeed the case. We can prove that the exponential moving average momentum after replacing \beta with \tilde{\beta}=1 - \frac{1}{k}(1-\beta) is approximately equal to the moving average momentum of k-step accumulated gradients with the original \beta: \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} Since \tilde{\beta} is usually very close to 1, we can approximately assume \tilde{\beta}^{i-1}\approx 1, \forall i=1,2,\cdots,k, and: \tilde{\beta}^k = (1 - (1 - \tilde{\beta}))^k \approx 1 - k(1 - \tilde{\beta}) = \beta Thus, we have the approximation: \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} This is exactly the form of gradient accumulation.
Therefore, if you are facing performance issues due to a small batch size but do not want to implement gradient accumulation yourself or modify the optimizer, you can try the following:
Change all \beta parameters to 1-\frac{1}{k}(1-\beta);
Multiply the learning rate by \chi_{t/k}, so that parameters are only updated every k steps.
Interestingly, I found that the second point is more important than the first: even if we don’t modify \beta and simply change the parameter update frequency to every k steps, the performance improves (provided that the small batch size is the current bottleneck of the model). This is the "counter-intuitive" phenomenon mentioned in the title:
Without changing anything else, just updating fewer times actually leads to better results.
Summary
This article introduced a discovery I made while tuning models—that gradient accumulation is hidden within the momentum of the optimizer. I then performed further simple analysis and experiments, finding that adjusting the smoothing coefficients and learning rate can achieve an effect similar to gradient accumulation. This led to the counter-intuitive observation that "updating fewer times can lead to better results."
Original Address: https://kexue.fm/archives/8634
For more details on reprinting, please refer to: "Scientific Space FAQ"