技術的サマリー：Muon は Adam で事前学習されたモデルを微調整できるか？

問題提起

Muon（Newton-Schulz による直交化された MomentUm）は、大規模言語モデル（LLM）の事前学習において極めて効率的なオプティマイザとして登場し、Adam に対して約 2 倍の計算効率と、2 次モーメント推定を排除することによるメモリ使用量の削減を提供しています。しかし、その微調整への実用的な応用は依然として限られています。主な障壁はオプティマイザの不一致問題です：大多数の公開されているオープンモデルは Adam で事前学習されており、微調整に Muon を単純に切り替えると性能が低下します。逆に、Muon で事前学習されたモデルを Adam で微調整しても、最適ではない結果となります。この非互換性は、事前学習での成功にもかかわらず、Muon の支配的な微調整パラダイムへの採用を厳しく制限しています。

手法と分析

著者らは、Adam と Muon の固有の暗黙的バイアスに焦点を当て、制御された実験と理論的分析を通じてこの不一致を調査しました。

1. 暗黙的バイアスの特性：

   • Adam は、1 次モーメントと 2 次モーメントに基づく要素ごとの前処理を使用し、最大ノルママージン（$\|W\|_{\max}$）を最大化する解に収束します。
   • Muon は、Newton-Schulz 直交化による行列レベルの前処理を適用し、スペクトルノルママージン（$\|W\|_2$）を最大化する解に収束します。
   • 単純化された線形回帰問題に対する理論的分析は、これらのオプティマイザが本質的に異なる解（$W^*_{\max}$ 対 $W^*_2$）に収束することを示しています。実証的には、Muon 学習済みの重みは、Adam 学習済みの重みに比べて、より高い安定ランクとより均一な特異値分布を示します。

2. 劣化のメカニズム：

   • この不一致は事前学習された知識を破壊します。異なる暗黙的バイアスを持つオプティマイザが適用されると、モデルの重みが事前学習フェーズと互換性のない構造的な方向へ強制されます。
   • 実験により、不一致による微調整は更新強度に対する感度を高めることが示されました。より強い更新は、事前学習された構造に対するより大きな破壊を引き起こし、より高いパープレキシティとより深刻な破滅的忘却をもたらします。

3. 提案する緩和策：LoRA：

   • 著者らは、更新の範囲を制限することがこの不一致を緩和できると仮説を立てました。彼らは、事前学習済みの重み（$W_0$）を固定し、更新を低ランク部分空間（$\Delta W = BA$）に制限する**低ランク適応（LoRA）**を採用しました。
   • このアプローチは重み変化の大きさを制限し、オプティマイザが事前学習された構造を攻撃的に変更するのを防ぎます。理論的分析は、「不一致の増幅」（一致した場合と不一致の場合の性能の差）が LoRA のランク $r$ によって有界であり、$r \to 1$ として消滅することを示唆しています。

主要な貢献

• 再現と分析：本論文は、アクセス可能なスケール（NanoChat、5 億 6100 万パラメータ）でオプティマイザの不一致現象を再現し、それを Adam と Muon の固有の暗黙的バイアス（最大ノルマ対スペクトルノルマ）に関連付けました。
• 解決策としての LoRA：著者らは、LoRA が不一致を効果的に緩和することを示しました。更新を制限することで、LoRA-Muon は、Adam で事前学習されたモデルを微調整する際、LoRA-Adam と同等かそれ以上の性能を発揮します。
• 包括的な検証：以下の分野で知見が検証されました：
  • 自然言語理解（NLU）：T5-Base による GLUE ベンチマーク。
  • 自然言語生成（NLG）：数学、コード、常識推論における Llama 2-7B/13B の指示微調整。
  • コンピュータビジョン：6 つの画像分類タスクにおける CLIP ViT-B/32 の微調整。
• ランクと忘却に関する研究：
  • ランク感度：LoRA-Muon の恩恵は、低～中程度のランクで最も顕著です。ランクが増加し（フル微調整に近づく）、不一致が再発し、Adam に対する性能が低下します。
  • 破滅的忘却：Adam で事前学習されたモデルを Muon でフル微調整すると、Adam よりも事前学習知識のより深刻な忘却を引き起こします。LoRA はこの忘却を大幅に軽減し、事前学習された構造を保持します。
• LoRA 変種の評価：本研究は既存の LoRA 変種（rsLoRA、PiSSA、LoRA-One など）を評価しました。LoRA 更新をフル微調整に近づけるように設計された変種（PiSSA など）は、必ずしも Muon に利益をもたらすわけではなく、むしろ不一致を悪化させる可能性があり、Muon 固有の適応が必要であることを示唆しています。

結果

• 性能ギャップの縮小：フル微調整において、Muon は Adam で事前学習されたモデルで一貫して Adam よりも性能が劣ります（WikiText-2 において +0.009 から +0.023 のパープレキシティギャップなど）。LoRA 下では、このギャップは大幅に縮小します（Muon 事前学習モデルで 39%、Adam 事前学習モデルで 78%）。
• LoRA における優位性：NLU、NLG、ビジョンタスク全体を通じて、LoRA-MuonはLoRA-Adamと同等かそれ以上の性能を発揮します。例えば、GLUE ベンチマークでは、LoRA-Muon-PE が最高平均精度（89.20%）を達成し、LoRA-Adam（88.93%）と Full-Adam（89.14%）の両方を上回りました。
• 効率性：Muon は Newton-Schulz 反復により 1 ステップあたりのオーバーヘッドがわずかに発生します（LoRA 設定では Adam より 1.1～1.2 倍遅い）が、メモリ効率（Adam の 2 つのモーメントバッファに対し、1 つのモーメントバッファのみを保存）は実用的な利点であり、オプティマイザ状態のメモリを約 50% 節約します。

意義と主張

本論文は、Adam と Muon の間のオプティマイザ不一致問題に関する最初の詳細な分析を提供すると主張しています。その主な意義は、広範な Adam 事前学習モデルのエコシステムに対して、Muon を実用的に微調整可能にすることにあります。

著者らは以下のように結論付けています：

1. オプティマイザの不一致は構造的である：更新強度が制限されていない場合、衝突する暗黙的バイアスが事前学習された知識を破壊することから生じます。
2. 制限された更新が鍵である：重み変化の大きさを制限する LoRA などの手法は、この不一致を効果的に無力化します。
3. 実用的な推奨事項：実務家は、フル微調整で観察される不一致のペナルティを被ることなく、メモリ効率と競争力のある性能を得るために、LoRA 微調整において Adam のドロップイン代替として Muon を使用できます。

本論文は控えめに、LoRA が問題を緩和する一方で、Adam 向けに最適化された既存の LoRA 変種が自動的に Muon へ転用されるわけではないと指摘し、Muon 固有の適応を開発するための将来の作業が必要であることを示唆しています。さらに、不一致の理論的性質の解明は未解決の課題のままです。