技術要約: Muonp – 分数次スペクトル冪乗を用いたMuon

問題提起

Muonオプティマイザは、深層ニューラルネットワーク、特に大規模言語モデル（LLM）の学習において強力なツールとして浮上しています。これは、勾配行列 $G = USV^\top$ をその極因子 $UV^\top$ に置き換えることで、特異値スペクトルを実質的に平坦化（flattening）します。この「スペクトルの平坦化」は、広範な特徴量の発見を促進し、支配的なモードが学習の更新を独占することを防ぎます。

しかし、著者らはMuonのアプローチにおける決定的な限界を指摘しています。それは、完全な平坦化はすべての特異値情報を破棄してしまうという点です。この一様性は事前学習においては有益ですが、適応タスク（ファインチューニングなど）においては、特異値の大きさがどの残差方向を変更すべきかをエンコードしているため、最適ではない可能性があります。本研究が扱う中心的な問いは、「Muonの行列レベルのスペクトル的利点を維持しつつ、効果的な適応に必要な大きさの構造をいかに保持するか？」というものです。

手法

理論的定式化: Muonp

著者らは、分数次スペクトル冪乗を用いて重みを更新する、Muonの一般化である Muonp を導入します：
$$\text{Update} \propto US^pV^\top$$
ここで、$p \in (0, 1)$ は有理数です。

• 幾何学的解釈: 著者らは、この更新規則が Schatten $q$-ノルム（ここで $q = 1 + 1/p$）の下での最急降下方向に対応することを証明しています（定理 2.1）。
• 補間:
  • $p = 0$: 標準的なMuon ($UV^\top$) に相当し、一様性を最大化します。
  • $p = 1$: 標準的な勾配降下法 ($USV^\top = G$) に相当します。
  • $0 < p < 1$: 両者の間を補間し、部分的な大きさの情報を保持します。

アルゴリズムの課題と解決策

$US^pV^\top$ を直接計算するには特異値分解（SVD）が必要であり、これは大規模モデルに対して計算コストが高すぎます。著者らは、明示的なSVDなしでこの更新を近似する方法を検討しています。

1. 一変数反復の不可能性: 著者らは、任意の有理数 $p$ に対して、固定された一変数多項式反復 $S_{n+1} = f(S_n)$ では $S^p$ に収束できないことを証明しています（補題 2.2）。極因子（$p=0$）とは異なり、分数次冪には元の行列を「記憶」するアルゴリズムが必要です。
2. 二変数多項式回帰: これを解決するために、著者らは低次奇関数二変数回帰を導出しました。アルゴリズムは、特異値が $S_{n+1} = f(S_n, S)$ を通じて進化する行列 $G_n = US_nV^\top$ を反復的に計算します。
   • 定理 2.3 は、任意の奇関数二変数多項式 $f(x, y)$ について、項 $Uf(S_n, S)V^\top$ が行列の積と転置のみを使用して計算可能であり、Muonの計算構造を維持できることを示しています。
   • 定理 2.4 は、任意の有理数 $p \in (0, 1)$ に対して $S^p$ に収束する多項式の存在を確立しています。
3. 明示的な構成: 本論文では、これらの多項式の具体的な公式（例：$p=1/3$ の場合は $f(x, y) = x + c(y - x^3)$）を提供しています。係数は、$S^p$ が吸引不動点となるように選択されています。

実装の詳細

• 複雑性: 計算複雑性は $O(\omega k N)$ のままです（$\omega$ は行列積の複雑さ、$k$ はステップあたりの乗算回数、$N$ は反復回数）。Muonpは、行列積の回数において定数倍の差があるのみです。
• 安定性: bfloat16 で実装した場合、数値的に安定しており、現代のハードウェアに適しています。
• カリキュラムへの適合性: Muon ($p=0$) はMuonpの特殊なケースであるため、学習中の切り替え（例：事前学習でMuonを使用し、ファインチューニングでMuonpを使用する）において、状態のリセットは不要であり、多項式近似の変更のみで済みます。

主な貢献

1. Muonp オプティマイザ: 明示的なSVDなしに、任意の有理数 $p \in (0, 1)$ に対して $US^pV^\top$ を計算する一般化されたオプティマイザの導入。
2. 理論的限界: 分数次スペクトル冪乗は固定された一変数多項式反復では計算できず、新しい二変数アプローチが必要であることの証明。
3. 構成的アルゴリズム: 行列積のみの構造を維持しながら分数次冪を近似する、低次奇関数二変数回帰の導出。
4. 実証的検証: 数学的推論、コーディング、およびダウンストリームタスクにおいて、MuonpがMuonおよびAdamWを大幅に上回る性能を示すことを実証。
5. スペクトル幾何学の分析: 最適なスペクトル幾何学はレジーム依存的であるという洞察（一様な更新 ($p=0$) は事前学習に、部分的な保持 ($0 < p < 1$) はファインチューニングに有益）。

実験結果

ファインチューニングの性能

Muonpは、Llama3-1BおよびLlama3-3Bモデルのファインチューニングにおいて、MuonおよびAdamWに対して一貫した改善を示しています。

• 数学およびコード推論: GSM8K、Numina、OpenCodeInstructにおいて、MuonpはMuonと比較してpass@8精度を向上させ、損失を減少させます。
• パープレキシティ: Muonpは、FinewebおよびPile of Lawデータセットにおいて、より低い検証パープレキシティを達成します。
• ダウンストリームタスク: 一連の推論ベンチマーク（PIQA、Winogrande、LAMBADAなど）において、MuonpはMuonおよびAdamWの両方よりも高い精度を実現します。
• オプティマイザのミスマッチ: 事前学習でMuon（一致するオプティマイザ）を使用していた場合でも、ファインチューニング中にMuonpを使用するとMuonを上回ります。これは、この利得が単なるオプティマイザの一貫性ではなく、スペクトル幾何学に根ざしたものであることを示唆しています。

事前学習 vs ファインチューニングの二分法

本論文は、学習フェーズに基づいた明確なパフォーマンスの二分法を明らかにしています。

• 事前学習: Muon ($p=0$) がMuonpを上回ります。著者らは、事前学習は広範な特徴を発見し、支配的な方向が更新を独占するのを防ぐための「等方化（isotropization）」から恩恵を受けると考えています。
• ファインチューニング: Muonp ($0 < p < 1$) がMuonを上回ります。このレジームでは、モデルの固有次元が低いため、特異値の情報を保持することが、最も関連性の高い残差方向への適応に集中するのに役立ちます。
• カリキュラム実験: 事前学習の最終ステップでMuonからMuonpへ切り替える（カリキュラム）と、即座に損失が減少します。これは、$p$ を増やすことで小さな特異方向に対する有効な学習率を適応させるという仮説を裏付けています。

実行時間

MuonpはMuonと同等の実行時間を維持します。例えば、$p=1/3$（3次回帰）の場合、MuonpはMuon（一部の実装では $p=0$ のために5次回帰を使用）よりもわずかに速く、H100 GPU上でのステップ時間は約100〜120msです。

意義と主張

本論文は、MuonpをMuonオプティマイザの原理的な一般化として位置づけています。その主な意義は以下の通りです。

1. ギャップの解消: Muonの一様性と、勾配降下法の大きさへの感度の間のトレードオフ（「スペクトル平坦化」と「大きさの保持」）を解決する体系的な方法を提供します。
2. レジーム依存の最適化: 「最適な」スペクトル幾何学は普遍的なものではない（事前学習には完全な平坦化が必要であり、ファインチューニングには部分的な平坦化が必要である）という証拠を提示します。
3. 実用的な効率性: 行列積のみの構造を維持することで、SVDのコストを負うことなく、数十億パラメータのモデルへのスケーリングに不可欠なハードウェアフレンドリーな特性を維持したまま、これらの理論的利点を実現します。
4. 将来の方向性: 特定のタスクに対する最適な $p$ の自動選択や、事前学習のカリキュラムへのMuonpの統合が、有望な今後の研究領域であることを示唆しています。

結論として、特異スペクトルを保持することは適応シナリオにおいて大きな利益をもたらす可能性があり、Muonpはこれを実現するための計算効率が高く、原理に基づいたメカニズムを提供します。