优化器 (3)（Gram Newton-Schulz）

2026-04-05 优化器 #优化器 #Muon #Newton-Schulz #优化理论 #大语言模型约 33 分钟 · 11353 字

📄 原文：Zifeng Mai · 阅读原文 · 代码

引言

在上一篇文章中，我们深入介绍了 Muon 优化器的数学原理。Muon 的核心创新在于使用 Newton-Schulz 迭代来近似矩阵的 msign 函数，也称为极分解（polar decomposition），从而实现对参数矩阵的正交化更新。

然而，标准的 Newton-Schulz 算法存在一个显著的计算瓶颈：它需要在每次迭代中对原始的矩形参数矩阵进行多次矩阵乘法运算。对于一个 $n \times m$ 的权重矩阵，每次 Newton-Schulz 迭代需要 $O(mn^2)$ 的时间复杂度，这在大规模训练中成为了不可忽视的开销。

近期，DAO Lab 提出了 Gram Newton-Schulz 算法，通过以下创新显著加速了 Muon 优化器：

数学等价性：在 Gram 矩阵 $\mathbf{X}\mathbf{X}^\top$ 上迭代而非原矩阵 $\mathbf{X}$ ，输出与标准 Newton-Schulz 数学等价
数值稳定性：通过重启策略和 float16 算术，实现与标准方法相当的稳定性
硬件感知：对称 GEMM 核充分利用 Hopper/Blackwell 架构特性

在 Kimi K2 这种万亿参数量级别的模型上，Gram Newton-Schulz 可以将优化器时间减少高达 50%。

一、Muon 与 Newton-Schulz 回顾

1.1. Muon 优化器更新规则

在第 $k$ 步训练时，令 $\mathbf{W}_k \in \mathbb{R}^{n \times m}$ 为权重矩阵， $\mathbf{G}_k$ 是 $\mathbf{W}_k$ 的梯度。Muon 的更新规则为：

(1)

其中 $\mu$ 是动量系数， $\eta$ 是学习率。 $\mathbf{M}_k$ 是动量矩阵， $\mathbf{M}_0 = 0$ 。

Definition 1（极分解）：若 $\mathbf{X} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top$ 是矩阵 $\mathbf{X}$ 的 SVD 分解，则 $\text{polar}(\mathbf{X}) = \mathbf{U}\mathbf{V}^\top$ 。

由于精确计算 $\text{polar}(\mathbf{X})$ 需要进行完整的 SVD 分解，计算开销较大，因此 Muon 使用 Newton-Schulz 迭代来近似它。

1.2. 标准 Newton-Schulz 迭代

Newton-Schulz 是一种基于矩阵多项式的迭代方法。从初始矩阵 $\mathbf{X}_0$ 开始，每次迭代按照以下规则更新近似值 $\mathbf{X}_t \approx \text{polar}(\mathbf{X})$ ：

(2)

其中 $(a_t, b_t, c_t)$ 是预先设定的多项式系数。

下面分析标准 Newton-Schulz 迭代的计算复杂度。

令 $T$ 表示迭代次数（在 Muon 中 $T=5$ ），并假设 $n \leq m$ ，定义纵横比 $\alpha = m/n \geq 1$ 。

每次迭代包含三次矩阵乘法：

$\mathbf{X}\mathbf{X}^\top$ ： $2mn^2$ FLOPs
$\mathbf{A}^2$ （其中 $\mathbf{A} = \mathbf{X}\mathbf{X}^\top$ ）： $2n^3$ FLOPs
$\mathbf{B}\mathbf{X}$ （其中 $\mathbf{B} = b_t\mathbf{A} + c_t\mathbf{A}^2$ ）： $2mn^2$ FLOPs

总计算量为：

(3)

当 $T=5$ 时，总计算量为 $(20\alpha + 10)n^3$ ，分布在 15 次 GEMM 操作上。

因此，标准的 Newton-Schulz 迭代存在以下问题：

未利用对称性： $\mathbf{A} = \mathbf{X}\mathbf{X}^\top$ 和 $\mathbf{B}$ 都是对称矩阵，但标准实现未利用这一结构。
强依赖 $\alpha$ ：当 $\alpha \gg 1$ 时，计算量主要由昂贵的 GEMM 操作主导。

二、Gram Newton-Schulz 的数学原理

2.1. 核心思想

Gram Newton-Schulz 的核心洞察在于下面的公式 (4)（证明见上一篇博客）。

(4)

基于公式 (4)，Gram Newton-Schulz 的策略是：

计算 $n \times n$ Gram 矩阵 $\mathbf{X}\mathbf{X}^\top$
使用迭代方法近似 $\mathbf{Q}_T \approx (\mathbf{X}\mathbf{X}^\top)^{-1/2}$
计算 $\mathbf{Q}_T \mathbf{X}$

Gram Newton-Schulz 的关键优势在于，第 2 步（占据绝大部分计算量）仅在小的 $n \times n$ 对称矩阵上操作，仅需两次矩形矩阵乘法（初始的 $\mathbf{X}\mathbf{X}^\top$ 和最终的 $\mathbf{Q}_T\mathbf{X}$ ）。

2.2. 从 Newton-Schulz 到 Gram Newton-Schulz

Proposition 1（奇异向量不变性）：设 $\mathbf{X}_0 = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top$ 是 $\mathbf{X}_0$ 的 SVD 分解。若 $\mathbf{X}_t$ 按照方程 (2) 迭代，则对任意 $t \geq 0$ ，存在奇次多项式 $p_t: \mathbb{R} \to \mathbb{R}$ 使得：

(5)

其中 $p_t(\mathbf{\Sigma}) = \text{diag}(p_t(\sigma_1), \ldots, p_t(\sigma_r))$ ， $r$ 为 $\mathbf{X}_0$ 的秩。

根据 Proposition 1，如果多项式序列 $(p_T \circ \cdots \circ p_1)(x) \approx 1$ 对所有奇异值成立，则 $\mathbf{X}_T \approx \mathbf{U}\mathbf{V}^\top = \text{polar}(\mathbf{X}_0)$ 。

那么，如何将迭代多项式方法 $(p_T \circ \cdots \circ p_1)(\mathbf{X}) \approx \text{polar}(\mathbf{X})$ 转换为近似 $\mathbf{Y} \mapsto \mathbf{Y}^{-1/2}$ 的迭代方法？

Lemma 1（奇次多项式的性质） 每个奇次多项式 $p(x) = ax + bx^3 + cx^5$ 可以重写为 $p(x) = x h(x^2)$ 的形式，其中 $h(y) = a + by + cy^2$ 是低一次的多项式。

因此，如果 $p(x) \approx 1$ ，则 $h(y) = p(y^{1/2})y^{-1/2} \approx y^{-1/2}$ ，因此 Newton-Schulz 多项式隐式地提供了近似逆平方根的方法。下面的定理展示了二者的对应关系。

Theorem 1（迭代求逆平方根）：设任取 $t \in \{1, \ldots, T\}$ 都有

(6)

则有

(7)

其中 $q_T$ 由以下迭代定义：

(8)

迭代的起始条件为 $r_0 = x^2$ ， $q_0 = 1$ 。

注意到当 $t\to+\infty$ 时，有 $q_t = x_t/x_0 \to 1/x_0 = (x_0^2)^{-1/2} = r_0^{-1/2}$ 。因此，我们就可以通过迭代来近似 $r_0^{-1/2}$ 。

2.3. Naive Gram Newton-Schulz 算法

将定理 1 的迭代提升到矩阵形式，作者提出了 Naive Gram Newton-Schulz 算法：

算法 1（Naive Gram Newton-Schulz）

输入： $\mathbf{X} \in \mathbb{R}^{n \times m}$ （ $n \leq m$ ），系数 $\{(a_t, b_t, c_t)\}_{t=1}^5$

$\mathbf{X} \gets \mathbf{X} / (\Vert\mathbf{X}\Vert_{\mathsf{F}} + \epsilon)$ （将奇异值归一化到 $[0, 1]$ ， $\epsilon = 10^{-7}$ ）
$\mathbf{R}_0 = \mathbf{X} \mathbf{X}^\top$
$\mathbf{Q}_0 = \mathbf{I}$
对于 $t = 1, \ldots, 5$ 按照下面的公式迭代

$(9)$
返回 $\mathbf{Q}_5 \mathbf{X}$

2.4. 计算复杂度分析

下面分析 Naive Gram Newton-Schulz 的 FLOPs 数量级。每次迭代包含四次矩阵乘法（使用对称 GEMM 核）：

$\mathbf{R}_{t-1}^2$ ： $n^3$ FLOPs
$\mathbf{Q}_{t-1} \mathbf{Z}_t$ ： $n^3$ FLOPs
$\mathbf{Z}_ {t} \mathbf{R}_ {t-1} \mathbf{Z}_ {t}$ ： $2n^3$ FLOPs

初始化和输出步骤：

$\mathbf{X}\mathbf{X}^\top$ ： $mn^2$ FLOPs
$\mathbf{Q}_5 \mathbf{X}$ ： $2mn^2$ FLOPs（非对称）

下面的计算量可以优化：

$\mathbf{Q}_ 1 = \mathbf{Q}_ {0}\mathbf{Z}_ 1=\mathbf{Z}_1$ 不需要计算：节省 $n^3$ FLOPs
$\mathbf{R}_5$ 不需要计算：节省 $2n^3$ FLOPs

因此 Naive Gram Newton-Schulz 的总 FLOPs 为：

(10)

当 $T=5$ 时，总计算量为 $(17 + 3\alpha)n^3$ ，分布在 19 次 GEMM 操作上。

使用对称 GEMM 的标准 Newton-Schulz 需要 $T(3\alpha + 1)n^3$ FLOPs。当 $\alpha > 1$ （非方阵）时，Gram Newton-Schulz 更优。

在 Muon 的场景下（ $T=5, \alpha=4$ ）：

相比使用对称 GEMM 的标准 Newton-Schulz 节省 55% FLOPs
相比典型实现（无对称 GEMM）节省 68% FLOPs

三、Stabilized Gram Newton-Schulz

3.1. Naive Gram Newton-Schulz 的数值稳定性问题s

尽管 Naive Gram Newton-Schulz 在精确算术下与标准 Newton-Schulz 等价，但在有限精度（尤其是半精度）下表现较差。实验发现，使用 bfloat16 训练 Transformer 时损失函数曲线会频繁出现尖峰，最终输出充满 Inf。

作者通过研究特征值和奇异值的演化，给出了矩阵发散的几个原因。

原因 1：伪负特征值（Spurious Negative Eigenvalues）

根据定义， $\mathbf{R}_t$ 应该是半正定矩阵（因为 $r_t = x_t^2 \geq 0$ ）。然而，在 bfloat16 下， $\mathbf{R}_0 = \mathbf{X}_0\mathbf{X}_0^\top$ 会引入微小的负特征值（如 $-10^{-5}$ ），这是由于浮点误差导致的。下面的命题指出，只要出现了负特征值，该特征值的幅度就会随着迭代而指数增长。

Proposition 2（负特征值指数增长）：设 $\mathbf{R}_0$ 为 Gram Newton-Schulz 迭代的初始 Gram 矩阵。若 $\mathbf{R}_0$ 存在负特征值 $\lambda < 0$ ，则该特征值的幅度随迭代步数指数增长。

Proof：回顾更新规则 $r_t = r_{t-1} z_t^2 = r_{t-1} h_t(r_{t-1})^2$ 。

代入 $h_t(x) = \frac{15}{8} - \frac{10}{8}x + \frac{3}{8}x^2$ ，当 $r_{t-1} < 0$ 时：

(11)

因此，若 $r_0 < 0$ ，伪特征值的幅度指数增长，引发链式反应：当 $r_t \to -\infty$ 时， $z_t \to \infty$ ，导致 $q_t \to \infty$ 和 $x_t \to \infty$ 。 $\blacksquare$

原因 2：特征向量漂移（Eigenvector Drift）

在精确算术下，所有中间矩阵的特征向量与 $\mathbf{X}_0$ 的左奇异向量 $\mathbf{U}$ 匹配，但在有限精度下会发生漂移。这导致 $\mathbf{Q}_t$ 和 $\mathbf{X}_t$ 的特征值偏离理论值。

3.2. 重启策略（Restarting）

为了缓解 Naive Gram Newton-Schulz 在低精度下的数值稳定性问题，作者提出了重启策略。重启的核心思想在于不直接计算 $\mathbf{X}_ T$ ，而是运行少量迭代（如 5 步）得到 $\mathbf{X}_ 5$ ，然后在 $\mathbf{X}_ 5$ 上再次应用 Gram Newton-Schulz 计算 $\mathbf{X}_ {10}$ ，依此类推。

每次重启时，由于重新计算了 Gram 矩阵 $\mathbf{R} = \mathbf{X}\mathbf{X}^\top$ ，这消除了之前积累的大幅度负特征值。另外，由于 $\mathbf{Q}_t$ 在每次重启时重置为单位矩阵，其特征值始终有界，因此 $\mathbf{X}_t = \mathbf{Q}_t \mathbf{X}$ 的特征值也能保持有界。

作者提出，重启的最佳时机取决于所使用的多项式系数序列。对于 Muon 来说，作者使用 Polar Express 的五次多项式：

$t$	$a_t$	$b_t$	$c_t$
1	8.123737	-22.232240	16.373715
2	4.026529	-2.776323	0.514551
3	3.870284	-2.739120	0.520999
4	3.253351	-2.343223	0.481420
5	2.300652	-1.668904	0.418807

通过数值模拟发现，在第 2 次迭代后重启可以最好地平衡稳定性和速度：

$\mathbf{R}_t$ 的最小特征值保持在 $-0.4$ 以上
$\mathbf{Q}_t$ 的条件数保持在 $\approx 100$ 以下

3.3. 进一步的安全措施

使用 float16 而非 bfloat16

float16 的范围（约 $6.1 \times 10^{-5}$ 到 $6.5 \times 10^4$ ）足够使用，且精度更高。在某些测试矩阵上，float16 给出更准确的 $\text{polar}(\mathbf{X})$ 近似。

引入安全因子

大多数 Newton-Schulz 多项式设计为在 $\Vert\mathbf{X}_0\Vert \leq 1$ 时收敛。由于数值误差可能导致奇异值略大于 1，因此作者建议引入安全因子：

(12)

这确保即使奇异值达到 1.05 也能收敛。

避免显式添加单位矩阵

计算矩阵二次型 $\mathbf{Z}_t \gets a_t\mathbf{I} + b_t \mathbf{R}_{t-1} + c_t \mathbf{R}_{t-1}^2$ 时，显式添加 $a_t\mathbf{I}$ 可能降低稳定性。更准确的方式是：

(13)

这样所有涉及 $a_t$ 的算术都在 float32 中进行，精度更高。

3.4. Stabilized Gram Newton-Schulz

综合以上分析，作者提出了完整的 Stabilized Gram Newton-Schulz 算法：

算法 2（Stabilized Gram Newton-Schulz）

输入： $\mathbf{X} \in \mathbb{R}^{n \times m}$ （ $n \leq m$ ），系数 $\{(a_t, b_t, c_t)\}_{t=1}^5$

$\mathbf{X} \gets \mathbf{X} / (\Vert\mathbf{X}\Vert_{\mathsf{F}} + \epsilon)$
$\mathbf{X} \gets \text{float16}(\mathbf{X})$ （转换为半精度）
第一次迭代：
- $\mathbf{R} \gets \mathbf{X}\mathbf{X}^\top$
- $\mathbf{Q} \gets \mathbf{I}$
- 对于：
  - $\mathbf{Z} \gets b_t \mathbf{R} + c_t \mathbf{R}^2$
  - $\mathbf{Q} \gets \mathbf{Q}\mathbf{Z} + a_t\mathbf{Q}$
  - $\mathbf{RZ} \gets \mathbf{R}\mathbf{Z} + a_t\mathbf{R}$
  - $\mathbf{R} \gets \mathbf{Z}(\mathbf{RZ}) + a_t(\mathbf{RZ})$
- $\mathbf{X} \gets \mathbf{Q}\mathbf{X}$ （重启）
第二次迭代：
- $\mathbf{R} \gets \mathbf{X}\mathbf{X}^\top$
- $\mathbf{Q} \gets \mathbf{I}$
- 对于：
  - $\mathbf{Z} \gets b_t \mathbf{R} + c_t \mathbf{R}^2$
  - $\mathbf{Q} \gets \mathbf{Q}\mathbf{Z} + a_t\mathbf{Q}$
  - $\mathbf{RZ} \gets \mathbf{R}\mathbf{Z} + a_t\mathbf{R}$
  - $\mathbf{R} \gets \mathbf{Z}(\mathbf{RZ}) + a_t(\mathbf{RZ})$
返回 $\mathbf{Q}\mathbf{X}$

下面来计算 Stabilized Gram Newton-Schulz 的计算复杂度（带一次重启）。

一次重启需要额外两次矩阵乘法：

$\mathbf{X} \gets \mathbf{Q}_2 \mathbf{X}$ ： $2mn^2$ FLOPs
$\mathbf{R}_2 \gets \mathbf{X}\mathbf{X}^\top$ ： $mn^2$ FLOPs

同时可以省略三次乘法：

$\mathbf{R}_2 \gets \mathbf{Z}_2 \mathbf{R}_1 \mathbf{Z}_2$ ：节省 $2n^3$ FLOPs
$\mathbf{Q}_3 \gets \mathbf{Q}_2 \mathbf{Z}_3$ ：节省 $n^3$ FLOPs

因此，带一次重启的 Stabilized Gram Newton-Schulz 的计算复杂度为：

(14)

对于 $T=5, \alpha=4$ ：

相比使用对称 GEMM 的标准 Newton-Schulz，减少 42% FLOPs
相比典型实现（无对称 GEMM），减少 58% FLOPs

此外，如果使用 $T-1$ 次重启，Gram Newton-Schulz 就完全退化为标准 Newton-Schulz。因此，添加重启可以视为在运行时间和稳定性之间权衡。

四、对称 GEMM 核优化

为了充分利用 Gram Newton-Schulz 带来的对称矩阵乘法机会，作者实现了专门的对称 GEMM 核。

三角调度器 (Triangular Scheduler)：作者仅将下三角部分的工作块分配给 warp group，上三角部分的值通过对称性获得（ $\mathbf{C}_{ij} = \mathbf{C}_{ji}$ ）。
融合二次型核：对称 GEMM 核可以融合矩阵二次型的计算，通过在寄存器级别添加 $\gamma$ 到所有对角线元素，完全避免了加载单位矩阵的 I/O 操作。
$(15)$

五、实验结果

5.1. 数值稳定性验证

在 Llama-430M 模型上的实验表明：

Naive Gram Newton-Schulz：出现损失尖峰，最终输出充满 Inf
Stabilized Gram Newton-Schulz：训练稳定，验证困惑度差异在 0.01 以内

5.2. 加速效果

在 Kimi K2 上：

Newton-Schulz 正交化步骤运行时间减少 40-50%
端到端训练时间减少 2-17%（取决于训练配置）

加速比随矩阵纵横比 $\alpha = m/n$ 增加而提高：

$\alpha = 1$ （方阵）：无加速（退化为标准 Newton-Schulz）
$\alpha = 4$ （典型 MLP）：42% FLOP 减少
$\alpha = 8$ （细粒度 MoE）：55% FLOP 减少

5.3. 对称 GEMM 核性能

在 Hopper (H100) 和 Blackwell (B200) GPU 架构上 benchmark：

相比 cuBLAS，对称 GEMM 核实现 1.5-2x 加速
在 Gram Newton-Schulz 中贡献了额外的 10-15% 加速
对于 $256 \times 256$ 以上矩阵，核效率超过 80%

Appendix

Apd.1. Proof of Proposition 1

(16)

其中 $p_t(\mathbf{\Sigma}) = \text{diag}(p_t(\sigma_1), \ldots, p_t(\sigma_r))$ ， $r$ 为 $\mathbf{X}_0$ 的秩。

Proof：我们使用数学归纳法。

当 $t=0$ 时，结论显然成立。

假设当 $t=k$ 时， $\mathbf{X}_k = \mathbf{U} p_k(\mathbf{\Sigma}) \mathbf{V}^\top$ ，其中 $p_k$ 是某个奇次多项式。

则当 $t=k+1$ 时：

(17)

其中 $p_{k+1}(x) = a_{k+1} p_k(x) + b_{k+1} p_k(x)^3 + c_{k+1} p_k(x)^5$ 。

因此结论对 $t=k+1$ 也成立。由数学归纳法，命题得证。 $\blacksquare$

Apd.2. Proof of Theorem 1

Theorem 1（迭代求逆平方根）：设任取 $t \in \{1, \ldots, T\}$ 都有

(18)

则有

(19)

其中 $q_T$ 由以下迭代定义：

(20)

迭代的起始条件为 $r_0 = x^2$ ， $q_0 = 1$ 。

Proof：任取对 $t \in \{1, \ldots, T\}$ ，定义 $x_0 = x$ 且 $x_t = p_t(x_{t-1})$ 。我们将通过归纳法证明 $r_t = x_t^2$ 且 $q_t = x_t/x_0$ 对所有 $t$ 成立。

当 $t=0$ 时，由定义 $r_0 = x^2$ ， $q_0 = 1$ ，结论成立。

假设结论对 $t-1$ 成立，即 $r_{t-1} = x_{t-1}^2$ 且 $q_{t-1} = x_{t-1}/x_0$ 。

由假设 $p_t(x) = x h_t(x^2)$ ，我们有：

(21)

注意到 $h_t(x_{t-1}^2) = h_t(r_{t-1}) = z_t$ ，因此 $x_t = x_{t-1} z_t$ 。

两边平方得：

(22)

两边除以 $x_0$ 得：

(23)

因此结论对 $t$ 也成立。由数学归纳法，定理得证。 $\blacksquare$