Abstract
Muonは、反復的な直交化によってニューラルネットワーク更新の行列構造を活用することで、大規模基盤モデルの事前学習における有望なオプティマイザとして登場しました。しかし、その実用効率は、各最適化ステップごとに複数回のNewton--Schulz（NS）反復を必要とすることにより制限されており、これが自明ではない計算および通信のオーバーヘッドを生みます。本研究では、直交化の前にAdamスタイルの適応的な2次モーメント前処理を適用するMuonの拡張であるMuon^2を提案します。私たちの重要な洞察は、Muonにおける極（polar）近似の根本的な難しさが、条件の悪い運動量行列にあり、Muon^2によってそのスペクトルが大幅に改善されることで、実務上十分な直交化に向けた収束がより速くなる点です。さらに、方向整合（directional alignment）によって実用的な直交化品質を特徴付けると、Muon^2は各極ステップにおいてMuonに比べて目覚ましい改善を示します。パラメータ数60Mから1.3Bの範囲におけるGPTおよびLLaMAの事前学習実験では、Muon^2は一貫してMuonおよび最近のMuonの変種を上回り、NS反復を40\%削減します。さらに、メモリ効率の高い因数分解版であるMuon^2-Fを導入します。これは、メモリオーバーヘッドが無視できるにもかかわらず、Muon^2の得られる利得の大部分を維持します。
