技術要約：Aurora – レバレッジを考慮したスペクトル最適化手法

問題提起

本論文は、大規模言語モデルの学習において、勾配モーメンタム行列の極分解因子（polar factor）を用いて重みの更新を近似する、普及しているスペクトル最適化手法である Muon における決定的な病理を特定している。具体的には、背の高い行列パラメータ（行数 $m$ が列数 $n$ を超えるもの、例えばMLP層のupおよびgate投影行列など）において、Muonの更新は極めて不均一な行ノルムを示す。

この不均一性は、以下の自己強化的なフィードバックループを生み出す：

1. 行ノルムが小さい行に対応するニューロンは、持続的に小さな更新を受ける。
2. これらのニューロンはネットワークの出力に意味のある寄与ができず、事実上の「死んだ」状態となる。
3. 本論文は、この現象を**レバレッジスコア（leverage scores）**に関連付けるメカニズム的な証拠を提示しており、モーメンタムバッファにおける低い行ノルムが低いレバレッジスコアと相関しており、Muonがこれを修正するのではなく維持してしまうことを示している。

著者らは、既存の行正規化手法（例：NorMuon）が、行を再スケーリングすることでこの問題を緩和しているものの、その過程でMuonの更新の幾何学的構造を歪め、真の極分解因子から遠ざけてしまうことを観察している。この「行ノルムの一様性を強制すること」と「正確な極分解因子の幾何学を保持すること」の間のトレードオフが、対処すべき中心的な問題である。

手法

1. 正規化と極分解精度の分析

著者らはまず、極分解精度（polar factor precision）と行正規化の間の緊張関係を分析している。

• 極分解近似誤差 (PAE): Newton-Schulz反復（CANS-12など）による極分解の計算精度を高めることが、検証損失の低下と相関していることを実証している。
• 幾何学的衝突: 背の高い行列の場合、列直交性（極分解因子の特性）と単位行ノルムを同時に満たすことは数学的に不可能であることを証明している。事後的な正規化は、更新を真の極分解因子から逸脱させ、重大な幾何学的誤差を導入する。
• U-NorMuon: 中間ステップとして、背の高い行列に対してのみ行正規化を適用し、正規化のためのステートフルなモーメンタムバッファを排除した U-NorMuon を提案している。これにより、統計量の陳腐化によるオーバーヘッドなしに、行正規化の利点のみを分離して示し、標準的なMuonやNorMuonよりも優れた性能を示すことを示した。

2. Aurora オプティマイザ

極分解の幾何学を保持することと、一様な行ノルムを強制することの間の衝突を解決するために、著者らは Aurora を提案している。

• 定式化: Auroraは、以下の2つの制約条件下で、勾配のドット積のトレースを最大化する制約付き最適化問題として更新を定式化している：
  1. 左半直交性: $U^\top U = \eta^2 I$ （極分解因子の構造を保持）。
  2. 等しい行ノルム: すべての行 $i$ に対して $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$。
• Riemannian-Aurora: 著者らは、これらの制約を結合多様体（スティフェル多様体と等行ノルム多様体の交わり）として扱う理論的な参照ソルバーである Riemannian-Aurora を導出している。これは理論的には精密であるが、レバレッジ制約を含む大規模な線形システムの解決を必要とするため、計算コストが極めて高い。
• Practical Aurora: 著者らは、以下のプロセスを交互に行う反復的かつ実用的なインスタンス化を提案している：
  1. 行正規化: ノルムを等しくするために行を再スケーリングする。
  2. 極分解近似: 行列を再直交化するためにNewton-Schulz反復を適用する。
     このプロセスは、収束を確実にするためにパラメータ $\beta$ によって減衰される。著者らは、局所的な正則性条件の下で、この交互投影が両方の制約を満たす解に線形に収束することを証明している。

主な貢献

1. 正規化されたスペクトル最適化器の分析: 本論文は、事後的な単位行正規化が、背の高い行列に対する極分解精度を低下させることを示している。また、統計量の陳腐化を伴わないステートレスなバリアントである U-NorMuon を導入し、行の一様性と正確な直交化が共に望ましいものであることを証明した。
2. ニューロン死の研究: 著者らは、Muonで訓練されたモデルにおけるニューロン死のメカニズム的な説明を提供し、それを背の高い行列パラメータにおける低いレバレッジスコアに関連付けた。これを独自の訓練ランおよびオープンウェイトモデル（Rnj-1, Laguna-XS.2）において経験的に検証し、高密度MLP層における死んだニューロンの顕著な集団を確認した。
3. Aurora と Riemannian-Aurora: 著者らは、Muonの理論的枠組みを厳密に遵守しながら、一様な行ノルムを強制する手法として Aurora を導出した。また、二重制約問題に対する高コストな参照ソルバーとして Riemannian-Aurora も提供している。
4. 広範な事前学習による検証: 著者らは、340Mおよび1.1BのパラメータスケールにおいてAuroraを検証し、Muon、NorMuon、U-NorMuonと比較して優れたダウンストリーム性能を示すことを実証した。

結果

ダウンストリーム性能

• ベンチマークスコア: 1.1Bスケールにおいて、Aurora はMuonのベースラインを大幅に上回り、MMLUで9.1ポイント高いスコアを達成した。また、340Mおよび1.1Bの両スケールにおいて、推論および知識のベンチマーク・スイート（ARC, HellaSwagなど）全体で最高の平均スコアを記録した。
• Speedrun: modded-nanoGPT speedrun（オプティマイザ・トラック）において、Auroraは既存の手法（Contra-Muon, flooring）と組み合わせることで、ターゲットの検証損失にNorMuonベースラインよりも25ステップ早く到達し、新たなSOTA（State-of-the-Art）を樹立した。

ニューロンの利用率

• レバレッジスコア: Muonで訓練されたモデルは、ゼロ付近に有意なモードを持つ二峰性のレバレッジスコア分布を示し、死んだニューロンが存在する。対照的に、Aurora と U-NorMuon は等方的なレバレッジスコア分布を生み出し、死んだニューロンを事実上排除している。
• オープンウェイトによる検証: Muonで訓練されたオープンウェイトモデル（Rnj-1 および Laguna-XS.2）の分析により、初期のデンス層における死んだニューロンの存在（Laguna-XS.2の第1層では最大約28%）が確認され、問題の汎用性が裏付けられた。

効率性とスケーリング

• スループット: 分散環境において、Auroraの追加的な計算オーバーヘッドは最小限である。なぜなら、追加の行列乗算は勾配通信とオーバーラップできるからである。単一GPUではスループットがわずかに低下するが、分散学習におけるウォールクロック・オーバーヘッドは緩やかである。
• MLP幅のスケーリング: AuroraのMuonに対する性能向上は、MLPの拡張係数（expansion factor）に対して単調にスケールする。MLP層が広くなるにつれてAuroraの優位性が増すことは、Auroraが非常に広いネットワークの訓練に特に効果的であることを示唆している。

意義と主張

本論文は、Aurora が、真の極分解因子の幾何学と一様な行ノルムという実用的な必要性の両方を尊重する最適化器を導出することによって、言語モデルの訓練における特定の経験的に観察された病理（背の高い行列パラメータにおけるニューロン死）を解決することを主張している。

著者らは、純粋な理論モデルではなく、観察された訓練ダイナミクス（ニューロン死、レバレッジスコアの不均衡）に基づいてアルゴリズムを導出するというボトムアップのデザイン哲学を強調している。彼らは、特定のパラメータクラス（例：背の高い行列 vs 幅の広い行列）に合わせて最適化器を調整する**アーキテクチャと最適化器の共同設計（co-design）**が、標準的な「ワンサイズ・フィッツ・オール（万能型）」戦略を凌駕する有望な方向性であると論じている。

最後に、本論文はAuroraが非常に広いMLP層の有効な訓練を可能にすることを提案している。これは、現在のMuonベースの訓練においては、死んだニューロンの蓄積によって制限される可能性がある能力である。この研究は、Auroraを単なる性能向上策としてではなく、大規模言語モデルにおける現在のスペクトル最適化の基礎となる幾何学的仮定に対する、必要な修正として位置づけている。