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


该研究采用以下技术来解决无限和稀疏搜索空间带来的挑战 。首先应用正则化,因为它简单,可扩展,并在许多搜索任务中很成功;其次是简化程序空间中的冗余;最后,为了降低搜索成本,该研究通过减少模型大小、训练示例的数量和距离目标任务的 step 使成本降低 。
左:显示了五次进化搜索实验的均值和标准误差 。右:随着搜索的进行,冗余语句的百分比和缓存命中率都在增加 。

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

文章插图
Lion 的推导与分析
研究者表示,优化器 Lion 具有简洁性、高内存效率以及在搜索和元验证方面的强大性能 。
推导
搜索和选择过程促成了程序 4,它通过从原始程序 8(附录)自动删除冗余语句获得 。研究者进一步简化并得到程序 1 中的最终算法(Lion) 。在简化过程中从程序 4 中删除了几个不必要的元素 。其中 cosh 函数被删除,因为 m 将在下一次迭代中重新分配(第 3 行) 。使用和 clip 的语句也被删除,因为研究者观察到没有它们质量不会下降 。三个红色语句转换为一个符号函数 。
尽管程序 4 中同时使用了 m 和 v,但 v 仅改变了动量的更新方式(两个具有常数~0.9 和 ~1.1 的插值函数等同于一个具有 ~0.99 的函数)并且不需要单独跟踪 。注意不再需要偏差校正,因为它不会改变方向 。
如下算法 2 显示了伪代码 。
【涨点神器!超越Adam!谷歌提出Lion炼丹优化器,内存更小、效率更高!】分析
符号更新和正则化 。Lion 算法通过符号操作在所有维度上产生了具有统一幅度的更新,这在原理上不同于各种自适应优化器 。直观来看,符号操作为更新添加了噪声,作为了一种正则化形式并有助于泛化 。下图 11(右)展示了一个证据 。
动量跟踪 。与 AdamW 和动量 SGD 中常用的 0.9 相比,Lion 中用于跟踪动量的默认 EMA 因子为 0.99 (β_2) 。EMA 因子和插值的这种选择,使得 Lion 在记住动量梯度的 10 倍和在更新中对当前梯度施加更多权重之间取得平衡 。
超参数和批大小选择 。与 AdamW 和相比,Lion 更简单并且超参数更少,因为它不需要 ? 和因式分解相关的参数 。Lion 需要更小的学习率,进而需要更大的解耦权重衰减才能达到类似的有效权重衰减强度 (lr * λ) 。
内存和运行时优势 。Lion 只保存动量,比 AdamW 等流行的自适应优化器占用的内存更小,这在训练大模型和 / 或使用大批量时非常有用 。例如,AdamW 需要至少 16 个 TPU V4 芯片来训练图像分辨率为 224 且批大小为 4,096 的 ViT-B/16,而 Lion 只需要 8 个(均具有动量) 。
Lion 评估结果
在实验部分,研究者在各种基准上对 Lion 进行了评估,主要将它与流行的 AdamW(或者当内存成为瓶颈时的 )进行比较 。
图像分类
研究者在图像分类任务上执行涵盖各种数据集和架构的实验 。除了在上从头开始训练之外,他们还在两个更大的成熟数据集 -21K 和 JFT 上进行预训练 。图片大小默认为 224 。
首先在上从头开始训练 。研究者训练 -50 达到 90 个 epoch,批大小为 1,024,其他模型训练 300 个 epoch,批大小为 4,096 。如下表 2 所示,Lion 在各种架构上显著优于 AdamW 。
其次在 -21K 上进行预训练 。研究者在 -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 在和V2 上训练的 ViT-H/14 性能相匹配,但预训练成本降低了 3 倍 。