The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm

要約

極座標分解とそれに関連する行列符号関数の計算は、数値解析において数十年にわたりよく研究されてきた問題である。最近では、特にミュオン最適化の枠組みの中で、ディープラーニングにおける重要なサブルーチンとして浮上している。しかし、この設定における要件は、従来の数値解析とは大きく異なる。ディープラーニングでは、手法は非常に効率的でGPUと互換性がなければならないが、高い精度は不要な場合が多い。その結果、Newton-Schulzのような古典的アルゴリズム（初期収束の遅さに悩まされる）や、有理関数に基づく手法（QR分解や行列反転に依存する）は、このコンテキストには適していない。この研究では、極座標分解を計算するためのGPUフレンドリーなアルゴリズムであるPolar Expressを紹介する。Newton-Schulzのような古典的な多項式手法と同様に、我々のアプローチは行列-行列の乗算のみを使用するため、GPUと互換性がある。Chen & ChowとNakatsukasa & Freundの先行研究に触発され、Polar Expressは最小最適化問題を解くことにより、各反復で多項式更新規則を適応させる。この性質により、早い収束と速い漸近収束の両方が保証される。また、有限精度の問題にも対処しており、実際にはbfloat16で安定する。我々はPolar ExpressをMuon最適化のフレームワークの中で適用し、GPT-2のような大規模モデルにおける検証損失において一貫した改善を示し、様々な学習率において最近の代替案を凌駕する。

要約(オリジナル)

Computing the polar decomposition and the related matrix sign function, has been a well-studied problem in numerical analysis for decades. More recently, it has emerged as an important subroutine in deep learning, particularly within the Muon optimization framework. However, the requirements in this setting differ significantly from those of traditional numerical analysis. In deep learning, methods must be highly efficient and GPU-compatible, but high accuracy is often unnecessary. As a result, classical algorithms like Newton-Schulz (which suffers from slow initial convergence) and methods based on rational functions (which rely on QR decompositions or matrix inverses) are poorly suited to this context. In this work, we introduce Polar Express, a GPU-friendly algorithm for computing the polar decomposition. Like classical polynomial methods such as Newton-Schulz, our approach uses only matrix-matrix multiplications, making it GPU-compatible. Motivated by earlier work of Chen & Chow and Nakatsukasa & Freund, Polar Express adapts the polynomial update rule at each iteration by solving a minimax optimization problem, and we prove that it enjoys a strong worst-case optimality guarantee. This property ensures both rapid early convergence and fast asymptotic convergence. We also address finite-precision issues, making it stable in bfloat16 in practice. We apply Polar Express within the Muon optimization framework and show consistent improvements in validation loss on large-scale models such as GPT-2, outperforming recent alternatives across a range of learning rates.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL