Gram Newton–Schulz#

主来源：Dao AI Lab, Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon（Jack Zhang, Noah Amsel, Berlin Chen, Tri Dao, 2026）。开源实现：Dao-AILab/gram-newton-schulz。

Muon 用 Newton–Schulz（NS）近似动量矩阵的 polar factor。标准 NS 每一轮都要对大矩形矩阵做乘法；Gram Newton–Schulz（Gram NS）通过代数重排，把大部分迭代搬到更小的方形 Gram 矩阵上，再配合对称 GEMM kernel 降低计算量。

一句话版本：

Gram NS = 只在开始、restart 和结束时碰大矩形矩阵 $X$，其余时间在小方阵 $R=XX^\top$ 上累计逆平方根多项式；再用 symmetric GEMM 只算半个输出。 它对长方形权重最有利，但必须用 restart、FP16 和安全缩放控制半精度误差。

1. 计算问题与 Gram 改写#

1.1 Muon 为什么要做 polar decomposition#

设某层权重为 $W_k$，梯度为 $G_k$。Muon 先积累矩阵动量：

$$M_k=\mu M_{k-1}+G_k,$$

然后用动量矩阵的 polar factor 作为更新方向：

$$W_{k+1}=W_k-\eta\,\mathrm{polar}(M_k).$$

如果

$$X=U_r\Sigma_r V_r^\top,$$

那么

$$\mathrm{polar}(X)=U_rV_r^\top.$$

这里取只包含非零奇异值的紧 SVD；秩亏时得到 canonical partial isometry，零奇异值不会被奇多项式变成 1。直觉上，这一步保留 $X$ 的左右奇异向量，但把奇异值大致压到 $1$。精确 SVD 太贵，Muon 因此用若干轮五次多项式 Newton–Schulz 近似它：

$$X_{t+1} =a_tX_t+b_tX_tX_t^\top X_t+c_t(X_tX_t^\top)^2X_t.$$

由于

$$X_t=U\Sigma_tV^\top$$

时有

$$X_tX_t^\top X_t=U\Sigma_t^3V^\top,$$

所以 NS 不改变奇异向量，只对每个奇异值应用标量多项式：

$$p_t(x)=a_tx+b_tx^3+c_tx^5.$$

只要多轮复合后的 $p(x)$ 把相关区间内的奇异值送到 $1$ 附近，就得到 $UV^\top$ 的近似。

1.2 标准 Newton–Schulz 慢在哪里#

实践中通常把 $X$ 转置到短边在前。设

$$X\in\mathbb R^{n\times m},\qquad n\le m,\qquad \alpha=\frac mn\ge1.$$

一轮标准 NS 可以写成：

A = X @ X.mT
B = b * A + c * (A @ A)
X = a * X + B @ X

对应的主计算量为：

$T$ 轮共需要：

$$T(4mn^2+2n^3)=2T(2\alpha+1)n^3.$$

Muon 常用 $T=5$，也就是：

$$(20\alpha+10)n^3.$$

真正的问题不是方阵 $A^2$，而是每一轮都有两次涉及 $n\times m$ 的大矩形乘法：

X @ Xᵀ       # 构造 Gram
B @ X        # 把多项式作用回 X

MLP、GQA/MLA 投影和细粒度 MoE expert 的权重往往很长方，$\alpha$ 越大，这两项越占主导。

1.3 核心改写：直接在 Gram 空间里迭代#

当 $X$ 满行秩时，polar factor 还有一个等价表达（秩亏时须在非零谱上取逆平方根，即用伪逆）：

$$\mathrm{polar}(X)=(XX^\top)^{-1/2}X.$$

这意味着不必每轮都显式构造新的长方形 $X_t$。可以改成：

1. R₀ = XXᵀ                         # 一次大矩形乘法
2. 在 n×n 的 R₀ 上近似 Q ≈ R₀⁻¹ᐟ²   # 主要迭代
3. X_out = QX                       # 一次大矩形乘法

如果 $m\gg n$，那么 $R,Q,Z$ 都只是 $n\times n$；相较于每轮都操作 $n\times m$ 的 $X$，计算量显著下降。

这也是 “Gram Newton–Schulz” 名字的来源：主要状态从长方形 $X$ 变成了 Gram 矩阵 $XX^\top$。

2. 等价性与计算成本#

2.1 为什么这种改写与标准 NS 数学等价#

五次奇多项式可以分解为：

$$p_t(x)=x h_t(x^2), \qquad h_t(y)=a_t+b_ty+c_ty^2.$$

对一个标量奇异值 $x$，定义：

$$r_0=x^2,\qquad q_0=1,$$

并迭代：

$$z_t=h_t(r_{t-1}),$$

$$r_t=r_{t-1}z_t^2,$$

$$q_t=q_{t-1}z_t.$$

那么有：

$$(p_T\circ\cdots\circ p_1)(x)=q_Tx.$$

其中：

• $r_t=x_t^2$，跟踪当前奇异值的平方；
• $q_t=x_t/x_0$（$x_0>0$），累计从输入到当前输出的缩放；
• 当 $x_t\to1$ 时，$q_t\to1/x_0=(x_0^2)^{-1/2}$。

把它提升到矩阵，就是 Naive Gram NS：

$$R_0=XX^\top,\qquad Q_0=I,$$

$$Z_t=a_tI+b_tR_{t-1}+c_tR_{t-1}^2,$$

$$Q_t=Q_{t-1}Z_t,$$

$$R_t=Z_tR_{t-1}Z_t,$$

最后返回：

$$X_T=Q_TX.$$

在精确算术中，它与逐轮生成 $X_1,\ldots,X_T$ 的标准 NS 等价；差别只是计算顺序。

2.2 标准 NS 与 Gram NS 的执行路径#

标准 NS#

X = normalize(X)
for a, b, c in coefficients:
    A = X @ X.mT
    B = b * A + c * (A @ A)
    X = a * X + B @ X
return X

特点：

• 每轮都更新 $n\times m$ 的 $X$；
• 每轮两次乘法涉及长方形 $X$；
• 中间误差不会在一个单独的逆平方根状态中长期积累，因此相对稳健。

Naive Gram NS#

X = normalize(X)
R = X @ X.mT
Q = I

for a, b, c in coefficients:
    Z = a * I + b * R + c * (R @ R)
    Q = Q @ Z
    R = Z @ R @ Z

return Q @ X

特点：

• 循环主体只操作 $n\times n$ 方阵；
• 开始和结束各有一次长方形乘法；
• 精确算术等价，但半精度下可能不稳定。

2.3 FLOP 节省来自哪里#

如果使用 symmetric GEMM，Naive Gram NS 在 $T=5$ 时的成本是：

$$(17+3\alpha)n^3.$$

实际的稳定版本需要一次 restart，成本变成：

$$(4T+6\alpha-6)n^3.$$

取一个常见例子：$T=5,\alpha=4$。

所以稳定 Gram NS 相比已经使用 symmetric GEMM 的标准 NS 仍少约 42% FLOP。

如果 $\alpha=1$，Gram NS 在 FLOP 上没有优势，而且会启动更多 GEMM；文章的实现对方阵回退到标准 NS。收益主要来自长方形权重，而不是所有形状。

3. 有限精度稳定性#

3.1 最大难点：精确等价不等于半精度稳定#

理论上：

$$R_t=X_tX_t^\top\succeq0,$$

所以 $R_t$ 的特征值都应非负。但在 FP16/BF16 中，构造 Gram 矩阵会引入很小的伪负特征值，例如 $-10^{-5}$。

Gram NS 本质上在近似：

$$R^{-1/2}.$$

对负特征值做实数逆平方根迭代没有意义。虽然这些负值最初只是舍入噪声，迭代映射却会不断放大其绝对值：

伪负特征值出现
    ↓
R_t 的负方向指数放大
    ↓
Z_t 变大
    ↓
Q_t 和最终 X_t 爆炸
    ↓
Inf / loss spike

文章还观察到第二种误差：eigenvector drift。精确算术中 $R,Q,Z$ 都是同一个 Gram 矩阵的多项式，应该共享特征向量；半精度乘法会逐渐破坏这种关系，使不同谱方向互相污染。

因此，Naive Gram NS 虽然代数上正确，却不能直接用于半精度训练。

3.2 Restart：用少量大矩形乘法换稳定性#

解决办法是在连续运行几轮 Gram 迭代后，把累计变换真正作用回 $X$，再从新的 $X$ 重建 Gram 状态：

# 先在 Gram 空间运行两轮
...

# restart
X = Q @ X
R = X @ X.mT
Q = I

# 再运行剩余轮次
...

对五轮 Polar Express 类系数，文章推荐在前两轮后 restart。它有三个作用：

1. 截断伪负特征值的连续放大；
2. 清除 $Q,R$ 之间累计的 eigenvector drift；
3. 重新建立与当前 $X$ 一致的 $R=XX^\top$。

restart 越多越稳定，但每次都会新增：

$$X\leftarrow QX,\qquad R\leftarrow XX^\top,$$

也就是两次涉及大矩形 $X$ 的乘法。极端情况下每轮都 restart，就退化成标准 NS。因此 restart 是速度与稳定性的调节旋钮。

3.3 其他稳定性措施#

FP16 优先于 BF16#

Gram NS 已把数值尺度控制在 $1$ 附近，不太需要 BF16 的大动态范围。FP16 尾数精度更高，文章因此默认使用 FP16；他们报告两者的训练 loss 没有明显差异，但 FP16 在部分测试矩阵上更准确。

更保守的 safety factor#

理论多项式通常只保证在 $\lVert X\rVert_2\le1$ 时稳定。舍入误差可能让奇异值略大于 $1$，因此 Polar Express 会做安全缩放。文章建议 Gram NS 使用更保守的值，例如从 1.02 调到 1.05。

不显式构造 $aI$#

朴素写法：

$$Z=aI+bR+cR^2$$

会先把整个 $Z$ 舍入回 FP16，然后再参与后续 GEMM，损失 $aI$ 的精度。文章改成：

$$Z=bR+cR^2,$$

$$Q_{\text{new}}=QZ+aQ,$$

$$RZ=RZ+aR,$$

$$R_{\text{new}}=Z(RZ)+a(RZ).$$

这样 $aQ$、$aR$ 可以在 GEMM epilogue 的 FP32 accumulation 中处理。

高精度数值计算要谨慎#

对满行秩矩阵、使用谱条件数时，显式形成 $XX^\top$ 会平方条件数：

$$\kappa(XX^\top)=\kappa(X)^2.$$

Muon 只需要近似正交化，通常可以接受；如果目标是高精度 polar decomposition，Gram NS 未必合适。

4. Kernel 实现与证据#

4.1 Symmetric GEMM kernel：只做下三角 MMA#

文章的第二层贡献是 Hopper/Blackwell 上的 CuTeDSL symmetric GEMM。它不局限于经典 $AA^\top$ SYRK，还支持在已知结果对称时计算：

$$AB,$$

或融合形式：

$$D=\alpha AB+\beta C.$$

执行方式：

普通 GEMM tile grid       symmetric GEMM tile grid

■ ■ ■ ■                  · · · ·
■ ■ ■ ■                  ■ · · ·
■ ■ ■ ■                  ■ ■ · ·
■ ■ ■ ■                  ■ ■ ■ ■

全矩阵执行 MMA            只给下三角分配工作
                          epilogue 同时写 (i,j) 和 (j,i)

工程要点：

• triangular scheduler 只调度下三角 tile，并均匀分配给 CTA cluster；
• 非对角 tile 算完后，同时转置写入上三角；
• 对角 tile 只写一次；
• 对称信息省掉约一半 MMA，但不会省掉最终完整矩阵的存储；
• 对标准 NS，$XX^\top$ 和 $bA+cA^2$ 走 symmetric GEMM，$aX+BX$ 仍是一般 GEMM+Add。

只换 kernel、不改算法时，文章报告标准 NS 在 Hopper 上已经能快约 26–29%。

4.2 实验结果应该怎样解读#

作者在 Llama-430M、Qwen-600M、Gemma-1B 和自定义 MoE-1B 上测试，并报告：

• 验证集 perplexity 与标准 Muon 的差异保持在 0.01 内；
• orthogonalization / Newton–Schulz 子步骤减少约 40–50% 时间；
• 长方形权重的 NS 最高约 2×；
• 对模拟 Kimi K2 的细粒度 MoE 权重配置，Hopper 和 Blackwell 上均报告约 2× NS speedup。

必须区分：“NS 子步骤快 2×”不等于“整个训练快 2×”。

如果 NS 原来只占训练 step 的比例 $f$，把它加速 $s$ 倍，Amdahl 定律给出的整体 speedup 是：

$$S_{\text{total}}=\frac{1}{(1-f)+f/s}.$$

例如 $f=0.02,s=2$ 时，整体只快约 1%；$f=0.17,s=2$ 时，整体约快 9%。这项优化在以下场景更重要：

• 前后向已使用 FP8/更低精度，而 NS 仍需 16-bit；
• SFT/RL 的全局 batch 较小，前后向工作较少；
• 大规模 pipeline 中 optimizer step 难以被 backward 完全隐藏；
• 权重高度长方形，尤其是细粒度 MoE expert。

4.3 它和 batched SYRK PR 不是同一个改动#

NVIDIA Emerging-Optimizers 的 batched SYRK PR #276 做的是：

A = batched_syrk(X)       # X @ Xᵀ，只算下三角并镜像
B = batched_syrk(A, A)    # bA + cA²
X = torch.baddbmm(X, B, X, beta=a)

它仍然是标准 NS：每轮都更新长方形 $X$，只是前两个对称乘法换成 batched SYRK。

Dao Lab 的 Gram NS 则是：

算法重排：不再每轮更新长方形 X
    +
restart：控制 Gram 空间的半精度误差
    +
symmetric GEMM：加速方形对称乘法
    +
fused GEMM epilogue：减少中间 I/O 和舍入

因此两者的包含关系更接近：

4.4 阅读时应保留的边界条件#

1. 等价只在精确算术中成立。 半精度下，计算顺序改变会产生不同结果。
2. 形状决定收益。 方阵应回退到标准 NS；长方形矩阵才有明显算法收益。
3. 性能数字依赖硬件和 batch。 文中的 kernel 主要针对 Hopper/Blackwell，不能直接外推到任意 GPU。
4. restart 位置不是普适常数。 换多项式系数或迭代次数后，应重新调 restart。
5. Gram 会平方条件数。 它适合 Muon 的近似正交化，不自动等价于通用高精度 polar solver。
6. 2× 是局部 speedup。 是否改善 end-to-end 训练，要看 NS 在总 step 中的占比与并行隐藏情况。

参考链接#

• Dao AI Lab：Gram Newton-Schulz
• 开源实现：Dao-AILab/gram-newton-schulz
• 对称 GEMM：Dao-AILab/quack
• NVIDIA batched SYRK PR #276
• NVIDIA 原始 Triton SYRK PR #47

4.1 要点回顾#

Gram Newton–Schulz 的核心不是“发现 $XX^\top$ 对称”，而是把 NS 的累计多项式从长方形 $X$ 搬到小方形 Gram 空间，再用 restart 修复这种重排在半精度下产生的谱误差。 Symmetric GEMM 是第二层加速；长方形越强、NS 越难被隐藏，收益越大。

延伸阅读#

• Triton Kernel 实现练习：怎样从逐元素计算走到归约、矩阵乘法和 Attention？