提出高效GPU矩阵极分解算法,提升大模型训练稳定性
The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm

- 基于最小最大优化设计迭代更新规则,加速收敛
- 支持bfloat16精度,实现在10亿参数模型上验证损失下降
- 适配深度学习需求,比现有方法在多种学习率下更优
计算极分解和相关矩阵符号函数是数值分析领域长期研究的问题。近年来,它成为训练深度神经网络的Muon优化器中的关键子程序。然而,该应用场景与传统设置差异显著:深度学习要求面向GPU的算法,优先考虑高吞吐而非高精度。本文提出Polar Express,一种新的极分解计算方法。与牛顿-舒尔茨等经典多项式方法类似,其仅使用矩阵-矩阵乘法,在GPU上极为高效。受陈与周及中田健人、弗伦德早期工作的启发,Polar Express 在每次迭代中通过求解一个极小极大优化问题来调整更新规则。我们证明该策略在最坏情况下最小化误差,使Polar Express在早期迭代和渐近阶段均实现最快收敛。同时解决有限精度问题,使其可在bfloat16下实用。集成到Muon后,该方法在基于FineWeb数据集(10亿至100亿词元)训练的GPT-2模型上,持续降低验证损失,优于近期各类替代方案,且在多种学习率下表现一致更优。
原文摘要 · Abstract (English)
Computing the polar decomposition and the related matrix sign function has been a well-studied problem in numerical analysis for decades. Recently, it has emerged as an important subroutine within the Muon optimizer for training deep neural networks. However, the requirements of this application differ sharply from classical settings: deep learning demands GPU-friendly algorithms that prioritize high throughput over high precision. We introduce Polar Express, a new method for computing the polar decomposition. Like Newton-Schulz and other classical polynomial methods, our approach uses only matrix-matrix multiplications, making it very efficient on GPUs. Inspired by earlier work of Chen & Chow and Nakatsukasa & Freund, Polar Express adapts the update rule at each iteration by solving a minimax optimization problem. We prove that this strategy minimizes error in a worst-case sense, allowing Polar Express to converge as rapidly as possible both in the early iterations and asymptotically. We also address finite-precision issues, making it practical to use in bfloat16. When integrated into Muon, our method yields consistent improvements in validation loss for a GPT-2 model trained on one to ten billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。