我们知道,梯度累积是在有限显存下实现大batch_size训练的常用技巧。在之前的文章《用时间换取效果:Keras梯度累积优化器》中,我们就简单介绍过梯度累积的实现,大致的思路是新增一组参数来缓存梯度,最后用缓存的梯度来更新模型。美中不足的是,新增一组参数会带来额外的显存占用。

这几天笔者在思考优化器的时候,突然意识到:梯度累积其实可以内置在带动量的优化器中!带着这个思路,笔者对优化了进行了一些推导和实验,最后还得到一个有意思但又有点反直觉的结论:少更新几步参数,模型最终效果可能会变好!

注:本文下面的结果,几乎原封不动且没有引用地出现在Google的论文《Combined Scaling for Zero-shot Transfer Learning》中,在此不做过多评价,请读者自行品评。

SGDM #

在正式讨论之前,我们定义函数
\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}
也就是说,$t$是一个整数,当它是$k$的倍数时,$\chi_{t/k}=1$,否则$\chi_{t/k}=0$,这其实就是一个$t$能否被$k$整除的示性函数。在后面的讨论中,我们将反复用到这个函数。

现在,我们来讨论带动量的梯度下降SGDM,它的一般形式为
\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}
这里$g_t$是第$t$步的梯度,$\theta$是参数,$\beta$是滑动系数。假设我们累积$k$步梯度,那么实际上就是每$k$步才更新一次,然后每次更新都使用$k$步的梯度平均,即
\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}
我们可以将动量$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}
或者写成一个通式:
\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}
同理,参数$\theta$的更新也可以分拆开:
\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}
或者写成一个通式:
\begin{equation}\theta_t = \theta_{t-1} - \chi_{t/k} \alpha_t m_t \end{equation}
所以,带动量的梯度下降,如果要累积$k$步梯度,那么只需要按照如下方式更新:
\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}
而不需要引入新的一组参数。

Adam #

对于Adam来说,它的更新公式如下:
\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}
动量$m$的处理方式同前述SGDM一致,所以关键在于二阶矩$v$。按照定义,在梯度累积$k$步的情况下,$v$的更新公式为
\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}
很遗憾,由于平方的存在,$v$无法做到像式$\eqref{eq:m-part}$那样的分拆,所以严格来说,不新增一组缓存参数是无法实现Adam的梯度累积的。

但是,如果我们假设“平均的平方”可以由“平方的平均”来估计,即假设
\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}
这里的$\sim$代表有比较紧密的线性相关关系,但不一定是近似相等,比如可以差一个倍数。那么我们依然可以像$m$那样对$v$的公式进行修改:
\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}
综合起来,就得到修改后的Adam公式:
\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}
笔者实验后发现,这样修改后的Adam,确实能起到跟梯度累积相似的效果。

结论反思 #

总的来说,在上面的两个修改版优化器中,主要包含两个改动:

1、修改$m$和$v$的更新公式;

2、参数改为每$k$步更新一次(或者学习率由$\alpha_t$改为$\chi_{t/k}\alpha_t$)。

其中$m,v$的更新公式改为了(不失一般性,以$m$为例):
\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}
如果要对应上原来的$m_t = \beta m_{t-1} + (1 - \beta) g_t$的格式,我们猜测上述迭代可能相当于将$\beta$换成了$\tilde{\beta}=1 - \frac{1}{k}(1-\beta)$。

事实上确实如此,我们可以证明将$\beta$换成$\tilde{\beta}=1 - \frac{1}{k}(1-\beta)$后的滑动平均动量,近似于原始$\beta$时累积$k$步梯度的滑动平均动量:
\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}
由于$\tilde{\beta}$通常相当接近于1,所以我们近似认为$\tilde{\beta}^{i-1}\approx 1, \forall i=1,2,\cdots,k$,并且
\begin{equation}\tilde{\beta}^k = (1 - (1 - \tilde{\beta}))^k \approx 1 - k(1 - \tilde{\beta}) = \beta\end{equation}
所以有近似
\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}
这就是梯度累积形式。

所以,如果你面临着batch_size比较小而导致性能不好的问题,但又不想自己实现梯度累积或者修改优化器,那么可以尝试如下操作:

1、将所有$\beta$参数改为$1-\frac{1}{k}(1-\beta)$;

2、学习率乘以$\chi_{t/k}$,使得参数每$k$步才更新一次。

有意思的是,笔者发现第2点比第1点更重要:哪怕我们不修改$\beta$,直接将参数更新频率改为每$k$步更新,效果也有所改善(前提是batch_size小是模型当前瓶颈)。这就是标题所提到的“反直觉”现象:

啥都不用改,少更新几次,效果反而更好了。

文章小结 #

本文介绍了笔者在调试模型中的一个发现——梯度累积隐藏在优化器的动量中,然后对此作了进一步的简单分析和实验,发现通过调整滑动系数和学习率也能达到近似的梯度累积效果,进而发现“少更新几次,效果反而更好”的反直觉现象。

转载到请包括本文地址:https://www.spaces.ac.cn/archives/8634

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Aug. 24, 2021). 《隐藏在动量中的梯度累积:少更新几步,效果反而更好? 》[Blog post]. Retrieved from https://www.spaces.ac.cn/archives/8634

@online{kexuefm-8634,
        title={隐藏在动量中的梯度累积:少更新几步,效果反而更好?},
        author={苏剑林},
        year={2021},
        month={Aug},
        url={\url{https://www.spaces.ac.cn/archives/8634}},
}