Lion : 超越 AdamW 的优化算法

Posted 2023-03-23 Zilliz Planet

出品人：Towhee 技术团队

优化算法，即优化器，在训练神经网络中起着基础作用。

近年来引入了大量手工优化器，其中大部分是自适应优化器。然而，具有解耦权重衰减的 Adam，也称为 AdamW，和具有因数二次矩的 Adafactor，仍然是训练大多数深度神经网络的事实上的标准优化器，尤其是最近最先进的语言、视觉和多模态模型。

另一个方向是自动发现这样的优化算法。学习优化 (L2O) 方法建议通过训练参数化模型（例如神经网络）来发现优化器以输出更新。然而，那些通常在有限数量的小任务上训练的黑盒优化器很难推广到最先进的设置，在这些设置中，更大的模型需要更多的训练步骤来训练。

另一类方法应用强化学习或蒙特卡洛采样来发现新的优化器，其中搜索空间由由预定义操作数和运算符组成的树定义。为了使搜索易于管理，他们通常通过使用固定操作数和限制树的大小来限制搜索空间，从而限制发现的可能性。因此，所发现的算法尚未达到最先进的水平。

Google 和 UCLA 提出 Lion (EvoLved Sign Momentum) 优化算法，以优越的性能和良好的效果超越了经典的优化算法。Lion 通过程序搜索来发现优化算法，并且用其改进深度神经网络训练。它能够解决在无限、稀疏空间的搜索挑战。

除了更具效率以外，Lion 还具有强大的泛化能力，跨越了体系结构、数据集和任务的限制。在图像分类方面，Lion 在 ImageNet 上将 ViT 的准确性提高了高达 2％，并在 JFT 上节省了高达 5 倍的预训练计算。

在视觉语言对比学习方面，Lion 在 ImageNet 上实现了 88.3％ 的zero-shot和 91.1％ 的微调准确度，分别超过了之前的最佳结果 2％ 和 0.1％。在扩散模型上，Lion 的表现优于 Adam，实现了更好的 FID 得分并将训练计算减少了高达 2.3 倍。对于自回归、掩码语言建模和微调，Lion 也表现出与 Adam 相当或更好的性能。

Lion 利用了高效的搜索技术来探索无限稀疏的程序空间。为了弥合代理和目标任务之间的广义间隙，它还引入了程序选择和简化策略。与自适应优化器不同，Lion 对于通过符号操作计算的每个参数都采用了相同的幅度的更新。它的性能增益随着训练批次大小的增加而增加。由于符号函数产生的更新的范数更大，它还需要较小的学习率。

AdamW vs. Lion optimizer.

上图为 AdamW 与 Lion 优化算法的对比。

可以看到，与 AdamW 和各种自适应优化器需要同时保存一阶和二阶矩相比，Lion 只需要动量，将额外的内存占用减半。这在训练大型模型和/或大批量时很有用。

例如，AdamW 需要至少 16 个 TPU V4 芯片来训练图像大小为 224、批量大小为 4,096 的 ViT-B/16，而 Lion 只需要八个。

另一个实际好处是，由于 Lion 的简单性，Lion 在我们的实验中具有更快的运行时间（步数/秒），通常比 AdamW 和 Adafactor 提速 2-15%，具体取决于任务、代码库和硬件。

因此，Lion 是一种表现出色的优化器，比AdamW 表现更强大。Lion 具有更快速和更节省显存的特点，有望成为未来主流优化器之一。如果你在训练时，感觉保存的训练过程中的状态太大，导致性能缓慢，不妨试试用 Lion。

自从 Adam 被提出以来，由于其快速收敛的特性，已成为许多模型的默认优化器。但是，一些学者担心这种现象可能导致所有的模型改进都朝着有利于 Adam 的方向发展。因此，发现比 Adam 更简单、更有效的优化器是很重要的。Lion 能够获得更好的泛化性能，这可能是因为它引入了额外的噪声，使得模型进入了更平坦但未必更小的损失区域。虽然 Lion 有更强的鲁棒性，但它并不是完美的，比如在小的 batch_size 时，Lion 的表现不如 AdamW，因为噪声必须适量才能更好地发挥作用。

相关资料：

• 代码地址：https://github.com/google/automl/tree/master/lion
• 论文链接：https://arxiv.org/abs/2302.06675v2