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

©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]，能不能进一步优化这部分参数量呢？暂时还不得而知。

文章小结

本文介绍了 Google 新提出的优化器 Lion，它通过大量算力搜索并结合人工干预得出，相比主流的 AdamW，有着速度更快且更省内存的特点，并且大量实验结果显示，它在多数任务上都有着不逊色于甚至优于 AdamW 的表现。

参考文献

[1] https://en.wikipedia.org/wiki/Sign_function

[2] https://arxiv.org/abs/1711.05101

[3] https://arxiv.org/abs/1802.04434

[4] https://parameterfree.com/2020/12/06/neural-network-maybe-evolved-to-make-adam-the-best-optimizer/

[5] https://kexue.fm/archives/7302

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·