Muon 的几何解释与实现约定#

TL;DR#

本页内容7 个部分

1. 背景与问题#

入门篇用理想 polar 方向解释 Muon，但实际训练只能做有限步迭代，还要决定矩阵边界和更新尺度。要把几何直觉变成可检查的方法，需要分别回答方向为什么自然、近似怎么算、尺度怎样校准。

2. 目标与定义#

Muon：方向、近似与尺度的分工

2.1 Muon 到底改变了什么#

设隐藏层权重与梯度为：

W_t,G_t\in\mathbb R^{m\times n}.

普通动量先形成：

B_t=\mu B_{t-1}+G_t.

若它的奇异值分解为：

B_t=U\Sigma V^\top,

Muon 使用的理想方向是：

O_t=UV^\top.

这一步有三个同时成立的解释：

解释	含义
polar factor	$UV^\top$ 是 $B_t$ 的矩形 polar factor
最近半正交矩阵	满秩时，在 Frobenius 距离下它是离 $B_t$ 最近的 row-orthogonal 或 column-orthogonal 矩阵
matrix sign / zeroth power	把每个非零奇异值从 $\sigma_i$ 映射到 $1$ ，记作 $B_t^0$

因此 Muon 不是把整个梯度除以一个标量，也不是逐元素做归一化。它保留 $U,V$ 所描述的方向结构，只改奇异值：

U\,\mathrm{diag}(\sigma_1,\ldots,\sigma_r)V^\top \quad\longrightarrow\quad U\,\mathrm{diag}(1,\ldots,1)V^\top.

如果原更新被少数大奇异值主导，polar 变换会压平大方向，同时显著抬升原来较小的方向。原始 Muon 说明把这些称为可能有学习价值的 “rare directions”；这是合理的经验动机，但不是已经证明完毕的因果机制。

3. 方法与推导#

3.1 为什么 polar factor 是一种自然的矩阵更新#

对当前位置做一阶近似：

L(W+\Delta W)\approx L(W)+\langle G,\Delta W\rangle_F.

如果约束更新的 Frobenius norm，最陡方向就是归一化梯度；如果把一层看作线性算子，并约束它的 spectral norm：

\|\Delta W\|_2\le 1,

那么线性化目标的最优方向满足：

\underset{\|A\|_2\le1}{\arg\min}\ \langle G,A\rangle_F =-UV^\top.

原因是 spectral norm 的对偶范数是 nuclear norm：

\max_{\|A\|_2\le1}\langle G,A\rangle_F =\|G\|_*=\sum_i\sigma_i(G).

所以可以把 Muon 理解为一种 spectral steepest descent：每个奇异方向最多走单位 operator-norm 步长，而不是让最大的奇异方向吃掉大部分预算。[1] 系统整理了这类“优化器对应哪一种范数几何”的解释。

不过要注意边界：实际 Muon 对动量矩阵而不是当前梯度取 polar factor，还会加入 Nesterov、有限步 NS、shape scaling 和 weight decay。严格的单步 steepest-descent 推导解释了方向的几何意义，不等于完整训练动力学的证明。

3.2 一次实际的 Muon step#

常见的 Nesterov 版本可以写成：

B_t=\mu B_{t-1}+G_t,

M_t=G_t+\mu B_t,

O_t\approx\mathrm{polar}(M_t),

W_{t+1}=(1-\eta\lambda)W_t-\eta\,s(m,n)O_t.

其中：

$B_t$ 是唯一的动量状态；
$O_t$ 由若干步 Newton–Schulz 得到；
$s(m,n)$ 处理矩阵形状带来的更新尺度差异；
$(1-\eta\lambda)W_t$ 是 decoupled weight decay；
不适合 Muon 的参数通常交给同一个 optimizer wrapper 内部的 AdamW 分支。

对应的简化伪代码是：

def muon_step(W, G, B, lr, momentum, weight_decay, ns_steps):
    B = momentum * B + G
    M = G + momentum * B                 # Nesterov-style
    O = newton_schulz_zeroth_power(M, ns_steps)
    O = adjust_update_scale(O, W.shape)
    W = (1 - lr * weight_decay) * W - lr * O
    return W, B

这里最重要的顺序是：先积累 momentum，再做 orthogonalization。如果先正交化每步梯度、之后才积累动量，非线性变换与时间平均不能交换，得到的是另一种优化器。

3.3 Newton–Schulz 如何避免显式 SVD#

对输入矩阵先做尺度归一化：

X_0=\frac{M}{\|M\|_F+\varepsilon}.

由于 $\|M\|_2\le\|M\|_F$ ，归一化后奇异值落在 $[0,1]$ 。若矩阵高大于宽，实践中通常先转置，让较短维位于 Gram 矩阵一侧。

一轮五次 NS 写成：

X_{k+1}=aX_k+bX_kX_k^\top X_k+c(X_kX_k^\top)^2X_k.

令：

X_k=U\Sigma_kV^\top,

就有：

X_{k+1}=U\left(a\Sigma_k+b\Sigma_k^3+c\Sigma_k^5\right)V^\top.

也就是说，矩阵乘法没有改变奇异向量，只是对每个奇异值反复应用标量多项式：

\phi(x)=ax+bx^3+cx^5.

精确收敛与快速谱整形不是同一个目标#

一组保守系数可以让 $\phi^k(x)$ 最终收敛到 $1$ 。Muon 常用的系数则是：

(a,b,c)=(3.4445,-4.7750,2.0315).

它特意增大 $\phi'(0)=a$ ，让很小的奇异值在前几步快速长大；代价是迭代不再严格收敛到全体奇异值都等于 $1$ ，而是在 $1$ 周围形成一个带宽。原始实验发现这种近似已经足够，通常使用 5 步。[2]

因此更准确的说法是：

实用 Muon 的 NS 是一个固定预算的谱整形器，而不一定是追求高精度 polar decomposition 的数值求解器。

PyTorch 当前实现也明确记录了这个区别：默认快速多项式得到的是类似 $US'V^\top$ 的结果，而不是精确 $UV^\top$ 。[3]

3.4 近似质量、尺度与参数分组#

怎样判断 NS 迭代是否“充分”#

“5 步够不够”至少有三个不同问题：

层次	问题	需要的证据
尺度充分	输出的总体能量是否已经达到目标？	RMS / Frobenius norm
数值充分	每个有效奇异值是否都接近目标带？	orthogonality residual 或抽样 SVD
优化充分	再加 NS 步数是否还能改善训练？	固定预算的训练 ablation

用 partial isometry 的 RMS 做第一道检查#

设理想输出：

X_*=U_{:,:r}V_{:r,:}\in\mathbb R^{m\times n},

其中 $r$ 是矩阵秩。它有 $r$ 个等于 $1$ 的奇异值，因此：

\|X_*\|_F^2=r,

\mathrm{RMS}(X_*) =\frac{\|X_*\|_F}{\sqrt{mn}} =\sqrt{\frac{r}{mn}}.

满秩时 $r=\min(m,n)$ ：

\mathrm{RMS}(X_*)=\frac{1}{\sqrt{\max(m,n)}}.

这个推导出现在 [4]，但它更一般的用途是给 NS 输出提供一个无量纲健康度指标：

\rho_k =\sqrt{\max(m,n)}\,\mathrm{RMS}(X_k) =\frac{\|X_k\|_F}{\sqrt{\min(m,n)}}.

对于精确、满秩 polar factor：

\rho_*=1.

满秩时 $r=\min(m,n)$ ，从奇异值看：

\rho_k^2 =\frac{1}{r}\sum_{i=1}^{r}\sigma_i(X_k)^2.

所以逐步记录 $\rho_0,\rho_1,\ldots,\rho_K$ ，可以检查小奇异值的总体能量是否已经被抬升、不同形状的矩阵是否在相同步数达到稳定区间，以及新 kernel 是否改变数值路径。

为什么 RMS 接近 1 仍然不够#

考虑两个 $2\times2$ 矩阵：

X_a=\begin{bmatrix}1&0\\0&1\end{bmatrix}, \qquad X_b=\begin{bmatrix}\sqrt2&0\\0&0\end{bmatrix}.

二者的 Frobenius norm 都是 $\sqrt2$ ，所以归一化 RMS 指标都是 $1$ ；但 $X_b$ 显然没有完成正交化。RMS 只看：

\frac1r\sum_i\sigma_i^2,

看不到奇异值之间的离散程度。

因此抽样诊断还要加 orthogonality residual。若 $m\ge n$ ：

e_{\mathrm{orth}} =\frac{\|X_k^\top X_k-I_n\|_F}{\sqrt n}.

若 $m<n$ ：

e_{\mathrm{orth}} =\frac{\|X_kX_k^\top-I_m\|_F}{\sqrt m}.

对于刻意不精确收敛的快速 Muon 多项式， $e_{\mathrm{orth}}$ 不会趋近机器精度， $\rho_k$ 也未必严格趋近 $1$ 。阈值应由所用系数、dtype 和步数的可信参考实现校准，不能把精确 polar 的理论值直接当生产 gate。

诊断	计算代价	能发现什么	不能证明什么
$\rho_k$ 随步数变化	一次平方和归约	尺度未饱和、异常 shape、溢出/下溢	奇异值是否集中在目标带
$e_{\mathrm{orth}}$	一次额外 Gram + 归约	行/列正交偏差	是否改善最终训练
fp32 SVD 抽样	高	完整奇异值分布、与 $UV^\top$ 的距离	全训练期间都安全
$K=3/5/8$ 训练 ablation	很高	优化意义上的“够用”	单个 kernel 的数值原因

同一个公式也解释了 shape scaling#

满秩 polar 输出满足：

\mathrm{RMS}(O_t)=\frac1{\sqrt{\max(m,n)}}.

所以不同形状的参数即使用同一个标量学习率，逐元素更新 RMS 也不同。若直接乘：

s(m,n)=\sqrt{\max(m,n)},

理想满秩输出就会有单位 RMS。这消除了一个由矩阵最大维度引入的 implicit learning-rate factor。

但“Muon scaling”并不是唯一约定。以权重形状 $A\times B$ 表示，PyTorch 当前暴露了三种选择：

名称	缩放因子	主要语义
`original`	$\sqrt{\max(1,A/B)}$	Keller 参考实现的宽深比修正
`match_rms_adamw`	$0.2\sqrt{\max(A,B)}$	先统一 update RMS，再经验性匹配 AdamW 尺度
`spectral_unclamped`	$\sqrt{A/B}$	按 spectral / modular 参数化使用未截断宽深比

这三种规则不能脱离对应的 base learning rate 比较；把 scaling 换掉却照搬学习率，等于同时改变优化器。0.2 也是实验校准值，不由 Appendix A 推出。

还有一个低秩边界。若有效秩 $r<\min(m,n)$ ，乘上 $\sqrt{\max(m,n)}$ 后：

\mathrm{RMS}\!\left(\sqrt{\max(m,n)}X_*\right) =\sqrt{\frac{r}{\min(m,n)}}<1.

因此“shape scaling 让所有层 RMS 一致”依赖满秩或近似满秩假设；低秩梯度、有限 NS 步数和低精度都会让实际尺度偏离。

参数怎样分组，本身就是算法的一部分#

Muon 的默认边界不是“所有二维 tensor 都自动使用”。更稳妥的路由是：

参数	常见优化器	原因
attention / MLP 隐藏层权重	Muon	有明确的二维线性算子结构
embedding	AdamW	输入层的参数化和训练动力学不同
LM head / classifier	AdamW	经验上单独处理更稳妥
bias、norm gain、标量/向量	AdamW	没有自然的矩阵 polar 几何
convolution kernel	视实现而定	可展平成 `[out_channels, -1]` 后使用 Muon

特别要注意：polar 变换不对任意 reshape、拼接或切分保持不变。

polar(concat(Q, K, V))  ≠  concat(polar(Q), polar(K), polar(V))

所以 fused QKV 是作为一张大矩阵正交化，还是拆成 Q、K、V 三张矩阵，会得到不同更新。原始 Muon 实验倾向于分别处理 Q/K/V；更一般地说，参数分组定义了 Muon 眼中的线性算子边界，不能只当成 kernel batching 的实现细节。

4. 计算过程#

用 $M=\operatorname{diag}(3,1)$ 手算一次有限精度之外的理想实数迭代，隔离多项式本身的作用。按 Frobenius 范数归一化，得到 $X_0=M/\sqrt{10}$ ，奇异值为 $(0.948683,0.316228)$ 。

取本篇前面定义的快速多项式 $\phi(x)=3.4445x-4.7750x^3+2.0315x^5$ ，逐个作用于奇异值：

步骤	第一奇异值	第二奇异值
归一化输入	0.948683	0.316228
一次多项式后	0.751846	0.944672
理想 polar 目标	1	1

本例 $X_1$ 仍为对角矩阵， $\rho_1=\|X_1\|_F/\sqrt2\approx0.853720$ 。一次迭代抬高了小奇异值，却没有得到精确正交矩阵；继续迭代的数值质量应按第 3 节协议检查。

若示意取 $W=I_2$ 、学习率 0.1、尺度因子 1 且无衰减，本步得到 $W'=\operatorname{diag}(0.924815,0.905533)$ 。这是一次实数多项式演算，不是 BF16 kernel 或训练实验；它只串起输入、近似方向、诊断量与更新。

5. 比较与边界#

5.1 Muon 与 SGD、AdamW、Shampoo 的关系#

优化器	主要状态	更新变换	保留的结构
SGD-momentum	一阶动量	标量学习率	全局欧氏几何
AdamW	一阶 + 对角二阶矩	逐元素预条件	坐标级尺度
Shampoo	每个 tensor 维度的二阶统计	Kronecker 预条件	tensor 轴结构
Muon	一阶矩阵动量	当前动量的近似 matrix sign	整张矩阵的奇异向量结构

Muon 与 Shampoo 有一个直接代数联系。若 Shampoo 不累计历史 preconditioner，只使用当前矩阵 $G=U\Sigma V^\top$ ：

(GG^\top)^{-1/4}G(G^\top G)^{-1/4} =UV^\top.

因此无动量 Muon 可以看作一种“instantaneous / accumulation-free Shampoo”；加上先于 polar 的 momentum 后，就得到实际 Muon 的核心形态。区别仍然很大：Shampoo 的重点是累计曲率代理与周期性矩阵逆根，Muon 的重点是对当前动量做固定预算的谱整形。[5]

还有一种后来的理论视角把 matrix sign 解释为 nuclear norm 的次梯度，并研究 decoupled weight decay 下与 spectral-norm constraint 的关系。这为隐式正则化提供了线索，但依赖相应假设，不能直接当作大模型训练中所有收益的完整解释。[6]

5.2 系统实现：真正昂贵的是矩阵乘法，不是 optimizer state#

相较 AdamW 的一阶矩与二阶矩，Muon 分支通常只保存一份 momentum buffer；额外计算来自每步对二维更新执行若干轮矩阵乘法。

标准五次 NS 的一轮核心是：

A = X @ X.mT
B = b * A + c * (A @ A)
X = a * X + B @ X

它带来几个系统问题：

单个小矩阵无法充分占满 GPU，需要按 shape batching 或 grouped scheduling；
$XX^\top$ 、 $A^2$ 是对称结果，可用 SYRK / symmetric GEMM 少算约一半 MMA；
$BX$ 是一般矩阵乘，不能利用输出对称性；
长方形参数每轮反复触碰大矩阵，适合改写到较小 Gram 空间；
分布式场景必须明确完整矩阵在哪个 rank 上完成正交化，以及更新/参数怎样同步。

这些优化不能改变上一节的数值诊断。替换 kernel 后至少要对齐：

\rho_k,\qquad e_{\mathrm{orth}},\qquad \text{singular-value quantiles},

而不是只比较最终输出的单个 max_abs_error。

标准 NS、Gram 空间重排、restart 稳定化和 batched SYRK 的详细 FLOP 分析见 Gram Newton–Schulz。

5.3 证据应该怎样读#

Muon 的证据来自不同层次，不能混成一个“快两倍”的结论：

证据	支持的结论	不支持的结论
小模型 speedrun 与受控对比	在这些设置下 sample efficiency / time-to-loss 更好	所有模型和任务都更好
大规模 LLM scaling 实验	Muon 可以扩展到较大训练，并可能降低达到目标 loss 的训练 FLOP	optimizer kernel 自身快 2×
NS 数值诊断	固定步数已达到某种谱整形质量	最终 validation loss 一定更好
norm geometry / 收敛分析	提供方向与隐式约束的理论解释	完整解释深网非凸训练中的收益来源

[7] 报告了接近同等 loss 时约 $0.52\times$ 训练 FLOP 的结果。它应被记录为特定模型、配方和拟合方法下的实验结果；这与 NS kernel 的局部加速比、optimizer step 占比和端到端 wall-clock 是三件不同的事。

当前仍值得单独验证的问题包括：

收益主要来自压低头部奇异值，还是抬升尾部奇异值？
不同层真正需要的 NS 步数是否一致？
快速多项式的目标带应怎样随 dtype、shape 和训练阶段变化？
Q/K/V、expert 和 tensor-parallel shard 应以哪种矩阵边界正交化？
预训练中的收益能否稳定迁移到 finetuning、持续训练和 RL？
shape scaling、parameterization 与 base learning rate 怎样联合标定？

5.4 实现与实验检查表#

开始一组 Muon 实验前，至少确认：

[ ] Muon 与 AdamW 参数组互斥且完整，embedding/head/bias/norm 没有误入；
[ ] Q/K/V 是按预期拆分或拼接，而不是被框架偶然决定；
[ ] momentum、Nesterov、weight decay 的顺序与参考实现一致；
[ ] NS 输入先缩放，长宽转置逻辑正确；
[ ] 明确记录 NS 系数、步数、计算 dtype；
[ ] 明确记录 shape scaling 规则和对应 base LR；
[ ] 对代表性 shape 记录每轮 $\rho_k$ ，抽样检查 $e_{\mathrm{orth}}$ 与奇异值分位数；
[ ] kernel 替换同时比较数值、optimizer 时间和端到端 step time；
[ ] 训练对比使用相同 token、scheduler、batch 和调参预算。

参考资料#

Old Optimizer, New Norm. 原文。
原始设计说明. 原文。
PyTorch _muon.py. 原文。
Muon is Scalable 的 Appendix A. 原文。
Shampoo 原论文. 原文。
Muon Optimizes Under Spectral Norm Constraints. 原文。
Muon is Scalable for LLM Training. 原文。