AdaFactor优化器浅析（附开源实现）

自从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：

Github地址：https://github.com/bojone/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://kexue.fm/archives/7302

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

打赏

微信打赏

支付宝打赏

因为网站后台对打赏并无记录，因此欢迎在打赏时候备注留言。
你还可以点击这里或在下方评论区留言来告知你的建议或需求。