训练一个 linear layer 时，参数是矩阵，但 SGD 和 AdamW 的更新规则并没有充分利用这个矩阵的结构。Muon (Jordan et al., 2024) 的做法是：先用 momentum 平滑梯度，再调整更新矩阵的奇异值，让不同方向的更新具有相近的尺度。名字里的 MomentUm Orthogonalized by Newton–Schulz，说的就是这个过程。
从逐元素缩放到矩阵更新
设一层的计算是 \(y=Wx\)，其中 \(W\in\mathbb{R}^{m\times n}\)，梯度为 \(G_t=\nabla_W\mathcal{L}_t\)。SGD 直接沿着 \(-G_t\) 更新。AdamW (Loshchilov & Hutter, 2019) 则保存梯度的一阶、二阶矩估计 \(M_t,V_t\)，用经过 bias correction 的 \(\widehat M_t,\widehat V_t\) 做逐元素缩放：
这里的除法和平方根都逐元素进行，\(\eta_t\) 是 learning rate，\(\lambda\) 是 decoupled weight decay 的系数。AdamW 能根据每个坐标的历史梯度调整步长，但矩阵还描述了输入方向到输出方向的映射。Muon 从这些方向入手构造更新。
先看不带 Nesterov 的版本。令 \(M_{-1}=0\)，定义
\(\operatorname{Polar}\) 是下面要解释的矩阵变换，\(s_{m,n}\) 用来控制不同形状矩阵的更新尺度。Momentum 先于这个变换计算：累积的是梯度的历史，再变换平滑后的矩阵。
把奇异值拉到相同尺度
暂时省略时间下标，把待处理的 momentum 矩阵写成 compact SVD：
\(v_i\) 是输入空间里的一个单位方向，\(u_i\) 是对应的输出方向，两组向量各自正交。等式 \(Mv_i=\sigma_i u_i\) 告诉我们：这个更新在方向 \(v_i\) 上的强度是 \(\sigma_i\)。如果奇异值相差很大，更新便主要集中在少数方向。
Muon 的理想变换是
也就是保留这些方向，把每个非零奇异值变成 \(1\)。这和把整个矩阵除以 Frobenius norm 不一样：整体归一化保留奇异值之间的比例，polar 变换会改变这个比例。
例如 \(M=\operatorname{diag}(8,2,0.5)\)，除以 \(\lVert M\rVert_F\) 以后，三个方向仍然相差 \(16:4:1\)；理想 polar 变换得到 \(I_3\)，三个方向的强度相同。下面也画出了实际五步迭代的结果，计算方法见下一节。
这里的 orthogonalization 指的是更新矩阵。加上这个更新之后，权重 \(W\) 本身不必正交，也不是在要求相邻两个训练 step 的更新互相垂直。对于满秩矩形矩阵，只能有 \(O^\top O=I_n\)（\(m\ge n\)）或 \(OO^\top=I_m\)（\(m\le n\)），所以更准确地说是 semi-orthogonal。
如果 \(M\) 秩亏，上面的 compact SVD 只保留非零奇异值，\(UV^\top\) 是一个 partial isometry，不能声称它的所有行或所有列都正交。多项式迭代也不会凭空生成缺失的方向：在精确算术下，零奇异值始终为零。
和 sign、向量归一化有什么关系？
苏剑林在《Muon 优化器赏析》(Su, 2024) 中将这个 polar 变换记作 \(\operatorname{msign}\)。对非零标量 \(z\)，\(\operatorname{sign}(z)=z/\sqrt{z^2}\)；矩阵也有类似的表达式：
这里的逆平方根作用于矩阵的特征值；遇到零特征值时采用伪逆约定，将对应值保留为零。因此它一般不同于逐元素取 \(\operatorname{sign}\)。两个特例能说明矩阵形状的意义：
同一组数值，作为对角矩阵中的独立缩放系数，结果是逐元素取符号；作为一个列向量，结果是整体归一化。参数在计算中扮演什么角色，会影响合适的更新方式。这里的 \(\operatorname{msign}\) 指基于奇异值的变换，不是一般数值线性代数中基于特征值定义的 matrix sign function。
为什么这种更新可能有用？如果小奇异值对应的方向也包含学习信号，拉平谱就能提高它们相对于主导方向的更新强度。不过，小方向也可能包含噪声；这种直觉本身并不保证泛化更好。Momentum 能平滑输入给变换的梯度，实际收益仍然需要实验验证。
另一个出发点是约束特征的变化。《为什么我们选择尝试 Muon？》(Su, 2025) 从这个角度解释了谱范数：固定这一层的输入 \(x\)，更新权重后有 \(\Delta y=\Delta W x\)，而
所以，限制谱范数就是限制这个线性层的输出变化相对于输入的最大幅度，再在这个范围内尽量降低 loss。“下降最快”需要先说明怎样衡量一步的大小；Frobenius norm 和 spectral norm 给出的最优方向不同。
为什么 spectral norm 约束会给出这个方向？
Bernstein 和 Newhouse (Bernstein & Newhouse, 2024) 从 norm 约束下的 steepest descent 解释了这种矩阵更新。先忽略 momentum，用一阶近似
如果限制 \(\lVert\Delta W\rVert_F\le\eta\)，最优方向是 \(-G/\lVert G\rVert_F\)。换成 spectral norm 约束
其中 \(\lVert A\rVert_2\) 是最大奇异值，就得到一个最优解 \(\Delta W=-\eta UV^\top\)。因为若 \(G=\sum_i\sigma_i u_i v_i^\top\)，则
而 \(-\eta UV^\top\) 能取到等号。这里 \(\sum_i\sigma_i\) 也叫 nuclear norm。
这解释了单层、固定输入下的更新方向。整个网络同时更新时，输入特征也会改变；带 momentum、有限步近似的实际 Muon，也不等于每一步都精确求解当前 loss 的这个问题。
用 Newton–Schulz 避免 SVD
SVD 适合解释方向，但每个 step 都对每个权重矩阵做 SVD，计算成本太高。Newton–Schulz 可以只用矩阵乘法近似这个变换，更适合 GPU。
先归一化
因为最大奇异值不超过 Frobenius norm，\(X_0\) 的奇异值都小于 \(1\)。经典的三次迭代是
将 \(X_k=U\Sigma_kV^\top\) 代入，得到
矩阵乘法在这里实现了对每个奇异值分别应用 \(f(x)=\frac32x-\frac12x^3\)：奇异向量不变，位于 \((0,1)\) 的奇异值逐渐靠近 \(1\)。整个过程不需要显式算出 \(U,\Sigma,V\)。
实际的 Muon 实现 使用更激进的五次多项式，常用 \(5\) 步：
对应的标量变换是 \(f(x)=ax+bx^3+cx^5\)。它在零附近的斜率更大，可以更快放大小奇异值；代价是它并不精确收敛到 \(1\)，例如 \(f(1)=0.701\)。所以实际输出是 \(U\widetilde\Sigma V^\top\)，把奇异值拉到相近尺度，而非严格得到 \(UV^\top\)。前图的第三栏正是这个有限步计算，不能把更多迭代简单理解为更精确的正交化。
极小的非零奇异值在有限步后也可能仍然很小。苏剑林在 Newton–Schulz 专文 (Su, 2025b) 中讨论了另一种改进：每一步采用不同的 \((a_k,b_k,c_k)\)，针对剩余的奇异值区间设计后续变换。同样的迭代步数不必共享系数；上面的固定系数五步法只是一个具体实现，计算近似更准确也不自动意味着训练效果更好。
示例每步的奇异值
从对角矩阵的奇异值 8、2、0.5 开始，先除以 Frobenius norm 加 1e-7，再使用正文的五次多项式。表格保留三位小数，各列跟踪原始方向。
| 迭代次数 | \(\sigma_1\) | \(\sigma_2\) | \(\sigma_3\) |
|---|---|---|---|
| 0 | 0.968 | 0.242 | 0.061 |
| 1 | 0.729 | 0.768 | 0.207 |
| 2 | 1.079 | 1.025 | 0.673 |
| 3 | 0.689 | 0.687 | 1.143 |
| 4 | 1.127 | 1.129 | 0.771 |
| 5 | 0.740 | 0.744 | 1.021 |
计算时可以先转置高矩阵，让参与 \(X_kX_k^\top\) 的行数为 \(\min(m,n)\)，最后再转置回来。这样临时方阵更小。矩阵乘法可以在 BF16 下进行；上图用 FP64 演示标量递推，实际低精度结果会有差异。
完整的一步与更新尺度
前面为了说明机制使用了普通 momentum。作者实现默认使用 Nesterov-style momentum，采用 EMA 记法时，完整顺序如下。这里的 \(t\) 是训练 step，\(k\) 是一次更新内部的 Newton–Schulz 迭代，二者不要混淆。
- 计算当前梯度 \(G_t=\nabla_W\mathcal{L}_t\)。
- 更新 \(M_t=\beta M_{t-1}+(1-\beta)G_t\)，初始 \(M_{-1}=0\)。
- 构造 \(H_t=\beta M_t+(1-\beta)G_t\)；不用 Nesterov 时取 \(H_t=M_t\)。
- 以 \(H_t\) 为输入，归一化后执行 \(K\) 次 Newton–Schulz，得到 \(O_t\)。
- 更新 \(W_{t+1}=(1-\eta_t\lambda)W_t-\eta_t s_{m,n}O_t\)。
常用起点是 \(\beta=0.95\)、\(K=5\)。有些实现用不带 \((1-\beta)\) 的 momentum 累加式；配套的 Nesterov 写法也要一起对应，不能只替换其中一行。正的整体缩放会被进入迭代前的归一化消去，忽略 \(\epsilon\) 和舍入时，两套配套写法可以等价。
还需要选择 \(s_{m,n}\)。对满秩矩阵的理想 polar 结果，令 \(r=\min(m,n)\)，因为共有 \(r\) 个值为 \(1\) 的奇异值，
如果直接使用 \(O\)，矩阵越大，平均每个元素的更新反而越小。Moonlight (Liu et al., 2025) 使用
让理想更新的 RMS 为 \(0.2\)，方便与 AdamW 的更新尺度对齐。\(0.2\) 来自该工作的经验选择，不是 AdamW 更新 RMS 的普适定理；有限步近似、秩亏也会使实际 RMS 偏离它。Weight decay 直接作用于 \(W_t\)，不放进正交化，也不乘这个形状系数。
苏剑林的《Muon 优化器指南》(Su, 2025c) 将常见缩放整理为四种约定。沿用本文 \(y=Wx\) 的记号，\(m\) 是输出维度，\(n\) 是输入维度：
| 约定 | \(s_{m,n}\) | 缩放方式 |
|---|---|---|
| 朴素版 | \(1\) | 直接使用近似 polar 结果 |
| KellerJordan 版 | \(\sqrt{\max(1,m/n)}\) | 输出维度大于输入维度时放大更新 |
| MuP 版 | \(\sqrt{m/n}\) | 按输出与输入的维度比缩放 |
| Moonlight 版 | \(0.2\sqrt{\max(m,n)}\) | 将理想更新的元素 RMS 对齐到 \(0.2\) |
同一个 learning rate 数值不能脱离缩放约定比较。PyTorch 的 Muon 文档 (PyTorch Contributors, n.d.) 将 KellerJordan 版和 Moonlight 版分别称为 original 和 match_rms_adamw。维度比要按线性层的实际输入、输出判断：如果代码写成行向量形式的 x @ W，权重的存储形状就是 \([\text{in},\text{out}]\)，不能照搬 PyTorch Linear 的轴顺序。为降低 Newton–Schulz 成本而临时转置矩阵，也不应改变缩放因子。
MuP 版与官方版的 max(1, ·) 分别在控制什么？
MuP 关注模型宽度改变时的更新尺度。对 \(v\in\mathbb{R}^d\)，\(\operatorname{RMS}(v)=\lVert v\rVert_2/\sqrt d\)。用这个量衡量每个特征分量的变化，有
对非零矩阵的理想 polar 更新，\(\Delta W=-\eta s_{m,n}O\) 且 \(\lVert O\rVert_2=1\)。取 \(s_{m,n}=\sqrt{m/n}\)，上式就等于 \(\eta\)，不随输入、输出维度改变。这给出了 MuP 版的尺度直觉；它控制特征扰动，Moonlight 版则对齐更新矩阵本身的元素 RMS。
苏剑林在 2026 年的文章 (Su, 2026) 中进一步用各向同性输入解释官方版的截断，下面假设待处理的矩阵满秩。当 \(m<n\) 时，理想 \(O\) 会丢掉输入在其行空间之外的分量；若固定 \(O\)，并假设输入满足 \(\mathbb{E}[xx^\top]=I_n\)，就有
按这种平均尺度，取 \(s_{m,n}=1\) 就够了；当 \(m\ge n\) 时，相同计算则给出 \(s_{m,n}=\sqrt{m/n}\)，合起来正是官方版的 \(\max(1,\cdot)\)。这是一种依赖输入分布的解释。训练中输入与更新方向可能相关，不能据此认定该缩放在所有阶段都更优。
Weight decay 为什么能限制权重谱范数？
《为什么我们选择尝试 Muon？》(Su, 2025a) 还给出了一个简单的范数界。设 \(\lambda>0\)、\(0\le\eta_t\lambda\le1\)，矩阵形状和缩放 \(s_{m,n}\) 固定。对理想 polar 更新，由三角不等式得到
逐步使用这个关系，就有 \(\lVert W_t\rVert_2\le\max(\lVert W_0\rVert_2,s_{m,n}/\lambda)\)。因此，限制每步更新的谱范数，再配合 weight decay，可以限制权重谱范数的累积增长。
实际近似若满足统一的界 \(\lVert O_t\rVert_2\le C\)，则将上界中的 \(s_{m,n}\) 换成 \(Cs_{m,n}\)。这个界通常很宽松，也没有同时约束整个网络的输入特征，不能把它当作 attention logits 必然稳定的保证。
实际训练中的取舍
Muon 通常用于 Transformer 隐藏层的 attention projection 和 MLP 权重；embedding、LM head、bias 和 normalization 参数交给 AdamW。Embedding 和 LM head 虽然也是二维矩阵，但参数分组还要考虑它们的用途，不能只检查 ndim == 2。作者还报告 (Jordan et al., 2024)，把合并存储的 Q、K、V 分开做 Muon 更新更好；矩阵如何分块会改变变换本身。
| 比较项 | AdamW | Muon 的矩阵参数部分 |
|---|---|---|
| 更新处理 | 用二阶矩估计逐元素缩放 | 用矩阵多项式调整 momentum 的奇异值 |
| 持久优化器状态 | 一阶、二阶矩 | 一个 momentum buffer |
| 每步额外计算 | 主要是逐元素运算 | 多次矩阵乘法和临时方阵 |
| 参数切片 | 更新可逐元素独立计算 | 任意切片后分别正交化通常不等价 |
假设两者的状态都用 FP32，Muon 的矩阵参数每个元素需 \(4\) bytes 的持久 momentum，AdamW 的两个矩估计需 \(8\) bytes。这只比较优化器状态，不包含参数、梯度、master weights、activations、Newton–Schulz workspace，以及仍由 AdamW 管理的参数，不能据此说训练总显存减半。
分布式训练的区别也来自矩阵结构。例如按行分成 \(M_1,M_2\)，完整 Gram matrix 包含 \(M_1M_2^\top\) 这样的跨分片项；各卡独立处理自己的切片会漏掉它们。结合 FSDP / ZeRO 或 Tensor Parallel 时，需要恢复相关矩阵，或使用数学上对应的分布式矩阵运算。Moonlight 的实现会 gather 相关分片再计算完整更新，并保留本卡需要的部分；这些临时状态和通信都要计入成本。
Moonlight (Liu et al., 2025) 在其 compute-optimal scaling 实验中报告了相对 AdamW 约 \(2\times\) 的计算效率，并用 Muon 训练了总参数 \(16\)B、激活参数约 \(3\)B 的 MoE。这个结果描述的是其训练设置下的 compute–loss 关系，不代表一个 optimizer step 快了两倍。Muon 每步增加的矩阵计算能否被达到目标 loss 所需 step 的减少抵消，要在相同硬件和训练设置下测 wall-clock time。
更大规模还会遇到稳定性问题。Kimi K2 的 MuonClip (Kimi Team, 2025) 在 Muon 更新之外加入 QK-Clip：根据每个 attention head 观测到的 logits，在更新后缩放相关 query / key 权重，抑制 logits 过大。它是额外的稳定化机制；仅把更新矩阵的奇异值拉平，并不能限制整个训练过程中权重和 attention logits 的增长。