技術的サマリー：スペクトル平坦化が Muon に必要不可欠である

問題定義

現代のニューラルネットワークの学習には、平坦な領域、鋭い方向、および鞍点によって特徴付けられる高度に非凸なランドスケープ上での最適化が伴う。オプティマイザの選択は、安定性と効率性にとって決定的である。Adam や AdamW のような座標ごとの適応法は頑健性を向上させたが、行列値パラメータがもたらす幾何学的な課題を完全に解決するものではない。最近、Muonは、更新前にモーメントバッファを直交化することで、より大きな学習率と高速な収束を可能にする高性能な代替手段として登場した。しかし、最大安定ステップサイズや収束率の観点から Muon が標準的な勾配降下法（SGD）を上回る理由を説明する理論的メカニズムは、未だ十分に探求されていない。具体的には、既存の理論は、直交化が学習率のボトルネックや最適化問題の有効な条件数をどのように定量的に変化させるかを説明していない。

手法

本論文は、Muon の更新規則をスペクトル平坦化の一種として定式化することで、Muon に対する理論的解析を提供する。著者は 2 つの主要な理論的結果を開発し、実験的に検証した。

1. 理論的枠組み:

   • Muon の定義: Muon は、勾配 $G_t$ を計算し、その後ニュートン・シュルツ反復を適用して極因子 $O_t \approx (G_t G_t^\top)^{-1/2} G_t = UV^\top$ を近似することで重み $W$ を更新する。この操作は勾配の特異ベクトルを保持しつつ、特異値を 1 に置き換える。
   • ヘッセ行列の近似: 解析には、ヘッセ行列を $H \approx X^\top X \otimes G_t G_t^\top$ と近似するクリネッカー積因数化された曲率モデル（K-FAC）が用いられる。ここで、$X$ は入力活性化を、$G_t$ は勾配を表す。
   • 学習率の解析: 著者は、SGD と Muon の両者に対して正確な 1 ステップ降下の閾値を導出した。両オプティマイザの最大安定学習率（$\eta_{max}$）を比較し、SGD は勾配の最大特異値（$\sigma_{max}$）によってボトルネックされるのに対し、Muon の安定性は平均特異値（$\frac{1}{m}\sum \sigma_i$）によって支配されることを示した。
   • 収束の解析: Muon を、前処理行列 $P = I_n \otimes (G_t G_t^\top)^{-1/2}$ を持つ前処理勾配法として再構成する。ポリアク・ロジャエヴィッチ（PL）条件と相対的な滑らかさの下で、著者は収束因子を解析する。勾配共分散が異方的である場合、Muon の有効な条件比（$\tilde{\alpha}/\tilde{\beta}$）は SGD のもの（$\alpha/\beta$）よりも厳密に優れていることを証明する。これは、勾配のスペクトル分散に関連する因子分だけ改善される。

2. 実験設計:

   • セットアップ: 実験は CIFAR-10 を用いて CifarNet で実施された。行列値パラメータに対するオプティマイザの影響を分離するため、デュアルオプティマイザ戦略が採用された。畳み込みカーネルは対象オプティマイザ（Muon または SGD）によって更新され、バイアスやその他のパラメータは固定された参照オプティマイザによって更新された。
   • 安定性テスト: バッチ正規化（BN）なしで学習率のスイープを行い、内在的な安定性を観察した。
   • 収束テスト: 両オプティマイザを同一の学習率（および個別にそれぞれの最良の学習率）で実行し、収束速度とステップごとの収束比（$r_t$）を測定した。
   • 正規化の原理: 入力スペクトルスケール（$\lambda_{max}(X^\top X)$）を低減することがより高い学習率を可能にするという理論的主張を検証するため、新しい正規化スキームであるFrobNorm（フロベニウスノルムによる正規化）が提案され、テストされた。

主要な貢献

1. 最大学習率の導出: 本論文は、Muon の最大安定ステップサイズが勾配の最大特異値ではなく、平均特異値に比例することを証明する。

   • SGD の場合: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)}$。
   • Muon の場合: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \frac{\sum \sigma_i}{m} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)} \frac{\sum \sigma_i}{m \sigma_{max}(G)}$。
   • これは、スペクトル平坦化が単一の支配的な特異方向がステップサイズの制約を決定することを防ぎ、Muon が著しく大きな学習率を許容できることを示している。

2. 収束率の加速: 著者は、Muon が有効な条件比を改善する前処理として機能することを示す。

   • K-FAC 下では、Muon の収束因子は $1 - \frac{\tilde{\alpha}}{\tilde{\beta}}$ であり、ここで $\frac{\tilde{\alpha}}{\tilde{\beta}} = \frac{\alpha}{\beta} \sqrt{\frac{\lambda_{min}(G_t G_t^\top)}{\lambda_{max}(G_t G_t^\top)}}$ である。
   • $\frac{\lambda_{min}}{\lambda_{max}} < 1$ であるため、項 $\sqrt{\dots}$ は 1 未満であるが、導出は前処理システムという文脈において $\frac{\tilde{\alpha}}{\tilde{\beta}} > \frac{\alpha}{\beta}$ であることを示しており、勾配共分散が条件悪化している場合、厳密に高速な線形収束率をもたらす。

3. 新しい正規化の原理: 解析により、最大安定学習率が $\lambda_{max}(X^\top X)$ に反比例することが明らかになった。著者は、正規化層はこのスペクトルスケールを低減することを目的とすべきであると提案する。彼らはFrobNormでこれを検証し、標準的な学習が失敗する高い学習率（例：0.25）での学習を可能にすることを示した。これは特に SGD において顕著であった。

実験結果

• 安定性: バッチ正規化がない場合、SGD は Muon が安定し損失を減少させ続けた学習率において、最初の数回の反復内で発散した。Muon はパラメータノルムと勾配ノルムにおいて制御された成長を示したのに対し、SGD は壊滅的な爆発を示した。
• 収束速度: 同一の学習率（例：$\eta=0.05$）において、Muon は SGD よりも数エポック早く精度のマイルストーン（70%–91%）に到達した。経験的な収束比 $r_t$ は Muon において一貫して低く、より高速な線形率という理論的予測を確認した。
• 正規化の影響: FrobNorm の適用により、両オプティマイザは高い学習率（CifarNet では 0.25、GPT-2 風トランスフォーマーでは 0.2）で安定して学習できるようになり、層入力のスペクトルスケールを制御することが、より大きなステップサイズを可能にする一般的な原理であるという仮説を検証した。

意義と主張

本論文は、Muon の経験的成功に対する原理的かつ幾何学的な説明を提供すると主張する。Muon の優位性は単なるヒューリスティックではなく、スペクトル平坦化の直接的な帰結であると仮定している。

1. 学習率の制約を最大特異値から切り離し、より大きなステップサイズを可能にする。
2. 最適化ランドスケープの有効な条件数を改善する片側前処理として機能し、より高速な収束をもたらす。

著者は限界を認め、理論的解析は決定論的なフルバッチ設定と正確な極因子を仮定しているが、実際の実装では確率的勾配、有限のニュートン・シュルツ反復、およびモーメントを使用することを指摘している。これらの結果を確率的設定や大規模アーキテクチャに拡張することが、将来の重要な研究方向であると述べている。論文は、これらの結果が Muon の経験的性能をより精密な理論的理解へと変換し、将来のオプティマイザや正規化戦略の設計を導く可能性があると結論付けている。