Muon需要多少正交化？

# Muon 需要多大的正交化？
来源：https://arxiv.org/html/2606.00371

###### 摘要

Muon 优化器通过将病态动量更新替换为近似半正交更新来改进神经网络训练。这引出一个实际问题：Muon 到底需要多大的正交化？我们使用一个直接从 Muon 低精度奇异值带导出的松弛三次牛顿-舒尔茨调度来研究这个问题。由此产生的五步三次构造使用了十次主导矩阵乘法，而五次五次牛顿-舒尔茨迭代则需要十五次。三次调度并非旨在作为更精确的极分解求解器；相反，它是一种原则性的低成本变体，使我们能够探究极分解精度、频谱塑形和训练质量之间的关系。通过综合诊断、NanoGPT 消融实验以及混合 MoE/Mamba 模型的训练实验，我们发现训练质量并非单调地由极分解精度决定：截断的 Polar Express、Muon-Jordan、三次牛顿-舒尔茨以及显式的 FP32 SVD 极分解因子在 GPT-2 Small 上可以达到几乎无法区分的最终损失，而 cubic5 在具有十亿到四十亿参数的混合 MoE/Mamba 模型上，其验证损失与 Muon-Jordan 五次更新的差距约在10^{-3}以内。这些结果支持 cubic5 作为一种实用的低成本 Muon 正交化变体，并在测试的设置中提供了训练质量一致性的经验证据。

## 1 引言

### 1.1 Muon 优化器

Muon 是一种新提出的用于神经网络中矩阵值隐藏层参数的优化器[9 (https://arxiv.org/html/2606.00371#bib.bib1)]。它可以被视为带有动量的随机梯度下降，随后进行正交化步骤。设W_t ∈ R^{m×n}为权重矩阵，G_t为其随机梯度，M_t为动量缓冲区。简化的 Muon 更新为

M_t = β M_{t-1} + (1-β) G_t, (1)
W_{t+1} = W_t - η polar(M_t). (2)

如果M_t = U Σ V^⊤是奇异值分解，那么

polar(M_t) = U V^⊤. (3)

因此，精确的 Muon 将保留动量更新的奇异向量，但将所有非零奇异值替换为1。Jordan 等人通过观察发现，变压器的动量更新通常是近似低秩的：少数方向主导更新，而许多“稀有”方向的奇异值较小，但可能对学习仍然重要，从而在经验上激发了这一操作[9 (https://arxiv.org/html/2606.00371#bib.bib1)]。因此，正交化可以理解为一种频谱重塑操作，用于放大这些被抑制的方向。

这种视角与计算高度精确极分解的标准数值线性代数目标不同。在神经网络训练中，更新在与动量、学习率调度、归一化、权重衰减、随机梯度和低精度算术交互后仍必须保持有用。精确的 SVD 极分解因子是自然的数学参考点，但它不必是训练损失最优的。因此，本工作的核心问题不仅仅是牛顿-舒尔茨迭代是否收敛到极分解因子，而是 Muon 必须正交化到何种精度才能保持训练质量。

这些迭代的关键观察在于，奇次矩阵多项式对角作用在奇异值上。如果X = U Σ V^⊤且

X^+ = a X + b (X X^⊤) X + c (X X^⊤)^2 X, (4)

那么

X^+ = U (a Σ + b Σ^3 + c Σ^5) V^⊤. (5)

因此，矩阵问题可以通过奇异值上的标量多项式映射来研究。经典极分解[6,5,2,11,7,13,12]算法以及最近提出的混合极分解方法[8]优化这些多项式以实现高精度收敛到1。然而，Muon 用于深度学习场景，其中低精度的近似方向可能已经足够。这为多项式调度创造了空间，这些调度作为极分解求解器可能精度较低，但作为频谱变换则更便宜或有不同的偏差。

### 1.2 用于 Muon 的多项式极分解

原始的 Muon 实现使用了五次固定的五次牛顿-舒尔茨多项式迭代，

p(x) = 3.4445 x - 4.7750 x^3 + 2.0315 x^5, (6)

在 Frobenius 归一化后以 bfloat16 精度运行[9 (https://arxiv.org/html/2606.00371#bib.bib1)]。我们将此基线称为 Muon-Jordan。其系数是根据松弛的 Muon 目标选择的：在经过多次复合后，奇异值允许位于1附近的一个带状区域内，而不是收敛到机器精度。Liu 等人[10 (https://arxiv.org/html/2606.00371#bib.bib2)]进一步证明，Muon 可以扩展到包含160亿参数的更大混合专家（MoE）模型。

Polar Express 从矩阵符号和极分解方法[1 (https://arxiv.org/html/2606.00371#bib.bib3)]的角度重新审视了多项式设计问题。它使用由最小化最大误差准则选择的自适应五次多项式序列。这提高了极分解近似的质量，并且在插入 Muon 后，可以在学习率扫描中改善验证损失。Polar Express 保持了与标准五次牛顿-舒尔茨相同的基本计算结构：每次迭代使用三次主导矩阵乘法应用一个五次奇次多项式。

一个互补的方向是通过重构计算来降低相同多项式迭代的成本。Gram 牛顿-舒尔茨观察到许多 Muon 矩阵是矩形的，而标准牛顿-舒尔茨会重复形成对称的 Gram 矩阵[15 (https://arxiv.org/html/2606.00371#bib.bib4)]。通过将更多工作转移到较小的 Gram 矩阵并使用硬件感知的对称内核，Gram 牛顿-舒尔茨在保持底层多项式迭代的同时减少了正交化步骤的运行时间。这与系数设计是正交的：可以改进多项式、内核，或者两者兼顾。

### 1.3 贡献

本工作使用一个松弛的三次构造来探测一个更广泛的问题：Muon 需要多大的正交化，以及什么样的频谱塑形对训练重要？我们做出以下贡献。

- 我们推导了一个自适应三次牛顿-舒尔茨调度，其系数根据当前最坏情况下的奇异值下界和一个松弛的目标带[0.7, 1.3]选择。
- 我们确定了一个实用的五步调度，使用 bfloat16 有效下界 l_0 = 7×10^{-3}。该调度使用10次主导矩阵乘法，而五次五次迭代需要15次。
- 我们在 NanoGPT 训练中比较了 cubic5、截断的 Polar Express、截断的 Muon-Jordan 以及显式的 FP32 SVD 极分解因子。这将极分解精度与优化器质量分开：在我们的 GPT-2 Small 运行中，精确 SVD 并未优于最强的近似牛顿-舒尔茨更新。
- 我们刻画了一个成本-质量前沿。Cubic5 在默认 GPT-2 Small 设置下接近 Polar Express 和 SVD，并且在具有1B到4B参数的混合 MoE/Mamba 模型上，其验证损失与 Muon-Jordan 五次更新的差距保持在约10^{-3}以内。这些结果支持 cubic5 在评估的机制中作为一种可行的低成本 Muon 正交化变体，而不是一个普遍更好的更新规则。
- 我们在合成矩阵上验证了标量构造，并通过微基准测试表明，降低多项式次数会降低正交化子例程的成本。

我们的声明是经过校准的。Cubic5 并非作为比 Polar Express 更精确的极分解算法提出，也不是一个普遍更好的 Muon 更新。相反，它是一种实用的低成本 Muon 正交化变体，其经验行为也有助于探究多项式诱导的频谱变换如何影响训练，而不仅仅是经典的极分解收敛性。

## 2 用于 Muon 优化器的松弛三次牛顿-舒尔茨方法

### 2.1 松弛的奇异值目标

经典极分解旨在将每个非零奇异值映射到1。Muon 放宽了这一要求。原始 Muon 系数搜索的指导观察是，训练可以容忍奇异值处于诸如[0.7, 1.3]的带状区域内[9 (https://arxiv.org/html/2606.00371#bib.bib1)]。这改变了多项式设计问题：松弛带是一个有用的设计目标，但并不保证训练质量仅由每个奇异值是否落在此区间内决定。

我们使用归一化迭代 X_0 = X / ||X||_F，因此所有奇异值至多为1。在精确算术中，最小的非零奇异值可能任意小。由于 Muon 用于 bfloat16 精度的神经网络训练场景，远低于相对精度的下界实际上没有意义。因此，我们使用

l_0 = 7×10^{-3} < ε_bf16 = 2^{-7} = 0.0078125 (7)

作为有效下界。该值产生一个调度，在五次三次迭代后达到松弛的下界目标0.7。Muon-Jordan 和 Polar Express 使用更保守的下界 l_0 = 10^{-3}。如果我们使用相同的下界，调度需要七次三次迭代。

### 2.2 三次系数的推导

我们使用与 Chen–Chow 缩放三次牛顿-舒尔茨方法相同的方法推导三次系数[3 (https://arxiv.org/html/2606.00371#bib.bib7)]。在迭代 t 时，假设所有受保护的奇异值位于区间 [l_t, r_t] 内。我们选择一个奇次三次多项式

f_t(x) = a_t x + b_t x^3, b_t < 0, (8)

其峰值为 u = 1.3。由于 b_t < 0，该多项式增加到其临界点 k_t，然后下降。我们施加三个条件：

f_t'(k_t) = 0, (9)
f_t(k_t) = u, (10)
f_t(l_t) = f_t(r_t). (11)

前两个条件一旦 k_t 已知就确定了形状。由 f_t'(x) = a_t + 3 b_t x^2 得

a_t = -3 b_t k_t^2. (12)

将其与 f_t(k_t) = u 结合得

b_t = -u / (2 k_t^3), a_t = 3u / (2 k_t). (13)

端点条件确定了峰值位置。将上述形式代入 f_t(l_t) = f_t(r_t) 得

3 l_t k_t^2 - l_t^3 = 3 r_t k_t^2 - r_t^3, (14)

因此

k_t^2 = (r_t^2 + r_t l_t + l_t^2) / 3. (15)

等价地，如果

α_t = 1/k_t = sqrt(3 / (r_t^2 + r_t l_t + l_t^2)), (16)

那么

f_t(x) = u/2 (3 α_t x - α_t^3 x^3). (17)

端点相等是重要的。因为三次多项式在受保护的区间上只有一个内部最大值，在 [l_t, r_t] 上的最小值出现在其中一个端点。强制 f_t(l_t) = f_t(r_t) 平衡了两个端点，并在选择的峰值约束下最大化新的最坏情况下的下界。

### 2.3 自适应最坏情况调度

在每一步根据当前下界重新计算系数。我们初始化

l_0 = 7×10^{-3}, r_0 = 1, (18)

并且在第一步后通过设置 r_t = u = 1.3 (t ≥ 1) 来保护松弛的上界范围。最坏情况下的下界演变为

l_{t+1} = f_t(l_t). (20)

对于 l_0 = 7×10^{-3}，此构造在五步后达到 l_t ≥ 0.7。由此产生的 cubic5 系数为

对应的矩阵迭代为

X_{t+1} = a_t X_t + b_t (X_t X_t^⊤) X_t, (21)

对于高矩阵使用常见的转置技巧，以便形成更小的 Gram 矩阵。这保留了奇异向量并将 f_t 应用于每个奇异值。

参阅图注图1：有效 bfloat16 下界 l_0 = 7×10^{-3} 的自适应松弛三次牛顿-舒尔茨调度。系数根据当前最坏情况下的下界重新计算，标量下界轨迹在五次三次迭代后达到松弛的目标带。
### 2.4 FLOP 与矩阵乘法计数

我们只计算主导矩阵乘法，忽略标量乘法、矩阵加法和归一化。一个三次步骤

X_{t+1} = a_t X_t + b_t (X_t X_t^⊤) X_t (22)

需要两次主导乘法：一次形成 Gram 矩阵，一次将其乘回 X_t。因此五次三次步骤需要十次主导矩阵乘法。

一个五次牛顿-舒尔茨步骤的形式为

X_{t+1} = a_t X_t + b_t (X_t X_t^⊤) X_t + c_t (X_t X_t^⊤)^2 X_t (23)

需要三次主导乘法：一次用于 Gram 矩阵，一次用于 Gram 平方或 Gram 中的多项式项，以及一次将结果乘以 X_t。五次五次步骤，包括 Muon-Jordan 和 Polar Express，因此需要十五次主导乘法。

对于 m×n 矩阵且 m ≥ n，每个三次步骤的主要浮点运算成本大约为

2 m n^2 + 2 m n^2 = 4 m n^2 (24)

当使用较小的 n×n Gram 矩阵时。五次步骤增加了一个 n×n 乘法，得到大约

4 m n^2 + 2 n^3. (25)

对于高度矩形的矩阵，矩形乘法占主导，三次调度节省了大约三分之一的主导乘法。对于接近方形的矩阵，n^3 Gram 平方项也很大，因此收益仍然有意义。实际墙钟加速取决于批处理、内核融合、对称内核使用、矩阵纵横比以及正交化子例程是否占优化器步骤的很大一部分[15 (https://arxiv.org/html/2606.00371#bib.bib4)]。

## 3 数值结果

实验围绕两个问题组织。首先，松弛的三次推导是否产生预期的频谱变换和子例程成本降低？其次，一旦将正交化例程插入 Muon，训练质量与极分解精度的关联程度如何？这第二个问题对于优化器设置至关重要：一个更新在数值上作为极分解因子可能不太精确，但对于训练仍然同样有用。

### 3.1 奇异值诊断与正交化微基准