技術要約: DOT-MoE: MoEficationのための微分可能な最適輸送

1. 問題提起

大規模言語モデル（LLM）の急速なスケーリングは、推論効率における重大なボトルネックを生み出しています。Mixture-of-Experts（MoE）アーキテクチャは、トークンごとにパラメータの一部のみを活性化することで、モデルサイズと推論コストを切り離しますが、MoEをゼロから学習させることは極めて不安定で、大量のデータを必要とし、計算コストも高いことが知られています。

その結果、事前学習済みの密な（dense）モデルを疎な（sparse）MoEへと変換する「MoEfication」が、有力な代替案として浮上しています。しかし、既存の密なフィードフォワードネットワーク（FFN）を疎なエキスパートに変換する手法は、ヒューリスティックな戦略（例：ランダムな分割、重みに基づくクラスタリング、または活性化に基づくクラスタリング）に大きく依存しています。これらのアプローチには、主に以下の3つの制限があります。

1. 関心の分離（Separation of Concerns）: これらは通常、ニューロンの割り当て（FFNの分割）とルーターの学習を、相互に適合させることができない別々の逐次的なステージとして扱います。
2. プロキシ最適化（Proxy Optimization）: ニューロンの実際の出力への寄与ではなく、中間的な統計量（重みや活性化）に基づいて最適化を行うため、密なモデルの挙動を再構成する際に最適ではない結果を招くことがよくあります。
3. 微分不可能性（Lack of Differentiability）: ニューロン割り当ての離散的な性質により、エンドツーエンドの勾配フローが妨げられ、構造的な分解とルーティング・ポリシーを共同で最適化することが困難になります。

2. 手法: DOT-MoE

著者らは、密なレイヤーをエキスパートへと分解するプロセスを**微分可能な最適輸送（Differentiable Optimal Transport: DOT）**問題として定式化するフレームワーク、DOT-MoEを提案します。このアプローチにより、エキスパートの分解とルーティング機構の両方をエンドツーエンドで学習することが可能になります。

2.1. 最適輸送の定式化

核心となる洞察は、ニューロンの割り当てを、厳格な容量制約の下で「質量」（ニューロン）をソース分布（密なFFNのニューロン）からターゲット分布（エキスパート）へと輸送する、バランスの取れた輸送問題として捉えることです。

• ソース: $dffn$ 個のニューロン。各ニューロンは単位質量を持つ。
• ターゲット: $E$ 個のエキスパート。各エキスパートは正確に $s$ 個のニューロンを必要とする（$s \cdot E = dffn$）。
• コスト関数: 固定されたコスト行列の代わりに、学習可能なアフィニティ行列 $A$ を使用します。ここで、$A_{i,e}$ はニューロン $i$ をエキスパート $e$ に割り当てる際のアフィニティ（親和性）を表します。目的は、周辺制約を満たしながら、総アフィニティを最大化することです。

2.2. Sinkhorn-Knoppによる微分可能な緩和

離散的な割り当て問題を微分可能にするために、DOT-MoEはエントロピー正則化を採用しています。

• 最適輸送問題は、輸送計画行列の行と列を反復的に正規化するSinkhorn-Knoppアルゴリズムを用いて緩和されます。
• これにより、期待値において容量制約を満たすソフトな割り当て行列 $M_{soft} \in [0, 1]^{dffn \times E}$ が得られます。
• 温度パラメータ $\tau \to 0$ となるにつれ、解は離散的な最適解に近づきます。

2.3. Straight-Through Estimator (STE) による共同最適化

デプロイメントには離散的な割り当てが必要であり、ルーティングには離散的な top-k 選択が必要ですが、これらは共に微分不可能であるという点が重要な課題となります。DOT-MoEは、Straight-Through Estimators (STE) を用いてこれに対処します。

• フォワードパス: ハードな決定を使用します（割り当てには $M_{soft}$ の貪欲な丸め、ルーティングには top-k 選択を使用）。
• バックワードパス: ソフトな対応物（$M_{soft}$ およびルーターのロジット $P$）を通じて勾配を伝播させます。
• これにより、訓練中にルーターの重み ($W_r$) と割り当てアフィニティ行列 ($A$) を共適応させることができます。

2.4. 学習目的関数

このフレームワークは、密なモデルの重みを凍結した状態でのアライメント・フェーズを通じて訓練されます。モデルは以下の組み合わせの損失を最小化します。

1. 出力再構成（Output Reconstruction）: 密なティーチャーモデルと疎なスチューデントモデルの出力間のKLダイバージェンスによって測定されます。
2. 言語モデリング（Language Modeling）: 言語モデリングの目的関数に対するクロスエントロピー損失。
3. 安定化補助損失（Stability Auxiliaries）: ルーターの $z$-loss（ロジットの爆発を防ぐため）およびロードバランシング損失（エキスパートの崩壊を防ぐため）。

3. 主な貢献

1. 最適輸送フレームワーク: 密なモデルからMoEへの変換を、厳格なエキスパート容量制約を強制する微分可能なSinkhorn反復を通じて解かれる、バランスの取れた輸送問題として定式化する新しい手法を導入しました。
2. 二段階の共適応（Dual-Level Co-Adaptation）: ルーターを訓練する前に割り当てを固定する従来の手法とは異なり、DOT-MoEは、補完的なSTEを用いて、ニューロンからエキスパートへの分解と、トークンからエキスパートへのルーティングを共同で最適化します。
3. 出力認識型の構築（Output-Aware Construction）: ニューロンの寄与を（重みや中間活性化などのプロキシではなく）最終的な出力に対して直接モデル化することで、より高い再構成忠実度を実現しました。

4. 実験結果

著者らは、3つのモデルファミリー（LLaMA-2, LLaMA-3, Qwen2.5）と6つのベンチマーク（ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ）にわたってDOT-MoEを評価しました。

• プルーニングとの比較: 50%のパラメータ予算を用いたLLaMA-2 7Bにおいて、DOT-MoEはWikiText-2で 7.99 のパープレキシティを達成し、最先端の構造化プルーニング手法であるDISP-LLM（9.84）を大幅に上回りました。
• 既存のMoEficationとの比較: ゼロショット設定において、DOT-MoEはCMoEやLLaMA-MoE-v2といった手法を大幅に上回りました。例えば、LLaMA-2 7Bにおいて、DOT-MoEはCMoEの44.5%に対し、平均精度**61.5%**を達成しました。
• 効率性: DOT-MoEは、活性パラメータを**50%削減しながら、元の密なモデルの性能の約90%**を維持しています。
• ファインチューニング: 1.2Bトークンの連続的なファインチューニングのみで、DOT-MoEモデルは密なモデルに近い性能を回復しました（例：LLaMA-2 7Bにおいて66.6% vs 76.6%）。これは、同等またはより多くのデータ予算を必要としたベースラインを上回る結果です。
• アブレーション研究:
  • 粒度（Granularity）: エキスパートの粒度を増やすことで、ある飽和点までは性能が向上しますが、その後は利得が減少します。
  • 推論速度: エキスパート数が増えても、融合されたMoEカーネル（例：vLLM）を使用する場合、総融合中間次元が一定であるため、推論スループットは安定しています。
  • スパース性: 高いスパース性レベルで訓練されたモデルは、推論時のさらに極端なスパース性レジームに対しても堅牢に汎化します。

5. 意義と主張

本論文は、DOT-MoEが、事前学習済みモデルを疎なMoEへと変換するための原理的かつ微分可能な手法を提供することにより、LLM推論の効率化における重要な進歩であることを主張しています。

• パラダイムシフト: ヒューリスティックで静的な分割から、エキスパートの機能的な一貫性を尊重する、動的で学習可能な分解へと移行します。
• 実用性: 本手法は、ゼロからのフル再学習を必要とせずに、事前学習済みモデルの事後的（post-hoc）なスパース化を可能にし、活性パラメータを劇的に削減しながら高い容量を維持します。
• 堅牢性: FFNの忠実度を維持するには、ニューロンの寄与を明示的にモデル化する必要があることを示しており、これは既存のプロキシベースの手法が捉えきれない原則です。

著者らは、DOT-MoEが疎なモデルに対する優れた初期化を提供し、最小限の追加学習データで、密なアーキテクチャと疎なアーキテクチャの間の性能ギャップを埋めるものであると結論付けています。