先记住一句话
Muon 是对 2D hidden-weight momentum 做 Newton–Schulz 近似的 polar/orthogonalized update；名字含 Newton，但它不计算 Hessian，通常仍归为 first-order optimizer。
1. Preconditioning 的共同语言
Δθ = −η Pₜ gₜP 改变参数空间几何。SGD 的 P=I；Adam 的 P 是由 squared gradients 得到的 diagonal matrix；full-matrix preconditioner 能旋转并缩放相关方向，却存储/计算昂贵。
2. Shampoo：按 tensor axes 分解
对 matrix gradient G，Shampoo 维护 left/right statistics（简化表示）：
L←βL+(1−β)GGᵀR←βR+(1−β)GᵀGG_pre≈L^−1/4 G R^−1/4它用 row/column structure 近似不可承受的 full parameter preconditioner。matrix inverse roots、额外 state 与 distributed communication 是主要成本，因此常 block 化、降低 preconditioner update 频率。
3. SOAP：在 Shampoo eigenbasis 里跑 Adam
SOAP 的核心解释是：周期性从 Shampoo statistics 取 eigenbasis，把 gradient 旋转到这个缓慢变化的 basis，在其中持续做 Adam-style first/second moments，再旋回参数坐标。这样即使不频繁更新 eigendecomposition，diagonal Adam statistics 仍每步更新。
它比 Adam 多 matrix statistics、basis 和分解频率超参；是否 wall-clock 更快取决于 batch、matrix size、hardware 与 implementation，而不只是减少 steps。
4. Muon 的三步
- 对 matrix gradient 积累 momentum（可用 Nesterov）；
- 把 momentum matrix 归一化后做若干 Newton–Schulz iterations，近似其 polar factor
UVᵀ； - 按适合 matrix shape 的 scale/LR 应用 update，并做 decoupled decay。
若 M=UΣVᵀ，理想正交化把不同 singular values 变得接近 1，只保留 singular directions。直觉上避免少数高-gain direction 支配整个 matrix update。
5. Newton–Schulz 在做什么
直接 SVD 昂贵。Newton–Schulz 用 matrix multiplications 迭代逼近 polar factor；这非常适合 accelerator，但系数、归一化、iteration 数、bf16/float32 与矩阵宽高处理影响稳定和 update scale。官方/框架实现经过特定数值设计，不应只看三行伪代码重写。
6. 为什么只用于 hidden 2D weights
- Linear/attention/MLP matrix 有明确 input/output axes，可定义 matrix singular directions；
- bias/norm 是 1D，没有相同 orthogonalization；
- embedding/output head 的巨大、语义特殊矩阵常需不同 LR/scale；
- Conv kernel 要 reshape，choice 会决定何种 geometry；
- 因此实践通常用 Muon + auxiliary AdamW/Adam，而非一个公式覆盖所有参数。
7. 三者不要混为“二阶”
| 统计/变换 | 是否显式 Hessian | |
|---|---|---|
| Shampoo | gradient covariance-like axis preconditioners | 否 |
| SOAP | Shampoo eigenbasis 中的 Adam | 否 |
| Muon | momentum matrix 的 polar/orthogonalized update | 否 |
| Newton method | H⁻¹g | 是/近似 loss Hessian |
Shampoo 常被宽泛称为 higher-order/preconditioned method，但它的 matrices 来自 gradient statistics；Muon 的 Newton–Schulz 是求 matrix function 的数值迭代，不是对 loss 做 Newton step。
8. 训练与系统指标
- Muon/Adam parameter numel 与每组 LR/decay；
- matrix update 的 spectral norm、Frobenius norm、
‖ΔW‖/‖W‖； - Newton–Schulz residual/nonfinite 与 iteration cost；
- optimizer step time、通信、tokens/s 和 memory；
- validation loss per token 与达到目标 loss 的 wall-clock；
- 从 AdamW checkpoint 切换时 optimizer state reset 的 transient。
9. 如何做替换实验
先锁 data/global batch/schedule budget，分别 sweep optimizer 推荐 LR 与 decay；核对 parameter routing；从相同 model weights 开始并明确 moments 是否重置；同时比较 steps、tokens、wall-clock 与 end quality。若只看 early training loss，matrix method 的额外 step time可能被隐藏。
10. 四个 matrix-aware 计算
例 1：diagonal 与 full geometry
Gradient g=(2,1)，Adam-like diagonal scale P=diag(0.5,2) 得 Pg=(1,2)，只缩放坐标。若 P=[[1,−0.5],[−0.5,1]]，Pg=(1.5,0)，还发生旋转，体现 cross-coordinate correlation。
例 2：Shampoo left statistic
G=[[1,2],[0,1]]，则 GGᵀ=[[1²+2²,1×0+2×1],[2,0²+1²]]=[[5,2],[2,1]]；off-diagonal 2 记录两行 gradient 的相关性。
例 3：Newton–Schulz 标量直觉
对已缩放 singular value x=0.5，用迭代 x_next=1.5x−0.5x³ 得 0.75−0.0625=0.6875，向 1 靠近；x=1 则保持 1。矩阵版本同时拉平多个 singular values。
例 4：parameter routing
一个 block 有四个 1024×1024 matrices（共 4,194,304 参数）和四条 1024D bias/norm vectors（4096 参数）。若矩阵走 Muon、向量走 Adam，Muon 覆盖约 4,194,304/4,198,400≈99.90% 参数，但两套 state/checkpoint 都仍要正确维护。
Muon 不是“把 Adam 的 v 换成 SVD”。它更接近 momentum SGD 的 matrix update 后处理；非矩阵参数仍需另一个 optimizer，且 LR scale 不能直接继承 AdamW。
自测
1. Adam 丢失了什么 matrix 信息？
它只保存每 coordinate 的平方统计，不建模 row/column directions 的相关性。
2. Muon 的 Newton–Schulz 为什么不是二阶 loss optimization？
迭代用于近似 momentum matrix 的 polar factor，没有计算 loss Hessian 或 Hessian-vector product。
3. embedding 为什么常不走 Muon？
它结构/尺度特殊且维度巨大；标准 Muon recipe聚焦 hidden-layer 2D weights，embedding/head 用 auxiliary Adam。
一手与官方资料
Shampoo提出 tensor-axis preconditioning；SOAP在 Shampoo eigenbasis 中连接 Adam；Muon 原始说明与 PyTorch Muon 文档给出 momentum、Newton–Schulz、Nesterov 与 decoupled decay 的实现定义。