跳转至

Optimistic Dual Averaging Unifies Modern Optimizers

会议: NeurIPS 2026
arXiv: 2605.11172
OpenReview: https://openreview.net/forum?id=fsvzA0o8vh
领域: 优化/理论
关键词: 乐观对偶平均、现代优化器统合、无超参权重衰减、Muon与Lion、原始对偶平均

一句话总结

本文提出将乐观对偶平均(SODA)作为现代一阶优化器的统合理论框架,不仅自然囊括了 Adam、Lion、Muon、NAdam 等主流算法,更推导出将权重衰减重新解释为原始迭代平均的免调参 \(1/(k+2)\) 调度包装器。

研究背景与动机

近年来深度学习优化器的发展主要沿着两条互为补充的技术主线推进:其一是几何适配与归一化,即针对现代大模型高维参数结构设计非欧几何更新,代表方案包括基于 \(\ell_\infty\) 几何坐标归一化的 Adam/Lion 以及基于矩阵谱范数收缩的 Muon 与 Scion;其二是算法组件的高效编排,即探索梯度动量、外推、多序列平均与调度器的协同机制。然而,这一快速演进导致不同优化器看似割裂,缺乏一个统一的理论基石来解释为什么特定的动量与更新规则能带来加速。

更为紧迫的工程痛点在于超参数调节的巨大开销。随着模型规模与训练步长激增,按网格搜索超参数在算力上不可承受。特别是在单周期(single-epoch)预训练中,尽管过拟合并不存在,权重衰减(weight decay)却对最终收敛具有关键作用。近期针对 Muon 等优化器提出的按模型维度倒数 \(1/d\) 缩放权重衰减的经验法则,不仅严重依赖优化器类型(例如 Lion 需要比 AdamW 大得多的衰减值),而且混淆了模型维度与训练步长(horizon)对衰减强度的真实影响。

面对现代优化器设计碎片化与权重衰减跨训练时域难以迁移的困境,本文从经典对偶平均(Dual Averaging)理论出发,构建能够跨越欧氏与非欧几何的统一分析视角。核心 idea:将现代优化器统一刻画为兼具对偶乐观外推与原始迭代平均的泛化乐观对偶平均(SODA),从而证明权重衰减本质上是原始空间的平均机制,并衍生出理论严密、无需针对不同时域调节超参数的 \(1/(k+2)\) 衰减封装器。

方法详解

整体框架

SODA 将经典乐观对偶平均(ODA)与原始外推序列深度耦合,形成对偶空间(梯度滑动平均与乐观预测)和原始空间(迭代点平均与外推)高度对称的双流动力学。算法在对偶端维护带有两级动量的梯度估计,在原始端通过 Fenchel 共轭将对偶方向映射为几何更新向量,并与历史迭代点进行加权平均。

在实际实现中,针对现代深度学习的“现代化”(MODA)简化版令原始外推参数 \(\bar{\lambda}_k=0\),此时梯度计算点与输出点完全重合,一举覆盖了 Adam、NAdam、Lion 与 Muon;进一步结合凸优化光滑性分析给出的最优参数配置,SODA 被实例化为一个通用的包装器(SODA Wrapper),能够直接包裹任意基优化器(如 Adam 或 Muon),将标准权重衰减转化为围绕初始点锚定的自适应平均。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["当前迭代点 x^k 与初始化锚点 z^0"] --> G["计算随机梯度 g^k = ∇f(x^k, ξ_k)"]
    G --> B["基优化器无权重衰减更新<br/>BaseUpdate(g^k)"]
    B --> Z["对偶映射外推点更新<br/>z^(k+1) = z^0 + (k+2)·BaseUpdate(g^k)"]
    Z --> X["原始空间迭代平均<br/>x^(k+1) = (1 - 1/(k+2))·x^k + 1/(k+2)·z^(k+1)"]
    X --> Out["展开为中心化与时间衰减更新<br/>x^(k+1) = z^0/(k+2) + (1 - 1/(k+2))·x^k + BaseUpdate(g^k)"]

关键设计

1. SODA统合框架:对偶与原始双流乐观动力学 传统对偶平均仅在对偶空间累加梯度,难以解释现代自适应优化器中的乐观动量与多序列交互。SODA 引入对称的五步迭代:在对偶空间,算法不仅维护基础梯度动量 \(m^{k+1}=(1-\alpha_k)m^k + \alpha_k \nabla f(y^k,\xi_k)\),还通过乐观步 \(\bar{m}^{k+1}=(1-\bar{\alpha}_k)m^{k+1} + \bar{\alpha}_k \nabla f(y^k,\xi_k)\) 提前预估下一步梯度;随后借助正则化子(镜面映射)\(h_k(x)\) 的 Fenchel 共轭子微分求出中间点 \(z^{k+1}\in \arg\min_x \{\gamma_k \langle \bar{m}^{k+1},x\rangle + h_k(x)\}\);最后在原始空间执行迭代平均 \(x^{k+1}=(1-\lambda_k)x^k + \lambda_k z^{k+1}\) 与外推 \(y^{k+1}=(1-\bar{\lambda}_k)x^{k+1} + \bar{\lambda}_k z^{k+1}\)。通过选择不同的 \(h_k(x)\),该框架直接统一了现有 SOTA 方法:\(\ell_\infty\) 示性函数对应 Lion,平滑坐标对角二次型对应 NAdam(AdEMAMix 的简化版),而谱范数示性函数则对应带有 Nesterov 动量的 Muon。

2. 权重衰减的原始平均本质:无需调参的 \(1/(k+2)\) 调度 现存实践通常将权重衰减视作防止过拟合的 \(\ell_2\) 正则项,但无法解释为何在无过拟合的单周期训练中其依然极其关键。SODA 的理论分析表明,权重衰减本质上是原始迭代平均的一种变体。在现代配置(\(\bar{\lambda}_k=0\))下,令迭代平均权重 \(\lambda_k = 1/(k+2)\) 并将步长设为 \(\gamma_k = (k+2)\eta_k\),将带有参考点 \(z^0\) 的正则项展开后,迭代更新化简为: $\(x^{k+1} = \frac{1}{k+2}z^0 + \left(1 - \frac{1}{k+2}\right)x^k + \operatorname{BaseUpdate}(g^k)\)$ 该式揭示出三个物理意义明确的组成部分:第一项是以初始权重 \(z^0\) 为锚点的中心化(centering)项;第二项是系数严格为 \(1 - 1/(k+2)\) 的时间衰减项;第三项则直接继承基优化器的非正则更新量。相比传统将衰减强度乘以学习率的解耦权重衰减(AdamW),SODA 的等效衰减率随步数按 \(1/(k+2)\) 自然退火,完全剥离了对人工超参数 \(\mu\) 的调参依赖。

3. 极简 SODA 包装器与 SODA† 架构 针对工程落地的便捷性,作者提出即插即用的通用 SODA 包装器(SODA Wrapper):算法不改变底层基优化器对特定几何的处理(如 Muon 的奇异值正交化或 Adam 的对角二阶矩),仅在每一轮外层迭代中按时间步长放大基更新量并执行凸组合。初始点 \(z^0\) 可以通过随机种子按需伪随机生成,不引入任何额外的显存常驻开销。在此基础上,论文进一步提出规范化的 SODA† 变体:对输入/输出嵌入层采用统一的 Lion 动量(\(\ell_\infty\) 几何),对内部隐藏层采用谱范数几何并施加统一的乐观动量参数 \(\alpha_k=\bar{\alpha}_k=0.05\),彻底摒弃了 Muon 原版中对不同网络层混合使用 Adam 与专用动量的繁琐拼凑,实现了全模型动量超参数的大一统。

损失函数 / 训练策略

在理论上,SODA 针对 \(L\)-光滑、凸且具有二阶梯度变化界(梯度方差 \(\sigma^2\) 与相关系数 \(\rho\))的目标函数建立收敛保证。当选用固定正则项 \(h\) 且其具备 \(\mu\)-强凸性时,配置 \(\alpha_k=1/(k+1)\)、\(\bar{\alpha}_k=\lambda_k=1/(k+2)\)、\(\bar{\lambda}_k\leq \lambda_k/10\) 及步长 \(\gamma_k=\eta(k+2)\),可达到 \(O\left(\frac{(\rho+1)LR_\star}{\mu n} + \frac{\sigma \sqrt{R_\star}}{\sqrt{\mu n}}\right)\) 的最优非渐近收敛率。在实际语言模型预训练策略中,采用 nanoGPT 架构在 FineWeb100 数据集上进行单周期训练,训练步长采用 Chinchilla 算力最优规模(1× 至 4× Chinchilla,即每参数对应 20 至 80 tokens),学习率配合余弦退火调度,彻底关闭基优化器内建的手动权重衰减。

实验关键数据

主实验

在 FineWeb100 预训练任务上,采用 124M 规格的 nanoGPT 模型,对比在 1× Chinchilla(标准算力最优)与 4× Chinchilla(过训练场景)下,经 SODA 包装的优化器与基线优化器(带最优网格搜索权重衰减)的验证集最终 Cross-Entropy Loss。SODA 包装器完全不调节权重衰减超参数,表现优于经细致网格搜索的最优基线。

基优化器 训练时域 (Chinchilla) 基线最优损失 (网格调优权重衰减) SODA 包装后损失 (零权重衰减调参) 优势提升 (验证损失降低)
Muon (124M) 1× Chinchilla 3.328 3.312 -0.016 (更优收敛)
Muon (124M) 4× Chinchilla 3.125 3.101 -0.024 (超越最优衰减)
Adam (124M) 1× Chinchilla 3.415 3.390 -0.025 (一致性提升)
Adam (124M) 4× Chinchilla 3.208 3.182 -0.026 (超越基线)
Scion (124M) 1× Chinchilla 3.342 3.326 -0.016 (消除调参)
Scion (124M) 4× Chinchilla 3.148 3.123 -0.025 (持续收益)
SODA† (统合架构) 1× Chinchilla — (基线混合架构 3.328) 3.308 -0.020 (达最佳SOTA)
SODA† (统合架构) 4× Chinchilla — (基线混合架构 3.125) 3.096 -0.029 (全局表现最优)

消融实验

针对 SODA 包装器中初始锚定点 \(z^0\) 的中心化作用,以及将权重衰减跨时域迁移时不同缩放法则的表现进行消融对比(以 124M 模型在 4× Chinchilla 训练时域上的验证集损失为例):

算法配置 / 策略 4× Chinchilla 验证损失 说明
SODA(Muon) 完整包装器 3.101 包含 \(z^0\) 锚定与 \(1/(k+2)\) 动态衰减
w/o \(z^0\) 锚定 (设 \(z^0=0\)) 3.142 初始点移回原点,远离真实最优解,性能显著恶化 (+0.041)
固定权重衰减 (沿用 1× 最优衰减) 3.155 权重衰减未随训练时域动态退火,过拟合严重 (+0.054)
仅按维度缩放 \(1/d\) (Xiao/Qiu 经验法则) 3.136 忽略时域变化对衰减强度的影响,次优于 SODA (+0.035)
SODA† 去除乐观项 (\(\bar{\alpha}_k=0\)) 3.124 禁用乐观动量退化为普通双重平均,损失上升 (+0.028)

关键发现

  • 权重衰减与时域直接耦合:实验证实固定网络宽度下,最优权重衰减强度随训练步长增长而单调递减;单纯采用 \(1/d\) 的无量纲缩放无法应对不同时域的泛化要求,而 SODA 的 \(1/(k+2)\) 自然衰减曲线完美吻合这一趋势。
  • 锚定初始化点(Centering)不可或缺:若将式 (4) 中的锚定项设为 \(0\)(即直接向原点衰减),损失出现大幅回退,证明在非凸大模型优化中,以初始化权重作为正则中心能够提供极其关键的隐式偏置。
  • 模型扩展性单调增强:在 64M 至 1B 参数量的 nanoGPT 弱扩展实验中,SODA† 相较于基线的优势随模型规模与训练步长同步放大,呈现更佳的 Scaling 表现。

亮点与洞察

  • 为割裂的优化器流派提供通用对偶透镜:过去 Lion 与 Muon 的成功常被归功于启发式的符号或矩阵正交化工程技巧,本文证明它们均是经典乐观对偶平均在 \(\ell_\infty\) 与谱几何下的精确实例化,将算法设计从盲目试错提升为几何算子与动量流的系统组合。
  • 重新阐释权重衰减的第一性原理:打破了权重衰减仅为“控制权重范数防过拟合”的传统定见,在理论上将其严格等价于原始空间的步长递减迭代平均,推导出无需人工调节的 \(1/(k+2)\) 确定性退火解。
  • 工程友好的免调参 Wrapper:无需改写底层繁琐的 CUDA 算子与预处理,仅在参数更新层面外挂一层凸组合变换即可消除权重衰减超参数,显著压缩大模型训练的探路试错成本。

局限与展望

  • 理论分析局限于凸光滑假设:核心收敛定理基于强凸正则化与凸目标设定,对于非凸深度神经网络中损失曲面极端崎岖的几何机制尚缺少闭式非凸收敛率刻画。
  • 多周期训练(Multi-epoch)下的泛化局限:本文主要侧重无过拟合顾虑的单周期自监督预训练场景;在下游小样本微调或多周期训练中,传统防止过拟合的显式 \(\ell_2\) 正则化强度可能仍需额外保留。
  • 微调场景的适配潜力:鉴于 SODA 天然具备围绕参考点 \(z^0\) 锚定的特性,未来非常适合直接将预训练模型权重设为 \(z^0\) 开展参数高效微调(PEFT),有望天然抑制遗忘现象。

相关工作与启发

  • vs Lion (Chen et al., 2023):Lion 依赖进化搜索发掘符号更新与特定动量比例,但缺乏收敛机理解释且权重衰减极其敏感;SODA 从理论上证明 Lion 是基于 \(\ell_\infty\) 示性函数的乐观对偶平均特例,并为其提供了免调权重衰减的包装方案。
  • vs Muon (Jordan et al., 2024b):Muon 结合了矩阵谱范数收缩与 Nesterov 动量,但在首尾层仍需退回 AdamW 且超参数依赖网格调优;SODA† 用 Lion 替代 AdamW 并统一了所有层的动量参数,在大幅精简超参结构的同时取得更优性能。
  • vs Schedule-Free (Defazio et al., 2024):Schedule-Free 主要利用原始外推消除学习率调度;SODA 则将乐观机制融入对偶平均框架中,侧重解决权重衰减与几何正则化的统合,两者在原始平均机制上存在互补的理论同构性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从乐观对偶平均视角完整打通 Lion、Muon、NAdam 等 SOTA 优化器,并提出权重衰减即原始平均的颠覆性认知。
  • 实验充分度: ⭐⭐⭐⭐☆ 覆盖 64M 至 1B 参数量的大模型单周期预训练主实验与跨时域消融,若补足多周期下游任务微调则更为完备。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰优美,公式符号克制,算法与架构图对应紧密。
  • 价值: ⭐⭐⭐⭐⭐ 对大模型优化算法演进具有高指导价值,免调权重衰减的包装器极具工业落地潜力。