Newton-Schulz 正交化,FLOPs 减少 3.4 倍
- 将 Newton-Schulz 正交化整合为单一运算,FLOPs 减少 3.4 倍
- 正交化误差减少 96 倍 (3.838 → 0.040)
- 通过可学习系数确保收敛稳定性
发生了什么事?
中国研究团队发布了 Newton-Schulz 正交化的新集成框架 UNSO。[arXiv] 它用单一多项式运算替代了现有的 NS 迭代。核心是“消除无意义的项并引入可学习的系数”。
以 128×512 矩阵为基准,FLOPs 从 2.533×10^8 降至 8.831×10^7。[arXiv]
为什么重要?
Muon 作为替代 AdamW 的优化器而备受关注。它以 175 美元的价格训练了 GPT-2 XL,并且被用于 Kimi K2。[Keller Jordan] 但 Newton-Schulz 迭代是一个瓶颈。
UNSO 打破了这个瓶颈。它不是减少迭代,而是完全消除迭代。正交化误差也减少了 96 倍。误差累积是学习不稳定性的原因,而 UNSO 解决了这个问题。
未来会怎样?
PyTorch 2.10 正式包含了 Muon。[PyTorch] NVIDIA NeMo 也支持它。[NVIDIA] UNSO 很有可能被迅速吸收。
常见问题 (FAQ)
Q: 现在可以使用 UNSO 吗?
A: 是的。作者们在 GitHub 上公开了代码。可以在 PyTorch 环境中以替代现有 Muon 的形式使用。但建议在生产应用之前进行自身基准测试。
Q: 应该用 Muon 代替 AdamW 吗?
A: 视情况而定。Muon 仅适用于隐藏层。嵌入或输出层需要 AdamW。组合使用两个优化器是标准做法。
Q: 实际学习时间会减少多少?
A: 在正交化阶段快 3.4 倍。整体学习时间取决于模型大小。正交化比重越大,UNSO 的优势就越大。
如果这篇文章对您有帮助,请订阅 AI Digester。
参考资料
- UNSO: Unified Newton Schulz Orthogonalization – arXiv (2026-02-04)
- Muon: An optimizer for hidden layers – Keller Jordan Blog (2025-01-15)
- Muon Optimizer – PyTorch Documentation (2026-01-20)