技術概要：粒状な混合専門家モデルのための適応的逆インデックスルーティング（AIR-MoE）

1. 問題定義

混合専門家（MoE）モデルは、トークンごとに専門家のサブセットのみを活性化することで、スケーラブルな大規模言語モデル（LLM）の基盤となっています。最近の実証的なスケーリング則は、少数の大きな専門家ではなく、多数の小さな専門家（粒状な専門家）を使用することで性能が向上することを示唆しています。しかし、この regime は「ルーティングコスト」という重要な計算上のボトルネックをもたらします。専門家数（$E$）が増加するにつれて、ルーターは、トップ-$K$の専門家を特定するために、すべてのトークン表現とすべての専門家の重心との内積を計算する必要があります。このオーバーヘッドは総計算量を支配し、高度に粒状なモデルのトレーニングを実質的に不可能にする可能性があります。

ルーティングコストを削減するための既存の解決策は、通常、モデルの表現力を制限する構造的制約を課します：

• 階層的 MoEは、専門家を固定された非重複のグループに分割し、ルーティングをそれらのグループ内でのみ制限します。
• プロダクトキー/構造的制約（例：He, 2024）は、正確な検索を可能にするために専門家の重心をコンポーネントに分割しますが、許容される専門家の表現の空間を制限します（例：デカルト積構造を強制するなど）。
• 固定ハッシュアプローチは学習可能性を欠き、本質的に限界があることが示されています。

課題は、専門家のパラメータに構造的制約を課すことなく、または全専門家のスコアリングを必要とせずに、トップ-$K$ルーティングを効率的に近似することです。

2. 手法：AIR-MoE

著者は、情報検索および近似最近傍検索で使用される逆ファイルインデックスに着想を得た、MoE 用の適応的逆インデックスルーティング（AIR-MoE）というルーティングアーキテクチャを提案します。この手法は、すべての専門家の完全なスコアリングを回避する、粗から細への 2 段階の検索プロセスを通じて、真のトップ-$K$専門家を近似します。

2.1 2 段階検索

各入力トークン $h$ に対して、ルーティングは並行して進行します：

1. 粗い候補抽出（Coarse Shortlisting）：トークンは、コサイン類似度を使用して、学習されたコードブック $\mathcal{C} = \{c_1, \dots, c_G\}$ 内の最も近いコードワード $c_g$ にマッピングされます。各コードワード $c_g$ は、$M$ 個の候補専門家の事前計算された短リスト（shortlist） $L_g$ と関連付けられています。これらの短リストは、逆インデックスにおける「ポストリスト」に相当し、コードワード重心に対する最も高い内積を持つ専門家のインデックスを格納します。
2. 詳細スコアリング（Fine Scoring）：ルーターは、トークン $h$ と検索された短リスト $L_g$ 内の専門家との間の正確な内積のみを計算します。その後、トップ-$K$の専門家が、この削減された候補セットから選択されます。

この階層構造により、アクティブな専門家の集合 $T$ が短リスト $L$ の部分集合であり、$L$ が完全な専門家集合 $[E]$ の部分集合であることが保証されます。

2.2 学習と最適化

中心的な課題は、トークンのコードワードへの割り当てが微分不可能であることです。AIR-MoE は、バイレベル最適化戦略を通じてこれを解決します：

• 勾配ベースの学習：トークン表現と専門家の重心は、下流の言語モデリング損失からの標準的な逆伝播によって更新されます。
• 勾配なしコードブック学習：コードブック $\mathcal{C}$ とその関連する短リストは、適応的球面 k-meansアルゴリズムを使用して独立して更新されます。このアルゴリズムは：
  • トークン表現の分布を追跡するために指数移動平均（EMA）を使用します。
  • 高次元のコサイン幾何学を効果的に処理するために単位球面上で動作します。
  • 「死んだ」コードワード（割り当て回数が少ないもの）の崩壊を防ぐための再初期化メカニズムを含みます。
  • 短リスト $L_g$ をオプティマイザのステップ後にのみ更新し、インデックスが専門家の重心の現在の状態と一貫性を保つようにします。

重要なのは、AIR-MoE はstraight-through 推定トリックを使用しないことです。勾配は量子化ステップを通じてコードブックに流れません。

2.3 理論的保証

論文は、マスリコール（短リストに保持された真のトップ-$K$専門家の総確率質量）の下限を導出します。この限界は以下の要素に依存します：

1. 量子化誤差 $\epsilon(h) = \|h - c(h)\|$：誤差が小さいほど保持が向上します。
2. 短リストマス $\rho_M(c)$：コードワードに対するトップ$M$の専門家によって捕捉される確率質量。
   理論的には、トークンが適切に量子化され、短リストサイズ $M$ が十分に大きい場合、AIR-MoE はルーティングマスの大部分を保持します。

3. 主要な貢献

1. AIR-MoE アーキテクチャ：専門家の表現に対して構造的制約（例：グループ化やパラメトリック制限）を課すことなく、粒状な MoE 向けのトップ-$K$ルーティングを近似する、逆インデックスに着想を得たルーティングメカニズム。
2. バイレベル最適化：モデルパラメータの勾配ベースの学習から、コードブックの勾配なし適応（適応的球面 k-means を通じて）を分離する学習戦略。straight-through 推定の必要性を回避します。
3. 理論的解析：マスリコールの下限を導出し、近似がルーティング品質を保持する条件に関する直観を提供します。
4. 実証的性能：既存の効率的な粒状ルーティング手法と比較して、AIR-MoE が優れたパープレキシティ–FLOPs トレードオフを達成することを示実証。

4. 実験結果

著者は、WikiText-103、C5、OpenWebText2 において、小・中・大規模（最大約 0.45B パラメータ）のモデルスケールで、65,536 個の専門家を用いて AIR-MoE を評価しました。

• ベースラインとの性能比較：AIR-MoE は、PEER（He, 2024）や階層的 MoEを含む既存の効率的な粒状ベースラインを常に上回りました。
  • 同程度の FLOP コストにおいて、PEER ベースラインに対して最大**10%**のパープレキシティ改善を達成しました。
  • 標準的な粒状ルーティング（すべての専門家のスコアリング）は最低のパープレキシティを達成しましたが、 prohibitive な計算コストを伴いました（例：WikiText-103 Small において、AIR-MoE が 324.1P FLOPs であるのに対し、約 467.5P FLOPs）。
• 効率性：AIR-MoE は、競争力のあるパープレキシティを維持しながら、正確な粒状ルーティングと比較してトレーニング FLOPs を大幅に削減しました。
• アブレーション研究：
  • 重心の正規化を除去するか、球面 k-means の代わりにユークリッド距離を使用すると、性能がわずかに低下しました。
  • 静的コードブック（非適応的）を使用すると最悪の性能となり、適応的コードブックの必要性を確認しました。
  • 専門家選択ゲーティングは高い重複をもたらしましたが、深刻な専門家枯渇（78.6% の死んだ専門家）を招き、重複のみがルーティング品質の十分な指標ではないことを示しました。
• 専門家使用：構造的制約によりルーティングのアーティファクトを引き起こす階層的または PEER ルーターとは異なり、AIR-MoE は構造的制約なしで柔軟な専門家使用を可能にしました。

5. 意義と主張

この論文は、AIR-MoE がスパース混合専門家アーキテクチャにおけるルーティングのための効果的なメカニズムを提供し、特に完全なルーティングの計算ペナルティなしに粒状な専門家の使用を可能にすると主張しています。

• ドロップイン交換：この手法は、基盤となるモデルアーキテクチャや損失関数への修正を必要とせず、標準的なルーターのドロップイン交換として機能します。
• 柔軟性：専門家のパラメータに対する構造的制約を回避することで、AIR-MoE は、固定されたグループ化やプロダクトキー構造を必要とする手法よりも、トークンと専門家の依存関係をより柔軟に学習することを可能にします。
• スケーラビリティ：粗い候補抽出のコストが償却され、ルーティングされるトークンの数に依存しないため、このアプローチは大きな有効バッチ領域で有利にスケーリングします。

著者は限界を認め、FLOP の削減は有意義であるものの、実際のウォールクロック速度向上は、ハードウェアのサポートと、粒状 MoE に固有のメモリトラフィックおよびインデックスオーバーヘッドを軽減するための最適化された実装に依存すると指摘しています。また、現在の実装は次元削減を採用していないため、それは将来の作業のための直交する方向性であると述べています。