谷歌的 Lion 优化器将成为训练大模型或大批量的「福音」。
优化器即优化算法，在神经网络训练中起着关键作用。近年来，研究者引入了大量的手工优化器，其中大部分是自适应优化器。Adam 以及 Adafactor 优化器仍然占据训练神经网络的主流，尤其在语言、视觉和多模态领域更是如此。
除了人工引入优化器外，还有一个方向是程序自动发现优化算法。此前有人提出过 L2O（learning to optimize），该方法通过训练神经网络来发现优化器。然而，这些黑盒优化器通常是在有限数量的小任务上训练而成，很难泛化到大模型。
还有人尝试了其他方法，他们应用强化学习或蒙特卡罗采样来发现新的优化器。然而，为了简化搜索，这些方法通常会限制搜索空间，从而限制发现其他优化器的可能性。因此，目前方法尚未达到 SOTA 水平。
最近几年，值得一提的是 AutoML-Zero，它试图在评估任务时搜索机器学习 pipeline 的每个组件，对优化器的发现很有借鉴意义。
在本文中，来自谷歌、 UCLA 的研究者提出了一种通过程序搜索发现深度神经网络训练的优化算法的方法，进而发现 Lion（EvoLved Sign Momentum）优化器。实现这一目标面临两个挑战：首先是在无限稀疏的程序空间中寻找高质量的算法，其次是选择可以从小型任务泛化到更大、SOTA 任务的算法。为了应对这些挑战，该研究采用了一系列技术，包括具有热启动和重启的进化搜索、抽象执行、funnel 选择和程序简化。
与 AdamW 和各种自适应优化器（这些优化器需要同时保存一阶和二阶矩）相比，Lion 只需要动量并利用符号操作来计算更新，并且将额外的内存占用减半。这在训练大模型或大批量时很有用。例如，AdamW 需要至少 16 个 TPU V4 来训练图像大小为 224、批大小为 4096 的 ViT-B/16，而 Lion 只需要 8 个。另一个实际好处是，由于 Lion 的简单性，Lion 在实验中具有更快的运行时间（steps /sec），通常比 AdamW 和 Adafactor 提速 2-15%，具体取决于任务、代码库和硬件。
论文一作 Xiangning Chen 表示：我们的符号程序搜索发现了一个只跟踪动量的有效优化器 ——Lion。与 Adam 相比，它获得了 88.3% 的零样本和 91.1% 的微调 ImageNet 准确率，以及高达 5 倍（与 ViT 相比）、2.3 倍（与扩散模型相比）和 2 倍（与 LM 想比）训练效率。
Lion 在一系列模型（Transformer、MLP、ResNet、U-Net 和 Hybrid）和任务（图像分类、视觉 - 语言对比学习、扩散、语言建模和微调）中展示了出色的性能。值得注意的是，通过在 BASIC 中用 Lion 替换 Adafactor，该研究在 ImageNet 上实现了 88.3% 的零样本和 91.1% 的微调精度（Pham 等人，2021），分别超过之前的 SOTA 结果 2% 和 0.1%。
此外，Lion 将 JFT 上的预训练计算量减少了多达 5 倍，将扩散模型的训练效率提高了 2.3 倍，并获得了更好的 FID 分数，并在语言建模方面提供了类似或更好的性能，节省了多达 2 倍的计算量。
推特用户 crumb 表示：自己使用谷歌的 Lion 优化器训练了一个 124M 参数 GPT2 模型，发现达到与 Adam 相同的损失所需的 step 数减少了 37.5%。
图源：https://twitter.com/aicrumb/status/1626053855329898496
算法的符号发现
本文使用程序形式的符号表示具有以下优点：（1）它符合算法必须作为程序来执行的事实；(2) 与神经网络等参数化模型相比，程序等符号表示更易于分析、理解和迁移到新任务；(3) 程序长度可用于估计不同程序的复杂性，从而更容易选择更简单、通常更具通用性的程序。这项工作侧重于深度神经网络训练的优化器，但该方法通常适用于其他任务。
下图中，程序 2 这个简化代码片段使用与 AdamW 相同的 signature ，以确保发现的算法具有更小或相等的内存占用空间；程序 3 中给出了 AdamW 的示例表示。
该研究采用以下技术来解决无限和稀疏搜索空间带来的挑战。首先应用正则化，因为它简单，可扩展，并在许多 AutoML 搜索任务中很成功；其次是简化程序空间中的冗余；最后，为了降低搜索成本，该研究通过减少模型大小、训练示例的数量和距离目标任务的 step 使成本降低。
左：显示了五次进化搜索实验的均值和标准误差。右：随着搜索的进行，冗余语句的百分比和缓存命中率都在增加。
Lion 的推导与分析
研究者表示，优化器 Lion 具有简洁性、高内存效率以及在搜索和元验证方面的强大性能。
推导
搜索和 funnel 选择过程促成了程序 4，它通过从原始程序 8（附录）自动删除冗余语句获得。研究者进一步简化并得到程序 1 中的最终算法（Lion）。在简化过程中从程序 4 中删除了几个不必要的元素。其中 cosh 函数被删除，因为 m 将在下一次迭代中重新分配（第 3 行）。使用 arcsin 和 clip 的语句也被删除，因为研究者观察到没有它们质量不会下降。三个红色语句转换为一个符号函数。
尽管程序 4 中同时使用了 m 和 v，但 v 仅改变了动量的更新方式（两个具有常数∼0.9 和 ∼1.1 的插值函数等同于一个具有 ∼0.99 的函数）并且不需要单独跟踪。注意不再需要偏差校正，因为它不会改变方向。
如下算法 2 显示了伪代码。
分析
符号更新和正则化。Lion 算法通过符号操作在所有维度上产生了具有统一幅度的更新，这在原理上不同于各种自适应优化器。直观来看，符号操作为更新添加了噪声，作为了一种正则化形式并有助于泛化。下图 11（右）展示了一个证据。
动量跟踪。与 AdamW 和动量 SGD 中常用的 0.9 相比，Lion 中用于跟踪动量的默认 EMA 因子为 0.99 (β_2)。EMA 因子和插值的这种选择，使得 Lion 在记住动量梯度的 10 倍 history 和在更新中对当前梯度施加更多权重之间取得平衡。
超参数和批大小选择。与 AdamW 和 Adafactor 相比，Lion 更简单并且超参数更少，因为它不需要 ϵ 和因式分解相关的参数。Lion 需要更小的学习率，进而需要更大的解耦权重衰减才能达到类似的有效权重衰减强度 (lr * λ)。
内存和运行时优势。Lion 只保存动量，比 AdamW 等流行的自适应优化器占用的内存更小，这在训练大模型和 / 或使用大批量时非常有用。例如，AdamW 需要至少 16 个 TPU V4 芯片来训练图像分辨率为 224 且批大小为 4,096 的 ViT-B/16，而 Lion 只需要 8 个（均具有 bfloat16 动量）。
Lion 评估结果
在实验部分，研究者在各种基准上对 Lion 进行了评估，主要将它与流行的 AdamW（或者当内存成为瓶颈时的 Adafactor）进行比较。
图像分类
研究者在图像分类任务上执行涵盖各种数据集和架构的实验。除了在 ImageNet 上从头开始训练之外，他们还在两个更大的成熟数据集 ImageNet-21K 和 JFT 上进行预训练。图片大小默认为 224。
首先在 ImageNet 上从头开始训练。研究者训练 ResNet-50 达到 90 个 epoch，批大小为 1,024，其他模型训练 300 个 epoch，批大小为 4,096。如下表 2 所示，Lion 在各种架构上显著优于 AdamW。
其次在 ImageNet-21K 上进行预训练。研究者在 ImageNet-21K 上对 ViT-B/16 和 ViT-L/16 进行了 90 个 epoch 的预训练，批大小为 4,096。下表 2 显示，即使训练集扩大 10 倍，Lion 仍然优于 AdamW。
最后在 JFT 上进行预训练。为了突破极限，研究者在 JFT 上进行大量实验。下图 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 倍。
下表 3 显示了微调结果，具有更高的分辨率和 Polyak 平均。研究者使用的 ViT-L/16 与之前由 AdamW 训练的 ViT-H/14 结果相匹配，同时参数减少了 2 倍。在将预训练数据集扩展到 JFT-3B 后，Lion 训练的 ViT-g/14 优于之前的 ViT-G/14 结果，参数减少了 1.8 倍。
视觉语言对比学习
本节重点介绍 CLIP 风格的视觉语言对比训练。研究者没有从头开始学习所有参数，而是使用强大的预训练模型初始化图像编码器。
对于锁定图像文本调整 (LiT) ，研究者通过使用相同的 frozen 预训练 ViT 以对比方式训练文本编码器，进而在 LiT 上对 Lion 和 AdamW 进行比较。下表 4 显示了 3 种模型尺度上的零样本图像分类结果，Lion 展示出了对 AdamW 的持续改进。
下图 5（左）展示了 LiT-B/16-B 的示例零样本学习曲线，并在其他两个数据集上得到了类似的结果。
扩散模型
最近，扩散模型在图像生成方面取得了巨大的成功。鉴于其巨大的潜力，研究者测试了 Lion 在无条件图像合成和多模式文本到图像生成方面的表现。
对于 ImageNet 上的图像合成，研究者利用 2021 年论文《 Diffusion models beat gans on image synthesis 》中引入的改进 U-Net 架构，在 ImageNet 上执行 64×64、128×128 和 256×256 图像生成。如上图 5（中和右）所示，Lion 可以在 FID 分数上实现更好的质量和更快的收敛速度。
对于文本到图像的生成，下图 6 展示了学习曲线。虽然 64 × 64 base 模型没有明显改进，但 Lion 在文本条件超分辨率模型上优于 AdamW。与 AdamW 相比，Lion 获得了更高的 CLIP 分数并且具有更小的噪声 FID 指标。
语言建模和微调
本节重点介绍语言建模和微调。在纯语言任务上，研究者发现调整 β_1 和 β_2 可以提高 AdamW 和 Lion 的质量。
对于自回归语言建模，下图 7 展示了 Wiki-40B 的 token 级困惑度和 PG-19 的 word 级困惑度。Lion 始终实现比 AdamW 更低的验证困惑度。在 Wiki-40B 和 PG-19 上训练中型模型时，它分别实现了 1.6 倍和 1.5 倍的加速。当模型增加到大尺寸时，PG-19 进一步实现了 2 倍加速。
对于掩蔽语言建模，研究者还在 C4 数据集上执行 BERT 训练，其中 Lion 在验证困惑度方面的表现略好于 AdamW。相关学习曲线如下图 11（左）所示。
对于微调，研究者在 GLUE 基准上微调 Base（220M）、Large（770M）和最大的 11B T5 模型。下表 6 显示了 GLUE 开发集上的结果，平均而言，Lion 在所有 3 个模型规模上都击败了 AdamW。
与其他流行优化器的比较
该研究还使用四种流行的优化器 RAdam、NAdam、AdaBelief 和 AMSGrad 在 ImageNet 上训练 ViT-S/16 和 ViT-B/16（使用 RandAug 和 Mixup）。如下表 7 所示，Lion 仍然是表现最好的。
更多技术细节请参阅原论文。
