TL;DR
Muon（MomentUm Orthogonalized by Newton-Schulz）是一个面向矩阵参数的优化器。它与 Adam 最本质的区别在于更新粒度：Adam 把参数当成一大群互不相干的标量、逐元素（element-wise）更新；Muon 则把权重矩阵当成一个整体（一个线性变换）来更新——先把 SGD-momentum 得到的动量矩阵正交化（丢弃奇异值、只保留”方向”，即 SVD 极分解的 $UV^\top$ 部分），再用它去更新权重。工程上用 Newton-Schulz（NS）迭代近似这个正交化，而不是昂贵的显式 SVD，且 NS 迭代只含矩阵乘法、可以稳定地在 bf16 下运行。
收益主要有两点：
- 效率：Kimi（Moonshot）团队在 Moonlight 论文中报告，compute-optimal 设定下 Muon 达到与 AdamW 同等性能只需约 52% 的训练 FLOPs（约 2× 计算效率）；
- 内存：Muon 每个参数只维护一份动量 buffer（AdamW 需要维护 $m$、$v$ 两份），且正交化过程不需要额外高精度状态。
Muon 的适用范围是 hidden layer 的矩阵参数：1D 参数（bias、LayerNorm 等）没有矩阵结构、无法正交化；embedding 与 lm_head 虽为 2D，实践中按惯例仍交给 AdamW；4D 卷积权重可以 reshape 成 2D 后再使用。因此工程上 Muon 几乎总是与 AdamW 组成混合优化器。此外，由于每次更新都依赖完整矩阵的奇异结构，Muon 无法像 Adam 那样在 ZeRO 分片下”随手切分”，需要专门的分布式实现。
核心思想
优化器要回答的问题是：给定梯度 $G_t$，参数 $W \in \mathbb{R}^{d_{out}\times d_{in}}$ 应该往哪走、走多远。朴素 SGD 的回答是顺着梯度走：$W_{t+1}=W_t-\eta G_t$。
Adam 把 $W$ 拍平成标量集合，对每个坐标独立做归一化：
$$\Delta_{ij}=\frac{\hat{m}{ij}}{\sqrt{\hat{v}{ij}}+\epsilon}$$
它给每个坐标自适应的学习率，代价是丢弃了矩阵结构：不同方向之间的相对关系（奇异值结构）完全不被利用——奇异值大的方向被当成一堆”大标量”处理，奇异值小的方向被当成一堆”小标量”处理，方向与方向之间没有协同。
Muon 选择尊重矩阵结构。它的更新分为三步。
第一步：动量。 维护动量矩阵 $M_t=\beta M_{t-1}+(1-\beta)G_t$（通常用 Nesterov 动量、$\beta\approx0.95$）。这一步与 SGD-momentum 相同，只需要一份 buffer。
第二步：正交化。 把动量矩阵替换为”最近的正交（半正交）矩阵”：
$$O_t=\mathrm{Ortho}(M_t)=UV^\top,\quad \text{其中 } M_t=USV^\top \text{ 是 SVD}$$
即只保留方向、丢弃奇异值 $S$。效果是让每个矩阵方向获得均等的更新强度，不再被奇异值大小左右；可以理解为谱范数意义下的最陡下降。
第三步：缩放与更新。 $W_{t+1}=W_t-\eta\cdot(\mathrm{scale}\cdot O_t)+$ weight decay。$O_t$ 是”纯方向”（谱范数约为 1），需要乘上形状相关的缩放系数才能获得合适步长——这一步正是各实现分歧最大的地方，见下文”变体”。
为什么用 NS 迭代而不是 SVD？ 每步都做完整 SVD 太贵。NS 迭代先用 Frobenius 范数把 $M$ 归一化到谱范数 ≤ 1，再做约 5 次的多项式迭代：
$$X\leftarrow aX+b(XX^\top)X+c(XX^\top)^2X$$
系数 $(a,b,c)$ 决定奇异值被压向 1 的收敛速度。NS 迭代只含矩阵乘法，因此可以在 bf16 上稳定运行——这正是它替代 SVD 的工程理由（参考实现对此有专门注释）。理论上也已有证明：Muon 用少数几步 NS 与理想 SVD 极分解的收敛率只差一个常数因子，且该误差随迭代步数双指数衰减（见 References）。
把 Adam 与 Muon 的分野收拢成一张速查表：
| 维度 | AdamW | Muon |
|---|---|---|
| 更新单元 | 逐元素标量 | 整个矩阵（正交化方向 $UV^\top$） |
| 优化器状态 | $m$、$v$ 两份 | 一份动量 buffer |
| 计算精度 | 状态与更新通常 fp32 | NS 迭代可在 bf16 下进行 |
| 更新强度 | 每个坐标自适应 | 每个方向均等，$\lVert\Delta W\rVert_2\approx\eta$ |
| 适用参数 | 任意形状 | hidden 层矩阵，其余配 AdamW |
工程实现
变体
Muon 没有唯一的”标准配方”，主流实现的差异主要在三处，且会互相影响 lr 的取值。
NS 迭代的多项式系数 $(a,b,c)$。 系数决定迭代把奇异值压向 1 的速度。参考实现采用调优过的 quintic 系数 $(3.4445,-4.7750,2.0315)$、迭代 5 步；更早的基线是 $(2,-1.5,0.5)$；后续社区还提出了 Polar-Express、Gram-Schmidt 版 NS 等系数与迭代格式。
正交化后的缩放系数 $\mathrm{scale}$。 它直接乘在更新上、等效于缩放 lr，是各实现差异最微妙的地方。$O_t$ 是谱范数约为 1 的”纯方向”，需按矩阵形状缩放成合适步长，常见取值有：
- 参考实现：$\sqrt{\max(1,,d_{out}/d_{in})}$——对 $d_{in}$、$d_{out}$ 不对称；
- Moonlight/Kimi 的 spectral 缩放：$\sqrt{\max(d_{out},d_{in})}$——关于两维对称，并额外乘约 0.2 使更新 RMS 对齐 AdamW，从而 AdamW 的 lr / weight decay 可以直接迁移（MuonClip 沿用此缩放）；
- Scion 系：$\sqrt{d_{out}/d_{in}}$。
scale 与 lr 是相乘关系，因此”是否对称”直接决定该用多大的 lr。网上各种 Muon lr 经验值（0.02、0.05 等）对不上号，很大程度就源于这个系数差异。
顺带解释一个常见困惑——为什么 Muon 的 lr 看起来比 AdamW 大这么多：Muon 的 lr 单位是”每步的谱范数变化”，$\Delta W=\eta O_t$ 且 $|O_t|_2\approx1$，所以 $\eta$ 几乎就等于 $|\Delta W|_2$；而 Adam 的 lr 是逐元素缩放，两者数值天然不在一个量级（实践中 Muon 常用 $10^{-2}$ 量级，AdamW 常在 $10^{-4}$ 量级）。这也正是 spectral 缩放要做 RMS 对齐的原因——对齐之后，AdamW 的 lr 与 weight decay 才能直接迁移给 Muon。
动量与正则细节。 例如是否启用 Nesterov（参考实现默认启用，经验上略优）、weight decay 是否解耦（参考实现省略了 weight decay，Moonlight 论文指出大规模训练下必须加上，否则权重/层输出 RMS 会涨出 bf16 的表示范围）；Kimi 还提出 QK-Clip 抑制 attention logits 爆炸，与 Muon 打包成 MuonClip 用于 K2 的训练。
分布式实现
先回忆 ZeRO-1 下 Adam 为什么好切：optimizer 状态（$m$、$v$）按 DP 组切分，backward 后用 reduce-scatter 把梯度归位到持有对应分片的 rank，各 rank 独立逐元素更新自己那份，最后 all-gather 把更新后的权重同步给所有 rank。这套流程成立的根本原因是逐元素更新中每个坐标只依赖自身状态——分片切在哪里都不改变计算结果。
Muon 行不通：正交化需要完整矩阵，奇异方向是全局量，NS 迭代在行/列碎片上无法进行。目前有两类解法：
- Distributed Muon（Moonlight 开源方案，ZeRO-1 风格，已落地）：optimizer 状态仍按 DP 切分、内存照常省；但在每个 step 中，对本地分片负责的参数，先从 DP group gather 出完整的梯度/动量矩阵（只传 bf16），在本地完成 NS 正交化，然后丢弃不属于本地分片的更新、只应用本地部分，最后 all-gather 权重。由于 gather 只发生在 DP group 内、只覆盖本地负责的参数、且以 bf16 传输，总通信量约为 AdamW 的 1~1.25 倍。
- 真正分布式的正交化（开放问题）：不重建完整矩阵，而是把矩阵按 Row/Column 做 TP 式切分，让 NS 迭代中的矩阵乘通过集合通信协同完成（类似分布式 SVD / Polar 分解）。这能进一步省掉重建矩阵的带宽与临时显存，但需要把正交化 kernel 与 ZeRO/FSDP 的通信原语深度融合——DeepSpeed、NeMo 等正在这个方向推进（例如实现通信高效的 Gram-Schmidt 版 NS kernel）。也可以另辟蹊径，设计能直接支持分布式正交化的 Muon 变体。
相关方法
把 Muon 放回优化器谱系，能更清楚地看出它的取舍：
- Orthogonal-SGDM：更早的想法是先把梯度做 SVD 正交化、再施加动量；Muon 把动量挪到正交化之前，经验上效果更好，并用廉价的 NS 迭代替代昂贵的 SVD。
- Shampoo：去掉 preconditioner 的累积后，其更新会退化为 $UV^\top$（与 Muon 的正交化同构）。从这个角度看，Muon 可以理解为”只保留旋转方向、不累积二阶矩”的 Shampoo，这也能解释它为什么比 AdamW 省内存却依然有效。
- Newton-Muon：对梯度右乘输入二阶矩的逆 $G(ZZ^\top)^{-1}$ 再做正交化，相当于补回标准 Muon 忽略的输入数据几何，可看作 Muon 的改进方向。
References
- Keller Jordan. Muon: An optimizer for hidden layers. https://kellerjordan.github.io/posts/muon/
- Keller Jordan. Muon（参考实现 muon.py）. https://github.com/KellerJordan/Muon/blob/master/muon.py
- Moonshot AI. Muon is Scalable for LLM Training（Moonlight）. https://arxiv.org/abs/2502.16982
- Moonshot AI. Kimi K2 Technical Report. https://arxiv.org/abs/2507.20534
- Vineet Gupta, Tomer Koren, Yoram Singer. Shampoo: Preconditioned Stochastic Tensor Optimization. https://arxiv.org/abs/1802.09568
- Zhehang Du, Weijie Su. Newton-Muon. https://arxiv.org/abs/2604.01472
- Gyu Yeol Kim, Min-hwan Oh. Convergence of Muon with Newton-Schulz. https://arxiv.org/abs/2601.19156
- PyTorch Blog. Using Muon Optimizer with DeepSpeed. https://pytorch.org/blog/using-muon-optimizer-with-deepspeed/