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. 核心分析模块
- 一步损失分解:在匹配的验证损失下比较两种优化器的单步损失下降量,发现 Muon 的一阶增益与 Adam 相近,但二阶曲率惩罚始终更小。
- 惩罚溯源:将曲率惩罚拆解为 更新范数的平方 与 归一化方向锐度(NDS),实验表明两者的更新范数相当,因此 Muon 的惩罚优势完全源于更低的 NDS,而不是更新尺度。
- 数据与结构驱动因素:
- 利用 Zipf-PCFG 控制数据不平衡程度,观察到数据不平衡显著放大 Muon 在 NDS 上的优势。
- 通过层内/跨层曲率分解,发现在训练中后期,Muon 较低 NDS 主要来自层内 Hessian 块的贡献,跨层部分贡献较小。
- 理论印证:构建一个结构化矩阵块二次模型,假设 Hessian 具有低秩 Kronecker 结构、可同时对角化以及显著的曲率异质性,证明在梯度对齐高曲率方向时,Muon 通过将更新能量在不同曲率组之间均衡分配,获得比 GD 更低的平均 NDS;当曲率异质性足够强时,经相同步数后局部二次损失也更低。
3. 输出与结论
实验与理论共同表明,Muon 在每步更新时能避开高曲率的尖锐方向,从而减少二阶损失惩罚,这种能力在数据分布失衡和模型层内曲率异质性强时更为突出。
跟同类方法的差异:不同于先前工作仅从经验上报告 Muon 的墙钟时间加速,本文首次从 curvature(曲率)与 NDS 角度提供可量化的几何解释,并设计了受控合成实验与理论二次模型来分离变量、验证因果。
实验
实验设计
论文在 语言模型训练 的设置下对比 Muon 与 Adam。主实验使用 FineWeb 数据集的主文本部分,在相同验证损失的前提下,通过 二阶 Taylor 展开 量化单步损失下降(分为一阶增益与二阶曲率惩罚)。曲率惩罚进一步分解为 更新范数平方 与 归一化方向锐度 (NDS)。为探究数据特性影响,作者构造了 Zipf‑PCFG 合成数据,通过调整 Zipf 指数控制数据不平衡程度,系统测量 NDS 差异。结构层面,将 Hessian 拆分为 层内块 与 跨层块,观察不同训练阶段 NDS 贡献变化。最后,在 结构化二次模型 上通过理论分析验证观察到的现象。
关键发现
- 在匹配验证损失时,Muon 的单步损失下降显著大于 Adam;两者一阶增益相近,但 Muon 的曲率惩罚一致更小。
- 更新范数对比显示二者尺度相当,因此 Muon 的优势完全来自 更低的 NDS——即更新方向更不“刺入”高曲率区域。
- 数据不平衡会放大 Muon 对 Adam 的 NDS 优势:Zipf 指数越大(不平衡越强),Muon 的曲率惩罚相对越低。
- 训练中后期,Muon 的低 NDS 主要由 层内曲率 贡献,表明其方向更适应层内 Hessian 结构。
- 理论分析证明:在曲率异质性足够强且梯度对齐高曲率模式时,Muon 通过在各曲率分组间平衡更新能量,获得低于 GD 的平均 NDS 和更低的局部二次损失。
与基线对比解读
Adam 通过自适应学习率调节更新尺度,但对 方向曲率 缺乏显式约束。Muon 本质是一种结构化更新(如利用矩阵低秩补全或 Newton‑Schulz 迭代),本文发现其 隐式降低了 NDS,在不牺牲更新幅度的前提下避开了高曲率方向。这与常见的“Muon 只是加速收敛”的表象不同:更低的 NDS 意味着更稳定的训练动力学,尤其在数据极度不平衡或模型层间曲率差异大的场景下优势更突出。这提示实际工程中除了关注损失曲线,还应监控 NDS 或类似曲率指标,并可能将方向正则化融入优化器设计。
行业影响
落地场景
Muon 优化器对大规模模型训练场景的加速效果(约 2 倍于 Adam)可直接应用于:
- 大语言模型(LLM)预训练与微调:例如通用对话模型、代码生成模型,在同等算力下可训练更多 tokens 或更大模型。
- 多模态模型训练:如文生图、视频生成模型的扩散过程或 VAE 训练。
- 推荐与广告排序模型:工业级推荐系统常需以周为单位滚动训练千万级稀疏特征的大模型,Muon 可缩短训练周期,更快响应数据分布变化。
商业价值
- 降本:GPU 小时数是 LLM 训练的最大支出。2 倍加速直接意味着 → 训练成本砍半,或同样预算下模型规模/数据量翻倍,边际效应显著。
- 加速迭代:更短的实验周期让算法团队能更快地验证想法、调参或进行消融实验,缩短产品上线时间,间接带来增收。
- 体验提升:论文指出 Muon 在匹配验证损失条件下能获得更大的一步损失下降,理论上可收敛到更优的泛化平坦区域,可能提升模型效果,但需业务方自行验证。
与现有产品/工作流的接口
Muon 是一个即插即用的替换项,无需改动模型架构或训练框架。集成要点:
- 替换优化器:在 PyTorch 或 JAX 训练脚本中将
AdamW替换为 Muon 实现(如 KellerJordan/Muon),并调整学习率等超参数(Muon 通常需要更小的 LR)。 - 保留现有流水线:梯度累积、混合精度训练、分布式策略(如 FSDP/ZeRO)均可直接兼容,仅需在更新参数时切换到 Muon 的
Newton-Schulz后处理。 - 监控与调优:论文揭示了 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 如何减小曲率惩罚,缺乏多步累积效应的讨论。实际训练中,优化路径的长期稳定性和泛化性能同样重要,单步局部优势能否转化为整体收敛加速仍待进一步研究。网络结构影响的层级分解也只停留在现象观察,未给出机制性理论。