Muon优化器的谱缩放定律

# Muon 的谱缩放定律

来源：https://arxiv.org/html/2606.04058

###### 摘要

正交归一化更新规则已迅速成为训练大型语言模型的主流优化器选择，近期多个开源最先进模型均采用了 Muon。为使这些更新在计算上可行，Muon 使用 Newton–Schulz（NS）迭代执行正交归一化。由于 NS 仅为近似方法，奇异值较小的方向无法被正确正交归一化。在 Muon 中，NS 在每一步都被应用于动量矩阵，然而这些动量矩阵的奇异值谱在训练过程中如何演变，以及该行为如何随模型规模变化，目前几乎无人研究。我们首次对这一问题进行了系统性研究。通过追踪参数量从 77M 到 2.8B 的模型各层动量缓冲区的奇异值分位数，我们观察到了一致的规律：经过短暂的预热期后，分位数稳定在由层类型和模型规模共同决定的值上。这些稳定值在模型规模上呈现出异常整洁的幂律关系，且指数因层而异。从浅层到中后期层的稳定值随模型规模变化非常缓慢（约为 $M^{-0.25}$），因此学术规模下常用的 5 步 NS 配置在更大规模下仍能对这些层进行正交归一化。然而，部分末尾层的衰减速度要快得多（最高达 $M^{-0.96}$），在前沿规模下将落入 NS 失效区间，除非使用更多 NS 迭代次数或更优的系数。NS 迭代在大规模下计算开销不小；我们的幂律为从业者提供了一种有原则、感知层差异的方法，用于选择能够正交归一化关键方向的最小 NS 配置——在不牺牲更新质量的前提下避免不必要的计算。

## 1 引言

预训练大型语言模型（LLM）是一项耗费数百万 GPU 小时的昂贵过程，使得优化器的选择成为核心设计决策：即便是优化器效率上的微小提升也能带来可观的成本节省。AdamW（Loshchilov and Hutter, 2019; Kingma and Ba, 2015）长期以来是训练 LLM 的标准优化器（DeepSeek-AI, 2024; Llama Team, 2024; Team OLMo et al., 2025）。近年来，以 Muon（Jordan et al., 2024b; Bernstein and Newhouse, 2025）为代表的正交归一化更新优化器开始取而代之，提供了更稳定的训练效果以及跨规模更好的超参数迁移性（Pethick et al., 2025）。在更大规模下，Liu et al.（2025）表明 Muon 的计算效率是 AdamW 的两倍。值得注意的是，近期最先进的模型 Kimi-K2、GLM-5 和 DeepSeek-V4（Kimi Team et al., 2026; GLM-5, 2026; DeepSeek-AI, 2026）均使用 Muon 进行训练。

Muon 使用 Newton–Schulz（NS）迭代（Higham, 2008; Kovarik, 1970; Björck and Bowie, 1971）对动量矩阵进行近似正交归一化，该迭代反复对矩阵施加奇次多项式，将每个奇异值推向 $1$。由于 NS 仅为近似，奇异值过小的方向无法被正确正交归一化（见图 2）。因此，给定 NS 配置是否足够精确，取决于训练过程中动量奇异值的实际分布：若奇异值较大，即便是廉价的 NS 配置也能正确正交归一化；若奇异值较小，则需要更精确的配置。学术界通常使用 Cesista et al.（2025）提出的 $5$ 次多项式 NS 系数，这一配置因 NanoGPT 速度竞赛（Jordan et al., 2024a）而广为流传。近期前沿规模的 DeepSeek-V4（DeepSeek-AI, 2026）则使用了精度更高的 $10$ 次多项式组合。两种情况下，相同的配置被均匀地应用于所有层。由于每个 NS 步骤在大规模下都具有不可忽视的计算成本（Essential AI, 2025; Ahn et al., 2025b），一个自然的问题是：在大规模下 $5$ 次多项式是否已经足够，还是必须使用 $10$ 次多项式——更关键的是，答案对于每一层是否相同。

为回答这一问题，我们首次系统研究了 Muon 动量矩阵奇异值在训练过程中的演变规律，追踪了参数量从 77M 到 2.8B 的 GPT-2 风格模型在多个深度处的分位数。各层和各模型规模下均呈现出一致的规律：经过短暂的预热期后，奇异值分位数稳定在某一值，该值取决于层类型，并随模型规模减小。对这些稳定值拟合幂律揭示了与模型规模之间异常整洁的对数线性关系，且指数因层而异（见图 1）。这使我们能够在每个深度处外推：要在前沿规模下正交归一化足够多的方向以保持更新质量，NS 需要达到什么精度。

参见图注图 1：不同深度层和模型规模的归一化动量矩阵奇异值分位数稳定值的缩放定律。各层的指数差异显著。从浅层到中后期层的稳定值随模型规模变化非常缓慢（约为 $M^{-0.25}$）；NanoGPT 实验所用的 NS 近似在更大规模下对这些层仍然足够精确。然而，部分末尾层的衰减速度要快得多（最高达 $M^{-0.96}$），在前沿规模下将落入 NS 失效区间，除非使用更多 NS 迭代次数或更优的系数。

综合这些发现，从业者可以逐层选择能够在任意目标规模下正交归一化关键方向的最廉价 NS 配置。我们的主要贡献如下：

- 首次系统研究了跨层和模型规模（77M–2.8B）的 Muon 动量缓冲区奇异值谱。
- 建立了稳定值与模型规模之间的谱幂律关系，指数因层而异。
- 基于拟合所得幂律，提出了一套在前沿规模下逐层选择 NS 配置的实用方法。

### 1.1 相关工作

多项工作致力于解决在权重矩阵跨设备分片时大规模运行 Muon 的开销问题。Ahn et al.（2025b, a）提出了 Dion，一种通过低秩近似实现通信高效正交归一化更新的分布式优化器；而 Zhao et al.（2026）则对每个分片独立应用 NS，并通过周期性全局同步来保证训练稳定性。

另一条并行研究路线探索了深度学习中的矩阵预条件优化器，包括 Shampoo（Gupta et al., 2018）、SOAP（Vyas et al., 2024）和 COSMOS（Chen et al., 2026）。Anil et al.（2020）和 Shi et al.（2023）使 Shampoo 在大规模下具有实用性，但它需要学习率嫁接（Agarwal et al., 2020）等启发式方法才能在实践中匹敌 Adam，且据我们所知尚未被应用于前沿规模。Eschenha​​gen et al.（2025）通过自适应更新预条件器缓解了部分启发式依赖。与 Muon 更为接近的是，Li et al.（2025）在正交归一化更新基础上引入了 Adam 风格的二阶矩，在 Muon 谱范数步骤之上叠加了自适应的逐坐标缩放。Wen et al.（2025）在不同模型规模和数据与模型规模比例下对许多这类优化器进行了基准测试。

缩放定律由 Kaplan et al.（2020）开创，他们表明语言模型损失在参数量、训练 token 数和计算量上遵循整洁的幂律。Hoffmann et al.（2022）将这些关系细化为计算最优的 token 与参数比例，确认了此前的大型模型存在显著的训练不足问题。一条互补的研究方向探讨了哪些优化器超参数能随模型规模可预测地变化。Yang et al.（2022）表明，在 $\mu$P 参数化下，最优学习率可以跨规模零样本迁移，证明优化器超参数遵循其自身的缩放结构。我们的工作为这些研究方向作出贡献：我们表明，经过短暂预热期后，Muon 动量缓冲区的奇异值分位数稳定在遵循模型规模幂律的值上，且指数因层而异。

## 2 背景：Muon 与 Newton-Schulz

本节为后续内容奠定背景。我们首先介绍 Muon 优化器（第 2.1 节），并回顾其用于近似正交归一化的 Newton-Schulz 迭代（第 2.2 节），重点指出其关键局限性：奇异值足够小的方向无法被正交归一化。随后，我们通过受控实验（第 2.3 节）确定需要正交归一化多大比例的奇异方向才能使 Muon 保留其优势，从而确定本文后续追踪的分位数范围。

### 2.1 Muon 优化器

Muon（Jordan et al., 2024b; Bernstein and Newhouse, 2025）将二维权重矩阵上的原始梯度更新替换为**正交归一化**更新。在每一步，Muon 维护一个动量缓冲区 $M_t$，并在更新参数前对其施加 Newton-Schulz（NS）迭代以近似正交归一化（见算法 1）。

算法 1 Muon 优化器
1：学习率 $\eta$，动量系数 $\mu$，初始参数 $\Theta_0$
2：初始化动量缓冲区 $M_0 \leftarrow 0$
3：**for** $t = 0, 1, 2, \ldots$ **do**
4：计算梯度 $G_t = \nabla_\Theta \mathcal{L}(\Theta_t)$
5：更新动量：$M_{t+1} \leftarrow \mu \cdot M_t + G_t$
6：正交归一化：$O_{t+1} \leftarrow \mathrm{NS}(M_{t+1})$
7：更新参数：$\Theta_{t+1} \leftarrow \Theta_t - \eta \cdot O_{t+1}$
8：**end for**

其中 $\mathrm{NS}(\cdot)$ 表示第 2.2 节描述的 Newton-Schulz 迭代。Muon 的动机来源于谱范数下的最速下降（Bernstein and Newhouse, 2025），并已被证明在语言模型训练任务中于规模上的计算效率是 AdamW 的两倍（Liu et al., 2025）。

### 2.2 用于近似正交归一化的 Newton-Schulz 迭代

我们现在将注意力集中在近似正交归一化问题上。对大矩阵进行精确正交归一化代价高昂，因此实用实现通常使用快速迭代近似方法。

这一方向上的标准近似方法是 Newton-Schulz（NS）迭代。设 $A$ 为待正交归一化的动量矩阵。NS 首先归一化

$$\widetilde{A}_0 = \frac{A}{\|A\|_F},$$

将所有奇异值变换到 $[0, 1]$ 区间，然后依次应用一组奇次多项式 $p_0, \ldots, p_{n-1}$。在实践中，每个 $p_k$ 取如下形式的 5 次奇次多项式：

$$\widetilde{A}_{k+1} = p_k(\widetilde{A}_k) = a_k \widetilde{A}_k + b_k \left(\widetilde{A}_k \widetilde{A}_k^\top\right) \widetilde{A}_k + c_k \left(\widetilde{A}_k \widetilde{A}_k^\top\right)^2 \widetilde{A}_k, \quad k = 0, 1, \ldots, n-1, \tag{1}$$

回顾对于任意具有 SVD $A = USV^\top$ 的矩阵，奇次多项式满足

$$p(A) = U\,p(S)\,V^\top,$$

其中 $p(S)$ 将 $p$ 逐元素作用于 $S$ 的对角线上。这意味着奇异向量在每一步都被**精确保留**，只有奇异值被修改。将方程 1 展开 $n$ 次，最终结果为

$$\widetilde{A}_n = U\,\underbrace{(p_n \circ \cdots \circ p_1)(S)}_{=:\, f(S)}\,V^\top,$$

因此，NS 过程归结为一个一维问题：找到标量复合函数 $f = p_n \circ \cdots \circ p_1$，使得对每个奇异值 $\sigma \in [0, 1]$ 都有 $f(\sigma) \approx 1$。换言之，$f$ 应在 $(0, 1]$ 上近似符号函数，将每个奇异值推向 $1$，无论其初始值如何。当该条件成立时，$\widetilde{A}_n \approx UV^\top$，从而恢复 $A$ 的极分解中的正交归一因子。然而，由于每个 $p_i$ 都是奇次多项式，$f(0) = 0$ 对任意多项式选择均成立，因此 $f$ 无法在零的邻域内近似符号函数——这是 NS 族方法的根本局限性。

作为具体示例，考虑 NS 中使用的标准多项式，即 Muon 引入时所采用的（Jordan et al., 2024b）：

$$p(x) = 2x - 1.5x^3 + 0.5x^5,$$

应用 $n = 5$ 次，即 $f = p^{\circ 5}$。图 2 绘制了 $f(\sigma)$ 关于 $\sigma \in [0, 1]$ 的函数图像。可以看到，该近似对 $\sigma > 0.05$ 是精确的，能将这些值推近 $1$。然而，复合函数在原点附近近似线性：数值验证表明，对 $\sigma \leq 0.003$，有 $f(\sigma) \leq 0.1$。换言之，任何奇异值低于约 $0.003$ 的方向，在经过五步 NS 后仍基本处于**未正交归一化**状态——相比奇异值较大的方向，其对更新的有效贡献被抑制了 $10\times$ 或更多。

参见图注图 2：标准多项式 $p(x) = 2x - 1.5x^3 + 0.5x^5$ 对应的 NS 映射 $f(\sigma) = p^{\circ 5}(\sigma)$。左图展示完整范围 $\sigma \in [0, 1]$，右图为 $\sigma \in [0, 0.05]$ 区域的放大视图。在实践中，不同实现使用不同的 NS 配置。NanoGPT 速度竞赛（Jordan et al., 2024a）使用