跳转至

Why Muon Outperforms Adam: A Curvature Perspective

会议: NeurIPS 2026
arXiv: 2606.04662
OpenReview: https://openreview.net/forum?id=KWe4ab5CCr
论文: https://arxiv.org/abs/2606.04662
代码: https://github.com/KellerJordan/modded-nanogpt
领域: 优化/理论
关键词: Muon优化器、Adam、曲率分析、归一化方向锐度、Hessian谱

一句话总结

本文首次从二阶曲率视角解密大模型优化器 Muon 相比 Adam 的性能优势,证明两者的步长与一阶梯度增益相当,但 Muon 的谱归一化更新方向遭遇显著更小的归一化方向锐度(NDS),从而大幅降低二阶曲率惩罚并在结构化二次曲面上实现跨模态能量平衡。

研究背景与动机

近年来在大型语言模型预训练中,针对矩阵参数设计的优化器 Muon(Momentum Orthogonalized by Newton-Schulz)展现出惊人的训练效率,在各类模型规模下均能实现相比传统 Adam 约 2 倍的端到端加速。Muon 通过对梯度动量矩阵进行谱归一化(将非零奇异值统一重缩放至同一尺度),充分利用了 Transformer 权重的矩阵几何结构。然而,学术界对 Muon 为何能大幅超越 Adam 的局部几何机理仍缺乏清晰认知。近期虽有研究尝试从联想记忆容量和数据长尾分布的宏观视角提供机制解释,但在优化动力学的微观视角下,参数空间局部地形究竟如何对 Muon 和 Adam 产生异构影响依然是一个未解之谜。

这一探索的核心挑战在于,深度神经网络尤其是 Transformer 的损失曲面具有极高维度且高度非凸,直接计算全参数 Hessian 矩阵并在动态轨迹中进行横向对比面临巨大的计算壁垒。过去关于自适应优化器的研究多聚焦于坐标维度的预条件缩放或全局锐度指标,无法有效解耦“更新步长大小”与“更新几何方向”对实际单步损失下降的贡献差异。同时,训练数据的重尾分布与 Transformer 的多层级联结构如何与优化器更新方向交互,亦缺乏系统性的局部曲率量化框架。

本文的切入角度是回到优化地形的局部二阶泰勒展开,在严格对齐验证损失的公平基准下,系统拆解单步优化收益中的一阶梯度对齐项与二阶曲率惩罚项。核心 idea:通过将局部曲率惩罚解耦为更新范数与归一化方向锐度(NDS),揭示 Muon 的加速优势并非来自更大的有效步长或一阶梯度收益,而是源于谱归一化使更新方向天然避开了极端尖锐方向;结合 Kronecker 结构化二次模型,在理论上证明 Muon 能够在高低曲率模态间均匀分配更新能量,从而避免了 Adam 与梯度下降因梯度对齐尖锐模态而导致的严重超调。

方法详解

整体框架

本文构建了一套结合实证局部几何解耦与结构化二次优化的综合分析框架。针对真实 LLM 预训练轨迹(124M NanoGPT 于 FineWeb 数据集),框架首先利用二阶泰勒展开将单步损失下降精确分解为一阶项与二阶曲率项;进而将二阶项拆解为更新范数与归一化方向锐度(NDS),在排除步长尺度干扰后定位方向优势。随后,框架从数据(受控 Zipf-PCFG 不平衡语料)和模型(层内与层间 Hessian 块)两个维度解析 NDS 优势的生成与演化规律。最后,提炼出符合 Transformer 注意力权重真实曲面特性的结构化二次模型,给出严格的收敛性与方向锐度理论证明。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["参数矩阵与小批量梯度输入"] --> B["单步损失二阶泰勒展开<br/>解耦一阶收益与二阶曲率惩罚"]
    B --> C["归一化方向锐度与步长解耦<br/>分离更新范数与方向锐度NDS"]
    C --> D["层内与层间曲率解耦<br/>定位层内主导与边界层聚集效应"]
    D --> E["结构化二次模型理论证明<br/>谱归一化平衡高低曲率模态能量"]
    E --> F["输出优化机理与多步收敛保证"]

关键设计

1. 单步损失二阶泰勒展开:解耦一阶梯度收益与二阶曲率惩罚 为了从局部优化地形厘清 Muon 相比 Adam 的优势来源,该设计通过对小批量损失函数 \(\mathcal{L}_D\) 在当前权重 \(W\) 处进行二阶泰勒展开,将单步实际损失下降量 \(\Delta_D(W, Z) = \mathcal{L}_D(W) - \mathcal{L}_D(W - Z)\) 严格分解为一阶线性下降增益与二阶曲率代价两部分: $\(\Delta_D(W, Z) \approx \langle G, Z \rangle - \frac{1}{2} \langle Z, \mathcal{H}[Z] \rangle = I_D^{(1)}(W, Z) - I_D^{(2)}(W, Z)\)$ 其中 \(G = \nabla_W \mathcal{L}_D(W)\) 为随机小批量梯度,\(\mathcal{H} = \nabla_W^2 \mathcal{L}_D(W)\) 为局部 Hessian 算子,\(I_D^{(1)}\) 代表沿着更新方向 \(Z\) 移动所获得的线性下降量,而 \(I_D^{(2)}\) 则代表损失曲面局部弯曲所导致的曲率惩罚。在 124M NanoGPT 预训练轨迹中,二阶近似展现出极高的预测保真度。实测对比表明,在严格对齐的验证损失下,Adam 与 Muon 的一阶收益 \(I_D^{(1)}\) 几乎处于相同水平;然而,Adam 的二阶曲率惩罚曲线全程显著高于 Muon。这证明 Muon 更大的单步损失下降完全由显著更低的二阶曲率损耗所驱动。

2. 归一化方向锐度与步长解耦:剥离步长尺度定位几何方向本质 曲率惩罚项 \(I_D^{(2)}(W, Z) = \frac{1}{2} \langle Z, \mathcal{H}[Z] \rangle\) 同时受更新步长的绝对尺度 \(\|Z\|_F^2\) 与更新向量所指方向的曲率影响。为彻底排除由于步长设置不同可能带来的假象,该设计定义了尺度无关的归一化方向锐度(Normalized Directional Sharpness, NDS): $\(\mathcal{S}_F(W; Z) = \frac{\langle Z, \mathcal{H}[Z] \rangle}{\|Z\|_F^2}\)$ 从而实现因式分解 \(I_D^{(2)}(W, Z) = \frac{1}{2} \|Z\|_F^2 \cdot \mathcal{S}_F(W; Z)\)。实验追踪表明,Adam 与 Muon 的更新 Frobenius 范数平方比值在整个训练进程中恒定徘徊在 1.0 附近,说明两者的有效步长尺度高度一致。相反,Adam 与 Muon 的曲率惩罚比率与 NDS 比率曲线几乎完全重合(对齐验证损失下平均 NDS 比率为 1.76,对齐训练步数下平均比率高达 2.94)。这确凿地证实:Muon 较小的曲率惩罚完全由其更新方向所诱导的极低 NDS 所决定,Muon 探索出的矩阵更新轨迹天然穿行于地形更平缓的低曲率峡谷。

3. 层内与层间曲率解耦:刻画多层深度架构下的锐度动态迁移 为探究 Transformer 架构各层参数如何交互并构筑全模型 NDS 优势,该设计将全网络权重更新 \(Z = (Z_1, \dots, Z_L)\) 及 Hessian 算子分解为主对角块(层内曲率 \(\mathcal{H}_{\ell\ell}\))与非对角块(跨层交互曲率 \(\mathcal{H}_{\ell\ell'}\)): $\(\mathcal{S}_F(W_t; Z_t) = \mathcal{S}_F^{\text{within}}(W_t; Z_t) + \mathcal{S}_F^{\text{cross}}(W_t; Z_t) = \sum_{\ell=1}^L \frac{\langle Z_{t,\ell}, \mathcal{H}_{\ell\ell}[Z_{t,\ell}] \rangle}{\|Z_t\|_F^2} + \sum_{\ell \neq \ell'} \frac{\langle Z_{t,\ell}, \mathcal{H}_{\ell\ell'}[Z_{t,\ell'}] \rangle}{\|Z_t\|_F^2}\)$ 量化追踪揭示出两项关键动态:首先,随着训练向后推进,Muon 的跨层交互曲率项衰减速度远快于层内项,使得层内曲率贡献比 \(\rho_t^{\text{within}} = \mathcal{S}_F^{\text{within}} / \mathcal{S}_F\) 从训练初期的约 14% 急剧上升至中后期的 44%(接近翻三倍),而 Adam 的层内占比则稳定保持在 30% 上下;其次,在层级空间分布上,Adam 与 Muon 的层内 NDS 差距呈现极强的两头聚集性,约 70% 的锐度差值集中于最前端的输入层 L1 与最末端的输出层 L12,中间层仅占 2% 左右。这说明 Muon 的全局曲率优势在预训练中后期主要由单层权重内部的矩阵谱结构所主导。

4. 结构化二次模型理论证明:谱归一化平衡高低曲率模态能量 基于层内块 Hessian 在真实 LLM 中展现出的低 Kronecker 秩特性(秩为 4 时解释超 85% 能量)、正交近似联合对角化特性(JADE 得分超 0.84)以及极端长尾曲率谱(正曲率比值 \(w_1 / w_{88} \approx 2.59 \times 10^6\)),该设计构建了典型二次问题 \(\mathcal{Q}(Y) = \mathcal{L}(W_0) - \langle G, Y \rangle + \frac{1}{2} \langle Y, \mathcal{H}[Y] \rangle\)。理论证明显示:梯度由于内生对齐机制主要投射在高曲率特征模态上(\(\sigma_H > \sigma_L\));梯度下降(GD)与 Adam 的更新方向直接正比于残差梯度,导致过多更新能量堆积在极少数极尖锐的方向上,诱发严重超调与往复振荡。与之相反,Muon 的谱归一化更新 \(Z_t^{\text{Muon}} = \sum_i \text{sgn}(r_{i,t}) M_i\) 强制抹平了奇异值模态间的幅度差异,将更新能量严格均匀分配给高低曲率方向,使得方向锐度恒定为无偏平均 \(\mathcal{S}_F(Z_t^{\text{Muon}}) = \alpha w_H + (1-\alpha) w_L\)。当曲率异质性充分强且尖锐组占比满足 \(\rho + 1 > 1/\alpha > 1 + \sigma_H / \sigma_L\) 时,Muon 的残差收缩因子严格优于 GD 的交替收缩率,在任意有限步数 \(T \ge 1\) 下均实现更低的局部损失。

实验关键数据

主实验

主实验在 124M NanoGPT(FineWeb-10B 语料)以及结构化二次曲面上对比了 Muon、Adam 与 GD 的各项局部几何与损失指标。

实验设置 / 优化问题 评价指标 Muon Adam / GD 基线 相对差异 / 优势
NanoGPT 124M (对齐验证损失) 一阶损失增益 \(I_D^{(1)}\) 约 \(1.4 \times 10^{-2}\) 约 \(1.4 \times 10^{-2}\) (Adam) 一阶增益基本持平 (差值 < 3%)
NanoGPT 124M (对齐验证损失) 二阶曲率惩罚 \(I_D^{(2)}\) 约 \(0.4 \times 10^{-2}\) 约 \(0.7 \times 10^{-2}\) (Adam) Muon 曲率惩罚降低约 42.8%
NanoGPT 124M (对齐验证损失) 更新范数平方 \(\|Z_t\|_F^2\) 基准 1.00 1.01 (Adam/Muon 比值) 步长尺度完全一致
NanoGPT 124M (对齐验证损失) 归一化方向锐度 NDS \(\mathcal{S}_F\) 基准 1.00 1.76 (Adam/Muon 比值) Muon 方向锐度低 43.2%
NanoGPT 124M (对齐训练步数) 归一化方向锐度 NDS \(\mathcal{S}_F\) 基准 1.00 2.94 (Adam/Muon 比值) 同步数下 Muon 锐度低 66.0%
结构化二次优化仿真 (Synthetic) 归一化方向锐度 NDS 比值 0.09 (vs GD) 0.50 (Adam vs GD) Muon 较 GD 锐度小 11 倍
结构化二次优化仿真 (Synthetic) 最优单步损失下降比值 5.60 (vs GD) 1.22 (Adam vs GD) Muon 单步损失下降达 GD 5.6 倍

消融实验

消融实验包含两部分:一是在受控 Zipf-PCFG 语言生成任务中检验数据不平衡度参数 \(s \in \{0, 0.5, 1.0\}\) 对 NDS 的影响;二是对 12 层 Transformer 展开层内与跨层锐度贡献分解。

实验配置 / 分解维度 测试状态 / 参数设置 关键指标取值 (Muon) 关键指标取值 (Adam) 说明与差异
Zipf-PCFG 数据不平衡 \(s = 0.0\) (均衡分布) 归一化 NDS = 1.00 归一化 NDS = 1.63 锐度差值 \(\Delta(0) = 0.63\)
Zipf-PCFG 数据不平衡 \(s = 0.5\) (中度偏斜) 归一化 NDS = 1.12 归一化 NDS = 1.95 锐度差值 \(\Delta(0.5) = 0.83\)
Zipf-PCFG 数据不平衡 \(s = 1.0\) (重度长尾) 归一化 NDS = 1.25 归一化 NDS = 2.38 锐度差值扩大至 \(\Delta(1) = 1.13\) (增幅 1.8x)
层级分解:层内贡献占比 训练初期阶段 \(\rho^{\text{within}} \approx 14\%\) \(\rho^{\text{within}} \approx 27\%\) 早期跨层交互占主导地位
层级分解:层内贡献占比 训练中后期阶段 \(\rho^{\text{within}} \approx 44\%\) \(\rho^{\text{within}} \approx 34\%\) Muon 层内占比激增,优势内化至单层矩阵
跨层空间分布差异定位 首尾边界层 (L1, L12) 贡献比 \(\approx 70\%\) - 直接耦合词表与输出 logits,曲率差最剧烈
跨层空间分布差异定位 中间隐藏层 (L2–L7) 贡献比 \(\approx 2\%\) - 中间特征层两者方向锐度表现接近

关键发现

  • 数据不平衡性是拉大优化器曲率鸿沟的核心放大器:当语料词频幂律指数从 0 提升到 1 时,Adam 的 NDS 暴增 46%(1.63 升至 2.38),而 Muon 仅轻微上升 25%(1.00 升至 1.25),两者的锐度差距单调扩张了 1.8 倍。这解释了为何在长尾效应更显著的海量自然语言真实语料中 Muon 的提速倍率更为突出。
  • 架构层面的锐度演化具有“阶段迁移”与“端点聚焦”两大特征:Muon 在前期的跨层协同曲率下降迅速,中后期主要依靠极低单层内部曲率压制全局 NDS;而在深层空间中,与数据输入输出直连的 L1 和 L12 承载了绝大部分曲率收益。
  • 二次结构化验证了 Hessian 的物理规律:注意力权重 Hessian 具备极其显著的低 Kronecker 秩结构与正交联合可对角化性,其高曲率方向呈现极端稀疏与尖锐(前 128 个模态中 88 个正模态,最大最小曲率比超两百万倍),且超过 87% 的梯度能量内生聚集在正曲率模态内。

亮点与洞察

  • 颠覆了“优化器靠更大有效步长或更高梯度信噪比提速”的传统直觉,明确指出 Muon 的根本性优势在于其更新方向诱导的极低归一化方向锐度(NDS),为大模型矩阵优化器研究开辟了局域二阶几何分析新范式。
  • 揭示了谱归一化(Newton-Schulz 迭代)的隐式平滑机制:其本质是剥离了梯度在外层极端特征方向上的过度能量富集,通过将所有有效特征模态幅度强制置为 1,实现了在高曲率阻尼与低曲率推进之间的无偏动力学平衡。
  • 解释了近期增强版 Muon 变体(如向平坦 Hessian 方向倾斜更新分量)的成功根源:任何能够进一步降低更新向量在尖锐曲率上投影权重的设计,都能直接缩减二阶曲率惩罚并扩大单步损失净收益。

局限与展望

  • 实验与理论模型当前主要立足于自回归因果 Transformer 架构,对扩散模型(Diffusion Models)、卷积骨干或多模态交叉架构的普适性验证有待拓展。
  • 理论分析基于简化的两级曲率分组(高曲率组与低曲率组)及无动量(\(\mu=0\))设定,未来可进一步拓展至连续幂律连续谱分布与非零动量累加的完整动态系统分析。
  • 实际训练中全量 Hessian 计算代价极大,未来可探索将轻量级 NDS 在线估计作为自适应超参数调节或动量混合比例的动态触发信号。

相关工作与启发

  • vs Adam (Kingma & Ba, 2014): Adam 采用坐标维度的二阶矩除法实现分量归一化,但忽略了权重矩阵固有的二维奇异谱结构,在高度各向异性的曲面中易沿尖锐主轴过度推进;Muon 在矩阵流形上通过奇异值谱归一化实现了几何自适应,在相同步长下遭遇的曲率惩罚大幅缩减。
  • vs 联想记忆机制解释 (Wang et al., 2025b): 现有理论多从表征学习、注意力头联想记忆均衡与长尾词表检索宏观切入;本文则深入最底层的参数空间微观几何,为 Muon 在长尾数据下性能卓越的现象提供了精准的二阶曲率(NDS 随不平衡度单调分离)硬核物理支撑。
  • vs K-FAC 与二阶预条件算法 (Martens & Grosse, 2015): K-FAC 显式利用 Kronecker 分解估算 Fisher 阵并求逆,系统开销与显存代价极其高昂;Muon 则利用廉价的 Newton-Schulz 多项式迭代完成矩阵符号正交化,在完全不计算逆矩阵的前提下,自然达成了在正交曲率模态上分配等量更新能量的拟二阶优化效果。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从二阶泰勒展开与局部曲率视角揭秘 Muon 超越 Adam 的微观动力学机制,立意极佳且视角独到。
  • 实验充分度: ⭐⭐⭐⭐⭐ 兼具 124M 真实 LLM 预训练对齐量化、受控 Zipf-PCFG 不平衡语料实验及精细的层内/跨层块 Hessian 谱分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑链条严丝合缝,符号系统规范统一,实证观察与二次模型数学证明高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为大模型训练下一代优化器的设计、理解与性能诊断提供了极具启发性的理论基石与量化工具。