Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

本論文は、Muon 最適化アルゴリズムが抱く等方的な制約の限界を克服し、Shampoo 由来の統計量を用いた曲率感知の事前条件付けを導入することで、大規模言語モデルの学習効率を大幅に向上させる新しい最適化手法「Mousse」を提案する。

要約

🍮 Mousse（ムース）：AI 学習の「地形」を正しく見る新しいコンパス

こんにちは！今日は、2026 年に発表された画期的な AI 研究論文「Mousse（ムース）」について、難しい数式を使わずに、誰でもわかるように解説します。

この論文は、**「AI が賢くなるための『歩き方（学習アルゴリズム）』を、もっと自然で効率的なものに変えた」**というお話です。

---

🌍 1. 問題：AI は「平坦な道」だと思い込んでいた

まず、AI（特に大規模言語モデル）を学習させるということは、**「山と谷が複雑に絡み合った、見えない地形を歩く」**ようなものです。

• 谷（Valley）： 正解に近い場所（損失が低い）。
• 山（Peak）： 正解から遠い場所（損失が高い）。

これまでの有名な AI の歩き方（Muonという名前）は、**「どんな道でも、すべての方向に均等に力を入れて歩けばいい」**という考え方をとっていました。

• 例え話： 雪原を歩くとき、北・南・東・西、すべての方向に「同じ強さ」で足跡を残そうとするようなイメージです。

しかし、現実の AI の学習地形はそうではありません。

• 急な崖（曲率が高い部分）もあれば、広大な平野（曲率が低い部分）もあります。
• 急な崖では、同じ強さで歩くと転んでしまいます（不安定になる）。
• 平野では、同じ強さだと進みが遅すぎて、いつまでたってもゴールにたどり着けません。

つまり、「すべての方向を平等に扱う」という考え方は、AI にとっては非効率で、時には危険だったのです。

---

🍮 2. 解決策：Mousse（ムース）の登場

ここで登場するのが、この論文で提案された新しい歩き方**「Mousse（ムース）」**です。
（※名前の由来は、Muon（ムオン）と Shampoo（シャンプー）を掛け合わせた「ムース」です。泡のように軽やかで、かつ形状に合わせて変形するイメージです）。

🧭 Mousse のすごいところ：「地形に合わせた靴」を履く

Mousse は、**「その場所の地形（曲がり具合）に合わせて、歩き方を調整する」**というアイデアを使います。

1. 地形を「白く」する（Whitening）：
   まず、AI が持っている「過去の歩行データ（勾配の統計）」を使って、複雑な地形を一旦**「平坦で均一な白紙」**のように見立てます。

   • 例え話： 歪んだ鏡で世界を見ていたのを、正しいレンズ（Shampoo という技術）を通して、歪みを補正してまっすぐな世界に見えるようにするイメージです。

2. その上で「ムオン」の歩き方をする：
   地形が補正された状態で、Muon が得意とする「安定した歩き方（直交化）」を行います。

3. 元の世界に戻す：
   補正された歩き方を、元の複雑な地形に合わせて変換し直して、実際に AI を動かします。

結果として：

• 急な崖では、慎重に、でも確実に進みます。
• 平野では、思い切って大きく進みます。
• すべての方向を平等に扱う必要がなくなり、ゴール（正解）に最短でたどり着けるようになりました。

---

🚀 3. 実際の効果：どれくらい速くなった？

研究者たちは、1 億 6000 万パラメータから 8 億パラメータまでのさまざまなサイズの AI で実験を行いました。

• 🏆 結果： Mousse を使った AI は、Muon を使った AI よりも約 12% 少ないステップ数で、同じレベルの性能に達しました。
  • 例え話： 100 歩でゴールするところを、88 歩でゴールできるようなものです。
• ⏱️ 計算コスト： 驚くべきことに、この「賢い歩き方」をするために、計算時間はほとんど増えませんでした（わずか 3% のオーバーヘッド）。
• 💾 メモリ： 従来の「地形を考慮する」方法（SOAP など）に比べて、メモリ使用量も大幅に少なく済みます。

つまり、「より速く、より安く、より賢く」AI を学習させることができるようになったのです。

---

🛠️ 4. 技術的な工夫（裏話）

Mousse がうまくいくためには、いくつかの「コツ」がありました。

• 🌊 揺れを鎮める（グラフティング）：
  学習の途中で、歩幅（更新の大きさ）が勝手に小さくなりすぎてしまうのを防ぐために、安定した「基準となる歩幅」を混ぜて調整しました。
• ⚖️ 強さを調整（スペクトル・テンパリング）：
  急な崖（曲率が高い部分）に対して、あまりにも強く補正しすぎると逆に転んでしまいます。そこで、補正の強さを少し抑えて（α=0.125 など）、バランスよく調整する「おだやかな補正」を採用しました。
• 🧱 片側だけの補正：
  計算を軽くするために、左右両方の補正をする代わりに、「片側だけ」の補正でも十分性能が出ることがわかりました。これにより、さらに高速化・省メモリ化が可能になりました。

---

🌟 まとめ

この論文「Mousse」は、**「AI の学習という旅において、地形を正しく理解し、それに合わせて歩き方を変える」**というシンプルな発想で、画期的な効率化を実現しました。

• Muon（旧来の方法）： 「どんな道でも均等に歩く」→ 急な崖で転んだり、平野で遅れたりする。
• Mousse（新方式）： 「地形に合わせて靴を履き替えて歩く」→ 最短距離でゴールに到達する。

これにより、今後さらに巨大で賢い AI を作る際、**「より少ない計算資源で、より高い性能」**が出せる道が開かれました。AI 開発の未来にとって、非常に重要な一歩と言えるでしょう！

技術概要

Mousse: 曲率感知前処理による Muon の幾何学的補正

技術的サマリー（日本語）

本論文は、大規模言語モデル（LLM）の学習において、スペクトル最適化手法である「Muon」と、第二階微分情報に基づく前処理手法（Shampoo など）の利点を統合した新しいオプティマイザ**「Mousse」**（Muon Optimization Utilizing Shampoo's Structural Estimation）を提案するものです。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

---

1. 背景と問題定義

近年、Muon オプティマイザは、更新ステップを Stiefel 多様体（直交行列の集合）に制限することで、学習の加速と汎化性能の向上を実現し、Moonshot-AI や DeepSeek-AI などの大手 AI 研究所で採用されるなど注目されています。Muon はニュートン・シュルツ（Newton-Schulz）反復を用いて勾配の符号行列（sign matrix）を計算し、スペクトルノルム制約を課します。

しかし、Muon には以下の幾何学的な限界が存在します：

• 等方性（Isotropic）の仮定: Muon は、すべての固有方向に対して均一なスペクトル更新ノルムを課す「平等主義的」な制約を前提としています。
• 実世界の損失地形との不一致: 深層ニューラルネットワークの損失地形は、極めて非等方的（アノイソトロピック）であり、ヘッシアン（曲率）のスペクトルは重たい裾（heavy-tailed）を持ち、条件数が悪い（ill-conditioned）ことが知られています。
• 結果: Muon の等方性制約は、高い曲率を持つ方向では不安定さを増幅させ、平坦な方向では必要な進展を制限するリスクがあります。

2. 提案手法：Mousse

Mousse は、スペクトル最適化の構造的安定性と、第二階微分前処理の幾何学的適応性を両立させることを目指します。

核心的なアイデア

Muons の等方性制約を、局所的な曲率情報を用いて「白化（Whitening）」された座標系で適用することで、損失曲面の実際の幾何学構造に整合させます。具体的には、Shampoo で用いられるKronecker 積分解されたヘッシアン近似（行と列の勾配相関行列 $L$ と $R$）を用いて勾配を前処理し、その変換された空間でニュートン・シュルツ直交化を実行します。

数学的定式化

Mousse は、アノイソトロピックな信頼領域（trust region）に制約されたスペクトル最急降下問題として定式化されます。

1. 白化変換: 勾配 $G$ に対して、Shampoo の統計量 $L$ と $R$ を用いて白化行列 $P=L^{1/4}, Q=R^{1/4}$ を定義し、変換された勾配 $\tilde{G} = P^{-1} G Q^{-1}$ を計算します。
2. スペクトル制約: 変換された空間で、標準的な Muon と同様にニュートン・シュルツ反復（$Y = -\text{msign}(\tilde{G})$）を適用し、直交行列を求めます。
3. 逆変換: 得られた更新を元の空間に戻します。
   $$\Delta W = -L^{-1/4} \cdot \text{msign}(L^{-1/4} G R^{-1/4}) \cdot R^{-1/4}$$

このプロセスにより、更新ステップは曲率の高い方向では抑制され、平坦な方向では促進されるように調整されます。

実装上の工夫（安定化技術）

第二階微分情報とスペクトル制約の組み合わせにおける安定性課題に対し、以下の技術を採用しています：

• Trace Normalization: 層ごとの $L, R$ のスケーリング差を補正するため、分解前に共分散行列のトレースを次元数で正規化し、ダミング係数 $\epsilon$ の影響を均一化します。
• Spectral Tempering: 曲率補正の強さを制御する指数 $\alpha$ を、標準的な Shampoo の $0.25$から$0.125$ に緩和することで、平坦な方向での過剰な勾配増幅を防ぎます。
• Gradient Grafting: 更新の大きさ（RMS ノルム）が学習中に減少する問題を防ぐため、AdamW などの安定したオプティマイザから更新の大きさを借用する「グラフティング」手法を適用し、学習率の調整を容易にしています。
• Single-Sided Preconditioner: 計算コストとメモリ使用量を削減するため、$L$ のみ（または $R$ のみ）を使用する片側前処理も有効であることを示しています。

3. 主要な貢献

1. 統一された幾何学的フレームワーク: Kronecker 積分解された曲率統計量に基づく白化座標系内でスペクトル最適化を再定式化し、第二階微分前処理とスペクトル手法のギャップを埋める理論的基盤を提供しました。
2. 堅牢なエンジニアリング知見: 重たい裾を持つ曲率推定とスペクトル制約の組み合わせにおける安定性課題（Trace Normalization や Spectral Tempering など）を分析し、大規模学習における実用的なガイドラインを提示しました。
3. パレート最適な効率性: 160M から 800M パラメータの言語モデルでの実験により、Muon よりも優れた性能を維持しつつ、計算コストの増加を最小限に抑えることを実証しました。

4. 実験結果

FineWeb データセット（200 億トークン）を用いた大規模な実験を行いました。

• 収束性能の向上: 800M パラメータモデルにおいて、Mousse は Muon と比較して最終検証損失を約 0.012 改善しました。
• サンプル効率: 同程度の損失レベルに到達するために必要な学習ステップ数を、Muon よりも約 12% 削減しました（図 1, 4 参照）。
• 計算コスト: Wall-clock 時間（実時間）のオーバーヘッドは Muon と比較して約 3% 程度と極めて小さく、SOAP（Shampoo の変種）のような第二階微分オプティマイザに見られるようなスループット低下は発生しませんでした。
• メモリ効率: 第二階微分情報を用いるにもかかわらず、AdamW 型の分散バッファを不要とするため、SOAP よりもメモリ使用量が約 88% まで削減され、Muon と同等の軽量さを維持しました（表 1 参照）。
• スケーラビリティ: 160M から 800M までのモデルサイズにおいて、一貫して Muon や SOAP、AdamW を凌駕する性能を示しました（図 3, 5 参照）。

5. 意義と結論

Mousse は、大規模言語モデルの前学習（Pre-training）における新しいパレートフロンティアを確立しました。

• 理論的意義: 等方性の仮定に依存していたスペクトル最適化を、実際の非等方的な損失地形に適応させることで、幾何学的に補正された最適化手法の枠組みを確立しました。
• 実用的意義: 第二階微分情報の恩恵を受けながら、計算リソースやメモリ制約を大幅に増大させないため、大規模モデルの効率的な学習に直結します。

本手法は、Muon の優れた収束特性を維持しつつ、Shampoo の曲率適応性を組み合わせた「両者の長所を兼ね備えた」オプティマイザとして、今後の大規模モデル学習の標準的な選択肢の一つとなり得る可能性があります。