极地快车：最优矩阵符号方法及其在Muon算法中的应用

计算极分解和相关的矩阵符号函数一直是数值分析中经过深入研究的问题。最近，它作为Muon优化器中的一个重要子程序出现在深度神经网络训练中。然而，此应用的要求与经典设置截然不同：深度学习需要GPU友好的算法，优先考虑高吞吐量而非高精度。我们提出了Polar Express，一种计算极分解的新方法。与Newton-Schulz和其他经典多项式方法一样，我们的方法仅使用矩阵-矩阵乘法，使其在GPU上非常高效。受Chen & Chow和Nakatsukasa & Freund早期工作的启发，Polar Express通过求解极小极大优化问题来调整每次迭代的更新规则。我们证明该策略在最坏情况下最小化误差，使Polar Express能够在早期迭代和渐近阶段尽可能快速地收敛。我们还解决了有限精度问题，使其在bfloat16中可实际使用。当集成到Muon中时，我们的方法在FineWeb数据集上训练的GPT-2模型（10亿至100亿token）上产生了验证损失的持续改进，在一系列学习率上优于最近的替代方案。