Google新搜出的优化器Lion：效率与效果兼得的“训练狮”

2023-02-17 03:02

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

昨天在 Arixv 上发现了 Google 新发的一篇论文《Symbolic Discovery of Optimization Algorithms》，主要是讲自动搜索优化器的，乍看上去没啥意思，因为类似的工作也有不少，大多数结果都索然无味。

然而，细读之下才发现别有洞天，原来作者们通过数千 TPU 小时的算力搜索并结合人工干预，得到了一个更省显存的优化器 Lion（EvoLved Sign Momentum，不得不吐槽这名字起得真勉强），并在图像分类、图文匹配、扩散模型、语言模型预训练和微调等诸多任务上做了充分的实验，多数任务都显示 Lion 比目前主流的 AdamW 等优化器有着更好的效果。

更省显存还更好效果，真可谓是鱼与熊掌都兼得了，什么样的优化器能有这么强悍的性能？本文一起来欣赏一下论文的成果。

论文标题：

Symbolic Discovery of Optimization Algorithms

论文链接：

https://arxiv.org/pdf/2302.06675.pdf

先说结果

本文主要关心搜索出来的优化器本身，所以关于搜索过程的细节就不讨论了，对此有兴趣读者自行看原论文就好。Lion 优化器的更新过程为：

其中是损失函数的梯度，是符号函数 [1]，即正数变为 1、负数变为 -1。我们可以对比一下目前的主流优化器 AdamW [2] 的更新过程：

对比很明显，Lion 相比 AdamW 参数更少（少了个），少缓存了一组参数（所以更省显存），并且去掉了 AdamW 更新过程中计算量最大的除法和开根号运算（所以更快）。

在此之前，跟 Lion 最相似的优化器应该是 SIGNUM [3]，其更新过程为：

跟 Lion 一样，SIGNUM 也用到了符号函数处理更新量，而且比 Lion 更加简化（等价于 Lion 在和的特例），但是很遗憾，SIGNUM 并没有取得更好的效果，它的设计初衷只是降低分布式计算中的传输成本。Lion 的更新规则有所不同，尤其是动量的更新放在了变量的更新之后，并且在充分的实验中显示出了它在效果上的优势。

论文实验

本文开头就说了，Lion 在相当多的任务上都做了实验，实验结果很多，下面罗列一些笔者认为比较关键的结果。

▲ Lion在NLU和NLG任务上的结果，大部分都比AdamW、Adafactor优秀

▲ 在视觉Transformer上Lion与众多优化器的对比

▲ 在视觉Transformer上Li

▲ 在NLP的自回归生成上，Lion的收敛速度更快

▲ 上右图是ImageNet上的训练曲线，显示Lion尽管验证集效果更好，但训练集上的效果未必会优于AdamW

超参设置

看到论文效果如此惊人，笔者也跃跃欲试。在跑实验之前，自然需要了解一下各个超参的设置。首先是，原论文自动搜索出来的结果是，并在大部分实验中复用了这个组合，但是在 NLP 的任务上则使用了这个组合（论文的详细实验配置在最后一页的 Table 12）。

比较关键的学习率和权重衰减率，由于 Lion 的更新量每个分量的绝对值都是 1，这通常比 AdamW 要大，所以学习率要缩小 10 倍以上，才能获得大致相同的更新幅度；而由于学习率降低了，那么为了使权重衰减的幅度保持不变，权重衰减率应该要放大相应的倍数。

原论文的最后一页给出了各个实验的超参数参考值，其中小模型（Base 级别）上使用的是和，大模型（参数 10 亿以上）则适当降低了学习率到甚至。

事实上，之前我们在《基于Amos优化器思想推导出来的一些“炼丹策略”》就推导过学习率和权重衰减率的一个组合方案，参考这个方案来设置是最方便的。在该方案中，更新量写为（记号跟前面的描述略有不同，但不至于混淆，应该就不强行统一了）：

其中：

最后的近似只保留了不超过的项。

其中是原本的更新量；是（初始阶段）参数变化的相对大小，一般是级别，表示每步更新后参数模长的变化幅度大致是千分之一；q 是一个超参数，没什么特殊情况可以设为 1；是控制学习率衰减速度的超参数，可以根据训练数据大小等设置。

由于经过了运算，因此，k 是参数的维度；，这我们在《基于Amos优化器思想推导出来的一些“炼丹策略”》已经推导过了，其中是参数的变化尺度，对于乘性矩阵，就是它的初始化方差。所以，经过一系列简化之后，有：

这里的就是前面的，而。按照 BERT base 的 d=768 来算，初始化方差的量级大致在 1/d 左右，于是，假设取（为了给结果凑个整），那么按照上式学习率大约是、衰减率大约是 0.015。在笔者自己的 MLM 预训练实验中，选取这两个组合效果比较好。

个人实现：

https://github.com/bojone/bert4keras

延伸思考

总体来看，Lion 表现可圈可点，不管是原论文还是笔者自己的实验中，跟 AdamW 相比都有一战之力，再加上 Lion 更快以及更省显存的特点，或者可以预见未来的主流优化器将有它的一席之地。

自 Adam 提出以来，由于其快速收敛的特性成为了很多模型的默认优化器。甚至有学者提出，这个现象将反过来导致一个进化效应：所有的模型改进都在往 Adam 有利的方向发展，换句话说，由于我们选择了 Adam 作为优化器，那么就有可能将很多实际有效、但是在 Adam 优化器上无效的改动都抛弃了，剩下的都是对 Adam 有利的改进，详细的评价可以参考《NEURAL NETWORKS (MAYBE) EVOLVED TO MAKE ADAM THE BEST OPTIMIZER》[4]。

所以，在此大背景之下，能够发现比 Adam 更简单且更有效的优化器，是一件很了不起的事情，哪怕它是借助大量算力搜索出来的。

可能读者会有疑问：Lion 凭啥可以取得更好的泛化性能呢？原论文的解释是这个操作引入了额外的噪声（相比于准确的浮点值），它使得模型进入了 Loss 更平坦（但未必更小）的区域，从而泛化性能更好。

为了验证这一点，作者比较了 AdamW 和 Lion 训练出来的模型权重的抗干扰能力，结果显示 Lion 的抗干扰能力更好。然而，理论上来说，这只能证明 Lion 确实进入到了更平坦的区域，但无法证明该结果是操作造成的。不过，Adam 发表这么多年了，关于它的机理也还没有彻底研究清楚，而 Lion 只是刚刚提出，就不必过于吹毛求疵了。

Lion 是不是足够完美呢？显然不是，比如原论文就指出它在小 batch_size（小于 64）的时候效果不如 AdamW，这也不难理解，本来已经带来了噪声，而小 batch_size 则进一步增加了噪声，噪声这个东西，必须适量才好，所以两者叠加之下，很可能有噪声过量导致效果恶化。

另外，也正因为加剧了优化过程的噪声，所以参数设置不当时容易出现损失变大等发散情况，这时候可以尝试引入 Warmup，或者增加 Warmup 步数。还有，Lion 依旧需要缓存动量参数，所以它的显存占用多于 AdaFactor [5]，能不能进一步优化这部分参数量呢？暂时还不得而知。