ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

技術要約：ProbMoE: Mixture-of-Expertsのための微分可能な確率的ルーティング

問題提起
Mixture-of-Experts (MoE) モデルは、トークンごとに少数のエキスパートのみを活性化させることで効率的にスケールアップします。しかし、標準的なルーティングは離散的で微分不可能な top-$k$ オペレータに依存しているため、これらのモデルの学習は依然として困難です。従来の学習手順では、ルーターのロジットに対する依存性を無視し、選択されたエキスパートのソフトマックス確率を通じてのみ勾配を伝播させることで、この問題を回避しています。この近似は、代替となるエキスパートの組み合わせに対する学習信号を制限し、結果としてルーティング分布の集中、エキスパート利用率の低下、および学習ダイナミクスの不安定化を招くことがよくあります。これを緩和するための既存の試み（高密度なバックプロパゲーションやヒューリスティックな勾配推定器など）は、実行可能なエキスパートの部分集合に対する分布を明示的にモデル化できていなかったり、学習中に代替の組み合わせを探索するメカニブズムが限定的であったりすることが多いのが現状です。

手法
著者らは、エキスパートのルーティングを、基数制約（cardinality-constrained）を持つエキスパートの部分集合に関する確率的推論問題として定式化するフレームワークである ProbMoE を提案します。ルーティングを決定論的なオペレータとして扱うのではなく、ProbMoE はエキスパートの選択を離散的な部分集合上の分布としてモデル化します。

1. 確率的定式化:

   • Exact-$k$ ルーティング: 本フレームワークは、$|S|=k$ である部分集合 $S$ 上の分布を定義します。これは、ルーターのロジット（シグモイド関数経由）から導出された独立なベルヌーイ選択確率を、正確に $k$ 個のエキスパートが選択されるという制約条件の下で条件付けたものです。この分布の正規化定数 $Z_k$ は動的計画法を用いて正確に計算されるため、この定式化は計算可能です。
   • Dynamic-$k$ ルーティング: 本フレームワークは、部分集合のサイズ $|S|$ が $[k_{min}, k_{max}]$ の範囲内に収まる制約付きの設定へと一般化します。この分布は範囲条件に基づいているため、モデルはエキスパートの識別と、各トークンに対するルーティングの基数（cardinality）を同時に推論することができます。

2. フォワードパス:

   • 学習時、ProbMoE は定義された分布（Exact-$k$ または範囲制約付き）からエキスパートの部分集合 $S$ をサンプリングします。
   • MoE レイヤは、サンプリングされたエキスパートのみを使用して出力を計算し、標準的な MoE モデルの疎な実行パターンを維持します。
   • 推論時には、最大事後確率（MAP）集合が選択され、推論コストが従来の Moemo ルーティングと一致するようにします。

3. バックワードパス:

   • 勾配ベースの最適化を可能にするため、ProbMoE は SIMPLE 勾配推定器 (Ahmed et al., 2023) を採用しています。
   • 離散的なサンプルを直接微分する代わりに、バックワードパスでは、制約付き分布の下で各エキスパートが選択される**周辺確率（marginal probabilities）**を使用します。
   • ルーティング重みは、ストレートスルー・エスティメータを用いて構築されます：$w = (\text{stopgrad}(z - m) + m) \odot \pi$。ここで、$z$ はサンプリングされたマスク、$m$ は正確な周辺確率のベクトル、$\pi$ はソフトマックス・ルーティング確率です。
   • これにより、勾配が周辺確率を通じて流れ、疎なエキスパート実行を維持しながら、選択確率の変化が期待損失にどのように影響するかを反映した、情報量の多いルーターへの更新が可能になります。

主な貢献

1. 確率的推論の定式化: 本論文は、MoE ルーティングを、ヒューリスティックな勾配近似を超えた、基数制約付きエキスパートの部分集合に関する確率的推論として再定義しました。
2. 扱いやすい周辺確率ベースの最適化: SIMPLE 推定器を活用することで、著者らは、疎なエキスパート側の計算を行わずに、確率的な部分集合サンプリングと正確な周辺確率ベースのルーター勾配を組み合わせ、微分可能な学習を実現しました。
3. 統一された Dynamic-$k$ フレームワーク: 本手法は、ヒューリスティックなゲーティング規則に頼ることなく、エキスパートの識別と、トークンあたりのアクティブなエキスパート数を事前定義された範囲内で同時に推論するダイナミック・ルーティングへと自然に拡張できます。
4. 実証的な改善:
   • Exact-$k$: ProbMoE は、OLMoE および Qwen バックボーンにおいて、GSM8K, Law, Translation, MBPP, MMLU などのベンチマークで強力な性能を示し、DenseMixer や従来の Top-$k$ と同等またはそれ以上の性能を発揮しています。
   • Dynamic-$k$: ProbMoE Dynamic-$k$ は、平均してより少ないエキスパートを活性化させながら、Exact-$k$ に匹敵する性能を達成しており、効果的な適応的割り当てを実証しています。
5. ルーティングの多様性: 分析により、ProbMoE は、ベースラインと比較して、より広いルーティング分布を誘導し、よりバランスの取れたエキスパート利用を実現することが示されており、エキスパートの崩壊（expert collapse）を抑制しています。

結果と意義
本論文は、ProbMoE がエキスパートの利用率とルーティングの多様性を向上させることを示しています。アブレーション研究において、確率的サンプリングと周辺確率ベースの最適化の組み合わせが、性能向上の主要な要因であることが示されており、一方のコンポーネントのみを使用する構成よりも優れた結果となっています。

著者らは、ProbMoE がエキスパート選択のための原理的な手法を提供し、学習の安定性と計算効率を向上させると主張しています。ルーティングを確率的にモデル化することで、本フレームワークは学習中に代替となるエキスパートの組み合わせの探索を可能にし、より適切に較正されたルーティング分布をもたらします。Dynamic-$k$ バリアントは、適応的な基数がタスクの複雑さに合わせて学習され、タスク性能を損なうことなく実効的なルーティングコストを削減できることをさらに示しています。本研究は、今後の方向性として、ダイナミックな設定における疎性を最大限に活用するための、事前学習への拡張やシステムレベルの最適化を挙げています。