OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

技術要約：「OptMuon: ゼロノイズ最適性を備えた、確率的最適化のためのクローズドループ直交化モーメンタム手法」

1. 問題設定

本論文は、行列値の非凸目的関数に対する確率的一次最適化の問題を扱います：
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
目標は、$T$ 回の確率的勾配呼び出しの範囲内で、期待フロベニウスノルム $\mathbb{E}[\|\nabla f(X_{out})\|_F]$ を最小化するランダムなイテレート $X_{out}$ を見つけることです。

Muon スタイルのオプティマイザは、極分解（polar factorization）を通じてモーメンタム行列を近似的な直交方向へと変換することで、経験的な安定性を示してきましたが、既存の手法は通常、これらの方向をオープンループまたは一定の大きさのルールと組み合わせています。これらのルールは、あらかじめ規定された平滑性定数（$L$）や分散レベル（$\sigma^2$）に依存することが多く、ハイパーパラメータのチューニングに敏感であり、観測された最適化の軌跡に基づいて更新の大きさを明示的に較正することができません。対照的に、STORM や AdaGrad-Norm のような適応型手法は、「クローズドループ」のスカラ適応とノイズ適応性を備えていますが、一般には生の勾配に対して動作し、直交化されたモーメンタム方向には作用しません。

中心となる課題は、幾何学的な安定性を維持しつつ、手動の再チューニングなしにノイズ適応性とゼロノイズ最適性を達成するために、クローズドループのスカラ適応を Muon スタイルの直交化モーメンタムと正常に組み合わせることができるかという点です。

2. 手法：OptMuon

著者らは、更新方向と更新の大きさを分離するフレームワークである OptMuon を提案しています。

• 方向（直交化）： 更新方向は、モーメンタム行列 $M_t$ の極因子（polar factor）であり、$\text{Orth}(M_t) = U W^\top$ （ここで $M_t = U \Sigma W^\top$ は SVD とする）と定義されます。これにより、方向は幾何学的に正規化され、方向（orientation）と特異値のスケーリングが分離されます。
• 大きさ（クローズドループ適応）： ステップサイズは、固定された問題定数ではなく、観測された軌跡（勾配およびモーメンタムの履歴）から導出されるスカラー係数 $\gamma_t$ によって決定されます。
  • 係数スケジュール： 大きさ制御の核となるのは、実行最大値補正付き AdaGrad-Norm スケジュールです。これは、ノイズの多い領域でステップを抑制するために分母に確率的勾配のエネルギーを蓄積しますが、孤立した勾配スパイクによって係数が過度に崩壊するのを防ぐために、分子に項（$\max g_i^2$）を含んでいます。
  • Lipschitz-free（リプシッツ定数不要）： このアルゴリズムは、パラメータ選択のために平滑性定数 $L$、分散 $\sigma^2$、または有界勾配定数 $\hat{G}$ の知識を必要としません（ただし、これらは理論的な定数の中に現れます）。

2つのバリアント：

1. OptMuon-A (平均平滑性): 平均平滑性の仮定向けに設計されています。1回の反復につき1回の確率的勾配を使用し、標準的なモーメンタム蓄積（$M_t = G_t + (1-\alpha_t)M_{t-1}$）を用います。スケジュールの指数は $q=1/2$ です。
2. OptMuon-I (個別平滑性): 個別平滑性の仮定向けに設計されています。分散を低減するために同じサンプルに対して2回の勾配評価を用いる、STORM 型の再帰的モーメンタム推定器（$M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)) + G_t$）を採用しています。指数は $q=2/3$ です。

両バリアントは、以下の式でパラメータを更新します：
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. 主な貢献

本論文は、以下の具体的な貢献を主張しています：

• クローズドループ・モーメンタム直交化： 行列の幾何学（極因子による）と、大きさの較正（自己正規化されたスカラスケジュールによる）を明確に分離するフレームワーク。
• 実行最大値補正スケジュール： AdaGrad-Norm の自己正規化による抑制効果を維持しつつ、単一の勾配スパイクによる係数の崩壊リスクを軽減する、新しい係数設計。
• ノイズ適応的な収束率： 確率的項が分散 $\sigma$ の減少に伴って自動的に改善されることを示す理論的保証。
• ゼロノイズ最適性： $\sigma = 0$ のとき、両方の境界が手動の再チューニングなしに、自動的にほぼ最適な決定論的レート $\tilde{O}(T^{-1/2})$ に減少することを実証。

4. 理論的結果

下限性、有界分散を持つ不偏確率的勾配、平滑性（平均または個別）、およびほとんど確実に有界な確率的勾配の仮定の下で、論文は以下を証明しています：

• OptMuon-A (平均平滑性): 以下の収束率を達成します：
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I (個別平滑性): より鋭い確率的レートを達成します：
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

結果の意義：

• ゼロノイズの回復： ゼロノイズ領域（$\sigma = 0$）において、両方の境界は $\tilde{O}(T^{-1/2})$ に減少し、最適な決定論的一次レートと一致します。これは、既存の Muon 解析（例：Chang et al., 2025）がしばしば $\tilde{O}(T^{-1/4})$ を示したり、固定ステップサイズを必要としたりすることとは対照的です。
• ノイズ適応性： 確率的項（$\sigma^{1/2}T^{-1/4}$ および $\sigma^{1/3}T^{-1/3}$）は適応的であり、ノイズレベルが低下するにつれてタイトになります。これは固定ステップサイズを持つ手法とは異なります。
• 比較： 本論文は、OptMuon-A が、低ノイズ領域において最近の分散低減型 Muon バリアント（例：Chang et al., 2025）で確立された $\tilde{O}(T^{-1/4})$ のレートを改善しつつ、ゼロノイズ最適性を維持していることを述べています。

5. 意義と限界

意義：
本論文は、幾何学的安定性を損なうことなく、クローズドループのスカラ適応を Muon スタイルの直交化モーメンタムと結合できることを確立しています。これは、以前はスカラ適応型ステップサイズ・ルールに見られた「ゼロノイズ適応性」の概念を、極因子モーメンタム更新へと拡張するものです。これにより、手動の介入なしに、ノイズレベルに自動的に適応し、最適な決定論的パフォーマンスを回復する堅牢な最適化が可能になります。

限界（著者による記載）：

• 正確な直交化： 理論は正確な極因子の計算を仮定しています。Newton-Schulz 反復を用いてこれを近似できますが、現在の解析では、不正確な直交化によって導入される誤差は定量化されていません。
• 有界勾配： 解析は、ほとんど確実に有界な確率的勾配の条件（仮定 4）に依存しています。本論文は、有界勾配を必要としない理論を主張しているわけではありません。
• 経験的検証： 本論文は理論的な収束保証に焦点を当てており、大規模な経験的評価や、合成/制約付き目的関数への拡張は含まれていません。これらは今後の研究課題として挙げられています。