How Much Orthogonalization Does Muon Need?

技術要約：Muonにはどの程度の直交化が必要か？

問題提起
Muonオプティマイザは、不良条件のモーメンタム更新を近似的な半直交更新に置き換えることで、ニューラルネットワークの学習を改善する。標準的な数値線形代数は高精度な極分解（すべての特異値を1に写像すること）を目指すが、Muonの文脈では、厳密な精度は必ずしも必要ではないことが示唆されている。本研究が取り組む中心的な問いは、「Muonは学習の質を維持するために、実際にどの程度の直交化精度を必要とするのか？」である。具体的には、既存の手法（Muon-JordanやPolar Expressなど）で使用されている高次多項式反復の計算コストが正当化されるのか、あるいは、より安価で低次の近似が同等の学習結果を実現できるのかを調査している。

手法
著者らは、Muonの低精度（bfloat16）の特異値バンドに特化して設計された、緩和された三次ニュートン・シュルツ・スケジュール（cubic5と命名）を提案している。

1. 緩和されたターゲット: 古典的な極分解がマシン精度を目指すのに対し、この手法は特異値に対して$[0.7, 1.3]$の緩和されたバンドをターゲットとする。これは、bfloat16において、実質的な下限値（$l_0 = 7 \times 10^{-3}$）を下回る特異値は実用的な意味を持たないことを認識しているためである。
2. 適応的係数の導出: 各反復において、適応的な三次多項式 $f_t(x) = a_t x + b_t x^3$ を導出する。係数は、現在のワーストケースの下限特異値の境界（$l_t$）と固定された上限ピーク（$u=1.3$）に基づいて計算される。これにより、多項式が保護された区間 $[l_t, r_t]$ をターゲットバンドへと写像しつつ、新しい下限を最大化するように端点をバランスさせる。
3. 計算効率:
   • Cubic5: 単一の三次ステップは、2つの主要な行列積（グラム行列の形成とその再乗算）を必要とする。5ステップでは合計10回の主要な行列積を必要とする。
   • ベースライン (Muon-Jordan/Polar Express): 五次ステップは3つの主要な行列積を必要とする。5ステップでは15回の主要な行列積を必要とする。
   • 本手法は、小さなグラム行列に対する演算を最小限にするため、長方形行列に対する「転置トリック」を利用している。

主な貢献

• 緩和された三次スケジュールの導出: bfloat16の実効的な下限値に適応した係数を持つ、5ステップの三次ニュートン・シュルツ・スケジュールを導出し、5回の反復でターゲットとなるスペクトルバンドを実現した。
• コスト・品質フロンティアの特性評価: 学習の質は、極分解の精度によって単調に支配されるものではないことを確立した。低コストの三次近似が、特定の領域において高コストの五次手法や明示的なFP32 SVDと同等の性能に達することを示した。
• 様々なスケールでの実証的検証: 著者らは以下の方法を通じて手法を検証している：
  • 合成診断: 低精度行列演算におけるスカラー多項式の転送が保持されることを確認。
  • NanoGPT アブレーション: GPT-2 Smallを用いた、cubic5、切断されたPolar Express、切断されたMuon-Jordan、および明示的なFP32 SVDの比較。
  • 大規模学習: Megatron-LMを用い、1Bから4Bパラメータに及ぶハイブリッドMoE/Mambaモデル（Nemotron-3-Nano）の学習。

結果

• 合成ベンチマーク: cubic5法は、保護範囲 $[7 \times 10^{-3}, 1]$ 内の特異値を、数値的な爆発を起こすことなくターゲットバンドへと写像することに成功した。NVIDIA H200 GPU上でのマイクロベンチマークでは、行列の形状（長方形か正方形か）に応じて、Polar Expressと比較して直交化サブルーチンのコストを1.16倍から1.52倍高速化した。
• GPT-2 Small アブレーション:
  • 5ステップのcubic5は、最終検証損失3.2811を達成した。
  • 5ステップのPolar Expressは3.2798を達成した。
  • 5ステップのMuon-Jordanは3.2822を達成した。
  • 明示的なFP32 SVDは3.2809を達成した。
  • 観察: 最も強力なバリアント間の差は、検証損失において0.003未満である。注目すべきは、厳密なSVDが近似的な多項式手法を上回らなかったこと、およびMuon-Jordanに5ステップ目を追加すると、4ステップの場合よりも性能がわずかに低下したことである。
• ハイブリッドMoE/Mamba スケーリング (1B–4B):
  • 1Bおよび2Bモデルにおいて、cubic5はMuon-Jordanをわずかに上回った（検証損失が約$10^{-3}$低い）。
  • 4Bモデルにおいて、cubic5はわずかに高かった（約$10^{-3}$高い）が、初期フェーズ以降の学習軌跡はほぼ区別がつかなかった。
  • 相対的な学習損失の差は急速に減衰し、残りの学習予算の間、ゼロ付近で変動した。

意義と主張
本論文は、cubic5を、普遍的に優れた更新ルールとしてではなく、実用的で低コストなMuon直交化のバリアントとして控えめに位置づけている。

• オプティマイザ vs 分解精度: 結果は、Muonが厳密な極分解の精度ではなく、「スペクトル整形（spectral reshaping）」の効果から恩恵を受けているという仮説を支持している。特異値の写像が学習を最適化するかどうかは、単なる数学的な1への収束ではなく、学習フェーズや学習率に依存する。
• 実用的な有用性: cubic5は、コスト・品質フロンティアにおける原理的な選択肢を提供する。テストされた領域において、標準的なMuon-Jordanの五次更新と同等の学習品質を維持しながら、主要な行列積の回数を約3分の1削減（15回に対し10回）している。
• 限界: 著者らは、これらの結果が学習品質の同等性を確立するものであり、エンドツーエンドのスループット向上を保証するものではないことを明示している。実際のウォーククロック（実時間）の高速化は、ハードウェアカーネル、バッチング、およびフルオプティマイザーループ内で直交化ステップが占める割合に依存する。本論文は、cubic5が数学的な意味でPolar Expressよりも優れた極分解アルゴリズムであると主張しているのではなく、特定のコンテキストであるMuonの学習において、より効率的なスペクトル変換であることを主張している。