一只小狐狸带你解锁炼丹术&NLP秘籍

作者：苏剑林（来自追一科技，人称“苏神”）

前言

自从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优化器的更新过程。设为迭代步数，为当前学习率，是损失函数，是待优化参数，则是防止溢出的小正数，那么Adam的更新过程为

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

AdaFactor

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

抛弃动量

我们知道，CV模型很多时候要靠“SGD+动量”来炼出最优效果来，自适应学习率优化器通常训练不出最好的效果。但对于NLP模型来说，情况有点相反，自适应学习率显得更重要一些，很少听到由纯靠SGD调NLP模型的案例。因此，作为省显存的第一步，我们可以抛弃Adam里边的动量，这样就少一组缓存参数了，自然也就省了显存：

这其实就是RMSProp的变种，比RMSProp多了这一步。

低秩分解

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

广义KL散度

在SGD中，所有参数都是共用一个标量学习率；在Adam中，则是每一个参数都有自己的学习率。我们知道通过精调学习率，SGD其实也能有不错的效果，这表明“每一个参数都有自己的学习率”这件事情都不是特别重要，或者换一种说法，就是“精调每一个参数自己的学习率”并不是特别重要。

这启发我们，将换一种参数更少的近似可能也就足够了。而“参数更少的近似”，我们就不难想到低秩分解了。对于的矩阵，我们希望找到的矩阵和的矩阵，使得

当足够小时，、的参数总量就小于的参数量。为了“省”到极致，AdaFactor直接让，即寻找和，使得

既然要近似，就要有一个度量的标准。很容易想到的标准是欧氏距离，即

但在这个距离之下，并没有解析解；此外，在优化过程中（即）是非负的，而通过上述目标优化出来的无法保证非负，因此很可能扰乱优化过程。原论文的作者们很机智地换了一个度量标准，使得有解析解。具体来说，它使用了“广义KL散度”，又称“I散度”，其形式为：

这个度量源自不等式，当且仅当时等号成立。所以代入，然后两端乘以，我们有

当且仅当成立，如果有多个分量，那么对多个分量的结果求和即可，这就得到了度量。显然，广义KL散度是概率的KL散度的自然推广，但它不要求和满足归一化，只要求它们非负，这正好对应了AdaFactor的场景。而且巧妙的是，这种情形配上这个目标，刚好有解析解：

其实这个解析解也很形象，就是行、列分别求和，然后相乘，再除以全体的和。

推导过程

直接对求偏导数并让偏导数等于0，得

整理得

注意到如果是一组最优解，那么也是，说白了，所有的乘以一个常数，所有的也除以这个常数，是不变的。那么我们就可以随意指定或，因为它们就只是一个缩放标量而已。不失一般性，我们指定，那么就解得。

直观理解

我们也可以从另一个角度理解结果。由于是非负的，我们可以将它归一化，变成具有概率分布的特性，即，然后我们试图完成分解，由于现在相当于一个二元联合概率分布，那么就相当于它们的边缘分布，即

现在到还需要乘上一个，我们可以把它乘到或中，不失一般性，我们假设乘到上，那么就得到。

AdaFactor雏形

有了结果后，我们就可以用它来构建更省内存的优化器了，这就是AdaFactor的雏形。简单来说，当参数是普通一维向量时，优化过程保持不变；但是的矩阵时，算出来的梯度也是矩阵，从而也是矩阵，这时候我们对做低秩分解，然后维护两组缓存变量，分别滑动平均低秩分解后的结果，最后用共同调整学习率：

（把加到上去而不是上去，这是AdaFactor整出来的形式，不是笔者的锅～）.

滑动权重

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

等价形式

为了认识到这一点，我们重写一下Adam的的更新过程：

所以如果设，那么更新公式就是

问题是这个够不够合理呢？答案是可能不大够。当时，这时候就是，也就是用实时梯度来校正学习率，这时候校正力度最大；当时，，这时候是累积梯度平方与当前梯度平方的加权平均，由于，所以意味着当前梯度的权重不为0，这可能导致训练不稳定，因为训练后期梯度变小，训练本身趋于稳定，校正学习率的意义就不大了，因此学习率的校正力度应该变小，并且，学习率最好恒定为常数（这时候相当于退化为SGD），这就要求时，。

新的衰减策略

为了达到这个目的，AdaFactor采用如下的衰减策略

它满足。但即便如此，也不是任何都适合，必须有。好理解，那为什么要呢？原论文包含了对它的分析，大家可以去读读，但笔者觉得原论文的推导过于晦涩，所以这里给出自己的理解。

首先，对于来说，一个很容易想到的方案是所有梯度平方的平均，即：

所以这等价于让。这个方案美中不足的一点是，每一步梯度都是平权的，这不符合直觉，因为正常来说越久远的梯度应该越不重要才对，所以应该适当降低历史部分权重，而当时，，因此一个简洁的方案是在式中取，AdaFactor默认的是。

层自适应

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

AdaFactor完整版

至此，我们终于可以写出完整版AdaFactor的更新过程了：

其中是模长的变种，这一步相当于做了个截断，即时才执行归一化。原论文中的默认参数为

如果参数是一维向量而不是矩阵，那么使用普通的更新公式就行了。此外，论文还提出如果没有传入学习率，那么可以使用为默认学习率，但笔者看源码的时候发现这个默认学习率很少使用，基本上还是需要自己传入学习率的。

开源实现

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

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的效果。