5 Muon 优化器 / 从 SGD 到几何更新

V4 的第三大创新不在模型结构里,而藏在优化器里。 1.6T 参数中 99.9% 用 Muon,只有 embedding、prediction head、RMSNorm、 mHC 静态偏置和门控因子(合计不到 20 亿参数)还留给 AdamW。 Muon 不是 AdamW 的微调版——它换了一个数学范式:把权重矩阵当成几何对象,对整矩阵做正交化更新。

5.1 优化器演化史:从逐参数 → 整矩阵

过去 70 年优化器一直在做一件事:在原始梯度上加越来越聪明的"修饰"。Muon 是第一个跳出"逐参数思路"的主流优化器。

SGD1951 · Robbins-Monro 逐参数
Wt = Wt-1 − η · gt
每个参数独立沿梯度走,遇到狭长峡谷就 zigzag。
Momentum1964 · Polyak + 历史方向
Mt = μ·Mt-1 + gt     Wt = Wt-1 − η·Mt
动量帮你冲过小坑,但学习率仍然全局共享。
Adam2014 · Kingma & Ba + 自适应 lr
Wt = Wt-1 − η · m̂t / (√v̂t + ε)
用梯度二阶矩给每个参数加个 lr 缩放——但仍然是逐元素的。
AdamW2017 · Loshchilov & Hutter + 解耦 wd
把 weight decay 从梯度里拆出来独立施加,2018 年起成为 LLM 训练事实标准。
Muon2024 · Jordan et al. 整矩阵
把权重矩阵 W ∈ ℝn×m 视为几何对象,对动量缓冲做近似正交化—— 让更新方向在所有奇异方向上均匀分散

5.2 AdamW vs Muon:一个关键的视角差

AdamW · 逐元素视角

把 W ∈ ℝn×m 看成 n×m 个独立标量,每个标量算自己的一阶/二阶矩,独立决定步长方向。

ΔWi,j = − η · mi,j / (√vi,j + ε)

问题:梯度矩阵 G 的奇异值往往差几个数量级—— 前几个大奇异方向把更新预算"吃光",其它方向几乎不动。

Muon · 整矩阵视角

把 W 视为矩阵流形上的一个点,更新 ΔW 是矩阵之间的"几何位移"。

ΔW = − η · orthogonalize(μM + G)

把动量矩阵 SVD 后的奇异值全部砸成 1—— 所有方向的"推进力度"被强制均衡,相当于自适应给每个奇异方向独立 lr。

5.3 Muon 算法(V4 论文 Algorithm 1)

# 每个训练步、对每个逻辑独立权重矩阵 W:
G_t  = ∇_W L(W_{t-1})                       # ① 当前梯度
M_t  = μ · M_{t-1} + G_t                    # ② 动量(μ=0.95)
O'_t = HybridNewtonSchulz(μ · M_t + G_t)    # ③ Nesterov + 混合 NS 正交化
O_t  = O'_t · sqrt(max(n, m)) · γ           # ④ 重新缩放更新的 RMS(V4: γ 使 RMS≈0.18)
W_t  = W_{t-1} · (1 − η·λ) − η · O_t        # ⑤ 解耦 weight decay + 应用更新

几个工程细节值得记:

5.4 核心数学:HybridNewtonSchulz 怎么把矩阵正交化?

目标:给定动量矩阵 M = U Σ VT,把它近似变成 U VT(也就是奇异值 Σ 全部归 1)。 这正是 M 的"极分解"中的正交因子。

为什么不直接 SVD?SVD 复杂度 O(min(n,m)·n·m) 且数值不稳,跑在每个 GPU 的每个权重矩阵上完全不可承受。 Newton-Schulz(NS)迭代用纯矩阵乘法逼近,BF16 也能跑、张量核能加速

初始归一化让 ‖M₀‖₂ ≤ 1(最大奇异值不超过 1):

M0 = M / ‖M‖F

然后每次迭代用一个三次多项式 p(x) = a·x + b·x³ + c·x⁵ 作用在每个奇异值上:

Mk+1 = a · Mk + b · (Mk MkT) Mk + c · (Mk MkT)² Mk

系数 (a, b, c) 决定 p(x) 的形状——好的系数让 p(p(p(...(σ)))) 快速收敛到 1。 V4 用 10 次迭代 + 两阶段系数

前 8 步:激进收敛

(a, b, c) = (3.4445, −4.7750, 2.0315)

p(x) 在 [0,1] 上把小奇异值大幅放大,把大奇异值压回 1 附近。代价:在 σ≈1 处有过冲。

后 2 步:精细稳定

(a, b, c) = (2, −1.5, 0.5)

p(x) 在 σ=1 是不动点且导数小,把残留误差精确抹平到 1。

5.5 可视化:奇异值在 NS 迭代下的"压平"过程

下图展示一个 64×64 随机矩阵的奇异值分布——初始参差不齐,经过 10 次混合 Newton-Schulz 后全部收敛到 1 附近:

0 / 10(初始)
当前奇异值分布 目标值 σ = 1

前 8 步(橙色系数)把奇异值快速拽到 1 附近但有过冲;后 2 步(绿色系数)精细收敛。 最大/最小奇异值之比从 压到

5.6 几何直觉:为什么"奇异值砸平"能让训练更稳?

AdamW 的隐患

梯度矩阵 G 的奇异值谱通常重尾——头部几个方向的能量是尾部的 1000×。 逐元素自适应只能让"每个标量"步长合理,无法让奇异方向均衡

结果:训练前期参数大幅旋转,后期"卡死"在前几个奇异方向,影响 loss landscape 探索。

Muon 的修正

正交化后的 O = U VT所有奇异方向上模长相等—— 小奇异方向被"放大"以匹配大奇异方向,整个矩阵各向同性地更新。

类比:AdamW 像一格格爬梯子;Muon 是调整整面身体的姿态同步迈步。

5.7 V4 的实际配置一览

V4 取值说明
覆盖参数比例99.9%除 embedding / head / RMSNorm / mHC 静态偏置 / 门控因子外,全部用 Muon
AdamW 兜底参数量< 2B这些参数要么不是矩阵,要么形状特殊
NS 迭代数10(8 + 2)两阶段系数:激进收敛 → 精细稳定
动量系数 μ0.95使用 Nesterov 变体:NS 输入用 (μM + G)
更新 RMS 缩放≈ 0.18故意贴近 AdamW 的典型 RMS,复用 AdamW 的 lr 调度
Weight decay解耦施加Liu et al. 2025 的改进,原版 Muon 没做
QK-Clip不使用Q/K 已有 RMSNorm,logits 不会爆

5.8 收益:V4 训练为什么能跑下来

Muon 在 V4 的实际作用 1.6T 参数 × 33T token 的训练规模,AdamW 在中后期容易出现 loss spike + 长 plateau。 Muon 提供了更稳的收敛速度更小的方差——配合 mHC 的稳定性约束,让 V4 的训练成本压到了约 520 万美元(与 V3 同等量级,但模型规模大 2 倍以上)。
一句话总结 Muon 的关键不是"更聪明的自适应 lr",而是用 Newton-Schulz 把动量正交化—— 让权重矩阵的每个奇异方向均匀获得更新预算。这是优化器 70 年来第一次跳出"逐参数视角"。