技術的サマリー：MONA – ニュートン加速を備えたムオン最適化器

1. 問題定義

大規模言語モデル（LLM）の学習は、ますます優れたサンプル効率を持つ最適化器を要求しています。AdamW は 10 年以上にわたり標準でしたが、最近Muonのような代替手法が登場しています。Muon は、重み行列を単一の幾何学的単位として扱い、モーメントバッファに行列直交化（ニュートン・シュルツ反復による）を適用することで、一階の手法を改善します。このアプローチはスペクトルノルムにおける最急降下法と整合性があり、大規模なエキスパート混合（MoE）学習において AdamW より約$2\times$の計算効率を提供します。

しかし、Muon はすべての一階勾配法と根本的な限界を共有しています。すなわち、局所的な勾配情報を超えて損失地形を探索する明示的なメカニズムを欠いていることです。その結果、Muon は特に大規模バッチ学習のシナリオにおいて、鋭い局所最小値に陥りやすくなります。鋭い最小値は一般化性能の低下と関連付けられる一方、平坦な最小値はより良い下流タスクのパフォーマンスをもたらす傾向があります。ALTO最適化器は、鋭い最小値から脱出するために勾配差の指数移動平均（EMA）に基づく加速項を導入しましたが、その設計が Adam と LAMB 向けの汎用アダプターであるため、Muon の高度に非線形な直交化プロセスとの統合に課題を呈しています。

2. 手法：MONA

著者らは、曲率を考慮した加速項を Muon の直交化フレームワークにシームレスに統合するMONA（ニュートン加速を備えたムオン最適化器）を提案します。

核心アルゴリズム

MONA は、モーメント蓄積と直交化の前に加速項を導入することで、標準的な Muon パイプラインを変更します。反復$k$における行列パラメータ$W$の更新プロセスは以下の通りです：

1. 勾配差の計算: 現在の確率的勾配$G_k$と前回の勾配$G_{k-1}$の差を計算します：
   $$D_k = G_k - G_{k-1}$$
2. 加速バッファの更新: これらの差の指数移動平均を維持します：
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   ここで、$\beta_a$は加速の記憶を制御します。
3. 曲率を考慮した勾配: スケールされた加速項を生の勾配に加えることで、拡張勾配$\tilde{G}_k$を構築します：
   $$\tilde{G}_k = G_k + \alpha A_k$$
   係数$\alpha$は加速をスケールします。理論的動機付けによれば、$G_k - G_{k-1}$は$H_k(\theta_k - \theta_{k-1})$（ここで$H_k$はヘッセ行列）を近似します。したがって、この項は勾配変化が大きい鋭い最小値から離れる方向を指し示し、平坦な領域では小さく保たれます。
4. Muon 直交化: 拡張勾配は、標準的な Muon のモーメントおよび直交化ステップに入力されます：
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

設計の根拠

直交化の前に加速項を適用することが不可欠です。後に適用した場合、Muon がスペクトルノルム正則化のために依存する直交構造が乱されてしまいます。生の勾配をまず曲率を考慮した方向に変換することで、モーメントバッファはこれらの豊かになった方向を捕捉し、その後 Muon の標準パイプラインを通じて処理されます。

ベクトル値パラメータ（埋め込みやバイアスなど）の場合、MONA は Muon の既存の慣習に従い、AdamW をデフォルトとして使用します。

3. 理論的解析

本論文は、標準的な仮定（L-滑らかさ、分散が有界な不偏勾配、期待される方向整合性）の下での MONA の収束解析を提供します。

• 収束率: 著者らは、MONA が非凸設定において$O(K^{-1/2})$の収束率を達成することを証明しており、SGD および AdamW の漸近率と一致します。加速項は境界内の定数に影響を与えますが、漸近率には影響しません。
• 鋭い最小値からの脱出: この解析は、加速項がどのように選択的に鋭い最小値に抵抗するかを形式化します。定常点の近くでは、期待される加速$E[A_k]$は、過去の方向のヘッセ行列重み付き和に比例します。大きな固有値を持つ方向（鋭い最小値）に対しては加速の大きさが大きくなり、最適化器を押しやります。小さな固有値（平坦な方向）に対しては加速が小さく、安定した収束を可能にします。このメカニズムは、曲率への意識を追加しながら Muon のスペクトルノルム正則化を維持します。

4. 実証結果

MONA は、3 つの規模（1B、6B、68B パラメータ）の MoE 事前学習モデルおよび教師あり微調整（SFT）段階で評価されました。

事前学習のパフォーマンス

• モデル: MOE-1B-A0d2B（400B トークン）、MOE-6B-A0d5B（1.2T トークン）、および MOE-68B-A3B（700B トークン）。
• 指標: コード、数学的推論、一般的な英語、中国語の学術テキストの各ドメインにおける検証損失が監視されました。
• 知見: MONA は、3 つのすべてのモデル規模およびすべてのドメインにおいて、Muon および AdamW よりも一貫して低い検証損失を達成しました。特に 68B の大規模モデルにおいて性能差が顕著で、MONA はより安定した収束を示しました。

下流ベンチマーク（MOE-68B-A3B）

MONA で学習された 68B モデルは、包括的なベンチマークスイートで評価されました：

• 一般能力: MMLU、MMLU-Pro、CMMLU、CEVAL、BBH などのベンチマークにおいて、MONA は最高平均スコア（0.4557）を達成し、Muon（0.4478）および AdamW（0.4382）を上回りました。
• 数学的推論とコード生成: MONA は GSM8K、GPQA、および各種コード生成ベンチマーク（HumanEval+、MBPP+、BigCodeBench、LiveCodeBench）で優れたパフォーマンスを示しました。特に、MONA は転移可能なコード表現を必要とするタスクで顕著な改善をもたらしました。
• 教師あり微調整（SFT）: MONA で事前学習されたチェックポイントがコードデータで SFT を受けた際、Muon で事前学習されたベースラインと比較して、プロセス全体を通じて低い学習損失を維持しました。BigCode スイートでの SFT 後の評価により、MONA 初期化モデルが Muon 初期化モデルと比較して 7 つのタスクのうち 6 つで高いスコアを達成したことが確認されました。

効率性と展開

前回の勾配（$G_{k-1}$）と加速バッファ（$A_k$）の保存に伴うメモリオーバーヘッドに対処するため、著者らはMONA-Liteを導入しました：

• 量子化: 補助状態を float32 ではなく bfloat16（BF16）で保存します。
• ストリーミング計算: 専用の$G_{k-1}$バッファの必要性を排除するために、勾配差をその場で計算します。
• 結果: これらの戦略により、学習品質を犠牲にすることなく、追加のメモリオーバーヘッドを約 75% 削減しました。計算オーバーヘッドの分析により、MONA は最適化器ステップレベルで約 1% 遅いのみであり、Muon と比較してエンドツーエンドの反復時間には実質的な差がないことが示されました。

5. 意義と主張

本論文は、Muon の幾何学的効率と曲率を考慮した加速の必要性との間のギャップを MONA が成功裏に埋めたと主張しています。その主な意義は以下の点にあります：

1. アルゴリズム的統合: 幾何学的性質を乱すことなく、加速項を行列直交化フレームワークに統合できることを実証しました。
2. スケーラビリティ: 幅広いモデルサイズ（1B から 68B パラメータ）およびトークン数（最大 1 兆トークン） across 手法を検証し、AdamW および Muon 両方に対する一貫した改善を示しました。
3. 一般化: 事前学習中に鋭い最小値から脱出することが、特に複雑な推論やコード生成において、下流タスクでのより良い一般化につながるという実証的証拠を提供しました。
4. 実用性: 制約のあるリソース環境でも最適化器を viable にするメモリ効率の高い変種（MONA-Lite）を提供しました。

著者らは、MONA が大規模 LLM 学習のための最先端の最適化器を表しており、Muon フレームワークの計算効率を維持しつつ、優れた収束と下流パフォーマンスを達成していると結論付けています。