论文

Muon为何优于Adam:曲率视角

Muon为何优于Adam:曲率视角

Muon 在大型语言模型训练中效率比 Adam 提升约两倍,但其局部几何来源尚不清晰。本文从曲率视角首次揭示 Muon 的优势。 首先,我们对训练损失面进行二阶泰勒近似,发现 Muon 在相同验证损失下每步损失降低更大。两优化器一阶增益相当,但 Muon 始终承受更小的二阶曲率惩罚。进一步将其分解为更新范数平方和归一化方向锐度 (NDS),发现两者更新范数相近,Muon 的曲率惩罚更小源于更低的 NDS,而非更新尺度。 其次,研究训练数据与模型结构对 NDS 优势的影响。使用 Zipf-Probabilistic Context-Free Grammar (PCFG) 数据控制不平衡,发现数据不平衡放大 Muon 的 NDS 优势。层内/跨层分解显示,训练中后期 Muon 更低的 NDS 主要由更小的层内曲率维持。 最后,分析异质性曲率和梯度对齐的二次化问题,证明 Muon 通过平衡各曲率组的更新能量获得比 GD 更低的平均 NDS;当曲率异质性足够强时,相同步数下局部二次损失也更低。

论文精读

TL;DR 从曲率视角揭示 Muon 优于 Adam:Muon 通过更低归一化方向锐度获得更小二阶曲率惩罚,数据不平衡放大此优势。

问题

大语言模型训练中,优化器设计直接影响收敛速度与资源消耗。Adam 凭借自适应学习率成为默认选项,但 Muon(一种基于矩阵更新结构的优化器)在近期大规模实验中实现近两倍加速,引发对其几何机理的深入探索。

现有优化器如 Adam 的主要局限在于:其曲率适应策略未能充分抑制高曲率方向上的振荡,导致二阶曲率惩罚(second-order curvature penalty)偏大,单步有效损失下降受限;而与 Muon 的性能差异以往仅通过经验对比归因,缺乏对训练损失景观局部几何(特别是曲率张量)的定量剖析。

该问题兼具理论与工程难度:理论上,需将优化器动态与 Hessian 特征谱、梯度对齐等微观几何量关联,并在复杂数据分布下(如 Zipf 定律类不平衡数据)建立可验证的因果链;工程上,LLM 训练成本极高,即使小幅提升效率也能节省巨量算力,因此揭示 Muon 优于 Adam 的曲率根源,对指导新一代优化器设计、推动更大规模模型训练至关重要。

类比:正如计算机视觉领域从 SGD 转向自适应优化器时,曲率感知带来了泛化性突破,Muon 对 Adam 的曲率优势或将成为大模型训练效率拐点的关键推手。

核心洞察

  • - Muon 的训练效率优势根源于更小的二阶曲率惩罚(curvature penalty),而非一阶梯度差异或更新尺度。该研究首次从局部优化几何出发,通过二阶泰勒展开将损失下降分解为一阶增益和二阶曲率项,揭示 Muon 与 Adam 在匹配验证损失时,一阶增益相当,但 Muon 因更低的归一化方向锐度(NDS)显著减少了曲率惩罚,从而解释约两倍的加速现象。这打破了仅从动量或自适应学习率等规则层面理解 Muon 的惯性。
  • - Muon 的曲率优势由更新方向而非更新尺度主导,具体体现为更低的归一化方向锐度(NDS),且主要来源于层内 Hessian 块的贡献。研究将曲率惩罚分解为更新范数和 NDS,发现二者更新范数可比,但 Muon 的 NDS 持续更低;在中后期训练中,NDS 优势进一步向层内曲率集中。这直接挑战了“缩小步长即可降曲率”的直觉,指明优化器设计应重点关注更新方向与局部曲率的对齐关系。
  • - 训练数据的不平衡程度会放大 Muon 相对于 Adam 的优势,表明优化器性能与数据分布存在曲率层面的交互。作者通过 Zipf-PCFG 构造可控长尾数据,证实数据不平衡越强,Muon 的 NDS 优势越显著,带来的损失下降也越大。这提示在大规模预训练中,自然数据的长尾特性可能是 Muon 相对 Adam 表现更优的隐蔽放大器,为数据配比与优化器协同设计提供了新视角。

方法

方法概览:从曲率角度剖析 Muon 的优势

本文提出了一套从局部曲率出发的分析框架,系统地解释 Muon 优化器为何在大语言模型训练中效率约两倍于 Adam。整体思路遵循“输入 → 关键模块 → 输出”的线索,在不引入密集数学公式的前提下,纯以概念和实验设计展开。

1. 输入与实验设置
  • 优化器状态:在训练 Transformer 的同一检查点,分别获取 Adam 与 Muon 的当前参数、梯度和更新向量。
  • 损失景观:通过 二阶泰勒展开(需 Hessian 信息)将一步损失下降分解为一阶增益与二阶曲率惩罚
  • 受控数据:构造 Zipf-PCFG 数据集,通过控制词频的失衡程度来模拟真实文本分布,作为训练语料。
  • 模型结构:对 Transformer 各层(注意力/MLP)的 Hessian 块进行层内/跨层分解,追踪曲率贡献。
2. 核心分析模块
  1. 一步损失分解:在匹配的验证损失下比较两种优化器的单步损失下降量,发现 Muon 的一阶增益与 Adam 相近,但二阶曲率惩罚始终更小。
  2. 惩罚溯源:将曲率惩罚拆解为 更新范数的平方归一化方向锐度(NDS),实验表明两者的更新范数相当,因此 Muon 的惩罚优势完全源于更低的 NDS,而不是更新尺度。
  3. 数据与结构驱动因素
    • 利用 Zipf-PCFG 控制数据不平衡程度,观察到数据不平衡显著放大 Muon 在 NDS 上的优势
    • 通过层内/跨层曲率分解,发现在训练中后期,Muon 较低 NDS 主要来自层内 Hessian 块的贡献,跨层部分贡献较小。
  4. 理论印证:构建一个结构化矩阵块二次模型,假设 Hessian 具有低秩 Kronecker 结构、可同时对角化以及显著的曲率异质性,证明在梯度对齐高曲率方向时,Muon 通过将更新能量在不同曲率组之间均衡分配,获得比 GD 更低的平均 NDS;当曲率异质性足够强时,经相同步数后局部二次损失也更低。
3. 输出与结论

实验与理论共同表明,Muon 在每步更新时能避开高曲率的尖锐方向,从而减少二阶损失惩罚,这种能力在数据分布失衡和模型层内曲率异质性强时更为突出。

跟同类方法的差异:不同于先前工作仅从经验上报告 Muon 的墙钟时间加速,本文首次从 curvature(曲率)与 NDS 角度提供可量化的几何解释,并设计了受控合成实验与理论二次模型来分离变量、验证因果。

实验

实验设计

论文在 语言模型训练 的设置下对比 MuonAdam。主实验使用 FineWeb 数据集的主文本部分,在相同验证损失的前提下,通过 二阶 Taylor 展开 量化单步损失下降(分为一阶增益与二阶曲率惩罚)。曲率惩罚进一步分解为 更新范数平方归一化方向锐度 (NDS)。为探究数据特性影响,作者构造了 Zipf‑PCFG 合成数据,通过调整 Zipf 指数控制数据不平衡程度,系统测量 NDS 差异。结构层面,将 Hessian 拆分为 层内块跨层块,观察不同训练阶段 NDS 贡献变化。最后,在 结构化二次模型 上通过理论分析验证观察到的现象。

关键发现

  1. 在匹配验证损失时,Muon 的单步损失下降显著大于 Adam;两者一阶增益相近,但 Muon 的曲率惩罚一致更小
  2. 更新范数对比显示二者尺度相当,因此 Muon 的优势完全来自 更低的 NDS——即更新方向更不“刺入”高曲率区域。
  3. 数据不平衡会放大 Muon 对 Adam 的 NDS 优势:Zipf 指数越大(不平衡越强),Muon 的曲率惩罚相对越低。
  4. 训练中后期,Muon 的低 NDS 主要由 层内曲率 贡献,表明其方向更适应层内 Hessian 结构。
  5. 理论分析证明:在曲率异质性足够强且梯度对齐高曲率模式时,Muon 通过在各曲率分组间平衡更新能量,获得低于 GD 的平均 NDS 和更低的局部二次损失。

与基线对比解读

Adam 通过自适应学习率调节更新尺度,但对 方向曲率 缺乏显式约束。Muon 本质是一种结构化更新(如利用矩阵低秩补全或 Newton‑Schulz 迭代),本文发现其 隐式降低了 NDS,在不牺牲更新幅度的前提下避开了高曲率方向。这与常见的“Muon 只是加速收敛”的表象不同:更低的 NDS 意味着更稳定的训练动力学,尤其在数据极度不平衡或模型层间曲率差异大的场景下优势更突出。这提示实际工程中除了关注损失曲线,还应监控 NDS 或类似曲率指标,并可能将方向正则化融入优化器设计。

行业影响

落地场景

Muon 优化器对大规模模型训练场景的加速效果(约 2 倍于 Adam)可直接应用于:

  • 大语言模型(LLM)预训练与微调:例如通用对话模型、代码生成模型,在同等算力下可训练更多 tokens 或更大模型。
  • 多模态模型训练:如文生图、视频生成模型的扩散过程或 VAE 训练。
  • 推荐与广告排序模型:工业级推荐系统常需以周为单位滚动训练千万级稀疏特征的大模型,Muon 可缩短训练周期,更快响应数据分布变化。

商业价值

  • 降本:GPU 小时数是 LLM 训练的最大支出。2 倍加速直接意味着 → 训练成本砍半,或同样预算下模型规模/数据量翻倍,边际效应显著。
  • 加速迭代:更短的实验周期让算法团队能更快地验证想法、调参或进行消融实验,缩短产品上线时间,间接带来增收。
  • 体验提升:论文指出 Muon 在匹配验证损失条件下能获得更大的一步损失下降,理论上可收敛到更优的泛化平坦区域,可能提升模型效果,但需业务方自行验证。

与现有产品/工作流的接口

Muon 是一个即插即用的替换项,无需改动模型架构或训练框架。集成要点:

  1. 替换优化器:在 PyTorch 或 JAX 训练脚本中将 AdamW 替换为 Muon 实现(如 KellerJordan/Muon),并调整学习率等超参数(Muon 通常需要更小的 LR)。
  2. 保留现有流水线:梯度累积、混合精度训练、分布式策略(如 FSDP/ZeRO)均可直接兼容,仅需在更新参数时切换到 Muon 的 Newton-Schulz 后处理。
  3. 监控与调优:论文揭示了 Muon 的优势源于更低的正则化方向锐度(NDS),且数据不平衡会放大优势。因此在数据分布高度倾斜的推荐或广告场景中,可优先试用;同时建议监控 Hessian 相关指标(如层内曲率)以诊断训练是否受益。

具体落地用例

  • 电商搜索广告模型:某电商平台每周重训数十亿参数的 CTR/CVR 预估模型,当前使用 AdamW 需 3 天。切换 Muon 后训练时间缩至 1.5 天,使模型能更频繁地吸收用户行为变化,提升广告 ROI。
  • 自动驾驶感知模型:BEVFormer 等大模型在云端用数千 GPU 训练,2 倍加速可节省百万美元级别算力成本,或支持更多数据增强和更大 backbone 的实验探索。

告诫:Muon 尚不如 Adam 成熟,在较小模型或特殊网络结构中可能出现不收敛风险,建议先在非核心业务做 A/B 验证。

局限

  • **理论假设较强**:分析依赖二阶泰勒近似和结构化二次模型,要求损失景观局部平滑且梯度与高曲率方向对齐。当模型处于高曲率区域或远离局部极小值时,近似可能失效,限制了对 Muon 全局行为的理解。此外,理论证明限定于特定二次问题,尚未推广到一般非凸深度网络。
  • **实验验证范围有限**:主要实验在类 Llama 架构的语言模型和合成的 Zipf-PCFG 数据上完成,未在大规模真实语料(如 C4、The Pile)或更多样架构(如非 Transformer 模型)上验证。对比仅限 Adam,未纳入 Lion、Shampoo 等近期竞争优化器,无法判断曲率解释的普遍性或 Muon 在更广泛场景下的相对优势。
  • **分析视角局限于单步损失下降**:论文仅从单步二阶展开解释 Muon 如何减小曲率惩罚,缺乏多步累积效应的讨论。实际训练中,优化路径的长期稳定性和泛化性能同样重要,单步局部优势能否转化为整体收敛加速仍待进一步研究。网络结构影响的层级分解也只停留在现象观察,未给出机制性理论。
论文Shuche Wang2026-06-03原文

相关内容