AI学習を速める新手法「Muon」の謎を解明、計算を省く方が速い理由も

学習して、人間のように文章を理解・生成できるAI。ChatGPTやGeminiの中身。">大規模言語モデルの学習を従来のAdamWよりも高速化する新手法「Muon」について、その中核技術である直交化更新（PolarGD）がなぜ機能するのかを理論的に解明した研究が発表された。

研究では、最適化手法を閉形式で解析しやすい行列二次形式を用い、学習率調整の影響を除去した状態で更新方向を検証した。その結果、勾配が低ランクからかけ離れ、最大の特異方向が高い曲率を持つ場合に直交化が有効に働くことが判明した。誤差の行が直交しノルムが等しければ、任意のヘッセ行列に対して1ステップで最小値に到達する。一方で、現在の反復値にある微小な特異値が原因となり、通常の勾配降下法より大幅に遅くなる状況が存在することも明らかになった。

さらに、Muonで採用されている少数のNewton-Schulz反復は、この微小な特異値を勾配降下法と同様に扱うため、厳密な直交化を行うよりも反復回数を抑えた方が速く収束することが示された。実際にMuonで学習させたTransformerでは、アテンション行列の特異値の約4分の1がこの条件に該当するという。