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

要約

極性分解と関連するマトリックス記号関数を計算することは、数十年にわたって数値分析でよく研究された問題でした。
最近では、特にMuon Optimization Framework内で、深い学習において重要なサブルーチンとして浮上しています。
ただし、この設定の要件は、従来の数値分析の要件とは大きく異なります。
深い学習では、方法は非常に効率的でGPU互換性がなければなりませんが、多くの場合、高精度は不要です。
その結果、Newton-Schulz（初期収束が遅い）や合理的な機能（QR分解またはマトリックスの逆に依存する）に基づく方法などの古典的なアルゴリズムは、このコンテキストにはあまり適していません。
この作業では、極性分解を計算するためのGPUに優しいアルゴリズムであるPolar Expressを紹介します。
Newton-Schulzなどの古典的な多項式方法と同様に、私たちのアプローチでは、Matrix-Matrix乗算のみを使用して、GPU互換性があります。
Chen＆ChowとNakatsukasa＆Freundの以前の研究に動機付けられたPolar Expressは、Minimaxの最適化問題を解決することにより、各反復で多項式更新ルールを適応させます。
この特性は、急速な早期収束と速い漸近収束の両方を保証します。
また、有限精度の問題に対処し、実際にはBFLOAT16で安定しています。
Muon Optimization Framework内にPolar Expressを適用し、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, Google