Muon 优化器学习笔记：原理、推导与实践

Muon 是近期备受关注的一个自适应学习率优化器，Kimi 2 的训练也采用了它。本文是对 Muon 优化器的原理、数学推导以及实际工程应用中细节的梳理。

1. 核心算法与定义

Muon 的更新规则非常简洁，其核心在于使用了动量矩阵的矩阵符号函数（Matrix Sign, msign） 来进行更新。

更新公式

$$\begin{array}{c}\begin{array}{rl}{\mathbf{M}}_t=& \beta {\mathbf{M}}_{t-1}+{\mathbf{G}}_t\\ {\mathbf{W}}_t=& {\mathbf{W}}_{t-1}-{\eta }_t[\text{msign}({\mathbf{M}}_t)+\lambda {\mathbf{W}}_{t-1}]\end{array}\end{array}$$

理解 msign (矩阵符号函数)

$\text{msign}$ 与奇异值分解（SVD）紧密相关。定义如下：

$$\begin{array}{c}\mathbf{U},\mathbf{\Sigma },{\mathbf{V}}^{\top }=\text{SVD}(\mathbf{M})\Rightarrow \text{msign}(\mathbf{M})={\mathbf{U}}_{[:,:r]}{\mathbf{V}}_{[:,:r]}^{\top }\end{array}$$

虽然定义依赖 SVD，但通过代数推导可以证明：

$$\begin{array}{c}\text{msign}(\mathbf{M})=(\mathbf{M}{\mathbf{M}}^{\top }{)}^{-1/2}\mathbf{M}=\mathbf{M}({\mathbf{M}}^{\top }\mathbf{M}{)}^{-1/2}\end{array}$$

为了直观理解 msign，我们可以看几个特例：

1. 标量 $x$：退化为普通的符号函数 $\text{sign}(x)=x(x^2{)}^{-1/2}$。
2. 对角阵 $\mathbf{M}$：$\text{msign}(\mathbf{M})=\text{diag}(\text{sign}(\mathbf{m}))=\text{sign}(\mathbf{M})$。
3. 向量 $\mathbf{m}$ ($n\times 1$)：$\text{msign}(\mathbf{m})=\mathbf{m}/\Vert \mathbf{m}{\Vert }_2$，即归一化。
4. 正方阵：最重要的性质是最优正交近似：

$$\begin{array}{c}\text{msign}(\mathbf{M})={\mathrm{argmin}}_{{\mathbf{O}}^{\top }\mathbf{O}=\mathbf{I}}\Vert \mathbf{M}-\mathbf{O}{\Vert }_F^2\end{array}$$

注：

• Muon 和 Signum、Tiger 可以视作是同一思路下的优化器，它们都以动量 $M$ 为出发点，只是为更新量选择了不同的规整化（Normalization）方法。
• 实际上 2015 年就已经有大致相同的算法被提出：Stochastic Spectral Descent for Restricted Boltzmann Machines。
• 存在一个 Nesterov 版本：将 $\text{msign}(M_t)$ 替换为 $\text{msign}(\beta M_t+G_t)$，效果通常稍好。

---

2. 理论视角：谱范数下的最速下降

为什么 Muon 有效？我们可以从矩阵范数和最速下降的角度来解释。

一般地，参数更新可以看作是一个约束优化问题：

$$\begin{array}{c}\mathbf{W}\ast t+1=\mathrm{argmin}\ast \mathbf{W}\frac{\Vert \mathbf{W}-\mathbf{W}\ast t{\Vert }^2}{2{\eta }_t}+\mathcal{L}(\mathbf{W})\end{array}$$

• 当 $\Vert \cdot \Vert$ 取 F 范数时，得到的结果就是 SGD。
• 当 $\Vert \cdot \Vert$ 取 谱范数时，得到的结果就是 $\beta =0$ 的 Muon。

当 $\beta >0$ 时，滑动平均生效，我们可以将它视为梯度的一种更精准的估计，所以改为对动量取 $\text{msign}$。

为什么选择谱范数？

在训练模型时，我们希望模型“又稳又快”：即更新幅度受控（稳），但 Loss 下降尽可能快。这可以转化为：

$$\begin{array}{c}\mathrm{argmin}\ast \Delta \mathbf{W}\text{Tr}({\mathbf{G}}^{\top }\Delta \mathbf{W})\text{s.t.}\rho (\Delta \mathbf{W})\le \eta \end{array}$$

其中 $\text{Tr}({\mathbf{G}}^{\top }\Delta \mathbf{W})$ 是 Loss 的变化量，$\rho$ 是某种衡量“稳”的指标。

考虑前向传播 $y=xW$，输出的变化量为：

$$\begin{array}{c}\Vert \Delta \mathbf{y}\Vert =\Vert \mathbf{x}\Delta \mathbf{W}\Vert \le \rho (\Delta \mathbf{W})\Vert \mathbf{x}\Vert \end{array}$$

如果我们希望通过 $x$ 的范数和 $\rho (\Delta W)$ 来严格控制 $\Vert y\Vert$（即控制激活值的变化幅度，避免爆炸），最准确的 $\rho$ 就是谱范数。

因此，Muon 本质上是谱范数约束下的最速下降算法。

---

3. 工程实现与超参数迁移

Muon 在实际应用中不仅理论优美，而且在显存和计算上也有独特权衡。

性能与显存

• 显存优势：相比 Adam，Muon 少存储了一组向量（二阶矩），显存占用更少。
• 计算开销：会有 2%~5% 的时间开销。虽然直接 SVD 昂贵，但可以使用 Newton-Schulz 迭代进行近似。
  • 计算量：对 $n\times m$ 矩阵，单次迭代开销 $O(6n{m}^2)$。对于大 Batch size，额外开销可忽略（如 Llama 405B 上仅占 0.5%）。
  • 精度：NS 迭代在 bfloat16 下即可运行（相比之下 Shampoo 需要更高精度）。

从 Adam 到 Muon 的超参迁移

Muon 的超参数调节可以通过“对齐 RMS”的方式，直接复用 Adam 的经验值。

统计发现，Adam 更新量的 RMS 约在 $0.2\sim 0.4$ 之间。而 Muon 的更新量 $\Phi =U{V}^{\top }$ 的 RMS 解析解为 $\sqrt{r/nm}$。考虑到实际矩阵往往不是低秩的，取 $\text{RMS}({\mathbf{\Phi }}_t)\approx \sqrt{1/\max (n,m)}$。

为了让 Muon 的更新幅度与 Adam 保持一致，我们引入缩放因子：

$$\begin{array}{c}\mathbf{W}\ast t=\mathbf{W}\ast t-1-{\eta }_t(0.2\mathbf{\Phi }\ast t\sqrt{\max (n,m)}+\lambda \mathbf{W}\ast t-1)\end{array}$$

这样，我们就可以直接使用 Adam 的学习率 $\eta$ 和权重衰减 $\lambda$。

Weight Decay 的安全性

公式中的 Weight Decay 项 $-\eta \lambda W$ 实际上隐含了一个很好的性质：

$$\begin{array}{c}\Vert \mathbf{W}\ast t\Vert \le \max (\Vert \mathbf{W}\ast t-1\Vert ,\Vert {\mathbf{\Phi }}_t/\lambda \Vert )\end{array}$$

这说明 $\Vert W\Vert$ 是有上界的，进而 $\Vert xW\Vert$ 也被控制住了，减少了训练发散的风险。

---

4. 实际使用中的“坑”与改进

1. Embedding 层不能用 Muon

Muon 最初设计用于 dense 的线性层（输入输出向量 RMS 接近 1）。Embedding 层的输入是 one-hot 的（极度稀疏），这会导致 RMS 估计失效。

• 建议：Embedding 层继续使用 AdamW。
• 参考：Modular Duality in Deep Learning 提出对不同层使用不同优化器。

2. Transformer 中的 QKV 处理

在 Transformer 中，将 Q、K、V 的权重矩阵拆开分别使用 Muon，比作为一个整体的大矩阵训练效果更好。

3. MaxLogit 爆炸与 QK-Clip

Muon 容易导致某些 Head 的 Attention Logit 过大（MaxLogit 爆炸）。解决方案是引入 QK-Clip，在更新前对权重进行缩放：

$$\text{MuonClip}\{\begin{array}{rl}& {\mathbf{O}}_t=\text{msign}(\mathbf{M}\ast t)\times \sqrt{\max (n,m)}\times 0.2\\ & \mathbf{W}\ast t=\mathbf{W}\ast t-1-{\eta }_t(\mathbf{O}\ast t+\lambda \mathbf{W}\ast t-1)\\ & \text{Apply QK-Clip if }S\ast >\tau \end{array}$$

4. 奇异值分布与微调 (SFT)

• Muon 训练出的矩阵奇异值分布较均匀。
• Adam 训练出的矩阵通常是 Low Rank 的。 这解释了为什么 Pre-train (Muon) + SFT (Adam) 效果可能不佳，或者依赖低秩假设的 LoRA 在 Muon 模型上效果打折。
• 结论：SFT 中直接换用 Muon 不一定比 AdamW 好；Muon 的优势主要在于 Pre-training 阶段。

---

5. 总结与展望

Muon 展示了非逐元素（element-wise）优化器的巨大潜力。

未来的改进方向：

1. 高维张量支持：目前 Muon 仅适用于 2D 矩阵。对于 Conv 层的 4D 张量，能否设计对应的正交优化器？
2. 分布式通信优化：Muon 需要跨节点 gather 动量矩阵，通信开销较大。需要设计通信效率更高的变体。
3. 二阶优化器：Muon 是谱范数下的一阶最速下降。能否设计出真正意义上的矩阵级二阶优化器，对矩阵空间的曲率进行建模？

参考资料

1. [2410.21265] Modular Duality in Deep Learning
2. GitHub - KellerJordan/modded-nanogpt: NanoGPT (124M) in 2 minutes
3. Muon: An optimizer for hidden layers in neural networks
4. Deriving Muon
5. [2502.16982] Muon is Scalable for LLM Training
6. AdaMuon：能否进一步助力Muon替代Adam？
7. 你有没有把 AdamW 换成 Muon？为什么？ - 涮月亮的谪仙人的回答 - 知乎
8. 大模型控制学——谱球优化器
9. Muon优化器科普，但从最速下降的本质出发：这篇综合了很多 muon 的知识，还不错