一只小狐狸带你解锁炼丹术&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的更新过程为
要省显存,就首先得知道显存花在哪里的。首先,计算量和显存的大头肯定都是 ,也就是说,计算梯度是很费资源的,这也是为啥“ALBERT相比BERT参数量虽然少了那么多,但训练速度也没见快多少”的原因了;除此之外,显存的消耗主要是 了,我们要维护两组缓存变量,来滑动计算梯度的前两阶矩(也就是 和 ),用以计算参数的更新量。这两组变量每一组都跟训练参数本身一样大,因此对于参数比较多的模型,两组缓存变量所消耗的显存也不少。
在这一节中,我们会相对详细地介绍一些AdaFactor优化器,介绍中会设计比较多的公式和推导。如果只求一个大致了解的读者,可以自行跳过部分数学内容~
我们知道,CV模型很多时候要靠“SGD+动量”来炼出最优效果来,自适应学习率优化器通常训练不出最好的效果。但对于NLP模型来说,情况有点相反,自适应学习率显得更重要一些,很少听到由纯靠SGD调NLP模型的案例。因此,作为省显存的第一步,我们可以抛弃Adam里边的动量,这样就少一组缓存参数了,自然也就省了显存:
这其实就是RMSProp的变种,比RMSProp多了 这一步。
去掉 之后,缓存变量直接减少了一半,但AdaFactor还不满意,它希望保留自适应学习率功能,但把缓存变量 的参数量再压一压。这一次,它用到了矩阵的低秩分解。
在SGD中,所有参数都是共用一个标量学习率;在Adam中,则是每一个参数都有自己的学习率 。我们知道通过精调学习率,SGD其实也能有不错的效果,这表明“每一个参数都有自己的学习率”这件事情都不是特别重要,或者换一种说法,就是“精调每一个参数自己的学习率”并不是特别重要。
这启发我们,将 换一种参数更少的近似可能也就足够了。而“参数更少的近似”,我们就不难想到低秩分解了。对于 的矩阵 ,我们希望找到 的矩阵 和 的矩阵 ,使得
当 足够小时, 、 的参数总量就小于 的参数量。为了“省”到极致,AdaFactor直接让 ,即寻找 和 ,使得
既然要近似,就要有一个度量的标准。很容易想到的标准是欧氏距离,即
但在这个距离之下, 并没有解析解;此外,在优化过程中 (即 )是非负的,而通过上述目标优化出来的 无法保证非负,因此很可能扰乱优化过程。原论文的作者们很机智地换了一个度量标准,使得 有解析解。具体来说,它使用了“广义KL散度”,又称“I散度”,其形式为:
这个度量源自不等式 ,当且仅当 时等号成立。所以代入 ,然后两端乘以 ,我们有
当且仅当 成立,如果 有多个分量,那么对多个分量的结果求和即可,这就得到了度量。显然,广义KL散度是概率的KL散度的自然推广,但它不要求 和 满足归一化,只要求它们非负,这正好对应了AdaFactor的场景。而且巧妙的是,这种情形配上这个目标,刚好有解析解:
其实这个解析解也很形象,就是行、列分别求和,然后相乘,再除以全体的和。
直接对求偏导数并让偏导数等于0,得
整理得
注意到如果 是一组最优解,那么 也是,说白了,所有的 乘以一个常数,所有的 也除以这个常数, 是不变的。那么我们就可以随意指定 或 ,因为它们就只是一个缩放标量而已。不失一般性,我们指定 ,那么就解得 。
我们也可以从另一个角度理解结果。由于 是非负的,我们可以将它归一化,变成具有概率分布的特性,即 ,然后我们试图完成分解 ,由于 现在相当于一个二元联合概率分布,那么 就相当于它们的边缘分布,即
现在 到 还需要乘上一个 ,我们可以把它乘到 或 中,不失一般性,我们假设乘到 上,那么就得到 。
有了结果 后,我们就可以用它来构建更省内存的优化器了,这就是AdaFactor的雏形。简单来说,当参数 是普通一维向量时,优化过程保持不变;但 是 的矩阵时,算出来的梯度 也是矩阵,从而 也是矩阵,这时候我们对 做低秩分解,然后维护两组缓存变量 ,分别滑动平均低秩分解后的结果,最后用 共同调整学习率:
(把 加到 上去而不是 上去,这是AdaFactor整出来的形式,不是笔者的锅~).
在Adam以及上述AdaFactor雏形中,滑动权重 都是恒为常数,AdaFactor指出这是不科学的,并提出新的策略。
为了认识到这一点,我们重写一下Adam的 的更新过程:
所以如果设 ,那么更新公式就是
问题是这个 够不够合理呢?答案是可能不大够。当 时 ,这时候 就是 ,也就是用实时梯度来校正学习率,这时候校正力度最大;当 时, ,这时候 是累积梯度平方与当前梯度平方的加权平均,由于 ,所以意味着当前梯度的权重 不为0,这可能导致训练不稳定,因为训练后期梯度变小,训练本身趋于稳定,校正学习率的意义就不大了,因此学习率的校正力度应该变小,并且 ,学习率最好恒定为常数(这时候相当于退化为SGD),这就要求 时, 。
为了达到这个目的,AdaFactor采用如下的衰减策略
它满 足 。但即便如此,也不是任何 都适合,必须有 。 好理解,那为什么要 呢?原论文包含了对它的分析,大家可以去读读,但笔者觉得原论文的推导过于晦涩,所以这里给出自己的理解。
首先,对于 来说,一个很容易想到的方案是所有梯度平方的平均,即:
所以这等价于让 。这个方案美中不足的一点是,每一步梯度都是平权的,这不符合直觉,因为正常来说越久远的梯度应该越不重要才对,所以应该适当降低历史部分权重,而当 时, ,因此一个简洁的方案是在式中取 ,AdaFactor默认的 是 。
最后,我们还可以进一步根据参数的模长来校正更新量,这个思路来自LAMB优化器,在之前的文章《6个派生优化器的简单介绍及其实现》中也介绍过。简单来说,它就是将最后的更新量标准化,然后乘以参数的模长,说白了,就是不管你怎么折腾,最后的更新量我只要你的方向,而大小由参数本身的模长和预先设置学习率共同决定,使得所有层所有参数的相对变化程度保持一致。
至此,我们终于可以写出完整版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所做的分析相当经典,值得我们认真琢磨体味,对有兴趣研究优化问题的读者来说,更是一个不可多得的分析案例。
当然,没有什么绝对能有效的方法,有的只是方法虽好,要想实际有效,依然要用心炼丹。
可
能
喜
欢
卖萌屋上线Arxiv论文速刷神器,直达学术最前沿!
13个offer,8家SSP,谈谈我的秋招经验
Google|突破瓶颈,打造更强大的Transformer
推荐系统的发展与简单回顾
ACL2020|FastBERT:放飞BERT的推理速度
LayerNorm是Transformer的最优解吗?
夕小瑶的卖萌屋
_
关注&星标小夕,带你解锁AI秘籍
订阅号主页下方「撩一下」有惊喜哦