Mixture of Experts (MoE) を1から解説 — Switch Transformerのルーティングと負荷分散

巨大LLMを支える疎なアーキテクチャ「Mixture of Experts」を、Switch Transformer論文に沿って前提知識ゼロから解説。ルーターの数式、エキスパート容量、負荷分散損失、学習を安定させる3つの工夫まで、一次資料の数値だけで組み立てる。

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

パラメータは増やしたい、計算は増やしたくない

ふつうのニューラルネットは、どんな入力にも同じパラメータ全部を使います。賢くするにはパラメータを増やすしかなく、増やせば1回の計算(FLOPs)も比例して重くなる — これが密（dense）モデルの宿命でした。

Mixture of Experts（MoE、専門家の混合）はこの前提を壊し、入力ごとに使うパラメータを変えます。結果は「パラメータ数は桁外れなのに、1トークンあたりの計算量は一定」という疎（sparse）なモデルです。

2021年の論文 Switch Transformer（Fedus, Zoph, Shazeer）はMoEを極限まで単純化し、総パラメータ1.6兆のモデルを安定して学習させ、同じ計算資源でT5比最大7倍の事前学習速度を報告しました（§Abstract）。本記事はこの論文を一次資料に、MoEを土台から組み立てます。

比喩: 総合病院の受付

• 患者 = トークン（文中の単語）
• 専門医 = エキスパート（それぞれ別の重みを持つ小さなネットワーク）
• 受付 = ルーター（どの患者をどの医師に回すか決める）

密モデルは「全患者を万能医ひとりが診る」病院で、医師を賢くするほど診察が遅くなります。MoEでは受付が患者を専門医に振り分けるので、医師を100人に増やしても、患者ひとりの診察は1回のまま。ここが核心です。

直感: 知識の置き場と計算を切り離す

Transformerで最もパラメータが多いのは各ブロックのFFN（フィードフォワード網）です。Switch TransformerはこれをN個のコピー（エキスパート）に置き換え、各トークンを1個だけに通します（Figure 2）。

スケーリング則はモデル・データ・計算量の3軸を扱ってきましたが、この論文は第4の軸を提案します: 「FLOPs/トークンを固定したまま、パラメータ数だけ増やす」（§2）。増える計算はルーターの $O(d_{model} \times \text{エキスパート数})$ だけで、軽量です（§3）。

仕組み(1): ルーターはsoftmaxで専門家を選ぶ

Switch以前の標準MoE（Shazeer et al. 2017、§2.1）では、トークンの表現 $x$ から重み $W_r$ でロジット $h(x) = W_r \cdot x$ を作り、softmaxで確率に変えます。

$p_i(x)$ は「トークン $x$ をエキスパート $i$ に送りたい度合い」で、N個の合計が1。Attentionでも使ったsoftmaxそのものです。つまりこの式は「各エキスパートの点数を指数に直し、全員分の合計で割って取り分にする」と言っているだけ。点数がほんの少し高いだけでも指数を通すと差が開くので、受付の「推し」ははっきり出ます。上位 $k$ 個（top-k）を選び、出力を確率で重み付けして混ぜます。

$\mathcal{T}$ は選ばれたエキスパートの集合、$E_i(x)$ はエキスパート $i$ の出力。「選ばれた専門医の診断を、受付の信頼度で加重平均する」と読めます。$p_i(x)$ を掛けるおかげでルーターにも勾配が流れ、学習できます（§2.1）。要するに、和に登場するのは選ばれたエキスパートだけで、残りの N−k 人は式に一度も出てきません。呼ばれなかった医師の人件費がかからないのと同じで、計算しない相手はコストにならない — 疎（sparse）と呼ぶ理由が、この1行に書かれています。

仕組み(2): Switchの答え — 1人だけに送る

2017年当時は「ルーターの学習には2人以上の比較が必要（$k>1$ が必須）」と考えられていました。Switch論文はこれに反して $k=1$、最上位の1エキスパートだけに送ります（§2.1）。利点は3つ: (1) ルーティング計算が減る、(2) 各エキスパートの受け持ち（後述の容量）を半分以下にできる、(3) 実装が単純になり通信コストが減る。式(2)のゲート値は残るので微分可能性も保たれます。

Table 1では同一条件（128エキスパート、TPUv3 32コア）でtop-2 MoEと直接対決し、Switchの方が速く、容量係数が小さい設定では品質でも上回りました（§2.3）。「比較なしでは学べない」という直感は実験で覆されたのです。