Mixture of Experts (MoE) を1から解説 — Switch Transformerのルーティングと負荷分散
巨大LLMを支える疎なアーキテクチャ「Mixture of Experts」を、Switch Transformer論文に沿って前提知識ゼロから解説。ルーターの数式、エキスパート容量、負荷分散損失、学習を安定させる3つの工夫まで、一次資料の数値だけで組み立てる。
要点
- Mixture of Experts(MoE)は入力ごとに使うパラメータを変え、トークンあたりの計算量を一定に保つ疎なモデルです。
- Switch Transformerは総パラメータ1.6兆のモデルを安定して学習させ、T5比で最大7倍の事前学習速度を報告しました。
- Switchはk=1で各トークンを最上位の1エキスパートだけに送り、top-2 MoEより速く、容量係数が小さい設定では品質でも上回りました。
- エキスパート容量はバッチのトークン数÷エキスパート数×容量係数で決まり、定員を超えたトークンは残差接続で素通りします。
- 学習の安定化は、ルーター内部のfloat32化、初期化スケールを1.0から0.1へ下げること、エキスパートdropoutの3つです。
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
一次資料 — この記事の根拠
論文の発表 2021-01-11→この解説の公開 2026-08-125年7か月後
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient SparsityWilliam Fedus, Barret Zoph, Noam Shazeer · 2021-01-11 · v3arXiv:2101.03961論文ページ·PDF原文の要旨(Abstract)を読む
In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers of parameters -- but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability -- we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over the T5-XXL model.
パラメータは増やしたい、計算は増やしたくない
ふつうのニューラルネットは、どんな入力にも同じパラメータ全部を使います。賢くするにはパラメータを増やすしかなく、増やせば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)。増える計算はルーターの だけで、軽量です(§3)。
仕組み(1): ルーターはsoftmaxで専門家を選ぶ
Switch以前の標準MoE(Shazeer et al. 2017、§2.1)では、トークンの表現 から重み でロジット を作り、softmaxで確率に変えます。
は「トークン をエキスパート に送りたい度合い」で、N個の合計が1。Attentionでも使ったsoftmaxそのものです。つまりこの式は「各エキスパートの点数を指数に直し、全員分の合計で割って取り分にする」と言っているだけ。点数がほんの少し高いだけでも指数を通すと差が開くので、受付の「推し」ははっきり出ます。上位 個(top-k)を選び、出力を確率で重み付けして混ぜます。
は選ばれたエキスパートの集合、 はエキスパート の出力。「選ばれた専門医の診断を、受付の信頼度で加重平均する」と読めます。 を掛けるおかげでルーターにも勾配が流れ、学習できます(§2.1)。要するに、和に登場するのは選ばれたエキスパートだけで、残りの N−k 人は式に一度も出てきません。呼ばれなかった医師の人件費がかからないのと同じで、計算しない相手はコストにならない — 疎(sparse)と呼ぶ理由が、この1行に書かれています。
仕組み(2): Switchの答え — 1人だけに送る
2017年当時は「ルーターの学習には2人以上の比較が必要( が必須)」と考えられていました。Switch論文はこれに反して 、最上位の1エキスパートだけに送ります(§2.1)。利点は3つ: (1) ルーティング計算が減る、(2) 各エキスパートの受け持ち(後述の容量)を半分以下にできる、(3) 実装が単純になり通信コストが減る。式(2)のゲート値は残るので微分可能性も保たれます。
Table 1では同一条件(128エキスパート、TPUv3 32コア)でtop-2 MoEと直接対決し、Switchの方が速く、容量係数が小さい設定では品質でも上回りました(§2.3)。「比較なしでは学べない」という直感は実験で覆されたのです。
エキスパートには「定員」がある
TPUではテンソルの形を静的に決める必要があるため、各エキスパートが処理するトークン数(エキスパート容量)を事前に固定します(§2.2)。
バッチのトークンを均等割りした数に、容量係数(capacity factor)という余裕率を掛けたものが定員です。1.0なら「完全均等ならぴったり」、1.25なら25%の空席を確保します。つまり右辺は「1バッチで来る患者数 ÷ 医師の人数 × 余裕率」でしかありません。ルーターの確率もトークンの中身も式に出てこない、ということは、定員は入力を見る前に決め打ちされる数字だということです。
ルーティングが偏って定員を超えたトークンは、計算をスキップされ残差接続で素通りします(dropped tokens、§2.2)。係数を上げれば素通りは減るが、空席分の計算とメモリが無駄になる。論文は素通り率を低く保つことがスケーリングに重要としつつ、Switchは小さい係数(1.0〜1.25)ほど速度面で有利と報告しています(§2.3)。
負荷分散損失: 人気の専門医問題を損失で解く
放っておくとルーターは、学習初期にたまたま良かった少数のエキスパートばかり選びます。これを防ぐため本来の損失に補助損失を足します(§2.2)。N個のエキスパート、T個のトークンのバッチに対して、
は「実際にエキスパート へ送られたトークンの割合」、 は「ルーターが に割り当てた確率の平均」。両方が均等()のとき最小になる、つまり偏るほど罰金が増える損失です。要するにこの式は、同じエキスパートについて「実際に来た患者の割合」と「受付が送りたがった割合」を掛け合わせ、全員分を足しているだけ。片方だけ大きい医師より、両方とも大きい医師がいるほうが積は跳ね上がる — 人気が一箇所に集まるほど支払いが重くなる、という設計です。(argmax)は微分できませんが が微分できるので勾配は流れます。 倍はエキスパート数を変えても損失のスケールを保つため。係数は が使われ、〜 の掃引から「負荷を素早く均しつつ本来の損失を邪魔しない」値として選ばれました(§2.2)。
numpy擬似コードなら10行です。
import numpy as np
def switch_route(x, W_r): # x: (T, d_model)
logits = x @ W_r # (T, N) 各エキスパートの点数
probs = softmax(logits) # 式(1): 割り当て確率
expert = probs.argmax(-1) # 式(2)をk=1に: 送り先は1人
gate = probs.max(-1) # 出力に掛けるゲート値
f = np.bincount(expert, minlength=N) / len(x) # 実際の負荷
P = probs.mean(0) # 確率の平均
aux_loss = alpha * N * (f * P).sum() # 式(4)
return expert, gate, aux_loss
学習を安定させる3つの工夫(§2.4)
1. 選択的精度。 bfloat16のままルーターを回すとsoftmaxが暴れて発散します(Table 2: 全bfloat16は-3.780で発散、全float32は-1.718だが1160例/秒と低速)。ルーター内部だけfloat32に上げると、品質-1.716と速度1390例/秒を両立できました。数値形式は数値フォーマット入門参照。
2. 初期化を10分の1に。 初期化スケールを から に下げると、32エキスパートモデルの初期品質が-3.60から-2.72へ改善し、シード間のばらつきは0.68から0.01へ激減(Table 3)。223Mから1兆超まで同じ設定で通用したと報告されています。
3. エキスパートだけ強くdropout。 パラメータが多いぶん小規模タスクで過学習しやすいため、非エキスパート層0.1・エキスパート層0.4の「エキスパートdropout」でGLUE 85.2など4タスク最良を達成(Table 4)。
で、どれだけ速くなったのか
論文に書かれた数値だけを並べます。
- FLOPs一致のT5-Baseに対し、64エキスパートのSwitch-Baseはステップ基準7.5倍、実時間でも7分の1で同品質に到達(§3.1–3.2)
- 3.5倍のFLOPsを使うT5-Largeと比べても2.5倍の速度優位(§3.3)
- 101言語の多言語事前学習で全言語が改善、平均5倍・91%の言語で4倍以上の高速化(§4.3)
- Switch-C(1.6兆パラメータ・2048エキスパート)はT5-XXL比で4倍高速に同品質へ(§5.6)
- 蒸留でパラメータを95%以上削っても品質向上分の約30%を密モデルに残せる(§4.2)
- スパコンがなくてもエキスパート2個から改善が出る(§7、Appendix D)
論文の限界と、まだ分かっていないこと
- 最大級モデルの不安定性は未解決。 FLOPsの大きいSwitch-XXL(395B)は学習が時々不安定で、対策でも解決しきれず予定の100万ステップを完走していません(§5.6、§8)。
- 事前学習の良さが下流タスクへ満額で反映されない。 Switch-Cは1.6兆パラメータでもSQuAD 87.7と、小さいSwitch-XXLの89.6に劣ります。FLOPs・パラメータ数・ファインチューニング品質の関係は「よく理解されていない」と論文自身が明記(§8)。
- Attention層への適用は見送り。 QKV行列のSwitch化は品質が上がるもののbfloat16で発散(Appendix A)。
- 素通りトークンの救済策は効かず。 あふれた分を第2候補へ回すNo-Token-Left-Behindに経験的利得なし(Appendix B)。
- 実験はTPU+Mesh-TensorFlow前提の設計で、比較対象も当時のT5系である点は割り引いて読むべきです。
現場ではこう使う
ここからは論文を離れた実務の話です(筆者の解釈を含みます)。
誰がいつ触るか。 LLMの事前学習を担うMLエンジニアのアーキテクチャ選定時と、推論基盤エンジニアのサービング設計時。公開モデルの諸元で「総パラメータ数」と「アクティブパラメータ数」が別々に書かれていたら、それが本記事の 個のエキスパートと 個選択の話です。
触るパラメータ名。 論文Appendix Fの擬似コードの名前がそのまま実装の語彙です: num_experts、capacity_factor、expert_capacity、補助損失係数 、エキスパートdropout率。学習ログでは素通り率と負荷の偏りを必ず監視します。論文の既定値(、容量係数1.0〜1.25、ルーターfloat32)は今も出発点として引用されます。
事故になる落とし穴。
- ルーターを半精度のまま回す — Table 2の「発散」は他人事ではありません。混合精度でもルーターのsoftmaxだけはfloat32に。
- 「計算量一定」をメモリ一定と誤読する — FLOPsは増えなくても重みは全エキスパート分メモリに載り、サービングではデバイス間の全対全通信も発生します(§5.4)。前提知識は分散学習入門へ。
- 小データのファインチューニングに密モデルの感覚を持ち込む — 過学習が早いので、エキスパートdropoutのような強い正則化を先に検討します。
面接で問われたら。「MoEはなぜ計算が増えない?」→ トークンが通るエキスパートは1個(Switch)で、ルーターは だから。「なぜ普及が遅れた?」→ 複雑さ・通信コスト・学習不安定の3点で、Switchはその簡略化(k=1)・通信削減・安定化をワンセットで示した論文だから — 式(4)の負荷分散まで話せれば十分です。
まとめ
- MoEは入力ごとにパラメータを切り替え、計算量一定のままパラメータ数という第4の軸をスケールさせる
- Switchの核心はk=1ルーティング。比較対象なしでも学習でき、速度・単純さ・通信で勝る
- 実装を支えるのはエキスパート容量と負荷分散損失という2つの地味な仕組み
- 安定化の三種の神器: ルーターだけfloat32、初期化0.1倍、エキスパートdropout
- 限界も明確: 最大級モデルの不安定性と、事前学習→下流タスクの反映問題は未解決
Transformer本体をまだ読んでいない方はAttention機構を1から理解するから、巨大モデルを動かす側の事情はKVキャッシュから辿るのがおすすめです。
コメント
コメントにはログインが必要です