使用 Newton-Schulz 迭代的 Muon 收敛性分析

我们分析了最初提出并在实践中使用的 Muon 优化器——即使用动量正交化与少量 Newton-Schulz 迭代步骤。先前的理论结果将 Muon 中的这一关键步骤替换为基于精确 SVD 的极因子。我们证明了使用 Newton-Schulz 迭代的 Muon 以与 SVD-极理想化相同的速率收敛到稳定点，对于给定数量的 Newton-Schulz 迭代步数 $q$，仅相差一个常数因子。我们进一步分析了这个常数因子，并证明它以 $q$ 的双指数速度收敛到 1，并且随着 Newton-Schulz 中用于近似正交化方向的多项式次数的增加而改善。我们还证明了，与其基于向量的对应方法（带动量的 SGD）相比，Muon 消除了典型的秩的平方根损失。我们的结果解释了为什么使用少量低次 Newton-Schulz 迭代步骤的 Muon 能以更快的实际运行时间匹配精确极因子（SVD）的行为，并解释了通过 Newton-Schulz 进行的动量矩阵正交化相比基于向量的优化器有多大益处。总体而言，我们的理论证明了 Muon 的实际 Newton-Schulz 设计的合理性，缩小了其实践与理论之间的差距。