技术博客

Lion 与 Sophia:两个挑战 Adam 的新优化器在改什么

第 129 / 147 篇 · 2025-04-10 · ai · 优化器Adam二阶方法

Adam 霸占了深度学习优化器十年,对它的不满也积累了十年。2023 年几乎同时出现的 Lion(Google, Chen et al.)和 Sophia(Princeton, Liu et al.)代表了两条不同的改良路线:Lion 做减法——砍掉二阶矩,更新只保留符号信息;Sophia 做加法——把 Adam 对曲率的粗糙近似(梯度二阶矩)换成真正的二阶信息(Hessian 对角)。两者都声称在大模型预训练上快过 AdamW,但快的方式完全不同。

Lion:动量的符号化

Lion 的更新只有两行:

$$ c_t = \beta_1 m_{t-1} + (1-\beta_1) g_t, \qquad \theta_{t+1} = \theta_t - \eta\,\big(\mathrm{sign}(c_t) + \lambda\theta_t\big), \qquad m_t = \beta_2 m_{t-1} + (1-\beta_2) g_t . $$

和 Adam 对比,差异一目了然:没有 $v_t$,没有除法,没有 $\varepsilon$。$c_t$ 是"旧动量与新梯度的凸组合",取符号之后,每个坐标的步长恒为 $\eta$,更新向量范数恒为 $\eta\sqrt{d}$——步长信息从更新幅度里被完全抽离,优化器只负责给方向。这和 Adam $\varepsilon \to 0$ 极限下的符号下降形态一致,但符号作用的对象不是单步梯度而是动量插值,等效于符号版动量。

为什么敢扔掉幅度信息?作者的论证分两头。理论上,signSGD 在光滑非凸目标下有和 SGD 同阶的收敛率,符号操作引入的"截断偏差"(把不同幅度的坐标压成等权)在梯度噪声重尾时反而是优点——大幅度坐标往往是噪声尖峰,截断正好把它们削掉。工程上,幅度信息并没有真丢:每一步 $\mathrm{sign}(c_t)$ 的幅度恒为 1,而 Adam 的幅度信息($|m|/\sqrt{v}$)本来就被归一化到 $[0,1]$ 的窄区间里,携带的信息量有限,砍掉省掉的是一整个二阶矩的显存和带宽——Lion 的优化器状态只有 Adam 的一半,显存敏感场景(大 embedding、微调)这是实打实的收益。论文实测 ViT 预训练快 2 倍、精度持平。

另一个反直觉的设计:Lion 动量的衰减系数是 $\beta_1 = 0.9$(插值)、$\beta_2 = 0.99$(记忆),比一般训练更小的 $\beta_2$。符号操作本身就是最强的归一化,长窗口的记忆反而会让陈旧方向残留过久。

Sophia:给 Adam 换上真二阶矩

Adam 用 $v_t \approx E[g^2]$ 当"曲率"的代理,这个代理在梯度噪声大时偏离真实曲率很远:$E[g^2] = (\nabla f)^2 + \mathrm{Var}(g)$,噪声方差被当成了曲率。Sophia 的替换直截了当——分母用 Hessian 对角项 $h_t \approx \mathrm{diag}(\nabla^2 f)$ 的指数平均:

$$ \theta_{t+1} = \theta_t - \eta\, \mathrm{clip}\!\left(\frac{\hat m_t}{\max(\hat h_t, \rho)},\, -c,\, c\right). $$

$\max(\hat h_t, \rho)$ 防止负曲率方向($h<0$,凸性缺失处)的更新反向爆炸;外层的 clip 把单坐标步长钳在 $c\eta$ 以内。和 Adam 相比,分母从"梯度的二阶统计"换成"loss 曲面的二阶导数",这一步在数学上更接近牛顿法的对角近似。

代价是 Hessian 从哪来。直接算对角 Hessian 要 $O(d)$ 次前向或反向,不可行;Sophia 的答案是每 $k$ 步估一次($k=10$),用两种估计器之一:小批量 Hutchinson(随机向量投影 $v^\top H v$,一次额外反向),或对语言模型更实用的 Gauss-Newton 近似——把模型在预训练数据上额外跑 $k$ 步,用梯度二阶矩在该子集上重新校准 $h$。低频估计 + 指数平均的组合让开销摊到约 5% 的额外计算。论文报告 LLM 预训练在同等 step 数下 loss 低约 0.04 nats,等效加速约 2 倍。

两条路线对比:Lion 砍掉幅度信息,Sophia 换掉分母的信息来源 Lion(减法) c = β1·m + (1−β1)·g Δθ = −η·sign(c) − ηλθ 状态量:只有 m(Adam 的一半) 每步更新范数恒为 η√d 噪声尖峰被截断天然抑制 Sophia(加法) Δθ = −η·clip(m̂/max(ĥ,ρ)) h = EMA(diag Hessian) 分母是真曲率,不是噪声方差 每 k 步估一次 h,开销约 5% 负曲率方向由 max(h,ρ) 兜底

怎么选:一个粗略的决策框架

两个方法都不是无成本的普适替代。Lion 的代价在超参迁移:它默认配更大的 $\eta$(论文用 $\eta_{\text{Lion}} \approx 3\sim10 \times \eta_{\text{Adam}}$ 级别)和独立的 $\lambda$ 扫描,且 sign 更新对 batch size 的敏感性和 Adam 不同,照搬 Adam 的配置直接训大概率发散。Sophia 的代价在实现复杂度:Hessian 估计器要侵入训练循环(每 $k$ 步插入估计 pass),分布式训练里还要处理估计 pass 的通信同步,工程上远比换一行优化器类麻烦。

实践里的观察是:两者在中等规模 LLM 预训练上的加速证据相对扎实,在微调和小数据场景优势缩水——Lion 论文自己也承认在需要精确小步逼近的场景(如 diffusion 微调)表现不如 AdamW。截至目前的生态现实是:AdamW 仍是默认,Lion 在视觉和部分 LLM 训练中有落地(如部分版本 Switch Transformer 后时代实验),Sophia 则更多停留在研究代码里,但它的"低频二阶校准"思路已经渗进若干工业训练栈的定制优化器中。

小结

Lion 与 Sophia 的共同点是对 Adam 的 $v_t$ 不满,分歧在解法:Lion 认为 $E[g^2]$ 携带的幅度信息不值得维护,砍掉换显存和噪声鲁棒性;Sophia 认为它信息量不够,换成真 Hessian 换收敛速度。把两者放进 Adam 的统一视角里:Adam 是"方向平滑 + 幅度归一"的组合,Lion 强化前者、删除后者,Sophia 保持结构、精化后者。优化器设计的搜索空间,本质上就是在回答"梯度的哪些统计值得花显存记住"这一个问题。