不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

機器之心編輯部

在數(shù)值分析領(lǐng)域，Newton-Schulz 及其相關(guān)方法已被研究多年，但大多數(shù)工作關(guān)注的是高精度計算、CPU 優(yōu)化或方陣輸入。

就在昨天，普林斯頓大學、紐約大學的四位研究者提出 Gram Newton-Schulz，通過重構(gòu) Newton-Schulz，使其更適配 GPU 和大模型訓練場景，在萬億參數(shù) MoE 模型中可將優(yōu)化器時間降低 40–50%

我們用一句話來總結(jié) Gram Newton-Schulz 的核心思想：

不再直接在矩陣 X∈R^n×m 上迭代，而是在更小的 Gram 矩陣 X X^?∈R^n×n 上迭代，從而降低計算量并利用對稱矩陣計算優(yōu)化。

核心貢獻包括如下：

• 將標準 Newton-Schulz 重寫為數(shù)學等價形式，主要作用在 n×n 空間；
• 提出 Naive Gram Newton-Schulz（基礎(chǔ)版本）；
• 分析半精度下的不穩(wěn)定性，并提出重啟策略（Stabilized 版本）；
• 實現(xiàn)對稱矩陣乘法 GPU kernel；
• 構(gòu)建 GramMuon，在實際訓練中實現(xiàn) 40–50% 加速且精度無損。

下圖 1 為基于 NVIDIA B300 平臺，AdamW 與 Muon 在 LLaMA 不同模型規(guī)模下的優(yōu)化器步驟耗時（實際墻鐘時間）對比。

作者之一、普林斯頓大學本科生 Jack Zhang 表示，「我們讓 Muon 在幾乎不增加額外成本的情況下，運行速度最高提升 2 倍！Gram Newton-Schulz 可以作為 Muon 中 Newton-Schulz 的即插即用替代方案：我們觀察到驗證集困惑度幾乎沒有變化，誤差在 0.01 以內(nèi)。同時，我們也分享了在這一算法穩(wěn)定化過程中的大量探索，確保在任何情況下都不犧牲訓練質(zhì)量。」

作者之一、普林斯頓大學助理教授、Together AI 聯(lián)合創(chuàng)始人兼首席科學家 Tri Dao 表示，「這是我最喜歡的一類工作 —— 線性代數(shù)洞察 + 高性能算子。我們花了幾個月時間，深入研究 Muon 中間過程里這些矩陣的特征值和特征向量，最終提出了一種簡單而優(yōu)雅的算法，讓這一思路真正落地。」

接下來我們來看 Gram Newton-Schulz 相較于標準 Newton-Schulz 方法的優(yōu)勢。在此之前，有必要先回顧一下 Muon 優(yōu)化器。

Muon 可以理解為在譜范數(shù)下的最速下降法：

其中：μ 表示動量系數(shù)，η 表示學習率，polar 表示極分解。

極分解定義：若 X=UΣV^?，則：

由于精確計算昂貴，Muon 使用 Newton-Schulz 進行近似。

標準 Newton-Schulz 的迭代形式如下：

本質(zhì)上是對奇異值進行逐步歸一化，使其趨向 1。不過，標準 Newton-Schulz 存在計算瓶頸。每輪迭代包含三次矩陣乘法：

• XX^?：2mn^2
• A^2：2n^3
• BX：2mn^2

總 FLOPs 如下：

這就導致了一些問題，包括大量矩形矩陣乘法（GPU 效率低） 、未利用對稱性，并成為優(yōu)化器瓶頸。

因此，本文 Gram Newton-Schulz 核心思想是：將迭代從 X 轉(zhuǎn)移到 Gram 矩陣

關(guān)鍵變換：任何奇數(shù)多項式

可寫為

這樣一來，就將問題轉(zhuǎn)為了 1）對 R 進行矩陣多項式迭代、 本質(zhì)是近似 Y^-1/2。優(yōu)勢在于：維度從 n × m → n × n、可用對稱矩陣乘法 kernel 以及 減少矩形 GEMM 次數(shù)。

不過，Naive Gram Newton-Schulz在 float16 下不穩(wěn)定，Gram 矩陣可能產(chǎn)生負特征值，也會出現(xiàn) loss spike 和 Inf。

下圖為 Llama-430M 上的 Naive Gram Newton-Schulz：

因此有了穩(wěn)定版的 Stabilized Gram Newton-Schulz

核心改進在于使用float16 加速

并在第 2 步后重啟：

以及重新初始化 Gram。

結(jié)果顯示，在保持穩(wěn)定性的同時，仍然比原算法更快。

標準 Newton-Schulz、Naive Gram 與 Stabilized Gram 的復雜度對比如下：

結(jié)果表明，當 α>1（常見情況），Gram 方法明顯更便宜，F(xiàn)LOPs 降低最高約 42%–58%。

最后來看下 Gram Newton-Schulz 在 Kimi K2 中的耗時：

Kimi K2 是一個萬億參數(shù)級的稀疏、細粒度 MoE 模型，每一層包含 384 個專家，隱藏層維度為 7168，專家的中間層維度為 2048。由于模型正朝著更細粒度的 MoE 架構(gòu)發(fā)展，且 Kimi K2 使用 Muon 進行訓練，因此這是一個評測 Gram Newton-Schulz 的理想場景。

團隊將 Kimi K2 一次全局訓練步驟中暴露出來的 Newton-Schulz 的墻鐘時間，近似為以下部分的總和：

• 216 個專家的 up/gate/down 權(quán)重，形狀為 2048 × 7168；
• 1 個 dense 的 up/gate/down 權(quán)重，形狀為 7168 × 18432。

下圖為 NVIDIA H100 Hopper 硬件上，在 Kimi K2 的流水線并行配置中，Gram Newton-Schulz 的速度是標準 Newton-Schulz 的 2 倍

下圖為 NVIDIA B300 Blackwell 硬件上，在 Kimi K2 的流水線并行配置中，Gram Newton-Schulz 的速度同樣是標準 Newton-Schulz 的 2 倍。

更多細節(jié)內(nèi)容請參閱原博客。

博客地址：https://dao-lab.ai/blog/2026/gram-newton-schulz/