为什么 Muon 优于 Adam：曲率视角的解释 Why Muon Outperforms Adam: A Curvature Perspective
Muon 因归一化方向锐度更低而获得更小的曲率惩罚
前置知识
归一化方向锐度 (NDS)
衡量优化更新方向上的局部曲率，定义为 SF(W; Z) = ⟨Z, H[Z]⟩ / ‖Z‖_F^2，其中 H 是 Hessian 算子，Z 是更新方向。这个指标归一化了更新规模，只保留方向信息，数值越小说明更新方向越平滑。
本文的核心发现是 Muon 的优势来自更低的 NDS，理解这个概念才能读懂论文的曲率分析框架。
谱归一化
Muon 优化器的核心技术，对梯度动量矩阵 B_t = U_t Σ_t V_t^⊤ 进行奇异值分解，然后将所有非零奇异值归一化为 1，得到更新方向 Z_t^Muon = η_t U_t V_t^⊤。这使得更新在所有奇异值方向上具有相同幅度。
谱归一化是 Muon 与 Adam 的根本区别，也是其能够平衡高曲率和低曲率方向更新能量的关键机制。
曲率惩罚
在二阶泰勒展开中，一步损失下降 Δ_D(W, Z) ≈ ⟨G, Z⟩ - (1/2)⟨Z, H[Z]⟩，其中 I^(2)_D(W, Z) = (1/2)⟨Z, H[Z]⟩ 就是曲率惩罚。它表示因为更新方向上的局部曲率而损失的一阶收益。
论文证明 Muon 的优势主要来自更小的曲率惩罚，而不是更大的一阶增益，这是理解论文贡献的基础。
Hessian 算子
二阶导数算子 H = ∇²_W L_D(W)，作用于矩阵扰动 Z 得到 H[Z]，其矩阵表示在向量化后满足 vec(H[Z]) = mat(H)vec(Z)。Hessian 描述了损失函数在参数空间中的局部曲率结构。
本文通过 Hessian 的 Kronecker 结构和对角化性质构建理论模型，理解 Hessian 是跟随理论部分的关键。
研究动机
Muon 优化器在大语言模型预训练中相比 Adam 优化器可以实现约 2 倍的训练效率提升，这一优势在多种模型规模和配置下都得到了验证。然而，这种优势的几何来源仍然不清楚。现有研究从联想记忆和数据长尾分布的角度尝试解释 Muon 的优势，但这些解释无法说明为什么在相同的参数点上，Muon 和 Adam 会产生不同的优化效果。具体来说，在 124M 参数的 NanoGPT 模型上训练 FineWeb 数据集时，Muon 和 Adam 在相同的验证损失水平下表现出一阶损失下降相当，但总损失下降存在明显差异，这意味着存在一阶分析无法捕捉的机制。
本文的目标是本文的目标是从曲率视角理解 Muon 优于 Adam 的根本原因。具体而言，作者希望回答两个核心问题：第一，什么景观属性支持 Muon 相对于 Adam 的优势；第二，训练数据和模型结构如何影响这一属性。通过二阶泰勒展开分析训练损失的一步下降，作者试图识别出 Muon 优势的曲率来源，并量化数据不平衡和层间交互对这个来源的影响程度。
与已有工作不同的是，本文的独特切入角度是采用局部、依赖于优化器的曲率分析视角。与现有研究关注全局锐度、平坦度或曲率感知训练规则不同，作者认识到 Muon 和 Adam 在相同的参数点会诱导不同的更新矩阵，因此需要比较每个优化器更新方向上的 Hessian 曲率。这种视角的差异使得作者能够发现优化器差距出现在二阶项中，即使一阶增益相似，从而为理解优化器性能差异提供了新的理论框架。
核心方法
方法整体思路是通过二阶泰勒展开分解一步损失下降，识别出 Muon 优势的曲率来源。具体而言，作者首先对训练损失进行局部二次近似，将损失下降分解为一阶项（梯度对齐）和二阶项（曲率惩罚）。然后通过实验验证 Muon 和 Adam 的一阶下降相当，但曲率惩罚存在显著差异。接着作者将曲率惩罚分解为更新范数和归一化方向锐度（NDS），证明 Muon 的更小曲率惩罚来自更低的 NDS 而非更小的更新规模。最后，作者通过合成数据实验和层内/层外分解分析数据不平衡和模型结构对 NDS 的影响，并通过结构化二次模型的理论分析支持经验发现。
核心创新点是识别出 Muon 的谱归一化更新方向具有更低的归一化方向锐度（NDS），从而产生更小的二阶曲率惩罚。与 Adam 的坐标归一化不同，Muon 的谱归一化将更新能量均匀分配到所有奇异值方向上，避免了过度集中在高曲率方向。在异质曲率条件下，当梯度与高曲率方向对齐时，梯度下降会继承这种偏向，在高曲率方向投入过多更新能量，导致更大的方向锐度和曲率惩罚。Muon 的谱归一化消除了这种幅度不平衡，通过平衡高曲率和低曲率方向的更新能量，实现了更小的平均 NDS 和更大的累积损失下降。
方法步骤详情
方法步骤包括四个主要部分。第一步是实验设置，在 124M 参数 NanoGPT 模型上训练 FineWeb 数据集，通过网格搜索选择 Adam 和 Muon 的最优学习率。第二步是损失下降分解，对每个优化器的更新方向 Z_t，计算实际损失下降 Δ_D(W_t, Z_t)、预测下降 I^(1)_D(W_t, Z_t) - I^(2)_D(W_t, Z_t)、一阶下降 I^(1)_D(W_t, Z_t) = ⟨G, Z_t⟩ 和曲率惩罚 I^(2)_D(W_t, Z_t) = (1/2)⟨Z_t, H[Z_t]⟩，并在匹配的验证损失水平上比较两个优化器。第三步是 NDS 分析，计算归一化方向锐度 SF(W_t; Z_t) = ⟨Z_t, H[Z_t]⟩ / ‖Z_t‖_F^2 和更新范数 ‖Z_t‖_F，验证曲率惩罚的来源。第四步是影响因素分析，通过 Zipf-PCFG 合成数据在不同不平衡水平 s ∈ {0, 0.5, 1} 下训练，计算轨迹平均 NDS，并通过层内和层外 Hessian 块分解分析不同层对 NDS 的贡献。
技术新颖性
技术新颖性体现在三个方面。第一是首次从局部曲率角度系统分析 Muon 相对于 Adam 的优势，引入 NDS 作为关键指标，将优化器的损失下降分解为一阶项和曲率惩罚。第二是将曲率惩罚进一步分解为更新范数和 NDS，精确识别出 Muon 优势的来源是方向而非规模。第三是通过合成数据实验和层内/层外分解，量化了数据不平衡和模型结构对 NDS 的影响，揭示了 Muon 优势在训练中后期主要由层内曲率驱动。此外，作者提出了结构化二次模型，通过四个可验证的假设，在理论上证明了 Muon 在任何有限水平上具有更小的平均 NDS，在异质曲率足够强时实现更大的损失下降。
实验结果
核心发现包括四个方面。第一，Muon 在匹配验证损失下实现了比 Adam 更大的一步损失下降，这个差距主要由更小的二阶曲率惩罚驱动，而不是更大的一阶增益。在 124M 参数 NanoGPT 模型上，Muon 的预测损失下降与实际损失下降接近，而 Adam 的预测值与实际值匹配更好。第二，Muon 和 Adam 的更新范数相当，但 Muon 的 NDS 明显更低。实验显示 Adam 对 Muon 的 NDS 比率平均为 1.76，而更新范数的 ‖Z_t‖_F^2 比率接近 1，证明曲率惩罚差距几乎完全由 NDS 差距解释。第三，数据不平衡加剧了 Muon 相对于 Adam 的 NDS 优势。在 Zipf 指数 s 从 0 增加到 1 时，Adam 的归一化 NDS 从 1.63 上升到 2.38，Muon 只从 1.00 上升到 1.25，NDS 差距从 0.63 扩大到 1.13，增加了 1.8 倍。第四，Muon 的 NDS 优势在训练中后期主要由层内曲率驱动。层内贡献比例从早期约 14% 上升到后期约 44%，而 Adam 的层内比例从 27% 稳定增长到 34%。二次模型实验显示 GD 和 Adam 在 NDS 和损失下降上行为相似，而 Muon 在 NDS 比率为 0.89 时实现了 1.22 倍的损失下降比率，验证了理论预测。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LLM 预训练 | 一步损失下降 | Muon 比 Adam 实现更大的损失下降 | Adam | 由更小的曲率惩罚驱动，平均 NDS 比率为 1.76 |
| 数据不平衡影响 | NDS 差距 | 从 s=0 时的 0.63 扩大到 s=1 时的 1.13 | 平衡数据 (s=0) | 1.8 倍增长 |
| 二次模型优化 | 平均 NDS | Muon 比 GD 更小 | GD | 任何有限水平 T ≥ 1 都成立 |
局限与改进
局限性包括作者承认和观察到的几个方面。首先，作者关注的是因果语言模型（Causal LLMs），没有在扩散模型等其他模型类上验证这些洞察。其次，理论分析基于结构化二次模型的四个假设，包括低 Kronecker 秩 Hessian、同时对角化、曲率异质性和梯度对齐，这些假设在真实预训练中只是近似成立，可能会限制理论结果的普适性。第三，本文只分析了 Muon 和 Adam 两个一阶优化器，没有考虑高阶优化器或优化器变体。第四，实验主要关注预训练阶段，没有探讨微调阶段的行为。第五，Hessian 计算的二次复杂度限制了实验规模到 124M 参数，无法在更大的模型上验证结果。第六，层内/层外分解假设 Hessian 块之间的交互是可分离的，这可能忽略了更复杂的跨层动态。
独立分析的弱点
独立分析的弱点包括五个方面。第一，理论假设的简化可能限制结论的普适性，例如 Hessian 的低 Kronecker 秩假设在深层网络中可能不成立，同时对齐条件在实际预训练中可能只近似满足。改进方向是构建更逼真的理论模型，考虑非对齐情况和高阶交互。第二，实验规模受限于 Hessian 计算的二次复杂度，124M 参数的模型规模远小于当前主流 LLM 的数十亿参数。改进方向是开发 Hessian 的近似计算方法或随机估计技术，在更大模型上验证结果。第三，只分析了一步损失下降的曲率分解，没有考虑多步动态和长期累积效应。改进方向是扩展分析到有限视界的优化轨迹，考虑曲率的时间演化。第四，数据不平衡分析使用合成的 Zipf-PCFG 数据，可能无法捕捉真实预训练数据的复杂结构。改进方向是在真实数据集上通过采样和数据增强控制不平衡水平，验证合成数据的结论。第五，层内/层外分解没有考虑注意力机制和残差连接的具体结构，可能忽略了架构特定的曲率模式。改进方向是设计架构感知的曲率分解，区分注意力层和前馈层的不同贡献。
未来方向
未来研究方向包括作者提出的和基于成果可延伸的几个方向。作者提出在扩散模型等其他模型类上验证曲率洞察，因为不同模型类的优化景观可能有不同的曲率特征。基于本文成果，可以探索多个方向。第一，研究高阶优化器的曲率行为，例如二阶牛顿方法或拟牛顿方法是否也表现出类似的 NDS 优势。第二，设计曲率感知的优化器变体，例如通过估计 Hessian 的奇异值方向，将更多更新能量分配到低曲率方向。第三，研究不同架构的曲率特性，例如 Transformer 的层归一化和残差连接如何影响 Hessian 结构，以及是否可以设计架构特定的优化策略。第四，探索数据分布与优化器性能的交互，例如通过数据重采样或课程学习控制训练数据的曲率特性，改善优化效率。第五，研究预训练和微调阶段的曲率差异，因为微调通常使用较小的数据集和不同的目标函数，可能产生不同的曲率景观。第六，开发更高效的曲率估计方法，使大规模模型的曲率分析成为可能，为优化器设计和超参数调优提供指导。
复现评估
复现评估方面，论文没有明确提供代码开源信息，但实验描述较为详细。数据方面使用了 FineWeb 数据集和合成的 Zipf-PCFG 数据，FineWeb 是公开数据集，Zipf-PCFG 的生成参数在附录中给出。算力方面，124M 参数模型的 Hessian 计算需要相当的内存和计算资源，因为 Hessian 的矩阵表示规模为 (d1 × d2) × (d1 × d2)，即使使用低秩近似也需要存储多个 d1 × d1 和 d2 × d2 矩阵。难度方面，复现主要实验需要实现 Hessian 计算、同时对角化算法（JADE Jacobi 扫掠）、NDS 分解和层内/层外分解，这些技术实现有一定复杂度。此外，理论证明的复现需要仔细推导和验证，特别是定理 5.5 的条件和证明概要。总体而言，复现难度中等偏高，需要较强的机器学习和优化理论基础。
论文图表
