Muon

分类: 训练优化

Muon

定义

Muon（Momentum + Orthogonalization Update Normalization）是一种基于 Nesterov 动量 + Newtonite 正交化的矩阵优化器，通过对梯度矩阵做 Schulz 迭代求近似 orthogonal 极坐标分解，实现接近二阶方法的更新效果，同时保持一阶方法的计算成本。

数学形式

Gt=Nesterov(∇ft)∈Rm×nG_t = \text{Nesterov}(\nabla f_t) \in \mathbb{R}^{m \times n} Ut=NewtonSchulz5(Gt)≈arg⁡min⁡U⊤U=I∥U−Gt∥FU_t = \text{NewtonSchulz5}(G_t) \approx \arg\min_{U^\top U = I} \|U - G_t\|_F

更新规则：Wt+1=Wt−η⋅UtW_{t+1} = W_t - \eta \cdot U_t

核心要点

动机：标准 Adam 逐元素更新，忽视参数矩阵的几何结构；Muon 以矩阵整体为单位做正交化

优势：在大型语言模型预训练中多次被报告比 AdamW 收敛更快（token效率更高）

局限：近期理论分析（MOGA论文）指出 Muon 在最坏情况下平滑常数随 width ww 以 O(w)\mathcal{O}(\sqrt{w}) 增长，缺乏 width-independent 的理论保证

与 AdamW 的对比：Muon 更适合宽网络；AdamW + 适当的 μP 缩放在窄网络到宽网络迁移时更稳定

实现：Newton-Schulz 5 次迭代（GPU上约 2ms 开销），已有 PyTorch 实现

代表工作

Kosson et al., 2024 (原始 Muon 提出)

On the Width Scaling of Neural Optimizers: 分析 Muon 的宽度缩放行为，提出 MOGA 作为改进替代

相关概念

μP

NTK