Lion Optimizer/zh
| Article | |
|---|---|
| Topic area | optimization |
| Prerequisites | Stochastic gradient descent, Momentum (optimization), Adam optimizer |
概述
Lion 优化器(全称 EvoLved Sign Momentum,演化式符号动量)是一种用于训练神经网络的一阶随机优化算法,由 Chen 等人于 2023 年提出。[1] 与大多数广泛使用、基于理论原理手工设计的优化器不同,Lion 是通过对候选优化器程序空间进行符号程序搜索而发现的。其更新规则在概念上比 Adam 更简单,但在大规模 深度学习 工作负载上的准确率通常能够与之相当或更优,同时所占用的优化器状态内存大约只有一半。
Lion 的定义性特征是参数更新方向是插值动量项的逐元素符号。这使得更新的每个坐标都具有相同的大小(由学习率缩放的固定步长),与 Adam 不同——后者的逐坐标自适应缩放由梯度平方的指数移动平均决定。Lion 已被应用于多种大规模 语言模型 和 视觉训练流水线,并成为自动发现优化器这一研究方向的标杆。
背景与动机
到 2020 年代初,AdamW 已成为训练 Transformer 的事实标准优化器,尽管存在众所周知的局限:每个参数需要存储两个矩缓冲区(一阶矩和二阶矩),使优化器内存达到参数本身的两倍,并且其超参数对模型规模十分敏感。一长串研究工作试图通过手工设计来取代 Adam,经验结果褒贬不一:LAMB、Adafactor、Sophia、Shampoo 等各自解决了特定问题,但都未能撼动 Adam 的默认地位。
Lion 源于一个不同的前提:与其设计一个优化器,不如去搜索一个优化器。作者将优化器形式化为一个对梯度和状态进行操作的短程序,定义了一个由原始操作(加法、乘法、符号、指数移动平均、裁剪等)构成的搜索空间,并使用 演化式程序搜索 配合 元学习风格的代理训练任务来评估候选项。Lion 是在对更大的发现候选项进行大量剪枝后涌现出的、最简洁的高性能程序。
算法
设 $ g_t $ 为损失函数在第 $ t $ 步关于参数 $ \theta_t $ 的梯度,设 $ m_t $ 为初始化为零的单个动量缓冲区,设 $ \eta $ 为学习率。Lion 有两个插值系数 $ \beta_1, \beta_2 \in [0,1) $ 和一个权重衰减系数 $ \lambda \ge 0 $。第 $ t $ 步的更新为:
$ {\displaystyle c_t = \beta_1 \, m_{t-1} + (1 - \beta_1) \, g_t} $
$ {\displaystyle \theta_t = \theta_{t-1} - \eta \, \big( \operatorname{sign}(c_t) + \lambda \, \theta_{t-1} \big)} $
$ {\displaystyle m_t = \beta_2 \, m_{t-1} + (1 - \beta_2) \, g_t} $
第一行通过在先前的动量与当前梯度之间进行插值,得到一个更新方向 $ c_t $。第二行仅使用 $ c_t $ 的逐元素符号来执行参数更新,并附加一个采用 AdamW 风格的解耦权重衰减。第三行则使用另一个不同的插值系数来更新持久动量缓冲区 $ m_t $。关键之处在于,用于下一步方向的缓冲区($ m_t $)与用于当前方向的缓冲区(涉及 $ m_{t-1} $ 的混合)由不同的 $ \beta $ 值控制。
Chen 等人报告的默认超参数为 $ \beta_1 = 0.9 $、$ \beta_2 = 0.99 $,学习率大致比 AdamW 的取值小一个数量级,而权重衰减则大致大一个数量级。
直观理解
三点观察有助于解释 Lion 为何有效。
首先,每次更新在每个坐标上都具有单位幅度。由于更新为 $ \eta \cdot \operatorname{sign}(c_t) $,任何一个坐标的步长都不会超过 $ \eta $。这种隐式裁剪类似于梯度裁剪,但它是在逐坐标层面、且在每一步都进行,从而对 语言模型 训练中常见的异常梯度提供鲁棒性。
其次,双 $ \beta $ 结构将方向平滑与缓冲区累积解耦。当 $ \beta_1 < \beta_2 $ 时,方向 $ c_t $ 对当前梯度的反应比持久缓冲区 $ m_t $ 更加灵敏。这使 Lion 内建了一种类似 Nesterov 的前瞻特性:在时刻 $ t $ 所走的一步比单独由动量缓冲区给出的步骤更多地反映 $ g_t $。
第三,取符号会丢弃梯度的幅度信息,而 Adam 是通过其二阶矩来利用这一信息的。Lion 在丢弃幅度信息的情况下仍能良好训练这一经验事实表明,对许多大型神经网络而言,方向比精确的幅度缩放更重要,前提是学习率和权重衰减得到相互兼容的调优。
内存与计算
Lion 每个参数仅存储一个动量缓冲区 $ m_t $;而 Adam 和 AdamW 则存储两个(一阶矩和二阶矩)。对于一个拥有 $ P $ 个参数、以 bfloat16 或 32 位精度训练的模型而言,这将优化器状态内存减少了一半:从 $ 2P $ 个值降到 $ P $ 个值。对于参数规模达到数百亿乃至数千亿的模型,在加速器内存预算固定的情况下,这种节省对分布式训练而言意义重大。
Lion 每步的计算也略微更便宜:仅需一次符号运算和两次插值,无需对二阶矩估计取平方根并进行除法。在实践中,每步的时间主要由前向传播和反向传播主导,因此优化器算术更便宜带来的实际加速幅度不大,但确实非零。
超参数调优
Lion 的超参数并不能直接从 Adam 沿用过来。作者经验观察到、并在独立复现中得到佐证的推荐调整方式如下:
- 学习率:大致比同一模型下适用的 AdamW 取值小 $ 3\times $ 到 $ 10\times $。由于 Lion 的更新在每个坐标上具有单位幅度,有效步长更可预测,但在条件良好的坐标上其绝对值小于 Adam 的自适应缩放。
- 权重衰减:大致比 AdamW 取值大 $ 3\times $ 到 $ 10\times $,以在学习率较小的同时,保持有效的正则化强度 $ \eta \lambda $。
- 批量大小:在某些情形下,Lion 比 Adam 更能从更大的批量中受益,可能是因为基于符号的更新每步方差更高,在更大批量上能更干净地求平均。
- $ \beta_1, \beta_2 $:默认值 $ (0.9, 0.99) $ 较为稳健;原论文发现进行调优可获得些许提升,但建议从默认值开始。
经验性能
在原始论文中,Lion 在 ImageNet 上的图像分类、图文对比学习(CLIP 风格)、扩散模型训练以及最高达数十亿参数的语言模型预训练上进行了评估。在所有这些设置下,Lion 在验证准确率或困惑度上与 AdamW 持平或略有超越,且通常具有可比的实际训练时间。
在较小模型上的独立复现总体证实 Lion 具有竞争力,但并非在所有情形下都更优。其结果取决于模型架构、批量大小以及调优投入。在某些微调场景下,尤其是数据集非常小的情形,Lion 对学习率调度的选择比 AdamW 更敏感。
变体与后续工作
自 2023 年以来,出现了多种变体和分析:
- Tiger 采用了类似的动量符号思想,但仅使用单个 $ \beta $,据报道在略微牺牲准确率的代价下进一步提升了内存效率。
- Lion-K 将逐元素符号替换为更一般的非线性,并通过理论分析将 Lion 与正则化镜像下降的解释联系起来。
- 研究者已经发展出 Lion 的连续时间和收敛性分析,将该算法视为带符号非线性的特定常微分方程的离散时间近似,为所观察到的鲁棒性提供了部分理论依据。
这些后续工作普遍认同,Lion 的行为是由符号非线性与动量的指数移动平均之间的相互作用所支配,而非单独由其中任一组件决定。
与相关优化器的比较
- 与 signSGD 对比:signSGD 直接对原始梯度取符号,而不使用动量缓冲区。Lion 可视为采用了特定双 $ \beta $ 动量方案的 signSGD;该动量对 Lion 的经验性能至关重要。
- 与 Adam / AdamW 对比:Lion 使用一个缓冲区而非两个,并丢弃了幅度信息。Adam 通过二阶矩进行逐坐标尺度自适应,而 Lion 则不会。
- 与带动量的 SGD 对比:带动量的 SGD 直接将原始动量向量用作方向;Lion 使用其符号,从而对每个坐标的更新幅度设了上限。
- 与 Adafactor 对比:Adafactor 通过对二阶矩进行因子分解来降低优化器内存;Lion 则完全消除了二阶矩。
局限性
Lion 并非 Adam 的通用替代品。已报告的局限性包括:
- 超参数敏感性:由于 Lion 的更新基于符号,在学习率、权重衰减和预热上的细微改动可能比 Adam 更剧烈地改变其行为。
- 小批量场景:基于符号的更新每步方差更大;在批量非常小时,Adam 的逐坐标自适应缩放可起到稳定作用,而 Lion 缺乏这种机制。
- 稀疏梯度:在内在具有稀疏梯度结构的模型中(部分推荐系统、嵌入表),取符号的操作可能放大那些很少被更新的坐标上的噪声。
- 理论保证有限:尽管已有后续分析,Lion 在凸收敛理论方面的深度仍不及 SGD 或 Adam 变体。
历史
Lion 见于 Google 的 Chen 等人在 2023 年发表的论文《Symbolic Discovery of Optimization Algorithms》。该优化器是更广泛的符号搜索方法学的代表性成果,作者认为这一方法可以推广到机器学习训练流水线的其他组件。在论文发布后的数月内,PyTorch、JAX 和 Optax 上就出现了开源实现,该算法被纳入多个大规模训练栈,用于语言模型和图像生成的预训练等任务。