技術要約：Muonオプティマイザの収束境界とクリティカル・バッチサイズ

問題提起

Muonオプティマイザは、AdamWのような標準的なオプティマイザの後継となり得る強力な経験的性能を示しているが、その理論的根拠は未だ発展途上である。既存の文献では、特にネステロフ・モメンタム（Nesterov momentum）や重み減衰（weight decay）を含む実用的な設定における、Muonの厳密な収束証明が欠如している。さらに、Muonのハイパーパラメータと、クリティカル・バッチサイズ（学習の確率的一次オーダー（SFO）複雑性を最小化するバッチサイズ）との関係も理論的に特徴付けられていない。本論文は、厳密な収束解析を提供し、Muonのクリティカル・バッチサイズの低次境界を導出することで、これらのギャップを埋めることを目的とする。

手法

著者らは、以下の4つの実用的な構成におけるMuonの理論的解析を行う：

1. ネステロフ・モメンタムなし、かつ重み減衰なし。
2. ネステロフ・モメンタムあり、かつ重み減衰なし。
3. ネステロフ・モメンタなし、かつ重み減衰あり。
4. ネステロフ・モメンタあり、かつ重み減衰あり。

理論的枠組み：

• 収束解析： 著者らは、平均期待核ノルム $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$ の上界を導出する。これには、$L$-平滑性（$L$-smoothness）および有界な勾配分散という標準的な仮定を用いている。主要な革新は、モメンタム行列 $C_t$ を直交化することによって得られるMuonの探索方向の解析である。著者らは、全パラメータ次元ではなく、モメンタム誤差行列のランク（$r_1$）および探索方向のランク（$r_2$）を考慮に入れている。
• Newton-Schulz近似： Muonは実際には、正確なSVD（特異値分解）ではなくNewton-Schulz反復によって直交化を近似するため、著者らはこの近似誤差（$\delta(k)$）が収束に与える影響を分析し、それが無視できる程度の項しか導入しないことを示している。
• クリティカル・バッチサイズ（CBS）の導出： ターゲットとなる精度 $\epsilon$ に到達するために必要な学習ステップ数 $T$ を、バッチサイズ $b$ の関数としてモデル化する。SFO複雑性（$T \times b$）を最小化することにより、クリティカル・バッチサイズ $b^*$ の下限を導出する。

実験的検証：
理論的な予測は、二段階の戦略を通じて検証される：

1. 制御されたフルMuon実験： Muonがすべてのパラメータを更新し、理論的仮定に厳密に従う合成的なTeacher-Studentタスク。
2. 実用規模のハイブリッド実験： CIFARデータセットを用いたResNet-18/VGG-16の学習、およびC4コーパスを用いたLlama3.1 (320M) の学習。これらの設定では、行列形状のパラメータにはMuonが適用され、それ以外にはAdamWが適用されるという、標準的なデプロイメントを反映している。

主な貢献

1. 4つのバリアントに対する収束証明

本論文は、4つのすべてのMuonバリアントに対して収束保証を確立している。導出された平均期待勾配ノルムの上界は、以下のオーダーとなる：
$$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$$
（ネステロフおよび重み減衰の項については若干の差異がある）。

• 意義： ベクトルベースのオプティマイザでは分散が全パラメータ数に比例してスケールするのに対し、Muonの境界はモメンタム誤差行列のランク $r_1$ に依存してスケールする。経験的な証拠は、$r_1 \ll n$ となる過剰パラメータ化された領域において、よりタイトな理論的境界を示唆している。
• 核ノルム： 解析には核ノルムが使用されている。これは多くのベクトルベースの解析で使用されるフロベニウスノルムよりも厳しい基準であり、Muonがニューラルネットワークのパラメータの固有の構造を尊重していることを示唆している。

2. 重み減衰の役割と安定性

著者らは、重み減衰を組み込むことで、有界な勾配という制限的な仮定に頼ることなく、パラメータノルム（$\|W_t\|_F$）と勾配ノルム（$\|\nabla f(W_t)\|_F$）の両方の**ほとんど確実な有界性（almost-sure boundedness）**が保証されることを証明している。

• 安定性条件： 彼らは安定性のための必要条件を導出した。すなわち、学習率 $\eta$ と重み減衰係数 $\lambda$ は $\eta \leq 1/\lambda$ を満たす必要がある。実験結果は、$\eta$ がこの閾値を超えると学習が不安定になることを裏付けている。
• 暗黙的な正則化： 重み減衰は漸近的な収束レートの次数自体は改善しないものの、ノルムを制限することで暗黙的な正則化を提供しており、これが優れた経験的性能の理由となっている。

3. クリティカル・バッチサイズ（CBS）の下限

本論文は、ハイパーパラメータに依存するクリティカル・バッチサイズ $b^*$ の下限を導出している：
$$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$$

• ハイパーパラメータへの依存性： この式は以下を明らかにしている：
  • モメンタム ($\beta$): $\beta \to 1$ となると、項 $(1-\beta)$ が消失し、クリティカル・バッチサイズが減少することを示唆している。
  • 重み減衰 ($\lambda$): 因数 $1/(1-\lambda)^2$ は、重み減衰を大きくするとクリティカル・バッチサイズが増加することを示している。
  • ランク ($r_1$): 境界はモメンタム誤差のランクに比例し、CBSが勾配の低ランク構造に関連していることを結びつけている。
• 定性的スケーリング： 著者らは、この式は問題固有の定数（$\sigma^2, r_1, \epsilon$）に依存するため、絶対的な予測ではなく「定性的なスケーリング則」を提供するものであると強調している。しかし、異なるMuonバリアント間のCBSの「比」は、$\beta$ と $\lambda$ のみに依存する。

結果

理論的知見：

• 収束レートは、特定のハイパーパラメータ・スケジュール下で標準的な非凸SGD型の手法（$O(T^{-1/4})$）と一致するが、ランク依存性により、よりタイトな分散項を持つ。
• Newton-Schulz近似誤差は収束に限定的な影響しか与えず、反復回数とともに指数関数的に減衰する。

実験的知見：

• 安定性： ResNet-18/CIFAR-10における実験は、安定性条件 $\eta \leq 1/\lambda$ を検証している。勾配ノルムと学習損失はこの境界付近で最小化される。
• 収束速度： ネステロフ・モメンタムと重み減衰を備えたMuonは、AdamW、Momentum SGD、および他のMuonバリアントと比較して、一貫して最も速い損失および勾配ノルムの減少を達成した。
• クリティカル・バッチサイズの検証：
  • ハイパーパラメータの影響： 実験により、ネステロフ・モメンタムを加えるとCBSが大きな値へとシフトし、重み減衰を増やすとさらにCBSが増加することが確認された。例えば、$\lambda$ を0から0.5に増やすと、CBSは $\approx 1024$ から $\geq 4096$ へとシフトしたが、これは予測された $1/(1-\lambda)^2$ スケーリング（比率4.0）と一致している。
  • ランクのスケーリング： モデルの幅を変化させた実験では、有効ランク $r_1$ は幅とともにスケールするが、$r_1 \cdot \sigma^2$ の積は比較的一定に保たれることが示された。これが、理論的なランク依存性があるにもかかわらず、経験的なCBSの値が幅に対して平坦に見える理由である。
  • 効率性： Muonは、特に小規模から中規模のバッチサイズにおいて、AdamWよりも少ないステップ数と低いSFO複雑度を必要とする傾向がある。

意義と主張

本論文は、出現しつつあるこのオプティマイザの理論的理解の空白を埋めるべく、ネステロフ・モメンタムと重み減衰を同時に組み込んだ、Muonに関する初の厳密な収束解析を提供すると主張している。

• 理論的洞察： 重み減衰は単なる正則化ツールではなく、有界性と安定性を確保するためのメカニズムであり、有界な勾配という仮定を排除できることを示している。
• 実践的なガイダンス： 導出されたCBSの下限は、実務者に実行可能な指針を与える。ハイパーパラメータ $\beta$ と $\lambda$ の調整が、最適なバッチサイズにどのように定性的な影響を与えるかを明確にしており、高モメンタムの設定は小さなバッチで効率的に学習でき、高い重み減衰は効率を維持するために大きなバッチを必要とすることを示唆している。
• 検証： 本研究は、ビジョンおよび言語モデリングのワークロードを通じてこれらの予測を検証することで、理論と実践の架け橋となっており、理論的なスケーリング則が、Muonが行列形状のパラメータにのみ適用されるハイブリッド最適化設定においても成立することを裏付けている。

著者らは、単一行列の解析では層ごとの不均一性を明示的にモデル化していないことを謙虚に認めているが、勾配ノイズのスケールの高い相関性が層間で存在することが、彼らの知見の汎用性を支持していると主張している。また、実験においてShampooよりも優れた性能を示したものの、Shampooに類似した境界を導出することは依然として未解決の課題であることも認めている。