Why Muon Outperforms Adam: A Curvature Perspective

发表时间: 2026-06 · arXiv:2606.04662

作者/机构: Shuche Wang (National University of Singapore), Fengzhuo Zhang (Yale University), Jiaxiang Li (University of Minnesota), Dirk Bergemann (Yale University), Zhuoran Yang (Yale University)

A1 主要贡献

本文旨在从优化曲率的视角，揭示Muon优化器在大语言模型（LLM）训练中优于Adam的原因。现有的研究已从关联记忆和数据长尾性等角度进行了解释，但本文首次从局部几何景观的特性出发进行探究。

核心研究问题：
1. 是什么样的优化景观特性支撑了Muon相对于Adam的优势？
2. 预训练的因素，如训练数据和模型结构，如何影响这一特性？

研究目标与创新点：
为了回答上述问题，本文进行了系统的实证和理论分析，主要贡献如下：
1. 识别核心优势来源：通过对训练损失进行二阶泰勒展开，本文发现Muon和Adam的一阶收益（梯度对齐）相当，但Muon的二阶曲率惩罚显著更小。这表明Muon实现更大单步损失下降的主要原因是其更低的二阶曲率成本。
2. 分解曲率惩罚：进一步将曲率惩罚分解为更新范数和归一化方向锐度（Normalized Directional Sharpness, NDS）。研究发现，两种优化器的更新范数相当，因此Muon的优势来源于其更新方向具有更低的NDS，而非更小的更新步长。
3. 探究影响因素：
* 数据层面：使用可控不平衡度的Zipf-概率上下文无关文法（PCFG）生成的合成数据进行训练，发现数据越不平衡，Muon在NDS上的优势越明显。
* 模型结构层面：将NDS分解为层内和跨层贡献，发现在训练中后期，Muon的NDS优势主要由更小的层内曲率贡献维持，其跨层贡献会迅速下降。

  1. 提供理论支撑：构建了一个能够反映LLM训练景观特征的风格化二次型优化问题。在该模型上，本文证明了Muon通过在不同曲率组之间平衡更新能量，获得了比梯度下降（GD）更小的平均NDS。当曲率异质性足够强时，Muon在相同步数后也能达到更低的局部二次损失。

综上所述，本文揭示了Muon优于Adam的一个具体机制：Muon通过其谱归一化的更新方向获得了更低的NDS，从而承受了更小的二阶曲率成本，最终实现了更大的单步损失下降。

A3 背景知识

本节介绍Adam和Muon优化器的细节，并建立本文使用的符号体系。

Adam优化器。Adam在过去十年中一直是训练LLM的默认优化器【Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. 2014】。它通过使用随机梯度的一阶和二阶矩的指数移动平均来逐坐标地归一化参数更新。对于一个矩阵参数 $W Adam_t \in R^{m \times n}$，令 $GAdam_t = \nabla_W L_{D_t}(W Adam_t)$ 表示在小批量 $D_t \subseteq D$ 上的梯度，其中 $L_D$ 是在 $D_t$ 上的经验训练损失。Adam维护以下状态：

它们分别估计了随机梯度的一阶矩和逐元素的二阶矩。这里，$β_1, β_2 \in [0, 1)$ 是超参数，⊙ 表示哈达玛积。经过偏差修正后，
$M'_t = M_t/(1 − β^t_1)$ 和 $V'_t = V_t/(1 − β^t_2)$，Adam使用更新方向 $ZAdam_t = η_tM'_t/(\sqrt{V'_t} + \epsilon)$ 并通过 $W Adam_{t+1} = W Adam_t − ZAdam_t$ 来更新参数，其中平方根和除法是逐元素应用的。

Muon优化器。Muon是一个矩阵参数优化器，它明确利用了矩阵梯度的谱结构【Jordan, K., et al. Muon: An optimizer for hidden layers in neural networks. 2024b】。这种对矩阵结构感知的设​​计已被证明在大规模LLM预训练中优于Adam【Liu, J., et al. Muon is scalable for llm training. 2025b】。对于矩阵参数 $W Muon_t \in R^{m \times n}$，令 $GMuon_t = \nabla_W L_{D_t}(W Muon_t)$ 表示在小批量 $D_t \subseteq D$ 上的梯度。Muon维护一个动量累加器 $B_t = µB_{t−1} + GMuon_t$，其中 $B_0 = 0$ 且 $µ \in [0, 1)$。给定动量矩阵的奇异值分解 $B_t = U_tS_tV^\top_t$，Muon通过设置 $O_t = U_tV^\top_t$ 来谱归一化动量矩阵，并通过以下方式更新参数：

在实践中，$O_t$ 可以通过少量牛顿-舒尔茨迭代来高效近似，而不是通过精确的奇异值分解来计算。由此产生的更新是尺度不变的：将 $B_t$ 乘以任何正标量都不会改变更新方向 $ZMuon_t$。

符号。对于正整数 $N$，令 $[N] = \{1, . . . , N\}$。对于矩阵 $A, B \in R^{m \times n}$，定义弗罗贝尼乌斯内积和相应的范数为 $\langle A, B \rangle = \text{tr}(A^\top B)$，$\|A\|_F = \sqrt{\langle A, A \rangle}$。在全文中，$L_D$ 表示在 $D$ 上的经验训练损失，它在感兴趣的区域内是二次连续可微的。$W$ 表示矩阵参数。$G = \nabla_W L_D(W)$ 表示梯度，而 $H = \nabla^2_W L_D(W) : R^{m \times n} \rightarrow R^{m \times n}$ 表示作用于矩阵扰动的Hessian算子。对于任何矩阵扰动 $Z \in R^{m \times n}$，我们定义 $H[Z] = \frac{d}{d\epsilon} \nabla_W L_D(W + \epsilon Z)|_{\epsilon=0}$。我们用 $\text{mat}(H) \in R^{mn \times mn}$ 表示 $H$ 在向量化下的矩阵表示，即 $\text{vec}(H[Z]) = \text{mat}(H) \text{vec}(Z)$。对于向量 $x = (x_1, . . . , x_d)$，$\text{Diag}(x_1, . . . , x_d)$ 表示对角线元素为 $x_1, . . . , x_d$ 的 $d \times d$ 对角矩阵。

A2 方法细节

4 主要结果

本节展示了从曲率角度表征Muon优于Adam的主要实证发现。

4.1 Muon比Adam承受更小的二阶曲率惩罚

局部优化进展分解。为了从曲率角度研究Muon相对于Adam的优越性，我们从单步优化进展的局部分解开始。具体来说，对于参数矩阵 $W$ 和一个更新 $Z$，在小批量 $D$ 上的经验损失下降 $\Delta_D(W, Z) = L_D(W) - L_D(W - Z)$ 可以通过二阶泰勒展开近似为：

这个方程将损失下降 $\Delta_D(W, Z)$ 分解为一阶下降 $I^{(1)}_D(W, Z) = \langle G, Z \rangle$ 和曲率惩罚 $I^{(2)}_D(W, Z) = 1/2 \cdot \langle Z, H[Z] \rangle$。一阶项衡量了沿更新方向移动所带来的损失减少，而曲率惩罚则捕捉了抵消这种减少的二阶损失增加。我们将公式(4.1)的右侧称为预测损失下降，左侧 $\Delta_D(W, Z)$ 称为实现损失下降。

实验设置。为了评估公式(4.1)中的近似是否能解释Muon相对于Adam的优势，我们在以下设置中计算 $\Delta_D(W, Z)$、$I^{(1)}_D$ 和 $I^{(2)}_D$。我们使用一个124M参数的NanoGPT模型在FineWeb数据集上进行训练【Penedo, G., et al. The fineweb datasets. 2024】。我们强调这个规模，因为124M参数是现有LLM预训练曲率研究中考虑的最大模型尺寸之一【Zhang, Y., et al. Why transformers need adam: A hessian perspective. 2024a; Dong, Z., et al. Towards quantifying the hessian structure of neural networks. 2025】，因为基于Hessian的计算随参数维度呈二次方增长。对于Adam和Muon，我们都通过网格搜索选择最优学习率。完整的实验细节在附录A中提供。在Muon轨迹的第t步，给定训练批次 $D_t$、当前参数 $W Muon_t$ 和更新方向 $ZMuon_t$，我们计算 $\Delta_{D_t}(W Muon_t, ZMuon_t)$、$I^{(1)}_{D_t}(W Muon_t, ZMuon_t)$ 和 $I^{(2)}_{D_t}(W Muon_t, ZMuon_t)$。我们对Adam在其自身的优化轨迹上也使用其相应的参数更新计算相同的量。由于Muon在相同的训练步数下比Adam达到更低的损失，为了公平比较，我们在匹配的验证损失下而不是匹配的训练步数下比较这两个优化器。

实验发现。在图1(a)中，我们绘制了在匹配的验证损失水平下，opt ∈ {Muon, Adam}的预测损失下降 $I^{(1)}_{D_t}(W^{opt}_t, Z^{opt}_t) - I^{(2)}_{D_t}(W^{opt}_t, Z^{opt}_t)$ 和实现损失下降 $\Delta_{D_t}(W^{opt}_t, Z^{opt}_t)$。在所有验证损失水平上，Muon都比Adam实现了更大的实现损失下降，这与其优越的训练效率一致。对于Adam，预测的损失下降与实现的损失下降非常吻合。对于Mu
on，预测的损失下降略小于实现的损失下降。因此，接下来的分析在公式(4.1)的二阶近似内解释了Muon相对于Adam的优势，同时也为高阶效应可能进一步贡献于Muon优势的可能性留下了空间。在图1(b)和1(c)中，我们报告了在匹配的验证损失水平下Adam和Muon的一阶下降和曲率惩罚。如图1(b)所示，Adam和Muon在整个优化过程中具有可比的一阶下降，因为两条曲线在所有验证损失值上都保持在相似的水平，尽管Adam在不同随机种子之间表现出更高的可变性。相比之下，图1(c)显示了曲率项的明显差距：Adam的曲线（蓝色）始终位于Muon的曲线（橙色）之上，表明Muon在其更新方向上承受的Hessian二次型惩罚要小得多。这些结果与图1(a)一起表明，Mu-on的更大单步损失下降主要由较小的二阶曲率惩罚驱动，而不是由较大的一阶增益驱动。这引出了我们的第一个观察。
观察1：在验证损失对齐的情况下，Muon比Adam实现了更大的单步损失下降。这一差距主要归因于更小的二阶曲率成本。


图1：Muon和Adam更新方向上单步优化进展的分解。面板(a)比较了预测的单步损失下降 $I^{(1)}_{D_t}(W_t, Z_t) - I^{(2)}_{D_t}(W_t, Z_t)$ 与实现的单步损失下降 $\Delta_{D_t}(W_t, Z_t)$。面板(b)报告了一阶下降 $I^{(1)}_{D_t}(W_t, Z_t)$，面板(c)报告了曲率惩罚 $I^{(2)}_{D_t}(W_t, Z_t)$。结果显示，Muon和Adam实现了一阶下降相似，而Muon的曲率惩罚更小。

4.2 Muon更小的曲率惩罚来自其更新方向

NDS分解。观察1表明，Muon在单步损失下降上的优势主要由其更小的二阶曲率成本解释。注意，曲率惩罚 $I^{(2)}_D(W, Z) = 1/2 \cdot \langle Z, H[Z] \rangle$ 同时取决于更新的尺度和其方向上的曲率。我们现在要问：Muon更小的曲率成本是因为它采取了更小的步长，还是因为其更新方向遇到的曲率更小？为了解耦更新方向和更新尺度的影响，我们通过归一化尺度来衡量局部二次型的曲率。具体来说，我们遵循【Pan, Y. and Li, Y. Toward understanding why adam converges faster than sgd for transformers. 2023】的方法，定义沿非零更新 $Z$ 的归一化方向锐度（NDS）为：

根据这个定义，曲率惩罚可以分解为 $I^{(2)}_D(W, Z) = 1/2 \cdot \|Z\|^2_F \cdot S_F(W; Z)$，因此更小的曲率成本只能来自更小的更新范数或更低的NDS。我们因此检验Muon更小的曲率惩罚是由更小的更新范数还是更小的NDS驱动的。

实验方法。与4.1节类似，我们在匹配的验证损失水平上绘制更新范数和NDS。为了计算Adam和Muon之间的比率，我们对相邻检查点之间的值进行线性插值，因为两个优化器可能在记录的步骤中没有达到完全相同的验证损失。我们也在附录B中报告了在匹配的训练步数下的相应比率，这也得出了相同的结论。

实验发现。在图2(a)和2(b)中，我们根据验证损失对齐绘制了 $\|Z_t\|_F$ 和 $S_F$ 的值。如图2(a)所示，Adam的曲线（蓝色）始终位于Muon的曲线（橙色）之上，表明Muon在整个训练过程中比Adam具有更低的NDS。相比之下，图2(b)显示两个优化器具有可比的更新范数：两条曲线在所有验证损失水平上都保持接近平坦且彼此靠近。图2(c)通过绘制Adam与Muon的比率在数量上证实了这种分解：更新尺度 $\|Z_t\|^2_F$ 的比率（绿色虚线）保持在接近1的水平，而曲率惩罚 $I^{(2)}_{D_t}$ 的比率（深红色实线）和NDS $S_F$ 的比率（青色虚线）紧密地相互追踪，平均NDS比率为1.76。这表明曲率惩罚的差距几乎完全由NDS的差距解释，而不是由更新尺度的差异解释。因此，我们得出以下结论。
观察2：Muon和Adam具有可比的更新范数，因此Muon更小的曲率惩罚是由显著更小的NDS驱动的。


图2：Muon和Adam的NDS与更新范数比较。面板(a)绘制了Muon和Adam的NDS。面板(b)绘制了Muon和Adam的更新范数。面板(c)报告了曲率惩罚、NDS以及更新的弗罗贝尼乌斯范数平方的Adam与Muon之比。Muon和Adam具有相似的更新范数，而Muon的NDS比Adam小。此外，NDS的Adam与Muon之比与曲率惩罚之比紧密相随。

4.3 数据集不平衡扩大了Adam和Muon之间的NDS差距

动机与实验设计。先前的工作表明，训练数据的尾部结构和不平衡性可以与优化器行为产生强烈的相互作用。首先，神经网络的Hessian分析表明，Hessian谱对数据混合非常敏感【Sagun, L., et al. Empirical analysis of the hessian of over-parametrized neural networks. 2017; Papyan, V. The full spectrum of deepnet hessians at scale. 2018】。其次，最近的研究发现，Muon在重尾数据上尤其能胜过Adam【Wang, S., et al. Muon outperforms adam in tail-end associative memory learning. 2025b; Vasudeva, B., et al. How muon’s spectral design benefits generalization. 2025b】。受这些发现的启发，我们研究数据集不平衡是否会放大上面确定的归一化方向锐度差距。

实验设置：数据生成。为了研究数据不平衡如何影响NDS，我们使用Zipf-PCFG构建合成训练数据，这使我们能够明确控制不平衡程度。具体来说，我们实例化一个带有主题 $k \in [K]$ 的Zipf-PCFG，并将词汇表 $V$ 划分为 $C$ 个词类 $V_1, . . . , V_C$，使得 $V = \cup^C_{c=1}V_c$。不同的词类对应不同的句子成分，如名词、形容词和动词。每个主题 $k$ 在每个词类 $c$ 中都有自己的偏好分布。例如，主题 $k = “\text{food}”$ 会为与饮食相关的动词分配更高的概率。给定主题 $k$ 和词类 $c$，令 $\phi_{k,c}(j)$ 表示词元 $j \in V_c$ 的基础概率。为了引入可控的不平衡，我们按 $\phi_{k,c}(j)$ 的降序对 $V_c$ 中的词元进行排序，并用 $r(j, c, k)$ 表示词元 $j$ 的排名，其中最可能的词元排名为1，且 $r(j, c, k) \in \{1, . . . , |V_c|\}$。在不平衡水平为 $s$ 时，我们从与 $r(j, c, k)^{-s}\phi_{k,c}(j)$ 成比例的重加权分布中采样词元。然后，我们使用Adam和Muon，在不平衡水平 $s \in \{0, 0.5, 1\}$ 下生成的数据集上训练一个9M参数的NanoGPT模型（4层，4个注意力头，模型维度256），训练10,000步，学习率通过网格搜索选择。更多实验细节见附录A。

度量：轨迹平均NDS。为了评估整个训练轨迹上的NDS，对于每个优化器opt ∈ {Muon, Adam}，我们将不平衡水平s下的轨迹平均NDS定义为：

其中 $T$ 表示训练步数集合，$s \in \{0, 0.5, 1\}$ 是不平衡水平，$W^{opt,s}_t$ 和 $Z^{opt,s}_t$ 分别表示优化器opt在不平衡水平s下第t步的参数和更新。为了突出Adam和Muon之间的差异，我们用Muon在 $s = 0$ 时的值对 $\bar{S}_{opt}(s)$ 进行归一化：$\tilde{S}_{opt}(s) = \bar{S}_{opt}(s)/\bar{S}_{Muon}(0)$。我们进一步将不平衡水平s下的归一化锐度差距定义为 $\Delta(s) = \tilde{S}_{Adam}(s) - \tilde{S}_{Muon}(s)$。

实验发现。在图3(a)中，我们绘制了归一化轨迹平均NDS $\tilde{S}_{opt}(s)$ 与Zipf指数 $s \in \{0, 0.5, 1\}$ 的关系。如图所示，两种优化器的NDS都随不平衡性单调增加，但对Adam的影响要大得多：随着s从0增加到1，Adam的归一化NDS从1.63上升到2.38，而Muon的仅从1.00增长到1.25。两条曲线之间不断扩大的阴影区域反映了日益增大的差距。此外，图3(b)直接量化了这一差距：$\Delta(s)$ 从 $s = 0$ 时的0.63单调扩大到 $s = 1$ 时的1.13，随着数据变得更加不平衡，差距增加了1.8倍。我们将此观察总结如下。
观察3：增加数据集的不平衡水平不仅放大了Muon和Adam的NDS，而且还扩大了它们之间的NDS差距。


图3：不同不平衡水平(s = 0, 0.5, 1)对NDS的影响。面板(a)报告了三种不平衡水平下(s = 0, 0.5, 1)的轨迹平均NDS，该值已通过Muon在s = 0时的值进行归一化。面板(b)报告了相同设置下Adam与Muon在NDS上的差距。结果表明，随着数据变得更加不平衡，Muon在NDS上相对于Adam的优势变得更大。

4.4 Muon的NDS越来越转向层内Hessian块

层内/跨层分解。观察2确立了Muon和Adam在所有模型参数上的归一化方向锐度差距。在本节中，我们研究不同层如何对这一差距做出贡献。

实验设置：层内/跨层分解。考虑一个有 $L$ 层的模型，其中第 $t$ 步的完整参数为 $W_t = (W_{t,1}, . . . , W_{t,L})$，其中 $W_{t,\ell} \in R^{m_\ell \times n_\ell}$ 表示第 $\ell$ 层的权重矩阵。相应的更新分解为 $Z_t = (Z_{t,1}, . . . , Z_{t,L})$，其中 $Z_{t,\ell} \in R^{m_\ell \times n_\ell}$ 是优化器对第 $\ell$ 层产生的更新。Hessian算子 $H$ 同样可以分解为逐层的块：对于层 $\ell, \ell' \in [L]$，令 $H_{\ell\ell'}$ 表示将第 $\ell'$ 层的扰动映射到第 $\ell$ 层产生的二阶效应的块。根据定义，对角块 $H_{\ell\ell}$ 捕捉层内曲率，而非对角块 $H_{\ell\ell'}$（$\ell \neq \ell'$）捕捉跨层交互。这种块结构使我们能够将NDS分解为层内和跨层贡献：$S_F(W_t; Z_t) = S^{within}_F(W_t; Z_t) + S^{cross}_F(W_t; Z_t)$，其中

直观地说，$S^{within}_F$ 衡量了当每层的更新仅与其自身的Hessian块相互作用时遇到的曲率，而 $S^{cross}_F$ 则捕捉了不同层更新之间相互作用产生的额外曲率。我们进一步定义相对层内贡献为 $\rho^{within}_t = S^{within}_F(W_t; Z_t)/S_F(W_t; Z_t)$。

实验发现。图4(a)绘制了在与4.1节相同的实验设置下，Adam（蓝色）和Muon（橙色）的 $S^{within}_F$（实线）和 $S^{cross}_F$（虚线）。所有四条曲线都随着训练的进行而下降，并且Muon的两个分量始终小于Adam：Muon的曲线（橙色）始终位于Adam的曲线（蓝色）之下。对于Adam，蓝色实线和虚线曲线以可比的速率下降，维持了两个分量之间大致稳定的比例。然而，对于Muon，橙色虚线曲线（$S^{cross}_F$）的下降速度远快于橙色实线曲线（$S^{within}_F$），因此两条Muon曲线在训练过程中趋于收敛。这表明随着训练的进行，层内分量在Muon的NDS中变得越来越占主导地位。图4(b)通过绘制层内分数 $\rho^{within}_t$ 来量化这一趋势。Muon的曲线（橙色）从训练早期的约14%急剧上升到训练后期的约44%，其层内份额几乎增加了三倍。相比之下，Adam的曲线（蓝色）在30%左右波动，仅从约27%温和增加到约34%。这凸显了Muon的小 $S^{within}_F$ 在训练中后期保持整个模型NDS较低的重要性，而Adam则在层内和跨层贡献之间保持相对稳定的平衡。此外，附录C显示，层内Adam-Muon锐度差距并非均匀分布在各层之间，几乎所有的差距都集中在第一层和最深层。
观察4：在训练过程中，Muon的方向锐度越来越偏向于层内Hessian块，而Adam的锐度构成则相对稳定。Muon的层内和跨层分量都始终小于Adam。


图4：训练过程中方向锐度的层内和跨层分解。面板(a)报告了Muon和Adam的 $S_F(W_t; Z_t)$ 的层内和跨层分量。面板(b)报告了层内分数。结果显示，对于Muon，层内分量在训练过程中占据了方向锐度越来越大的份额。

5 结构化矩阵块二次模型的案例研究

本节为经验发现提供理论依据。鉴于观察4显示层内曲率成为Muon的NDS优势主导部分，我们隔离单个权重矩阵，在一个二次模型上研究该块内的局部曲率，比较不同优化器引起的更新方向所遇到的曲率。

5.1 结构化二次模型

模型定义。受公式(4.1)中二阶泰勒近似的启发，我们关注LLM预训练中一个优化步骤周围的局部二次景观。给定一个固定的参数 $W_0 \in R^{d_1 \times d_2}$、其梯度 $G = \nabla L(W_0)$ 和Hessian算子 $H = \nabla^2 L(W_0)$，我们考虑以下针对更新 $Y \in R^{d_1 \times d_2}$ 的二次模型：

为了使该模型能代表LLM预训练，我们对梯度和Hessian结构施加了四个假设，每个假设都在真实的预训练动态中得到了经验验证。

假设5.1（Hessian低克罗内克秩）。矩阵块上的局部Hessian算子H具有小的克罗内克秩。具体来说，令 $\text{mat}(H) \in R^{d_1d_2 \times d_1d_2}$ 为H在向量化下的矩阵表示。则存在一个整数 $r \ll \min\{d_1^2, d_2^2\}$，以及对称矩阵 $A_k \in R^{d_1 \times d_1}$ 和 $B_k \in R^{d_2 \times d_2}$，使得 $\text{mat}(H) = \sum_{k=1}^r B_k^\top \otimes A_k$，其中 $\otimes$ 表示克罗内克积。
* 验证：该假设受K-FAC等工作的启发【Martens, J. and Grosse, R. Optimizing neural networks with kronecker-factored approximate curvature. 2015】。图5(a)显示，对于四个注意力矩阵（WQ, WK, WV, WO），低秩克罗内克近似能够解释大部分Hessian的弗罗贝尼乌斯能量。例如，秩为3-5的近似就能解释WK和WV超过80%的能量。图5(b)可视化了WV的Hessian、其秩4克罗内克近似以及残差，证实了近似捕获了主要结构。


图5：假设5.1的经验支持。面板(a)报告了四个注意力矩阵Hessian的低秩克罗内克近似所解释的弗罗贝尼乌斯能量比例。面板(b)可视化了WV的Hessian、其秩4克罗内克近似以及残差。结果表明，注意力矩阵的Hessian可以很好地被低秩克罗内克积近似。

假设5.2（同时对角化）。考虑假设5.1中的秩-r克罗内克分解 $\{(A_k, B_k)\}_{k=1}^r$。假设 $\{A_k\}_{k=1}^r$ 和 $\{B_k\}_{k=1}^r$ 分别被正交矩阵 $U \in R^{d_1 \times d_1}$ 和 $V \in R^{d_2 \times d_2}$ 同时正交对角化。即，对于任意 $k = 1, . . . , r$，我们有 $A_k = U \text{Diag}(a_k^{(1)}, . . . , a_k^{(d_1)})U^\top$ 和 $B_k = V \text{Diag}(b_k^{(1)}, . . . , b_k^{(d_2)})V^\top$。
* 验证：该假设表明Hessian的克罗内克因子近似共享公共的正交特征基。图6(a)报告了使用JADE算法计算的同时对角化分数 $\eta_{sd}$，$\{A_k\}$ 和 $\{B_k\}$ 的得分分别为0.892和0.845，接近最大值1.0，表明单个共享正交基捕获了每个因子家族的大部分能量。

假设5.3（曲率异质性）。假设在 $i \in [d']$ 的 $w_i$ 值中，恰好有 $q$ 个是正的，并且这些正曲率具有两级结构：对于 $i \in [m]$，$w_i = w_H$，对于 $i \in \{m + 1, . . . , q\}$，$w_i = w_L$，其中 $w_H > w_L$ 且 $\alpha = m/q < 1/2$。
* 验证：该假设表明Hessian的正配对曲率是异质的。图6(b)绘制了正配对曲率 $w_i$ 的对数尺度图，揭示了一个强烈的长尾分布，曲率值跨越了六个数量级。高曲率模式的比例很小，这与 $\alpha < 1/2$ 的条件一致。

假设5.4（梯度对齐）。梯度 $G$ 位于由 $\{M_i\}_{i=1}^q$ 张成的子空间中，即 $G = \sum_{i=1}^q \sigma_i M_i$，其中 $\sigma_i = \langle G, M_i \rangle$。此外，系数 $\sigma_i$ 具有与假设5.3相同的两组结构：对于 $i \in [m]$，$\sigma_i = \sigma_H$，对于 $i \in \{m + 1, . . . , q\}$，$\sigma_i = \sigma_L$，其中 $\sigma_H > \sigma_L$。
* 验证：该假设表明梯度主要位于Hessian的顶层曲率子空间中。图6(c)绘制了累积梯度能量比率 $\zeta(i)$，显示梯度能量迅速集中在少数顶部曲率方向上，前约30个方向就捕获了超过80%的能量。


图6：假设5.2-5.4的经验支持。面板(a)显示了 $\{A_k\}^r_{k=1}$ 和 $\{B_k\}^r_{k=1}$ 的平均同时对角化得分 $\eta_{sd}$。面板(b)显示了正曲率的值。面板(c)显示了累积梯度能量比率 $\zeta(i)$。这些结果支持Hessian分解中矩阵的近似同时对角化，以及梯度与Hessian谱顶部特征值相关联的特征向量的对齐。

GD与Muon的比较。在这些风格化的二次问题上，Adam的行为与GD相似，并且与Muon有显著不同（如图7所示）。因此，为简化分析，理论部分重点关注GD和Muon的比较。
* GD 更新：$Y^{GD}_{t+1} = Y^{GD}_t - \eta^{GD}_t \nabla Q(Y^{GD}_t) = Y^{GD}_t + \eta^{GD}_t Z^{GD}_t$。
* Muon 更新（动量为0）：$Y^{Muon}_{t+1} = Y^{Muon}_t - \eta^{Muon}_t \text{spec}(\nabla Q(Y^{Muon}_t)) = Y^{Muon}_t + \eta^{Muon}_t Z^{Muon}_t$，其中 $\text{spec}(G) = UV^\top$ for $G = U\Sigma V^\top$。


图7：面板(a)报告了NDS比率，面板(b)报告了损失下降比率。结果显示，在满足假设5.1-5.4的二次问题上，GD和Adam在NDS和损失下降方面表现出相似的行为。

5.2 理论结果

定理5.5。令 $\alpha = m/q$ 为高曲率组的相对大小，$\rho = w_H/w_L > 1$ 为曲率比。在假设5.1–5.4下，以下结果成立：
* Muon的NDS更小：对于每个有限视界 $T \ge 1$，Muon的有限视界平均NDS小于GD，即 $\bar{S}^{Muon}_T < \bar{S}^{GD}_T$。
* Muon的损失下降更大：如果 $\rho + 1 > 1/\alpha > 1 + \sigma_H/\sigma_L$，对于每个有限视界 $T \ge 1$，Muon比GD达到更低的损失，即 $Q(Y^{Muon}_T) < Q(Y^{GD}_T)$。

定理诠释。定理5.5提供了两个结果。首先，它表明Muon的平均NDS总是小于GD，为观察2提供了理论支持。其次，在足够的曲率异质性条件下，Muon实现了更大的累积损失下降，支持了观察1。这些条件（$\rho > 1/\alpha - 1$ 和 $\alpha < \sigma_L/(\sigma_L + \sigma_H)$）要求最大曲率远大于其他曲率，且高曲率组占比较小，这与实践中观察到的Hessian谱的离群值现象一致。

核心机制。这两个结果背后的关键机制是，Muon的谱归一化在所有正交曲率特征模式上均衡了更新幅度，将能量均匀地分布在高曲率和低曲率方向之间。相比之下，GD的更新与梯度成正比，而梯度根据假设5.4将更多能量集中在高曲率方向上。这种集中导致GD承受更大的方向锐度，当曲率异质性足够强时，还会产生更大的曲率惩罚，从而抵消其一阶增益。

5.3 证明梗概

模型对角化。在假设5.1-5.4下，局部二次模型在正交秩一模式 $\{M_i\}_{i=1}^q$ 的张成空间上对角化。在此基底下，$G = \sum_{i=1}^q \sigma_i M_i$ 且 $H[M_i] = w_i M_i$。对于任何更新 $Y = \sum_i y_i M_i$，二次模型分解为标量分量：

动态完全由残差梯度系数 $r^{opt}_{i,t} = \sigma_i - w_i y^{opt}_{i,t}$ 表征。

Muon vs. GD 更新机制。由于 $M_i$ 是成对的奇异方向，谱归一化将残差 $\sum_i r^{Muon}_{i,t} M_i$ 映射为：

因此Muon为每个活动模式分配相同的幅度。而GD使用：

其更新能量与当前残差能量成正比。核心机制在于：Muon消除了模式间的幅度不平衡，而GD继承了这种不平衡。

NDS比较。Muon的等幅更新导致一个固定的曲率平均值：

而GD的NDS是一个残差能量加权的平均值：

由于 $\sigma_H > \sigma_L$，GD在起始时高曲率组的残差能量比例 $P^{GD}_0$ 大于组大小比例 $\alpha$。之后 $P^{GD}_t$ 在 $P^{GD}_0$ 和 $1-P^{GD}_0$ 之间交替，导致其时间平均值总是大于 $\alpha$。而Muon总是精确地将 $\alpha$ 比例的更新能量分配给高曲率组。由于 $w_H > w_L$，这导致 $\bar{S}^{Muon}_T < \bar{S}^{GD}_T$。

损失比较。通过分析次优性 $\Phi^{opt}_T = Q(Y^{opt}_T) - Q(Y^\star)$，可以发现Muon和GD的次优性都以线性速率收敛，但收敛因子不同。Muon的收敛因子为 $\Gamma^2$，而GD的为 $R$。在定理的条件下，可以证明Muon不仅第一步的次优性更小（$\Phi^{Muon}_1 < \Phi^{GD}_1$），而且后续的收缩也更快（$\Gamma^2 < \sqrt{R}$）。结合这两点，得出对于所有 $T \ge 1$，$Q(Y^{Muon}_T) < Q(Y^{GD}_T)$。

A4 实验环境

FineWeb 主体实验 (Section 4.1, 4.2, 4.4)

  • 模型: 124M参数的NanoGPT模型，包含12个Transformer层，12个注意力头，隐藏维度为768。词汇表大小为50,257（GPT-2分词器）。
  • 数据集: FineWeb-10B，序列长度为1024。
  • 硬件: 4块A100 GPU，每块80GB显存。
  • 软件: 基于Modded-NanoGPT代码库【Jordan, K., et al. modded-nanogpt. 2024a】。
  • 优化器配置:
    • Muon: 对除词嵌入和语言模型头之外的所有矩阵参数使用Muon，动量 $\mu=0.95$，无权重衰减，5次Newton-Schulz迭代。其余参数（嵌入层、lm头、标量或一维参数）使用Adam（$\beta_1=0.8, \beta_2=0.95$）。
    • Adam: 所有参数均使用Adam（$\beta_1=0.8, \beta_2=0.95$）。
    • 两种优化器的学习率都通过网格搜索从 $\{1, 2, 5\} \times \{10^{-1}, 10^{-2}, 10^{-3}, 10^{-4}\}$ 中选取。

Zipf-PCFG 数据不平衡实验 (Section 4.3)

  • 模型: 9M参数的NanoGPT模型，包含4个Transformer层，4个注意力头，模型维度256。词汇表大小为4,412。
  • 数据集: 使用Zipf-PCFG生成的合成语料库，序列长度1024，梯度累积8次，有效批大小为8,192个token。训练在不平衡指数 $s \in \{0, 0.5, 1\}$ 的数据上进行。
  • 优化器配置:
    • Muon: 对所有注意力矩阵(WQ, WK, WV, WO)和MLP矩阵(Win, Wout)使用Muon，学习率为 $5 \times 10^{-3}$，动量 $\mu=0.95$，5次Newton-Schulz迭代。其余参数使用Adam，学习率为 $1 \times 10^{-2}$，$\beta_1=0.8, \beta_2=0.95$。
    • Adam: 所有参数均使用Adam，学习率为 $1 \times 10^{-2}$，$\beta_1=0.8, \beta_2=0.95$。

A4 实验结果

本文通过一系列实证和理论分析，系统地揭示了Muon优化器优于Adam的内在机制。

  • 更大的单步损失下降: 实验表明，在匹配验证损失的情况下，Muon的单步损失下降显著大于Adam。通过二阶泰勒展开分析，发现这一优势并非源于更大的一阶收益（梯度对齐项），而是来自显著更小的二阶曲率惩罚（图1）。

  • 更低的归一化方向锐度 (NDS): 进一步将曲率惩罚分解为更新范数和NDS，实验证实Muon和Adam的更新范数大小相当。因此，Muon更小的曲率惩罚是由其更新方向具有更低的NDS所驱动的。Adam与Muon的曲率惩罚比率与它们的NDS比率高度吻合（图2）。

  • 数据不平衡的影响: 在使用Zipf-PCFG生成的可控不平衡数据集上的实验表明，随着数据不平衡程度（由Zipf指数s控制）的增加，Muon和Adam的NDS都会上升，但Adam的增幅远大于Muon。这导致两者之间的NDS差距随数据不平衡而扩大，表明Muon在处理不平衡数据时更具优势（图3）。

  • 层级贡献分析: 将NDS分解为层内和跨层贡献后发现，在训练过程中，Muon的跨层曲率贡献下降得非常快，使得其总NDS越来越由层内曲率主导。相比之下，Adam的层内与跨层贡献比例相对稳定。这表明Muun在训练中后期的NDS优势主要由更小的层内曲率维持（图4）。进一步分析（附录C）显示，层内NDS差距主要集中在模型的第一层和最后一层（图9）。

  • 理论模型验证: 为解释上述现象，本文构建了一个结构化的二次型优化问题，其Hessian和梯度结构（低克罗内克秩、曲率异质性、梯度与高曲率方向对齐）通过对真实LLM的实证分析得到验证（图5, 6）。在该模型上，理论证明了Muon比GD（其行为在此模型中与Adam类似，如图7）具有更小的平均NDS和更大的损失下降。其核心机制在于Muon的谱归一化能够均衡地在所有曲率模式上分配更新能量，而GD/Adam则倾向于将能量集中在梯度较大的高曲率方向上，从而导致更高的曲率惩罚。

A5 结论

本文首次从曲率的视角系统地探究了Muon优化器相较于Adam的优越性。研究发现，Muon之所以能取得更大的单步损失下降，其核心原因在于它的更新方向具有更低的归一化方向锐度（NDS），从而承受了更小的二阶曲率惩罚。进一步的研究表明，这种NDS优势会随着训练数据的不平衡性加剧而扩大，并且在训练中后期，这一优势主要由较小的层内曲率贡献所维持。为了从理论上解释这些实证观察，本文构建了一个结构化的二次型模型，并证明了当曲率存在异质性时，Muon确实能够获得更低的NDS和更大的损失下降。

本文的一个局限性是研究主要集中在因果语言模型（causal LLMs）上。将这些发现推广和验证到其他模型类别，如扩散模型，是未来值得探索的工作方向。

A6 附录

B 对4.2节的补充实验结果

NDS比较（按训练步数对齐）。本节补充了4.2节中按验证损失对齐的NDS比较结果，展示了按训练步数对齐的锐度比较。如图8(a)所示，在整个训练轨迹中，Muon始终保持比Adam更小的NDS，这一差距从早期到晚期训练一直存在。图8(b)报告了曲率惩罚、NDS和更新范数平方的Adam与Muon比率。更新范数比率保持在接近1，证实了两个优化器在每一步的更新范数相当。相比之下，NDS比率远高于1，轨迹上的平均值为2.94。曲率惩罚比率紧密跟随NDS比率，再次印证了观察2的结论：Muon更小的曲率惩罚是由更低的NDS而非更小的更新范数驱动的。值得注意的是，按训练步数对齐的平均NDS比率（2.94）大于按验证损失对齐的比率（1.76），这是因为在同一步数下，Muon已经达到了更低的验证损失，这种对齐方式实际上是将处于更优化阶段的Muon与处于较不优化阶段的Adam进行比较，从而放大了表观差距。


图8：Muon和Adam训练步骤中的NDS及相应比率比较。面板(a)：锐度比较。面板(b)：曲率惩罚、NDS以及弗罗贝尼乌斯更新范数平方的Adam/Muon比率。

C 层级NDS贡献

层级NDS差距的定位。为了补充观察4中的锐度比较，图9将方向锐度分解为12个Transformer层中每一层的贡献 $S_F^{(\ell, \text{opt})} = \langle Z^{\text{opt}}_{\ell}, H_{\ell\ell}[Z^{\text{opt}}_{\ell}] \rangle / \|Z^{\text{opt}}\|^2_F$。面板(a)显示，两种优化器在不同深度上积累层内曲率的方式并不均匀。面板(b)报告了总的Adam-Muon层内差距中每一层的份额，即 $\Delta S_F^{(\ell)} / \sum_{\ell=1}^{12} \Delta S_F^{(\ell)}$，其中 $\Delta S_F^{(\ell)} = S_F^{(\ell, \text{Adam})} - S_F^{(\ell, \text{Muon})}$。从图9可以看出，这个差距在深度上是高度局部化的。具体来说，大约70%的层内NDS差距来自两个边界层L1和L12，约28%来自深层L8-L11，而只有约2%来自中间层L2-L7。差距集中在L1和L12与这些边界层直接与词嵌入和输出logits交互的事实相符，在这些地方数据分布对局部曲率有最直接的影响。


图9：12个Transformer层中Adam-Muon层内锐度差距的逐层定位。面板(a)绘制了两种优化器更新的逐层NDS贡献 $S_F^{(\ell)}$。面板(b)报告了份额比率 $\Delta S_F^{(\ell)} / \sum_{\ell} \Delta S_F^{(\ell)}$。

D 第5节的经验验证细节

实验设置。本附录描述了第5节中每个验证图背后的实验设置、数据来源和计算流程。所有实验均使用在附录A所示的Zipf-PCFG数据集上训练的最后一个检查点，并采用Muon优化器。块Hessian是为四个注意力矩阵 $W_Q, W_K, W_V, W_O \in R^{256 \times 256}$ 计算的，每个密集Hessian $\text{mat}(H) \in R^{256^2 \times 256^2} = R^{65536 \times 65536}$。

D.1 有效的低克罗内克秩块Hessian（假设5.1）。
* 计算方法：对于一个注意力块的密集Hessian $\text{mat}(H)$，我们使用Van Loan重排【Van Loan, C. F. and Pitsianis, N. Approximation with kronecker products.