Everything You Need to Understand Muon Optimizer 理解 Muon 你所需要知道的一切
本文在 Codex 辅助下完成 (动画挺好的, 其他文字部分说实话写的不好, 为了完整性我还是保留了, SVD 不如看手写笔记)
SVD 将一个 linear map 分成 change of basis、stretching、再映射到 output directions。Muon 则利用这套分解，保留 update 的 directions，重新归一化它在这些 directions 上的强度。
1 SVD
1.1 Motivation
整个故事开始于我们想找一组相互垂直的向量, 在经过一个线性变换后, 仍然保持相互垂直. 即: 找一组特殊的 orthonormal basis \(v_1,\ldots,v_n\)，使得它们的 images 仍然两两 orthogonal:
\[ v_i \perp v_j\quad\Longrightarrow\quad Av_i\perp Av_j, \quad \forall i,j. \]
注意:
\(A\) 不一定 injective, 因此可能存在映射到零向量的可能性, 即 \(Av_i=0\).
并不要求 \(\{Av_i\}\) 归一化 (只需 orthogonal), 当然我们可以通过归一化来实现: \[Av_i=\sigma_i u_i,\qquad \|v_i\|=\|u_i\|=1,\qquad \sigma_i>0.\] 这里 \(v_i\) 是 right singular vector，\(u_i\) 是 left singular vector，\(\sigma_i\) 是 singular value。若 \(\operatorname{rank}(A)=r\)，只有前 \(r\) 个 directions 被保留下来；其余满足 \(Av_i=0\)，构成 \(\ker A\) 的 basis。非零 images 对应的 \(u_1,\ldots,u_r\) 只 span \(\operatorname{col}A\)，不一定 span 整个 \(\mathbb R^m\)。
把这些等式并排写出，再把 \(u_i\) 补成 output space 的 orthonormal basis，就得到
\[ AV=U\Sigma \quad\Longrightarrow\quad \boxed{A=U\Sigma V^\top}, \]
其中 \(U\in\mathbb R^{m\times m}\)、\(\Sigma\in\mathbb R^{m\times n}\)、\(V\in\mathbb R^{n\times n}\)，且 \(U^\top U=I_m\)、\(V^\top V=I_n\)。\(\Sigma\) 的 diagonal entries 按 \(\sigma_1\geq\cdots\geq\sigma_r>0\) 排列，其余为零。
与 \(Av=\lambda v\) 不同，SVD 不要求 input 和 output 属于同一个空间，也不要求 image 与原 vector 平行。因此 rectangular matrix 同样有 SVD。本文用 real matrices；complex case 将 transpose \(\top\) 换成 conjugate transpose \(H\)，orthogonal 换成 unitary 即可。
1.2 Visualization
先看 \(2\times2\)：unit circle 上的点写成 \(x(t)=(\cos t,\sin t)^\top\)，其 image 是
\[ Ax(t)=(Ae_1)\cos t+(Ae_2)\sin t. \]
两条 columns 独立时，这是一条 ellipse；不独立时退化成线段或点。为什么 ellipse 的两个 principal axes 能反推到 orthogonal input directions？令 \(v_1\) 为 unit circle 上使 \(\|Av_1\|^2\) 最大的方向，取 \(v_2\perp v_1\)。沿圆走到最大值点时 derivative 为零：
\[ \left.\frac{d}{dt}\left\|A(v_1\cos t+v_2\sin t)\right\|^2\right|_{t=0} =2(Av_1)^\top Av_2=0. \]
所以两条 images 确实 orthogonal。令 \(\sigma_i=\|Av_i\|\)、\(u_i=Av_i/\sigma_i\)，便有
\[ A(v_1\cos t+v_2\sin t) =\sigma_1u_1\cos t+\sigma_2u_2\sin t. \]
这同时给出了 ellipse 的方向和半轴长度，也避免了用 SVD 来证明 ellipse、再用 ellipse 来证明 SVD 的循环。
下面用 hollow dot 与 filled dot 分别标记 \(e_1,e_2\)，右端对应 \(Ae_1,Ae_2\)。拖动右端的两个标记，或直接编辑总箭头上的 \(A\)，就能改变这个 linear map。每个 factor 标在连接相邻视图的下方弧线下面。点击 \(A\) 中的数字可以直接编辑；它下方的绿色播放图标和进度条控制分步动画。绘图尺度固定，编辑范围限制为整条 ellipse 都能留在视图内。3D 视图可拖动空白区域改变观察角度；播放时会依次经过 \(V^\top,D,E,U\)。
从右往左读图。 先看最后的 ellipse：principal axes 给出 \(u_1,u_2\)，半轴长度给出 \(\sigma_1,\sigma_2\)。用 \(U^\top\) 将 axes 对齐 coordinates，再除以对应的 \(\sigma_i\)，就回到一个圆。此时 hollow dot 和 filled dot 的位置给出 \(V^\top e_1,V^\top e_2\)，也就是 \(V^\top\) 的两条 columns。只有 ellipse 的轮廓还不够：\(A\) 与 \(AR\)（\(R\) orthogonal）画出的轮廓相同，标记保留了 input points 与 output points 的对应关系。
从左向右，对任意 \(x\) 则是
\[ x\ \xrightarrow{V^\top}\ \text{coordinates in the }v_i\text{ basis} \ \xrightarrow{\Sigma}\ \text{stretch each coordinate} \ \xrightarrow{U}\ Ax. \]
\(U,V\) 可能包含 reflection，不能一概称为 rotation。两条 singular values 相等时 principal axes 也不唯一；拖动经过这样的点时，显示的 basis 可能跳变，但 \(A\) 和 ellipse 仍连续。
对于 \(3\times2\)，circle 变成嵌在 \(\mathbb R^3\) 某个平面内的 ellipse。笔记里把中间一步拆成 \(D\) 和 \(E\)，非常直观：
\[ D=\begin{pmatrix}\sigma_1&0\\0&\sigma_2\end{pmatrix},\quad E=\begin{pmatrix}1&0\\0&1\\0&0\end{pmatrix},\quad \Sigma=ED,\qquad A=UEDV^\top. \]
\(D\) 拉伸，\(E\) 补一个零 coordinate，\(U\) 把这个平面放到 \(\operatorname{col}A\)。Full SVD 的 shapes 是 \((3\times3)(3\times2)(2\times2)\)；第三条 \(u_3\) 垂直于 image plane，不代表第三个非零 singular value。图中是 3D 的屏幕 projection，屏幕上看起来不垂直的 arrows，在 3D 中仍可能 orthogonal。
1.3 Computation
不看图时，关键是 \(A^\top A\)：它 symmetric positive semidefinite，因此一定有 orthonormal eigenbasis。若
\[ A^\top A v_i=\lambda_i v_i, \qquad \lambda_i\geq0, \]
那么
\[ (Av_i)^\top(Av_j)=v_i^\top A^\top Av_j=\lambda_j\delta_{ij}. \]
所以计算过程是：先求 \(A^\top A\) 的 orthonormal eigenvectors 得到 \(V\)；再取 \(\sigma_i=\sqrt{\lambda_i}\)；对 \(\sigma_i>0\) 用 \(u_i=Av_i/\sigma_i\) 得到相互配对的 \(U\) columns，最后补全 basis。\(\sigma_i=0\) 时不能做除法。
例如
\[ A=\begin{pmatrix}2&1\\1&2\end{pmatrix},\qquad A^\top A=\begin{pmatrix}5&4\\4&5\end{pmatrix}. \]
后者的 eigenvalues 是 \(9,1\)，对应 \(v_1=(1,1)^\top/\sqrt2\)、\(v_2=(-1,1)^\top/\sqrt2\)。因此
\[ \Sigma=\begin{pmatrix}3&0\\0&1\end{pmatrix},\qquad U=V=\frac1{\sqrt2}\begin{pmatrix}1&-1\\1&1\end{pmatrix}. \]
circle 沿 \(45^\circ\) direction 拉长三倍，沿垂直方向保持长度。这里 \(U=V\) 是因为选了 symmetric positive definite 的例子，一般矩阵并非如此。
同理，\(AA^\top=U\Sigma\Sigma^\top U^\top\)。两边的非零 eigenvalues 都是 \(\sigma_i^2\)，但多出来的零 eigenvalues 数目可以不同。直接独立计算两组 eigenvectors 时还要对齐 signs 和重复 eigenvalues 的 bases，不能随意拼起来。
这是理解 SVD 的计算推导；实际 numerical library 通常不会先显式形成 \(A^\top A\)，因为它会把 condition number 平方。需要数值 SVD 时应使用成熟的 SVD routine。
2 Muon Optimizer
2.1 Motivation
对于 linear layer \(y=Wx\)，\(W\in\mathbb R^{m\times n}\) 的 entries 共同定义了一个 operator。Adam 的 moment estimates 与 normalization 主要逐 entry 进行，并不显式根据整个矩阵的 singular directions 调整 update。Muon 把这个 matrix structure 纳入更新规则。不过，利用 matrix structure 并非 Muon 首创，Shampoo 等方法已经走过这条路线。
Muon 的名称来自 MomentUm Orthogonalized by Newton–Schulz。它在 2024 年的 neural network training speedrunning 中发展起来；核心是先形成 momentum update，再做 approximate orthogonalization。通常用于 hidden weight matrices，embedding、output head 以及 vector parameters 则交给 AdamW。原始介绍
设 momentum matrix 的 compact SVD 为
\[ M=U_r\Sigma_rV_r^\top =\sum_{i=1}^r\sigma_i u_i v_i^\top. \]
理想化的操作是
\[ \boxed{\operatorname{Orth}(M)=U_rV_r^\top =\sum_{i=1}^r u_i v_i^\top.} \]
它保留 singular directions，把非零 singular values 换成 \(1\)。例如 \(M=\operatorname{diag}(10,1)\) 变成 \(I_2\)：原本相差十倍的两个 mode 得到相同强度。它不是把每条 column 各自除以长度，也不是把 parameter \(W\) 本身变成 orthogonal matrix。
Full-rank 的 \(m\times n\) 矩阵，\(m\geq n\) 时结果满足 \(O^\top O=I_n\)，\(m\leq n\) 时满足 \(OO^\top=I_m\)；rectangular case 通常称为 semi-orthogonal。Rank-deficient case 的 compact expression 是 partial isometry，零 directions 仍为零，不应凭空补成非零 update。
2.2 Algorithm
为了把 conventions 说清楚，本文使用 normalized momentum 与一种 Nesterov-style mixing：
\[ \begin{aligned} G_t&=\nabla_W\mathcal L_t,\\ M_t&=\beta M_{t-1}+(1-\beta)G_t,\quad M_0=0,\\ H_t&=(1-\beta)G_t+\beta M_t,\\ O_t&=\operatorname{NS}_K(H_t),\\ W_{t+1}&=(1-\eta_t\lambda)W_t-\eta_t s_{m,n}O_t. \end{aligned} \]
\(G_t,M_t,H_t,O_t,W_t\) 都是 \(m\times n\)；\(\beta\) 是 momentum coefficient，\(\eta_t\) 是 learning rate，\(\lambda\) 是 decoupled weight decay。关闭 Nesterov-style mixing 时令 \(H_t=M_t\)。这种 mixing 与 reference implementation 对齐；不同代码里的 momentum buffer 可能采用不同 normalization，不能只比较某一行公式。
\(s_{m,n}\) 是额外的 update scale，不属于 orthogonalization。本节选用 Moonshot 的 \(s_{m,n}=0.2\sqrt{\max(m,n)}\)。原因是 full-rank ideal \(O\) 有
\[ \|O\|_F^2=\min(m,n),\qquad \operatorname{RMS}(O)=\sqrt{\frac{\min(m,n)}{mn}} =\frac1{\sqrt{\max(m,n)}}. \]
这里 RMS 是 root mean square。乘上该 scale 后，ideal update 的 RMS 为 \(0.2\)，方便跨 shape 比较步长。有限 NS steps 下并非严格等于 \(0.2\)；这个数也是 empirical choice。Muon is Scalable for LLM Training, §2.2
2.3 NS iteration
每个 training step 都做一次精确 SVD 成本较高。我们只想改变 singular values，不需要显式拿到 \(U,V\)。设 \(X=U\operatorname{diag}(x_i)V^\top\)，那么
\[ XX^\top X=U\operatorname{diag}(x_i^3)V^\top. \]
因此矩阵 polynomial
\[ X_{k+1}=aX_k+bX_kX_k^\top X_k+c(X_kX_k^\top)^2X_k \]
等价于对每个 singular mode 的 coefficient 应用 \(f(x)=ax+bx^3+cx^5\)，同时保留两侧 directions。这就把 matrix orthogonalization 转成了一个 scalar iteration 的设计问题。
经典 cubic NS 使用 \(f(x)=\tfrac32x-\tfrac12x^3\)：\(f(1)=1\)、\(f'(1)=0\)，且 \(0<x<1\) 时 \(x<f(x)\leq1\)，于是逐步趋向 \(1\)。Muon 常用的 quintic coefficients 则是
\[ (a,b,c)=(3.4445,-4.7750,2.0315). \]
它更积极地增大小的 singular values，但 \(f(1)=0.701\)，所以 \(1\) 根本不是 fixed point。这组 coefficients 的目标是少量 steps 后得到有用的近似，不是无限迭代到精确 \(UV^\top\)。极小的初始 singular values 在五步后也可能仍很小。实现说明
下面给出单个 hidden matrix 的最小 update pseudocode；不包含 forward/backward。为减少乘法量，先转成 rows 不多于 columns 的形状。示例 \(m=2048,n=768\)，scalar 都标为 []。
# 将 H 的 singular spectrum 用 K 步 polynomial 调整到近似单位尺度。
# H: [m=2048, n=768]; K: scalar []=5; eps: scalar []=1e-7
def ns5(H, K, eps):
X = H.T / (norm_fro(H) + eps) # X: [n=768, m=2048]
a, b, c = 3.4445, -4.7750, 2.0315 # each: scalar []
for k in range(K): # k: scalar integer []
S = X @ X.T # S: [n=768, n=768]
P = b * S + c * (S @ S) # P: [n=768, n=768]
X = a * X + P @ X # X: [n=768, m=2048]
return X.T # O: [m=2048, n=768]
# 用 gradient G 与旧 momentum M 更新单个 W，返回新 W 和 M。
# W, G, M: each [m=2048, n=768]; M initially zero
# beta=0.95, eta=3e-4, wd=0.1, eps=1e-7: each scalar []
# K=5: scalar integer []; m=2048, n=768: shape dimensions
def muon_step(W, G, M, beta, eta, wd, K, eps):
M = beta * M + (1 - beta) * G # M: [m=2048, n=768]
H = (1 - beta) * G + beta * M # H: [m=2048, n=768]
O = ns5(H, K, eps) # O: [m=2048, n=768]
s = 0.2 * sqrt(max(W.shape)) # s: scalar [] ≈ 9.051
W = (1 - eta * wd) * W - eta * s * O # W: [m=2048, n=768]
return W, M # each: [m=2048, n=768]2.4 Questions
2.4.1 Why divide by Frobenius norm?
初始化为 \(X_0=H/(\|H\|_F+\epsilon)\)，因为
\[ \|X_0\|_2=\frac{\sigma_{\max}(H)}{\|H\|_F+\epsilon}\leq1. \]
这把初始 spectrum 放进设计 polynomial 时考虑的区间。Frobenius norm 只需要平方求和，而精确 spectral norm 又涉及最大 singular value 的计算。它只做统一缩放，不会消除各个 singular values 的比例差异；那是之后 NS iteration 的工作。\(H=0\) 时输出仍为零。
2.4.2 Why is it good to orthogonalize the update?
一个明确的解释来自 spectral-norm constrained steepest descent。暂时忽略 momentum，用当前 gradient 的 compact SVD \(G=U_r\Sigma_rV_r^\top\)，考虑
\[ \min_{\|\Delta W\|_2\leq\eta}\langle G,\Delta W\rangle_F. \]
因为 \(|u_i^\top\Delta Wv_i|\leq\eta\)，objective 至少为 \(-\eta\sum_i\sigma_i\)；而 \(\Delta W=-\eta U_rV_r^\top\) 恰好达到这个 bound。因此 ideal orthogonalized direction 是这个局部问题的一个最优解。Old Optimizer, New Norm
这个 constraint 对 linear layer 有直接含义：对所有 input \(x\)，都有 \(\|\Delta Wx\|_2\leq\eta\|x\|_2\)。例如 \(G=\operatorname{diag}(10,1)\)，同样限制 spectral norm 不超过 \(\eta\)，按 gradient 比例走的 \(-\eta\operatorname{diag}(1,0.1)\) 带来一阶变化 \(-10.1\eta\)，而 \(-\eta I\) 是 \(-11\eta\)。
这说明在指定的 norm constraint 下为何合理，不是说忽略 gradient magnitude 总有利。小 singular modes 也可能包含 noise；实际 momentum direction 不一定是当前 loss 的 descent direction，有限步 NS、finite learning rate 也不等于上述理想问题。Muon 并没有计算 Hessian，也不等同于 natural gradient。
2.4.3 How to choose the coefficients?
先选初始区间、允许的输出误差和 iteration budget，再比较 scalar polynomial 在该区间的行为。\(f'(0)=a\) 越大，小 modes 起步越快；但太激进会增加 overshoot 和不稳定性。经典 NS 关心在 \(1\) 附近的收敛；Muon 的常用 quintic 则偏向有限步的速度与 training 效果，两者优化目标不同。不能因为名字里有 Newton–Schulz，就假设任意 coefficients 都能收敛。
2.4.4 How is the performance?
2025 年的 Moonshot 工作在其 compute-optimal scaling experiments 中报告：达到相近效果，Muon 所需 training compute 约为 AdamW 的一半；并训练了 16B total / 3B activated parameters 的 Moonlight MoE，使用 5.7T tokens。这是特定实验下的 compute efficiency 结果，不意味着每个 step 快两倍。论文
Muon 增加了 matrix multiplications，也只需要一份 momentum state，而 AdamW 维护两份 moments。实际 wall-clock 收益还取决于 batch size、matrix shape、通信与实现。评估时应在各自调好 learning rate 和 weight decay 后，比较达到同一 validation loss 的总时间，而不只看相同步数的 loss。
SVD 告诉我们一个 matrix 在哪些 input directions 上、以多大强度、映射到哪些 output directions。Muon 的设计便是沿这套 directions 修改 update spectrum，并用 matrix polynomial 将这件事做得足够便宜。