Muon 原理#

TL;DR#

本页内容7 个部分

Momentum 保留整张梯度矩阵的尺度差异；Adam 则分别处理每个元素。如果把一层视为线性映射，这两种视角都没有直接控制矩阵在各个奇异方向上的步幅。

Muon 对隐藏层的动量矩阵作谱变换：保留左右奇异向量，压平非零奇异值，再用处理后的方向更新权重。它改变的是矩阵更新几何。[1]

设权重与 batch 梯度为 $W,G_t\in\mathbb R^{m\times n}$ 。先用 $M_t=\beta M_{t-1}+G_t$ 累积动量， $M_0=0$ ；这里暂不加入 Nesterov 等实现选项。

希望把“大奇异值方向走得很远、小奇异值方向走得很少”的更新，变成非零奇异方向上更均衡的方向。用只包含非零奇异值的紧致分解写成 $M_t=U_r\Sigma_rV_r^\top$ ，其中 $r$ 为秩。

将 $\Sigma_r$ 中每个正奇异值换成 1，得到 $Q_t^*=U_rV_r^\top$ 。例如 $M=\operatorname{diag}(100,1)$ ，理想方向为 $\operatorname{diag}(1,1)$ ；它保留两个方向，去掉 100 倍的幅度差。

这不是逐元素 RMS 归一化。秩不足时，上式也不会补出新的非零方向。为什么该方向对应算子范数约束下的最陡下降，由几何解释推导。

精确 SVD 的成本较高。实际实现先归一化动量，再用少量 Newton–Schulz 多项式迭代近似谱变换，得到 $Q_t$ 。有限步结果不保证奇异值严格为 1；迭代系数和步数必须随实现记录。[1]

把尺度因子记为 $s>0$ ，学习率为 $\eta$ ，可选衰减系数为 $\lambda$ ，核心更新为：

W_t=(1-\eta\lambda)W_{t-1}-\eta sQ_t.

$s$ 可能依赖矩阵形状；它控制处理后方向的总体量级。动量归一化与具体尺度规则需要联合确定，不能只凭这个示意式照搬学习率。

手算理想 polar 方向的效果。设当前 $W=\operatorname{diag}(2,2)$ 、动量 $M=\operatorname{diag}(100,1)$ ，取 $\eta=0.01$ 、 $s=1$ 、 $\lambda=0$ ：

这里隔离了谱变换的作用。它不是有限步 NS 的运行结果，也没有比较调参后最优的训练效果；实际近似算例在进阶篇计算。

原始参考实现主要将 Muon 用于隐藏层矩阵，embedding、输出头、bias 与归一化参数交给 AdamW 等方法。参数如何分组、矩阵如何切分，都会改变处理的算子。[1]

谱方向更均衡不自动意味着任务更好；还需比较近似质量、更新尺度、优化器耗时与训练结果。NS 的计算重排由 Gram Newton–Schulz维护。