MAPLE: MoE Adaptive Plug-and-play Layer-wise Expert allocation

技術要約：MAPLE – MoEにおける適応的プラグアンドプレイ層別エキスパート割り当て

問題提起
現在の疎結合型（Sparsely-activated）Mixture-of-Experts (MoE) Transformerは、一様に固定された数のエキスパートをすべての層で活性化させる均一なルーティング戦略を普遍的に採用しています。この慣習は、層ごとの冗長性とエキスパートの重要性の違いに関する、十分に文書化された不均一性を無視しています。著者らは、この一様性は系統的に最適ではないと主張しています。すなわち、エキスパート数による変動に敏感な層に対して計算能力を向けることができず、一方で冗長な層に対してはリソースを浪費しているという点です。MoEの効率性を向上させる既存のソリューション（非構造化プルーニング、エキスパートのマージ、または構造化圧縮など）は、通常、重みの再構成、追加のキャリブレーションデータ、再学習、または恒久的なアーキテクチャ変更を必要とするため、「プラグアンドプレイ」での事前学習済みモデルへのデプロイには適していません。

手法
本論文では、MAPLE (MoE Adaptive Plug-and-Play Layer-wise Expert Allocation) を提案します。これは、重みの変更、再学習、またはモデルアーキテクチャの変更を行うことなく、事前学習済みMoE大規模言語モデル（LLM）の各層に対して、ルーティングされるエキスパートの予算を不均一に再割り当てするフレームワークです。MAPLEは、以下の3つの明確なフェーズで動作します。

1. 層別感度モデリング (Layer-wise Sensitivity Modelling):
   本フレームワークは、他の層をベースラインの数（$k_0$）に固定したまま、各層のルーティングされるエキスパート数（$e$）を変化させることで、各層を個別に調査します。そして、各構成における検証精度 $A_l(e)$ を記録します。

   • 単一層の最適値 ($k^*_l$): 特定の層において精度を最大化するエキスパート数。
   • 感度指標: 層がエキスパート数の変化にどの程度強く反応するかを定量化するために、MAPLEは3つの軽量な指標を用います。
     • $s^{acc}_l$: 精度の変動範囲 ($\max A_l(e) - \min A_l(e)$)。
     • $s^w_l$: 層内のパラメータテンソルの分散。
     • $s^{out}_l$: $e$ 個のエキスパートを用いた層の出力と、$k_0$ 個のエキスパートを用いた層の出力との間の平均二乗誤差 (MSE)。
   • 著者らは、主要な感度指標として精度の変動範囲 ($s^{acc}_l$) を採用しています。

2. 感度ガイドによる閉形式の割り当て (Sensitivity-Guided Closed-Form Allocation):
   MAPLEは、エキスパートの再割り当てを制約付き最適化問題として定式化します。目標は、グローバルなエキスパート予算 $K$ の下で、各層の感度 ($s_l$) によって重み付けされた、各層の単一層最適値 ($k^*_l$) に近い割り当てベクトル $\mathbf{k} = [k_1, \dots, k_L]$ を見つけることです。

   • 目的関数は、重み付き二乗偏差を最小化します：$\min \sum w_l (k_l - k^*_l)^2$ （ここで $w_l = s_l$）。
   • ラグランジュ未定乗数法により、連続解 ($k^{cont}_l$) が得られます。予算に対する余剰または不足のエキスパートは、低感度の層に優先的に割り当てられます。
   • 連続解は実行可能な範囲 $[k_{min}, k_{max}]$ にクリップされ、最大剰余法を用いて整数割り当てに変換されます。

3. 感度制約付き遺伝的洗練 (Sensitivity-Constrained Genetic Refinement):
   丸めによるアーティファクトや層間の相互作用に対処するため、軽量な遺伝アルゴリズムが閉形式の解を洗練させます。

   • 突然変異 (Mutation): 候補は、ガウス突然変異を通じて生成されます。ここで、分散 $\sigma^2_l$ は感度に反比例します ($\sigma^2_l \propto 1/s_l$)。これにより、高感度の層には小さな摂動が与えられるようになります。
   • 選択 (Selection): 世代ごとに上位2つの候補（エリート）が保持され、交叉なしで集団が再生成されます。
   • このプロセスは、最終的な割り当て $\mathbf{k}_{best}$ に向かって急速に収束します（多くの場合、数世代以内）。

主な貢献

• プラグアンドプレイ・フレームワーク: MAPLEは、事前学習済みMoEモデルにおける不均一なエキスパート割り当てのための、完全なポストホック（事後的）、かつ重みを使用しない初のフレームワークです。再学習、重みの更新、またはアーキテクチャの変更を必要としません。
• 感度ガイド型の最適化: 本論文は、層ごとの感度に基づいた、遺伝的探索によって洗練された閉形式の解析解を導入しています。これにより、網羅的な探索や強化学習の計算コストを回避しています。
• 削減された予算下でのパフォーマンス: 本手法は、総ルーティング・エキスパート予算が減少している場合（例：75%）でも、不均一な割り当てが均一な割り当てよりも優れた性能を発揮できることを示しています。これは、元のフル予算の均一なベースラインさえも実質的に上回ります。

実験結果
著者らは、4つのMoEモデル（DeepSeek-MoE-16B, DeepSeek-V2-Lite, Moonlight-16B-A3B, OLMoE-1B-7B）を用い、7つのベンチマーク（ARC-E, ARC-C, BoolQ, PIQA, RTE, GSM8K, BBH）で評価を行いました。

• 精度の向上: 75%のルーティング・エキスパート予算の下で、MAPLEはすべてのモデルにおいて、均一なベースラインおよびプルーニングベースの手法（SparseGPT, Wanda, EPP, MoNE, LExI）を上回りました。
  • DeepSeek-MoE-16Bにおいて、エキスパートを75%のみ使用した場合でも、MAPLEはARC-E (65.09% $\to$ 71.40%)、ARC-C (48.49% $\to$ 51.50%)、BoolQ (80.03% $\to$ 82.38%) の精度を、元の100%均一ベースラインと比較して向上させました。
  • OLMoE-1B-7Bにおいて、MAPLEは最高の平均精度を達成し、フル予算のベースラインに対して4.13ポイント向上しました。
• 長文推論 (Long-Form Reasoning): GSM8KおよびBBHにおいて、MAPLEは競争力のある精度を維持しつつ（GSM8Kで0.61ポイント以内）、評価時間を大幅に短縮しました（GSM8Kで53.0%の削減）。
• デプロイ効率: SGLangに実装されたMAPLEは、単一アクティブエキスパート（$k=1$）が割り当てられた層に対してマルチエキスパート・ディスパッチ・ロジックをバイパスすることで、シングルGPUのエンドツーエンド・サービング・レイテンシを**32.2%削減し、スループットを47.4%**向上させました。
• 堅牢性: アブレーション研究により、精度の変動範囲（accuracy range）メトリックが最も効果的な感度信号であることが確認されました。本手法は、ランダムシードに対して安定しており、タスク間転移（cross-task transfer）においても汎用性を示しました。これは、特定された感度パターンが特定のキャリブレーションセットに過学習したものではなく、モデルの層に固有のものであることを示唆しています。

意義と主張
本論文は、MAPLEがMoEの効率化における「原理的かつ実践的な軸」を確立すると主張しています。その主な意義は、適切に設計された不均一な割り当ては、単に多くのエキスパートを活性化させることよりも効果的であることを証明した点にあります。再学習なしに、異なる層の不均一な要求に合わせて容量を再分配できることを証明することで、MAPLEは、より効率的なMoEモデルをデプロイするための費用対効果の高い経路を提供します。著者らは、本アプローチが既存のプルーニングや圧縮手法のオーバーヘッドを回避し、生産環境における既存の事前学習済みモデルに即座に適用可能であることを強調しています。

限界と今後の課題
著者らは、現在の評価が特定のGPU環境（NVIDIA RTX A6000）およびサービングエンジン（SGLang, vLLM）に限定されていることを認めています。今後の課題は、多様なハードウェアプラットフォーム（CPU、エッジデバイス）および、より広範な実世界の推論シナリオにおいてMAPLEを評価し、削減されたエキスパート計算が、異なるシステム制約下で実際のレイテンシ、メモリ、およびエネルギーのメリットにどのように変換されるかを検証することです。