Sparse MoE アーキテクチャにおいて、ゲーティングネットワークがトークンごとにスコア上位 K 個のエキスパートを選択して活性化するルーティング戦略。K の値がモデルの計算コストと表現力のトレードオフを決定する。
Top-K Routing は Sparse MoE モデルのゲーティング(ルーター)が採用する最も一般的な選択戦略です。各トークンの隠れ表現 x に対し、ルーターが N 個のエキスパートそれぞれのスコアを算出し、スコアの高い上位 K 個のエキスパートだけを活性化します。残りの N-K 個のエキスパートは計算をスキップするため、演算量は K/N に抑えられます。
K の選択はモデル設計の根幹を成すハイパーパラメータです。
| K 値 | 特徴 | 代表モデル |
|---|---|---|
| K=1(Top-1) | 最も疎。計算効率最大だが勾配が 1 経路しか流れず学習が不安定になりやすい | Switch Transformer, ST-MoE |
| K=2(Top-2) | 効率と安定性のバランスが良く、最も広く採用される | Mixtral 8x7B, Grok-1 |
| K=4 以上 | 表現力は向上するが Dense に近づき MoE の計算効率メリットが薄れる | 一部研究モデル |
Top-1 では各トークンが 1 つのエキスパートにしか送られないため、ルーティングの決定が誤っても補正が効きません。Top-2 では 2 つのエキスパートの出力を重み付き合成するため、ルーティングの曖昧さを吸収しやすく、学習が安定します。
典型的な実装では、ルーターは単一の線形変換 W_g ∈ R^(d×N) と Softmax で構成されます。
入力 x に対しスコアベクトル s = Softmax(W_g · x) を算出し、上位 K 個のインデックスを選択します。選択されたエキスパートの出力は、対応するスコアで重み付けされて合算されます。
近年は Softmax の代わりに Sigmoid ベースのルーティング(DeepSeek-V3 が採用)や、ノイズ付き Top-K(学習時にガウスノイズを加えて探索を促進)など、バリエーションが登場しています。
Top-K Routing だけではエキスパートの負荷が偏りやすいため、補助損失(Auxiliary Loss)を訓練目標に追加するのが一般的です。負荷分散損失は各エキスパートに割り当てられるトークン数の分散を最小化し、特定エキスパートへの集中(Expert Collapse)を防ぎます。
Switch Transformer では α × N × Σ(f_i × P_i) という形式の負荷分散損失を導入し、α のチューニングで学習安定性と専門化のバランスを調整します。
推論時は学習時と異なりバッチ内のトークン数が少ないため、エキスパート間の負荷分散は問題になりにくい一方、エキスパート切り替えのメモリアクセスパターンが性能に影響します。GPU 上では全エキスパートの重みを VRAM に保持し、Top-K で選択されたエキスパートのカーネルのみ実行する実装が一般的です。
A1: 理論上は Top-1 が Top-2 の約半分の FLOPs ですが、実際にはメモリアクセスのオーバーヘッドが支配的で、速度差は 20〜40% 程度に収まることが多いです。
A2: Expert Choice Routing では K を固定せず、各エキスパートが処理するトークン数を固定する逆方向のアプローチを取ります。これによりエキスパート間の負荷が完全に均一化されます。
A3: エキスパート利用率(各エキスパートが選択される頻度)とルーティングエントロピー(選択の均一性指標)を監視します。エントロピーが低すぎるとエキスパート崩壊、高すぎると専門化不足を示唆します。