Muon 优化器

← 中文导读 · ← 仓库首页（EN） · ← 演进总览 §3.7 · V4 梗概 · CSA/HCA · mHC · Hash MoE + FP4 论文：DeepSeek-V4 arXiv:2606.19348 · Algorithm 1 · §2.4 Muon Optimizer · §3.4.1 Efficient Implementation of Muon

一句话

Muon（Momentum Orthogonalized by Newton–Schulz）在 V4 中替换大部分矩阵参数的 AdamW：对动量矩阵做 Newton–Schulz 正交化得到更新方向 $UV^\top$（极因子），各奇异方向步长归一，从而更快收敛、训练更稳；嵌入、输出头、RMSNorm、mHC 门控等仍用 AdamW。

V4 中的位置

1. Muon 一步在做什么

V3 全系 AdamW：对每个参数做元素级二阶矩估计。V4 对 二维权重矩阵 $W\in\mathbb{R}^{n\times m}$ 改用 Algorithm 1：

$$ G_t=\nabla_W\mathcal{L}t,\quad M_t=\mu M{t-1}+G_t,\quad O'_t=\mathrm{HybridNewtonSchulz}(\mu M_t+G_t),\quad O_t=O't\cdot\sqrt{\max(n,m)}\cdot\gamma,\quad W_t=W{t-1}(1-\eta\lambda)-\eta O_t $$

其中 Nesterov 技巧体现在对 $\mu M_t+G_t$ 做正交化（而非仅 $M_t$）；$\gamma=0.18$ 用于 复用 AdamW 学习率超参（与 Kimi Muon 做法一致）。

图示详情

1.1 与 AdamW 的几何差异

设梯度（或动量）矩阵 $M=U\Sigma V^\top$（SVD）：

直观上：AdamW 沿「梯度自然拉伸」的方向走；Muon 先 去掉 $\Sigma$ 的缩放，只保留方向结构，再在正交方向上统一步长——论文报告在 万亿级 MoE 上收敛更快、对学习率 schedule 更不敏感。

1.2 V4 为何不用 QK-Clip

Kimi Muon 曾配合 QK-Clip 抑制 attention logit 爆炸。V4 在 query 与 KV entry 上直接做 RMSNorm（配合 CSA/HCA 架构），logit 已受控，故 不再使用 QK-Clip（论文 §2.4）。

2. Hybrid Newton–Schulz 正交化

显式 SVD 每步太贵；Muon 用 Newton–Schulz 迭代近似极因子 $UV^\top$。先将 $M$ 按 Frobenius 范数归一化 $M_0=M/|M|_F$，再重复：

$$ M_k=aM_{k-1}+b(M_{k-1}M_{k-1}^\top)M_{k-1}+c(M_{k-1}M_{k-1}^\top)^2M_{k-1} $$

V4 相对 Kimi Muon 的改动：混合两阶段、共 10 步——

下面代码与论文 Algorithm 1 / 式 (28) 对齐，便于对照 HybridNewtonSchulz 与单步 Muon 更新（教学用 NumPy，非 V4 生产实现）：

import numpy as np

# V4 Hybrid Newton-Schulz: 8 fast + 2 stable (paper §2.4)
NS_STAGES = [
    (8, (3.4445, -4.7750, 2.0315)),
    (2, (2.0, -1.5, 0.5)),
]


def newton_schulz_step(m: np.ndarray, abc: tuple[float, float, float]) -> np.ndarray:
    a, b, c = abc
    mm = m @ m.T
    return a * m + b * (mm @ m) + c * (mm @ mm @ m)


def hybrid_newton_schulz(m: np.ndarray) -> np.ndarray:
    """Approximate polar factor UV^T; input M is n×m."""
    m = m.astype(np.float64)
    norm = np.linalg.norm(m, ord="fro")
    if norm == 0:
        return m
    m = m / norm
    for n_steps, coeffs in NS_STAGES:
        for _ in range(n_steps):
            m = newton_schulz_step(m, coeffs)
    return m


def muon_step(
    w: np.ndarray,
    grad: np.ndarray,
    momentum: np.ndarray,
    *,
    lr: float = 2.7e-4,
    mu: float = 0.95,
    weight_decay: float = 0.1,
    gamma: float = 0.18,
) -> tuple[np.ndarray, np.ndarray]:
    """One Muon step for matrix W (Algorithm 1)."""
    momentum = mu * momentum + grad
    nesterov = mu * momentum + grad
    ortho = hybrid_newton_schulz(nesterov)
    scale = np.sqrt(max(w.shape)) * gamma
    update = ortho * scale
    w = w * (1 - lr * weight_decay) - lr * update
    return w, momentum

读法：hybrid_newton_schulz 输出近似正交矩阵；sqrt(max(n,m)) * gamma 把更新 RMS 缩放到与 AdamW 超参兼容的量级；最后 decoupled weight decay 与 AdamW 形式相同。

3. 哪些参数用 Muon，哪些仍用 AdamW

动机：正交化是 矩阵级 操作；embedding / LayerNorm 类 向量或标量 参数不适合 Muon，继续 AdamW 更自然。

V4-Flash / V4-Pro 共用 Muon 超参：momentum 0.95、weight decay 0.1、$\gamma=0.18$；AdamW 侧 $\beta_1=0.9$、$\beta_2=0.95$、$\varepsilon=10^{-20}$、weight decay 0.1（§4 训练设置）。

4. 训练 infra：Muon 与 ZeRO

Muon 需要 完整梯度矩阵 才能做 NS 迭代，与传统 ZeRO 按元素切分 AdamW 状态 冲突。V4 训练框架（§3.4.1）要点：

Muon 与 CSA/HCA、mHC、Hash MoE 同期打包进 V4，难以单独 ablation 优化器 vs 架构（见 V4 定位）。

5. 演进链小结

6. 上下游

参考

• 论文：arXiv:2606.19348 — Algorithm 1、§2.4、§3.4.1、§4 超参
• Muon 原论文：Momentum Orthogonalized by Newton-Schulz
• Kimi Muon 实践（V4 引用）：Moonshot Kimi 技术报告中的 Muon + QK-Clip（V4 省略 QK-Clip）