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

文章小结

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

当然，没有什么绝对能有效的方法，有的只是方法虽好，要想实际有效，依然要用心炼丹。 

推荐阅读

AINLP年度阅读收藏清单

百度PaddleHub NLP模型全面升级，推理性能提升50%以上

斯坦福大学NLP组Python深度学习自然语言处理工具Stanza试用

数学之美中盛赞的 Michael Collins 教授，他的NLP课程要不要收藏？

自动作诗机&藏头诗生成器：五言、七言、绝句、律诗全了

From Word Embeddings To Document Distances 阅读笔记

模型压缩实践系列之——bert-of-theseus，一个非常亲民的bert压缩方法

这门斯坦福大学自然语言处理经典入门课，我放到B站了

可解释性论文阅读笔记1-Tree Regularization

征稿启示 | 稿费+GPU算力+星球嘉宾一个都不少

关于AINLP

AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括文本摘要、智能问答、聊天机器人、机器翻译、自动生成、知识图谱、预训练模型、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLPer(id：ainlper)，备注工作/研究方向+加群目的。