涨点神器!超越Adam!谷歌提出Lion炼丹优化器,内存更小、效率更高!

谷歌的 Lion 优化器将成为训练大模型或大批量的「福音」 。
优化器即优化算法,在神经网络训练中起着关键作用 。近年来,研究者引入了大量的手工优化器,其中大部分是自适应优化器 。Adam 以及优化器仍然占据训练神经网络的主流,尤其在语言、视觉和多模态领域更是如此 。
除了人工引入优化器外,还有一个方向是程序自动发现优化算法 。此前有人提出过 L2O( to ),该方法通过训练神经网络来发现优化器 。然而,这些黑盒优化器通常是在有限数量的小任务上训练而成,很难泛化到大模型 。
还有人尝试了其他方法,他们应用强化学习或蒙特卡罗采样来发现新的优化器 。然而,为了简化搜索,这些方法通常会限制搜索空间,从而限制发现其他优化器的可能性 。因此,目前方法尚未达到 SOTA 水平 。
最近几年,值得一提的是 -Zero,它试图在评估任务时搜索机器学习的每个组件,对优化器的发现很有借鉴意义 。
在本文中,来自谷歌、 UCLA 的研究者提出了一种通过程序搜索发现深度神经网络训练的优化算法的方法,进而发现 Lion( Sign )优化器 。实现这一目标面临两个挑战:首先是在无限稀疏的程序空间中寻找高质量的算法,其次是选择可以从小型任务泛化到更大、SOTA 任务的算法 。为了应对这些挑战,该研究采用了一系列技术,包括具有热启动和重启的进化搜索、抽象执行、 选择和程序简化 。
PS:公众号后台回复:Lion,即可下载
与 AdamW 和各种自适应优化器(这些优化器需要同时保存一阶和二阶矩)相比,Lion 只需要动量并利用符号操作来计算更新,并且将额外的内存占用减半 。这在训练大模型或大批量时很有用 。例如,AdamW 需要至少 16 个 TPU V4 来训练图像大小为 224、批大小为 4096 的 ViT-B/16,而 Lion 只需要 8 个 。另一个实际好处是,由于 Lion 的简单性,Lion 在实验中具有更快的运行时间(steps /sec),通常比 AdamW 和提速 2-15%,具体取决于任务、代码库和硬件 。
论文一作Chen 表示:我们的符号程序搜索发现了一个只跟踪动量的有效优化器 ——Lion 。与 Adam 相比,它获得了 88.3% 的零样本和 91.1% 的微调准确率,以及高达 5 倍(与 ViT 相比)、2.3 倍(与扩散模型相比)和 2 倍(与 LM 想比)训练效率 。
Lion 在一系列模型(、MLP、、U-Net 和 )和任务(图像分类、视觉 - 语言对比学习、扩散、语言建模和微调)中展示了出色的性能 。值得注意的是,通过在 BASIC 中用 Lion 替换,该研究在上实现了 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% 。
图源:
算法的符号发现
本文使用程序形式的符号表示具有以下优点:(1)它符合算法必须作为程序来执行的事实;(2) 与神经网络等参数化模型相比,程序等符号表示更易于分析、理解和迁移到新任务;(3) 程序长度可用于估计不同程序的复杂性,从而更容易选择更简单、通常更具通用性的程序 。这项工作侧重于深度神经网络训练的优化器,但该方法通常适用于其他任务 。
下图中,程序 2 这个简化代码片段使用与 AdamW 相同的,以确保发现的算法具有更小或相等的内存占用空间;程序 3 中给出了 AdamW 的示例表示 。