Gram Newton-Schulz：一种用于Muon的快速、硬件感知的牛顿-舒尔茨算法

# Gram Newton-Schulz: 一种快速且硬件感知的 Newton-Schulz 算法，用于 Muon 来源：https://tridao.me/blog/2026/gram-newton-schulz/ ### 目录 - 标准 Newton-Schulz 的运行时间（https://tridao.me/blog/2026/gram-newton-schulz/#runtime-of-standard-newton-schulz） - 朴素 Gram Newton-Schulz 的运行时间（https://tridao.me/blog/2026/gram-newton-schulz/#runtime-of-naive-gram-newton-schulz） - 跟踪中间矩阵的特征值（https://tridao.me/blog/2026/gram-newton-schulz/#tracking-eigenvalues-of-intermediate-matrices） - 虚假负特征值（https://tridao.me/blog/2026/gram-newton-schulz/#spurious-negative-eigenvalues） - 特征向量漂移（https://tridao.me/blog/2026/gram-newton-schulz/#eigenvector-drift） - 通过重启稳定 Gram Newton-Schulz（https://tridao.me/blog/2026/gram-newton-schulz/#stabilizing-gram-newton-schulz-by-restarting） - 进一步预防措施（https://tridao.me/blog/2026/gram-newton-schulz/#further-precautions） - 关于稳定性的收获（https://tridao.me/blog/2026/gram-newton-schulz/#takeaways-on-stability） - 稳定版 Gram Newton-Schulz 的运行时间（https://tridao.me/blog/2026/gram-newton-schulz/#runtime-of-stabilized-gram-newton-schulz） - 布局工程与工作调度（https://tridao.me/blog/2026/gram-newton-schulz/#layout-engineering-and-work-scheduling） - 代码中的实现策略（https://tridao.me/blog/2026/gram-newton-schulz/#implementation-strategy-in-code） - 标准 Newton-Schulz 的内核优化（https://tridao.me/blog/2026/gram-newton-schulz/#kernel-optimizations-for-standard-newton-schulz） - 模型质量保持不变（https://tridao.me/blog/2026/gram-newton-schulz/#model-quality-is-preserved） - 我们的方法加速了优化器步骤（https://tridao.me/blog/2026/gram-newton-schulz/#our-method-speeds-up-the-optimizer-step） Muon 正成为训练最先进语言模型（如 Kimi K2 Thinking1 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:kimi) 和 GLM-5.2 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:GLM)）的首选优化器。与 AdamW 相比，Muon 达到给定损失所需的优化器步数更少，但每一步的计算代价更高。这一开销源于 Muon 的 Newton-Schulz 正交化过程，这是一种旧优化器中没有的立方时间复杂度矩阵操作。 icml_optimizer_plot_blackwell (2) *图 1：AdamW vs. Muon：在 B300 上基准测试的不同 Llama 模型规模的优化器步骤的挂钟时间。* Muon 卓越的优化质量使其更昂贵的优化器步骤物有所值。然而，随着模型规模扩大，计算每个 Muon 步骤的开销迅速增长。传统优化方法（SGD, AdamW）执行逐元素操作，例如更新动量或按二阶矩重新缩放动量。对于一个大小为 $n \times m$ 的权重矩阵，给定梯度矩阵作为输入，执行优化器步骤需要 $O(mn)$ 时间。相比之下，许多现代优化器（Muon,3 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:muon) Scion,4 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:scion) Dion,5 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:dion) SOAP,6 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:soap) Shampoo,7 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:shampoo) SPlus,8 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:splus) 等）在每一步训练中使用正交化或高阶预处理来计算权重的更新。这些方法需要代价为 $O(mn^2)$ 的矩阵乘法（假设 $n \leq m$）。因此，每次调用优化器的运行时间远大于 AdamW。根据训练设置（全局批次大小、集群规模和并行化设置），Newton-Schulz 占端到端挂钟时间的 2% 到 17% (https://tridao.me/blog/2026/gram-newton-schulz/#appendix)。虽然 $O(mn^2)$ 运行时间是这些算法不可避免的开销，但在 FLOPs 和挂钟时间方面仍有很大的改进空间。 按照通常的实现方式，Newton-Schulz 例程存在几个缺点： 1. 它使用不只是一次或两次，而是 **十次** 大小为 $n \times m$ 的矩阵乘法，每次花费 $2mn^2$ 个 FLOPs。大多数流行架构中的权重是矩形矩阵，且 $m \gg n$，而近期拥有许多细粒度专家的 MoE 架构的权重矩形程度甚至 **更高**。因此，矩形矩阵乘法主导了其他操作（如小的 $n \times n$ 矩阵乘法）的代价。 2. 它计算的许多中间矩阵是对称的，但未利用这一结构计算上的优势。用于计算这些矩阵的一半工作是冗余的。 3. 它使用 cuBLAS 进行批量矩阵乘法/加法 $\alpha \mathbf A \mathbf B + \beta \mathbf C$，而 cuBLAS 并未针对 Hopper GPU 架构进行完全优化。 先前的工作试图通过优化 Newton-Schulz 的多项式系数或其归一化步骤来改进它。虽然这可以减少 Newton-Schulz 收敛所需的迭代次数，但并未解决上述缺点。其他人 9 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:flashmuon) 使用专门设计的对称矩阵乘法例程实现了 Newton-Schulz，但由于矩形和非对称乘法数量较多，运行时收益有限。 虽然 Newton-Schulz 及相关方法已在数值分析文献中研究了数十年，但研究关注主要集中在需要高精度的场景、算法针对 CPU 而非 GPU 优化的场景、或者输入矩阵为方阵的场景。近年来，随机草图化被用于设计许多涉及高度矩形矩阵计算的复杂算法；然而（除了进一步优化系数 10 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:PRISM) 之外），这些方法似乎不适用于 Muon。 ## 我们的贡献 为了解决这些缺点，我们引入了 **Gram Newton-Schulz**，这是对 Newton-Schulz 例程的重新设计，**将优化器时间最多减少了 50%**，在像 Kimi K2 这样的万亿参数 MoE 模型中。Gram Newton-Schulz 不是在矩形输入矩阵 $\mathbf{X} \in \mathbb{R}^{n \times m}$ 上进行迭代，而是在小型方对称 Gram 矩阵 $\mathbf{XX}^\top \in \mathbb{R}^{n \times n}$ 上进行迭代，从而降低了 FLOPs 代价并能够更多使用对称 GEMM 内核。 我们的贡献如下。首先，我们展示了如何将标准 Newton-Schulz 重写为一种 **数学上完全相同** 的形式（输出与浮点误差范围内完全相同），但主要作用在 $n \times n$ 矩阵的空间上。因为这些矩阵更小且允许专门的对称矩阵乘法例程，每次迭代比标准 Newton-Schulz 更快。只有预处理步骤（形成 $\mathbf X \mathbf X^\top$）和后处理步骤（乘以 $\mathbf X$）需要矩形矩阵乘法。我们将这种新形式称为朴素 Gram Newton-Schulz (https://tridao.me/blog/2026/gram-newton-schulz/#alg-naive-gram-ns)。 其次，我们对朴素 Gram Newton-Schulz 的数值特性进行了全面研究。我们发现了使用半精度浮点运算时可能出现数值不稳定的潜在问题，尤其是 Gram 矩阵中出现虚假负特征值的情况。我们使用一种“重启”策略来修复这种不稳定性，即在算法中途重建 Gram 矩阵。我们将修改后的算法称为稳定版 Gram Newton-Schulz (https://tridao.me/blog/2026/gram-newton-schulz/#alg-stable-gram-ns)。 第三，为了充分利用最新一代 GPU 以及 Newton-Schulz 的数学结构，我们实现了 **对称** 矩阵乘法的自定义内核。这些内核使用 CuTeDSL 为 Hopper 和 Blackwell 架构实现，达到了最先进的性能。 最后，我们将 Muon 的 Newton-Schulz 例程替换为 Gram Newton-Schulz，这种优化器我们称为 **GramMuon**，并观察到正交化步骤的运行时间减少了 40-50%。实验证实，使用 GramMuon 训练语言模型是稳定的，并且在验证困惑度 $0.01$ 的范围内保留了标准版本的优化质量，使我们的算法成为罕见的“免费午餐”性能改进实例。 为了促进 Gram Newton-Schulz 的采用，我们发布了以下开源实现： 1. 一个即插即用替代品 (https://github.com/Dao-AILab/gram-newton-schulz)，用于替代 Muon 的 Newton-Schulz 例程，在数学上等价、数值稳定，且速度提高多达两倍。 2. 快速 GPU 内核 (https://github.com/Dao-AILab/quack/blob/main/quack/gemm_symmetric.py)，用于对称矩阵乘法（$AB$, $\alpha AB + \beta C$），使用 CuTeDSL 为 Hopper 和 Blackwell 编写，可能具有独立的研究价值。 我们将首先回顾 Muon，了解为什么首先需要 Newton-Schulz，描述标准 Newton-Schulz 的数学工作原理，并分析其性能瓶颈。 ## Muon 回顾 Muon 优化器 3 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:muon) 最好被描述为关于谱范数的最陡下降方向。11 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:deriving_muon) 在训练的第 $k$ 步，令 $\mathbf W_k \in \mathbb{R}^{n \times m}$ 为权重矩阵，$\mathbf G_k$ 为损失关于 $\mathbf W_k$ 的梯度。Muon 的更新规则为 \\\[ \begin{align*} \mathbf{M}_k &= \mu \mathbf{M}_{k-1} + \mathbf{G}_k \\ \mathbf{W}_{k+1} &= \mathbf{W}_k - \eta \operatorname{polar}(\mathbf{M}_k) \end{align*} \\\] 其中 $\mu$ 是动量系数，$\eta$ 是学习率，$\mathbf M_k$ 是动量矩阵（$\mathbf M_0 := 0$）。在大多数方面，Muon 类似于带有动量的基本随机梯度下降 (SGD)。其关键创新在于使用了 $\operatorname{polar}$ 操作，其定义如下： > **定义 1: 极分解** 如果 $\mathbf X = \mathbf U \mathbf \Sigma \mathbf V^\top$ 是矩阵的奇异值分解 (SVD)，则 $\operatorname{polar}(\mathbf X) = \mathbf U \mathbf V^\top$。 由于精确计算 $\operatorname{polar}(\mathbf X)$ 代价高昂，Muon 使用 Newton-Schulz 方法来近似它。Newton-Schulz 是一种基于矩阵多项式的迭代方法。从 $\mathbf X_0$ 开始，每次迭代根据更新规则改进近似 $\mathbf X_t \approx \operatorname{polar}(\mathbf X)$ \\\[ \mathbf{X}_{t+1} = a_t \mathbf X_t + b_t \mathbf X_t \mathbf X_t^\top \mathbf X_t + c_t \left(\mathbf X_t \mathbf X_t^\top\right)^2 \mathbf X_t. \\\] 我们可以通过理解它如何影响奇异值分解来解释 Newton-Schulz。令 $\mathbf X_0 = \mathbf U \mathbf \Sigma \mathbf V^\top$ 为 SVD。回想 $\mathbf U$ 和 $\mathbf V$ 具有标准正交列，使得 $\mathbf U^\top \mathbf U = \mathbf V^\top \mathbf V = \mathbf I$，并且 $\mathbf \Sigma$ 是对角矩阵，其元素称为奇异值。那么 \\\[ \mathbf X_0 \mathbf X_0^\top \mathbf X_0 = \left(\mathbf U \mathbf \Sigma \mathbf V^\top\right) \left(\mathbf U \mathbf \Sigma \mathbf V^\top\right)^\top \left(\mathbf U \mathbf \Sigma \mathbf V^\top\right) = \mathbf U \mathbf \Sigma \mathbf V^\top \mathbf V \mathbf \Sigma \mathbf U^\top \mathbf U \mathbf \Sigma \mathbf V^\top = \mathbf U \mathbf \Sigma^3 \mathbf V^\top \\] 同理， \\\[ \mathbf X_1 = \mathbf U \left(a_1 \mathbf \Sigma + b_1 \mathbf \Sigma^3 + c_1 \mathbf \Sigma^5 \right) \mathbf V^\top = \mathbf U p_1(\mathbf \Sigma) \mathbf V^\top \\] 其中我们定义了多项式 $p_1(x) = a_1 x + b_1 x^3 + c_1 x^5$。由于 $\mathbf U$ 和 $\mathbf V$ 具有标准正交列且 $p_1(\mathbf \Sigma)$ 是对角矩阵，该等式的右侧必定是 $\mathbf X_1$ 的 SVD！这表明 $\mathbf X_1$ 与 $\mathbf X_0$ 共享相同的奇异向量 $\mathbf U$ 和 $\mathbf V$，并且其奇异值是 $\mathbf X_0$ 的奇异值经过多项式 $p_1$ 变换后的结果。推广到更高步，$\mathbf X_T$ 也与 $\mathbf X_0$ 共享相同的奇异向量 $\mathbf U$ 和 $\mathbf V$，并且其奇异值已经过多项式复合 $(p_T \circ \cdots \circ p_1)(\mathbf \Sigma)$ 的变换。如果对于所有奇异值都有 $(p_T \circ \cdots \circ p_1)(x) \approx 1$，那么 $(p_T \circ \cdots \circ p_1)(\mathbf \Sigma) \approx \mathbf I$，因此 $\mathbf X_T \approx \mathbf U \mathbf V^\top = \operatorname{polar}(\mathbf X_0)$。 剩下的工作是找到一系列奇次多项式，使得 $(p_T \circ \cdots \circ p_1)(x) \approx 1$ 在奇异值上成立。为了简化，我们首先将矩阵 $\mathbf X_0 = \mathbf X / \|\mathbf X\|_{\mathsf F}$ 归一化。这确保了 $\mathbf X_0$ 的奇异值位于区间 $[0, 1]$ 内。Muon 的开发者确定了一个由五个 5 次奇次多项式组成的序列，它们在 $[0, 1]$ 区间内对每个输入都近似于 $1$，从而在仅五次迭代中就对典型输入 $\mathbf X$ 给出了 $\operatorname{polar}(\mathbf X)$ 的良好近似。3 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:muon) 标准 Newton-Schulz 的实现如下所示： > **算法 1: 标准 Newton-Schulz** > 输入：$\mathbf X \in \mathbb{R}^{n \times m}$，系数 $\{(a_t, b_t, c_t)\}_{t=1}^5$ > 1. $\mathbf X \gets \mathbf X \,/\, (\lVert\mathbf X\rVert_{\mathsf F} + \epsilon)$ // 将奇异值归一化到 $[0, 1]$。$\epsilon = 10^{-7}$ > 2. $\mathbf X \gets \texttt{bfloat16}(\mathbf X)$ // 转换为半精度以提高速度 > 3. 如果 $m < n$: $\mathbf X \gets \mathbf X^\top$ // 使 $\mathbf X \mathbf X^\top$ 计算更便宜的技巧 > 4. 对于 $t = 1, \ldots, 5$: // 应用 $p_t(\mathbf X)$ > 5. $\mathbf A \gets \mathbf X\mathbf X^\top$ > 6. $\mathbf B \gets b_t \mathbf A + c_t \mathbf A^2$ > 7. $\mathbf X \gets a_t \mathbf X + \mathbf B \mathbf X$ > 8. 如果 $m < n$: $\mathbf X \gets \mathbf X^\top$ // 撤销技巧 > 9. 返回 $\mathbf X$ 后续工作试图通过多种方式改进 Muon。大多数这些提议修改了 Muon 的更新规则，以便在更少的训练步骤中达到相同的损失；然而，它们使用了上述相同的 Newton-Schulz 例程。有些方法（如 Polar Express12 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:polar-express)）确实通过改变多项式序列13 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:grishina) 或归一化步骤来改进 Newton-Schulz。虽然它们提高了近似精度，但并未改变其挂钟运行时间。Dion 优化器5 (https://tridao.me/blog/2026/gram-newton-schulz/#fn:dion) 在分布式设置（权重和梯度分片到不同 GPU）中减少了运行时间。它使用 Muon 的低秩近似来降低通信成本和 $\mathbf X$ 的维度，但每一步仍然调用标准的 Newton-Schulz 例程。相比之下，我们的工作加速了 Newton-Schulz 本身。由于 Gram Newton-Schulz 在数学上与标准版本完全相同，因此它几乎与所有种类的 Muon 兼容。 ## 标准 Newton-Schulz 的运行时间 让我们以 FLOP 为单位分析 Newton-Schulz 的运行时间。