自从GPT、BERT等预训练模型流行起来后,其中一个明显的趋势是模型越做越大,因为更大的模型配合更充分的预训练通常能更有效地刷榜。不过,理想可以无限远,现实通常很局促,有时候模型太大了,大到哪怕你拥有了大显存的GPU甚至TPU,依然会感到很绝望。比如GPT2最大的版本有15亿参数,最大版本的T5模型参数量甚至去到了110亿,这等规模的模型,哪怕在TPU集群上也没法跑到多大的batch size。

这时候通常要往优化过程着手,比如使用混合精度训练(tensorflow下还可以使用一种叫做bfloat16的新型浮点格式),即省显存又加速训练;又或者使用更省显存的优化器,比如RMSProp就比Adam更省显存。本文则介绍AdaFactor,一个由Google提出来的新型优化器,首发论文为《Adafactor: Adaptive Learning Rates with Sublinear Memory Cost》AdaFactor具有自适应学习率的特性,但比RMSProp还要省显存,并且还针对性地解决了Adam的一些缺陷。

Adam #

首先我们来回顾一下常用的Adam优化器的更新过程。设$t$为迭代步数,$\alpha_t$为当前学习率,$L(\theta)$是损失函数,$\theta$是待优化参数,$\epsilon$则是防止溢出的小正数,那么Adam的更新过程为
\begin{equation}\left\{\begin{aligned}&g_t = \nabla_{\theta} L(\theta_{t-1})\\
&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/\left(\sqrt{\hat{v}_t} + \epsilon\right)\right.
\end{aligned}\right.\end{equation}

要省显存,就首先得知道显存花在哪里的。首先,计算量和显存的大头肯定都是$\nabla_{\theta} L(\theta_{t-1})$,也就是说,计算梯度是很费资源的,这也是为啥“ALBERT相比BERT参数量虽然少了那么多,但训练速度也没见快多少”的原因了;除此之外,显存的消耗主要是$m,v$了,我们要维护两组缓存变量,来滑动计算梯度的前两阶矩(也就是$m$和$v$),用以计算参数的更新量。这两组变量每一组都跟训练参数本身一样大,因此对于参数比较多的模型,两组缓存变量所消耗的显存也不少。

AdaFactor #

在这一节中,我们会相对详细地介绍一些AdaFactor优化器,介绍中会设计比较多的公式和推导。如果只求一个大致了解的读者,可以自行跳过部分数学内容~

抛弃动量 #

我们知道,CV模型很多时候要靠“SGD+动量”来炼出最优效果来,自适应学习率优化器通常训练不出最好的效果。但对于NLP模型来说,情况有点相反,自适应学习率显得更重要一些,很少听到由纯靠SGD调NLP模型的案例。因此,作为省显存的第一步,我们可以抛弃Adam里边的动量,这样就少一组缓存参数了,自然也就省了显存:
\begin{equation}\left\{\begin{aligned}&g_t = \nabla_{\theta} L(\theta_{t-1})\\
&v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\
&\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\
&\theta_t = \theta_{t-1} - \alpha_t g_t\left/\sqrt{\hat{v}_t + \epsilon}\right.
\end{aligned}\right.\end{equation}
这其实就是RMSProp的变种,比RMSProp多了$\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.$这一步。

低秩分解 #

去掉$m$之后,缓存变量直接减少了一半,但AdaFactor还不满意,它希望保留自适应学习率功能,但把缓存变量$v$的参数量再压一压。这一次,它用到了矩阵的低秩分解。

广义KL散度 #

在SGD中,所有参数都是共用一个标量学习率;在Adam中,则是每一个参数都有自己的学习率$\alpha_t\left/\sqrt{\hat{v}_t + \epsilon}\right.$。我们知道通过精调学习率,SGD其实也能有不错的效果,这表明“每一个参数都有自己的学习率”这件事情都不是特别重要,或者换一种说法,就是“精调每一个参数自己的学习率”并不是特别重要。

这启发我们,将$\hat{v}_t$换一种参数更少的近似可能也就足够了。而“参数更少的近似”,我们就不难想到低秩分解了。对于$m\times n$的矩阵$C$,我们希望找到$m\times k$的矩阵$A$和$k\times n$的矩阵$B$,使得
\begin{equation}AB \approx C\end{equation}
当$k$足够小时,$A$、$B$的参数总量就小于$C$的参数量。为了“省”到极致,AdaFactor直接让$k=1$,即寻找$\{a_i\}_{i=1}^m$和$\{b_j\}_{j=1}^n$,使得
\begin{equation}a_i b_j \approx c_{i,j}\end{equation}
既然要近似,就要有一个度量的标准。很容易想到的标准是欧氏距离,即
\begin{equation}\sum_{i,j} (a_i b_j - c_{i,j})^2\end{equation}
但在这个距离之下,$a_i,b_j$并没有解析解;此外,在优化过程中$c_{i,j}$(即$\hat{v}_t$)是非负的,而通过上述目标优化出来的$a_i b_j$无法保证非负,因此很可能扰乱优化过程。

原论文的作者们很机智地换了一个度量标准,使得$a_i,b_j$有解析解。具体来说,它使用了“广义KL散度”,又称“I散度I-Divergence)”,其形式为:
\begin{equation}l = \sum_{i,j} c_{i,j}\log \frac{c_{i,j}}{a_i b_j} - c_{i,j} + a_i b_j \label{eq:i-div}\end{equation}
这个度量源自不等式$x\log x\geq x - 1(\forall x > 0)$,当且仅当$x=1$时等号成立。所以代入$x = p / q\,(p,q > 0)$,然后两端乘以$q$,我们有
\begin{equation}p\log \frac{p}{q} - p + q \geq 0\end{equation}
当且仅当$p=q$成立,如果$p,q$有多个分量,那么对多个分量的结果求和即可,这就得到了度量$\eqref{eq:i-div}$

显然,广义KL散度是概率的KL散度的自然推广,但它不要求$c_{i,j}$和$a_i b_j$满足归一化,只要求它们非负,这正好对应了AdaFactor的场景。而且巧妙的是,这种情形配上这个目标,刚好有解析解:
\begin{equation}a_i = \sum\limits_{j}c_{i,j},\quad b_j = \frac{\sum\limits_{i}c_{i,j}}{\sum\limits_{i,j}c_{i,j}}\label{eq:aibj}\end{equation}
其实这个解析解也很形象,就是行、列分别求和,然后相乘,再除以全体的和。

推导过程 #

直接对$\eqref{eq:i-div}$求偏导数并让偏导数等于0,得
\begin{equation}\left\{\begin{aligned}
&\frac{\partial l}{\partial a_i}=\sum_j -\frac{c_{i,j}}{a_i} + b_j = 0\\
&\frac{\partial l}{\partial b_j}=\sum_i -\frac{c_{i,j}}{b_j} + a_i = 0
\end{aligned}\right.\end{equation}
整理得
\begin{equation}\left\{\begin{aligned}
&a_i \sum_{j} b_j = \sum_j c_{i,j}\\
&b_j \sum_{i} a_i = \sum_i c_{i,j}
\end{aligned}\right.\end{equation}
注意到如果$(a_i,b_j)$是一组最优解,那么$(\lambda a_i,b_j/\lambda)$也是,说白了,所有的$a_i$乘以一个常数,所有的$b_j$也除以这个常数,$a_i b_j$是不变的。那么我们就可以随意指定$\sum\limits_{i} a_i$或$\sum\limits_{j} b_j$,因为它们就只是一个缩放标量而已。不失一般性,我们指定$\sum\limits_{j} b_j=1$,那么就解得$\eqref{eq:aibj}$。

直观理解 #

我们也可以从另一个角度理解结果$\eqref{eq:aibj}$。由于$c_{i,j}$是非负的,我们可以将它归一化,变成具有概率分布的特性,即$\hat{c}_{i,j}=\frac{c_{i,j}}{\sum\limits_{i,j}c_{i,j}}$,然后我们试图完成分解$\hat{c}_{i,j}\approx \hat{a}_i \hat{b}_j$,由于$\hat{c}_{i,j}$现在相当于一个二元联合概率分布,那么$\hat{a}_i,\hat{b}_j$就相当于它们的边缘分布,即
\begin{equation}\hat{a}_i = \sum_j \hat{c}_{i,j} = \frac{\sum\limits_{j}c_{i,j}}{\sum\limits_{i,j} c_{i,j}},\quad \hat{b}_j = \sum_i \hat{c}_{i,j} = \frac{\sum\limits_{i}c_{i,j}}{\sum\limits_{i,j}c_{i,j}}\end{equation}
现在$\hat{c}_{i,j}$到$c_{i,j}$还需要乘上一个$\sum\limits_{i,j}c_{i,j}$,我们可以把它乘到$\hat{a}_i$或$\hat{b}_j$中,不失一般性,我们假设乘到$\hat{a}_i$上,那么就得到$\eqref{eq:aibj}$。

AdaFactor雏形 #

有了结果$\eqref{eq:aibj}$后,我们就可以用它来构建更省内存的优化器了,这就是AdaFactor的雏形。简单来说,当参数$\theta$是普通一维向量时,优化过程保持不变;但$\theta$是$m\times n$的矩阵时,算出来的梯度$g_t$也是矩阵,从而$g_t^2$也是矩阵,这时候我们对$g_t^2$做低秩分解,然后维护两组缓存变量$v^{(r)}_t\in \mathbb{R}^m,v^{(c)}_t\in\mathbb{R}^n$,分别滑动平均低秩分解后的结果,最后用$v^{(r)}_t,v^{(c)}_t$共同调整学习率:
\begin{equation}\left\{\begin{aligned}&g_{i,j;t} = \nabla_{\theta} L(\theta_{i,j;t-1})\\
&v^{(r)}_{i;t} = \beta_2 v^{(r)}_{t-1;i} + \left(1 - \beta_2\right) \sum\limits_{j}\left(g_{i,j;t}^2+\epsilon\right)\\
&v^{(c)}_{j;t} = \beta_2 v^{(c)}_{t-1;j} + \left(1 - \beta_2\right) \sum\limits_{i}\left(g_{i,j;t}^2+\epsilon\right)\\
&v_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t}\left/\sum\limits_{j}v^{(c)}_{j;t}\right.\\
&\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\
&\theta_t = \theta_{t-1} - \alpha_t g_t\left/\sqrt{\hat{v}_t}\right.
\end{aligned}\right.\end{equation}
(把$\epsilon$加到$g_t^2$上去而不是$\hat{v}_t$上去,这是AdaFactor整出来的形式,不是笔者的锅~)

滑动权重 #

在Adam以及上述AdaFactor雏形中,滑动权重$\beta_2$都是恒为常数,AdaFactor指出这是不科学的,并提出新的策略。

等价形式 #

为了认识到这一点,我们重写一下Adam的$\hat{v}_t$的更新过程:
\begin{equation}\begin{aligned}
\hat{v}_t =& v_t\left/\left(1 - \beta_2^t\right)\right.\\
=&\frac{\beta_2 v_{t-1} + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\
=&\frac{\beta_2 \hat{v}_{t-1}\left(1 - \beta_2^{t-1}\right) + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\
=&\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\hat{v}_{t-1} + \left(1 - \beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\right)g_t^2
\end{aligned}\end{equation}
所以如果设$\hat{\beta}_{2,t}=\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}$,那么更新公式就是
\begin{equation}\hat{v}_t =\hat{\beta}_{2,t}\hat{v}_{t-1} + \left(1 - \hat{\beta}_{2,t}\right)g_t^2\end{equation}
问题是这个$\hat{\beta}_{2,t}$够不够合理呢?答案是可能不大够。当$t=1$时$\hat{\beta}_{2,t}=0$,这时候$\hat{v}_t$就是$g_t^2$,也就是用实时梯度来校正学习率,这时候校正力度最大;当$t\to\infty$时,$\hat{\beta}_{2,t}\to \beta_2$,这时候$v_t$是累积梯度平方与当前梯度平方的加权平均,由于$\beta_2 < 1$,所以意味着当前梯度的权重$1 - \beta_2$不为0,这可能导致训练不稳定,因为训练后期梯度变小,训练本身趋于稳定,校正学习率的意义就不大了,因此学习率的校正力度应该变小,并且$t\to\infty$,学习率最好恒定为常数(这时候相当于退化为SGD),这就要求$t\to\infty$时,$\hat{\beta}_{2,t}\to 1$。

新的衰减策略 #

为了达到这个目的,AdaFactor采用如下的衰减策略
\begin{equation}\hat{\beta}_{2,t} =1 - \frac{1}{t^c}\label{eq:beta2}\end{equation}
它满足$\hat{\beta}_{2,1}=0,\lim\limits_{t\to\infty} \hat{\beta}_{2,t}=1$。但即便如此,也不是任何$c$都适合,必须有$0 < c <1$。$c > 0$好理解,那为什么要$c < 1$呢?原论文包含了对它的分析,大家可以去读读,但笔者觉得原论文的推导过于晦涩,所以这里给出自己的理解。

首先,对于$\hat{v}_t$来说,一个很容易想到的方案是所有梯度平方的平均,即:
\begin{equation}\hat{v}_t = \frac{1}{t}\sum_{i=1}^t g_i^2=\frac{t-1}{t}\hat{v}_{t-1} + \frac{1}{t}g_t^2\end{equation}
所以这等价于让$\hat{\beta}_{2,t} =1 - \frac{1}{t}$。这个方案美中不足的一点是,每一步梯度都是平权的,这不符合直觉,因为正常来说越久远的梯度应该越不重要才对,所以应该适当降低历史部分权重,而当$c < 1$时,$1 - \frac{1}{t^c} < 1 - \frac{1}{t}$,因此一个简洁的方案是在式$\eqref{eq:beta2}$中取$c < 1$,AdaFactor默认的$c$是$0.8$。

层自适应 #

最后,我们还可以进一步根据参数的模长来校正更新量,这个思路来自LAMB优化器,在之前的文章《6个派生优化器的简单介绍及其实现》中也介绍过。简单来说,它就是将最后的更新量标准化,然后乘以参数的模长,说白了,就是不管你怎么折腾,最后的更新量我只要你的方向,而大小由参数本身的模长和预先设置学习率共同决定,使得所有层所有参数的相对变化程度保持一致。

AdaFactor完整版 #

至此,我们终于可以写出完整版AdaFactor的更新过程了:
\begin{equation}\left\{\begin{aligned}&g_{i,j;t} = \nabla_{\theta} L(\theta_{i,j;t-1})\\
&\hat{\beta}_{2,t} =1 - t^{-c}\\
&v^{(r)}_{i;t} = \hat{\beta}_{2,t} v^{(r)}_{t-1;i} + \left(1 - \hat{\beta}_{2,t}\right) \sum\limits_{j}\left(g_{i,j;t}^2+\epsilon_1\right)\\
&v^{(c)}_{j;t} = \hat{\beta}_{2,t} v^{(c)}_{t-1;j} + \left(1 - \hat{\beta}_{2,t}\right) \sum\limits_{i}\left(g_{i,j;t}^2+\epsilon_1\right)\\
&\hat{v}_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t}\left/\sum\limits_{j}v^{(c)}_{j;t}\right.\\
&u_t = g_t\left/\sqrt{\hat{v}_t}\right.\\
&\hat{u}_t = u_t \left/\max\left(1, \left. RMS(u_t)\right/d\right)\right.\times \max\left(\epsilon_2, RMS(\theta_{t-1})\right)\\
&\theta_t = \theta_{t-1} - \alpha_t \hat{u}_t
\end{aligned}\right.\end{equation}
其中$RMS(x)=\sqrt{\frac{1}{n}\sum\limits_{i=1}^n x_i^2}$是模长的变种,$\max\left(1, \left. RMS(u_t)\right/d\right)$这一步相当于做了个截断,即$RMS(u_t) > d$时才执行归一化。原论文中的默认参数为
$$\begin{array}{c|c}
\hline
\epsilon_1 & 10^{-30}\\
\hline
\epsilon_2 & 10^{-3}\\
\hline
d & 1\\
\hline
\hat{\beta}_{2,t} & 1 - t^{-0.8}\\
\hline
\end{array}$$
如果参数是一维向量而不是矩阵,那么$\hat{v}_t$使用普通的更新公式$\hat{v}_t = \hat{\beta}_{2,t} v_{t-1} + \left(1 - \hat{\beta}_{2,t}\right) \left(g_t^2+\epsilon_1\right)$就行了。此外,论文还提出如果没有传入学习率,那么可以使用$a_t = \min\left(10^{-2},\frac{1}{\sqrt{t}}\right)$为默认学习率,但笔者看源码的时候发现这个默认学习率很少使用,基本上还是需要自己传入学习率的。

开源实现 #

为了方便大家使用,笔者开源了自己实现的AdaFactor:

开源包括纯keras版和tf.keras版,使用方法跟普通keras优化器一样,tf.keras版也可以当做一个普通的tensorflow优化器使用。开源实现参考了mesh_tensorflow版的源码,在此表示感谢。优化器也已经内置在bert4keras中,方便大家调用。

需要提醒的是,用AdaFactor的时候,batch_size最好大一些,因为本身低秩分解会带来误差,而如果batch_size过小,那么梯度估算本身也带来较大的误差,两者叠加优化过程可能还不收敛。对于预训练模型来说,batch_size通常还是很大的,所以现在不少预训练模型开始用AdaFactor优化器了;对于普通的下游任务来说,AdaFactor也可以尝试,但可能需要多炼炼丹,才能搞出优于无脑Adam的效果。对了,还要提醒一下,用AdaFactor的时候,学习率要设大一点,大概是$10^{-3}$级别为好,哪怕是finetune阶段也是如此。

文章小结 #

本文介绍了Google提出来的AdaFactor优化器,一个旨在减少显存占用的优化器,并且针对性地分析并解决了Adam的一些缺陷。笔者认为,AdaFactor针对Adam所做的分析相当经典,值得我们认真琢磨体味,对有兴趣研究优化问题的读者来说,更是一个不可多得的分析案例。

当然,没有什么绝对能有效的方法,有的只是

方法虽好,要想实际有效,依然要用心炼丹。

转载到请包括本文地址:https://www.kexue.fm/archives/7302

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

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

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

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

苏剑林. (Mar. 23, 2020). 《AdaFactor优化器浅析(附开源实现) 》[Blog post]. Retrieved from https://www.kexue.fm/archives/7302

@online{kexuefm-7302,
        title={AdaFactor优化器浅析(附开源实现)},
        author={苏剑林},
        year={2020},
        month={Mar},
        url={\url{https://www.kexue.fm/archives/7302}},
}