谷歌推出新优化器Lion：优化算法的符号发现

谷歌推出新优化器Lion：优化算法的符号发现

优化器即优化算法，优化器怎么选对模型的好坏起着重要的作用。优化器对梯度进行优化更新，不仅使模型找到最优点，而且能大大提升模型收敛速度。当然不同的优化器找到的最优点可能不同，有可能是局部最优点，有可能是全局最优点，这个是决定模型好坏的重要因素之一。

近年来提出了很多不同的优化算法，分别利用不同的公式来更新模型的参数。Adam（Kingma and Ba，2015）自从2015年被推出后，一直到今天仍然是最常用的优化算法之一。这表明从机器学习从业者的角度来看，深度学习的优化的最好的方法在很大程度上并没有多大改变。

然而，今年已经产生了一些新的想法来改进深度学习的优化方法，这可能会成为未来我们模型的优化方式。来自谷歌、 UCLA 的研究者提出了一种通过程序搜索发现深度神经网络训练的优化算法，来自论文《Symbolic Discovery of Optimization Algorithms》。通过数千TPU小时的算力搜索并结合人工干预，得到一个更省显存且速度更快的优化器 Lion（EvoLved Sign Momentum），其在图像分类、图文匹配、扩散模型、语言模型预训练和微调等诸多任务上做了充分的实验，多数任务都显示 Lion 比目前主流的 AdamW 等优化器有着更好的效果。

论文：Symbolic Discovery of Optimization Algorithms
组织：Google, UCLA
作者：Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, Quoc V. Le
论文地址：https://arxiv.org/abs/2302.06675
代码地址：https://github.com/google/automl/tree/master/lion

关于常见的优化器Optimizer/优化算法的介绍，可以查看之前的文章：
传送门：各种优化器SGD，AdaGrad，Adam，AdamW，LBFGS都做了什么？

Lion VS AdamW

先来看看优化器本身，Lion 的更新过程：
$$Lion：=\{\begin{array}{l}{\mu }_t=\text{sign}({\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t)+{\lambda }_t{\theta }_{t-1}\\ {\theta }_t={\theta }_{t-1}-{\eta }_t{\mu }_t\\ m_t={\beta }_2{m}_{t-1}+(1-{\beta }_2)g_t\end{array}$$

对比主流优化器 AdamW 的更新过程：
$$AdamW：=\{\begin{array}{l}{m}_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t\\ \nu ={\beta }_2{\nu }_{t-1}+(1-{\beta }_2)g_t^2\\ {\hat{m}}_t=m_t/(1-{\beta }_1^t)\\ {\hat{\nu }}_t={\nu }_t/(1-{\beta }_2^t)\\ {\mu }_t={\hat{m}}_t/(\sqrt{{\hat{\nu }}_t}+ϵ)+{\lambda }_t{\theta }_{t-1}\\ {\theta }_t={\theta }_{t-1}-{\eta }_t{\mu }_t\end{array}$$

两者对比，可以发现 Lion 比 AdamW 参数更少：

• Lion 少了一个$ϵ$，少缓存一组参数$\nu$，所以更省显存；
• Lion 去掉了 AdamW 更新过程中计算量最大的除法和开根号运算，所以速度更快；

另外，值得注意的是 Lion 更新规则有所不同，Lion 使用符号函数sign处理更新量，尤其是动量的更新放在了变量的更新之后，并且在充分的实验中显示出了它在效果上的优势。

与 AdamW 和各种自适应优化器（这些优化器需要同时保存一阶和二阶矩）相比，Lion 只需要动量并利用符号操作来计算更新，并且将额外的内存占用减半。这在训练大模型或大批量时很有用。例如，AdamW 需要至少 16 个 TPU V4 来训练图像大小为 224、批大小为 4096 的 ViT-B/16，而 Lion 只需要 8 个。另一个实际好处是，由于 Lion 的简单性，Lion 在实验中具有更快的运行时间（steps /sec），通常比 AdamW 和 Adafactor 提速 2-15%，具体取决于任务、代码库和硬件。

论文实验

论文将 Lion 与广泛使用的优化器（如Adam和Adafactor）进行比较，以训练不同任务的各种模型。

1.图像分类

在图像分类方面，Lion 将 ImageNet上 的ViT精度提高了2%，并将JFT上的预训练计算节省了5倍。

研究者在 ImageNet 上从头开始训练，在更大的数据集 ImageNet-21K 和 JFT 上进行预训练。图片大小默认为 224。如下表 2 所示，Lion 在各种架构上显著优于 AdamW，即使训练集扩大 10 倍，Lion 仍然优于 AdamW。下图4显示了 JFT-300M 上不同预训练预算下三种 ViT 模型（ViT-B/16、ViT-L/16 和 ViT-H/14）的准确率。Lion 使 ViT-L/16 能够与 AdamW 在 ImageNet 和 ImageNet V2 上训练的 ViT-H/14 性能相匹配，但预训练成本降低了 3 倍。

2.视觉语言对比学习

在视觉语言对比学习方面，Lion 在ImageNet上实现了88.3%的零样本和91.1%的微调精度，分别超过了SOTA 2%和0.1%。

在 CLIP 风格的视觉语言对比训练上，研究者没有从头开始学习所有参数，而是使用强大的预训练模型初始化图像编码器。

对于锁定图像文本调整 (LiT) ，研究者通过使用相同的 frozen 预训练 ViT 以对比方式训练文本编码器，进而在 LiT 上对 Lion 和 AdamW 进行比较。下表 4 显示了 3 种模型尺度上的零样本图像分类结果，Lion 展示出了对 AdamW 的持续改进。下图 5（左）展示了 LiT-B/16-B 的示例零样本学习曲线，并在其他两个数据集上得到了类似的结果。

3.扩散模型

在扩散模型上，Lion 通过获得更好的FID分数和将训练计算减少2.3倍而优于Adam。

因为扩散模型在图像生成上的成功出圈（如DALL-E 2），研究者测试了 Lion 在无条件图像合成（Image synthesis on ImageNet）和多模式文本到图像生成（Text-to-image generation）方面的表现。

研究者在 ImageNet 上执行 64x64、128x128、256x256图像生成。如上图 5（中和右）所示，Lion 可以在 FID 分数上实现更好的质量和更快的收敛速度。

在文本到图像生成上，下图 6 展示了学习曲线。虽然 64 × 64 base 模型没有明显改进，但 Lion 在文本条件超分辨率模型上优于 AdamW。与 AdamW 相比，Lion 获得了更高的 CLIP 分数并且具有更小的噪声 FID 指标。

4.语言建模和微调

对于自回归、掩码语言建模和微调，Lion 表现出与 Adam 相似或更好的性能。

研究者重点关注 Lion 在语言建模和微调的性能表现，特别地，在纯语言任务上，研究者发现调整 β_1 和 β_2 可以提高 AdamW 和 Lion 的质量。

对于自回归语言建模（Autoregressive language modeling），下图7展示了 Wiki-40B 的 token 级困惑度（Perplexity）和 PG-19 的 word 级困惑度。Lion 始终实现比 AdamW 更低的验证困惑度。在 Wiki-40B 和 PG-19 上训练中型模型时，它分别实现了 1.6 倍和 1.5 倍的加速。当模型增加到大尺寸时，PG-19 进一步实现了 2 倍加速。

对于掩蔽语言建模（Masked language modeling），研究者还在 C4 数据集上执行 BERT 训练，其中 Lion 在验证困惑度方面的表现略好于 AdamW，如下图 11（左）所示。右图是 ImageNet 上的训练曲线， Lion 尽管验证集效果更好，但训练集上的效果未必会优于AdamW。

对于微调（Fine-tuning），研究者在 GLUE 基准上微调 Base（220M）、Large（770M）和最大的 11B T5 模型。下表 6 显示了 GLUE 开发集上的结果，平均而言，Lion 在所有 3 个模型规模上都击败了 AdamW。

5.与其他流行优化器的比较

论文中还将 Lion 与四种流行的优化器 RAdam、NAdam、AdaBelief 和 AMSGrad 进行比较，分别在 ImageNet 上训练 ViT-S/16 和 ViT-B/16（使用 RandAug 和 Mixup）。如下表 7 所示，Lion 仍然是表现最好的。

超参数设置

通过对 Lion 的分析表明，它的性能增益随着训练批量的增加而增加。它还需要比Adam更小的学习率，因为符号函数产生的更新范数更大。

论文就 Lion 优化器推导出来一些“炼丹策略”：一些超参数的组合设置，比如关键的学习率$\eta$和权重衰减率$\lambda$，参数${\beta }_1$和${\beta }_2$。具体的细节可以看原论文。

优化器的另一个潜在限制——批量大小（batch size）。通过实验，论文指出 Lion 在小batch_size（小于64）的时候效果不如AdamW。如图8（左）所示，AdamW的最佳批量大小为256，而Lion的最佳批量为4096。这表明 Lion 适合更大的batch_size，但即使在batch_size=64的情况下，其性能仍保持稳定。此外，当批量大小扩大到32K（导致仅包含11K个训练步骤）时，Lion 比 AdamW 获得了显著的2.5%的准确率增益（77.9%对75.4%），这证明 Lion 在大批量训练环境中的有效性。

Lion 只保存动量，比 AdamW 等流行的自适应优化器占用的内存更小，这在训练大模型和 / 或使用大批量时非常有用。

小结

Google 新提出的一个简单有效的优化算法 Lion，利用高效的搜索技术来探索无限且稀疏的程序空间。相比主流的 AdamW，Lion 有着速度更快且更省内存的特点；Lion 还能够取得更好的泛化性能。该研究采用了一系列技术，包括具有热启动和重启的进化搜索、抽象执行、funnel 选择和程序简化策略。

参考：
https://arxiv.org/abs/2302.06675
https://github.com/google/automl/tree/master/lion
https://www.163.com/dy/article/HTUKVDAG0542QSJ0.html
https://spaces.ac.cn/archives/9473
https://www.jianshu.com/p/a2ce9a117f0a

欢迎各位关注我的个人公众号：HsuDan，我将分享更多自己的学习心得、避坑总结、面试经验、AI最新技术资讯。