在数值分析领域，Newton-Schulz及其相关方法已被研究多年，但大多数工作关注的是高精度计算、CPU优化或方阵输入。
近日，来自普林斯顿大学和纽约大学的四位研究者提出了Gram Newton-Schulz算法。该研究通过重构经典的Newton-Schulz方法，使其更适配GPU硬件和大规模模型训练场景。实验表明，该算法在训练万亿参数的混合专家模型时，可将优化器步骤的耗时降低40%至50%。
Gram Newton-Schulz的核心思想可以概括为：
不再直接在原始矩阵 ( X in mathbb{R}^{n times m} ) 上进行迭代，而是转移到其维度更小的Gram矩阵 ( X X^top in mathbb{R}^{n times n} ) 上进行迭代。此举显著降低了计算量，并充分利用了对称矩阵的计算优化潜力。
该研究的主要贡献包括：
* 将标准Newton-Schulz方法重写为数学等价形式，使其主要作用于 ( n times n ) 的空间。
* 提出了基础的 Naive Gram Newton-Schulz 版本。
* 分析了该算法在半精度（float16）下的不稳定性，并提出了相应的重启策略，从而得到稳定的 Stabilized Gram Newton-Schulz 版本。
* 实现了高效的对称矩阵乘法GPU内核。
* 构建了 GramMuon 优化器，在实际模型训练中实现了40-50%的加速，且未损失模型精度。
下图展示了在NVIDIA B300平台上，AdamW与Muon优化器在不同规模的LLaMA模型上执行优化器步骤的耗时对比。
论文作者之一、普林斯顿大学本科生Jack Zhang表示：“我们的工作让Muon优化器的运行速度最高提升了2倍，而几乎没有增加额外成本。Gram Newton-Schulz可以作为Muon中Newton-Schulz模块的即插即用替代方案。我们在验证集上观察到，模型的困惑度几乎没有变化，误差控制在0.01以内。”
另一位作者、普林斯顿大学助理教授、Together AI联合创始人兼首席科学家Tri Dao评论道：“这是我最喜欢的一类工作——深刻的线性代数洞察与高性能算子实现的结合。我们花费了数月时间，深入分析Muon中间过程中涉及的矩阵特征值与特征向量，最终提出了一种简洁而优雅的算法，使这一想法得以实现。”
算法优势：从标准Newton-Schulz到Gram Newton-Schulz
为了更好地理解Gram Newton-Schulz的优势，首先需要回顾Muon优化器的背景。Muon可以理解为在谱范数下的最速下降法：
其中，( mu ) 表示动量系数，( eta ) 表示学习率，( text{polar} ) 表示极分解。
极分解定义为：若 ( X = U Sigma V^top )，则：
[
text{polar}(X) = U V^top
]
由于精确计算极分解代价高昂，Muon使用Newton-Schulz迭代进行近似。
标准Newton-Schulz的迭代形式如下：
该迭代本质上是对矩阵的奇异值进行逐步归一化，使其趋近于1。然而，标准Newton-Schulz存在计算瓶颈：每轮迭代包含三次矩阵乘法，总计算量（FLOPs）为：
这导致了几个问题：大量低效的矩形矩阵乘法、未充分利用矩阵的对称性，最终使优化器步骤成为训练瓶颈。
因此，Gram Newton-Schulz的核心思想是将迭代从原始矩阵 ( X ) 转移到其Gram矩阵 ( R = X X^top ) 上：
关键的理论变换在于：任何奇数多项式 ( p(t) ) 都可以重写为 ( t cdot q(t^2) ) 的形式。通过这一变换，迭代过程被转化为在 ( R ) 上进行的矩阵多项式迭代，本质上是近似 ( Y^{-1/2} )。
这种方法带来了显著优势：
* 维度降低：从 ( n times m ) 降至 ( n times n )。
* 计算优化：可以使用专为对称矩阵设计的GPU内核。
* 减少低效运算：减少了矩形通用矩阵乘法（GEMM）的次数。
稳定性挑战与解决方案
然而，基础的Naive Gram Newton-Schulz在float16精度下不稳定，Gram矩阵可能产生负特征值，导致训练出现损失尖峰（loss spike）或无穷大（Inf）值。
（图为在Llama-430M模型上使用Naive Gram Newton-Schulz时出现的不稳定现象）
为此，研究者提出了Stabilized Gram Newton-Schulz版本。其核心改进包括：
1. 在关键步骤使用float16进行加速。
2. 引入重启机制。
3. 对Gram矩阵进行重新初始化。
稳定版算法的流程如下：
结果表明，稳定版算法在保持数值稳定性的同时，速度仍显著快于原算法。
性能对比与实际应用
标准Newton-Schulz、Naive Gram与Stabilized Gram三种方法的计算复杂度对比如下：
结果表明，当矩阵形状比 ( alpha = m/n > 1 )（在深度学习模型中极为常见）时，Gram方法在计算量上明显更优，FLOPs最高可降低约42%至58%。
最后，研究在真实的万亿参数模型Kimi K2上评估了该算法。Kimi K2是一个稀疏、细粒度的混合专家模型，每层包含384个专家，隐藏层维度为7168，专家中间层维度为2048。由于其采用更细粒度的MoE架构并使用Muon进行训练，是评测Gram Newton-Schulz的理想场景。
在NVIDIA H100和最新的B300硬件上，于Kimi K2的流水线并行配置中进行测试，Gram Newton-Schulz的速度均是标准Newton-Schulz的2倍。
（NVIDIA H100 Hopper平台上的性能对比）
（NVIDIA B300 Blackwell平台上的性能对比）
更多技术细节请参阅研究团队发布的原博客。
博客地址：https://dao-lab.ai/blog/2026/gram-newton-schulz/
关注“鲸栖”小程序，掌握最新AI资讯
本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/archives/28104
