今回の論文
今回取り上げるのは、William Fedus、Barret Zoph、Noam Shazeer による 2022 年の論文「Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity」です。公開元は Journal of Machine Learning Research（JMLR）で、研究分野は大規模言語モデル、Mixture of Experts、疎な条件付き計算です。URL は https://www.jmlr.org/papers/v23/21-0998.html です。
この論文を選んだ理由は、今の LLM 開発でも重要な「モデルを大きくしたいが、計算量と運用コストは増やしすぎたくない」という問題に、かなり本質的な答えを出しているからです。MoE はよく知られた考え方ですが、Switch Transformer はそれを実装しやすい形まで単純化し、後の大規模モデル設計にもつながる土台を作りました。
どんな技術か
Switch Transformer は、Transformer の FFN 層を 1 つの大きな層として使うのではなく、複数の「専門家ネットワーク」に分け、各トークンごとにそのうち 1 つだけを使う技術です。
普通の Transformer では、どの入力トークンも毎回同じ重み群を通ります。そのため、モデル容量を増やしたければ、すべてのトークンに対して毎回使う重みを増やすしかありません。これは性能向上には効きますが、計算量もメモリも一緒に膨らみます。
Switch Transformer はここを変えます。モデル全体としては大量のパラメータを持ちながら、1 トークンが実際に使うのはその一部だけです。つまり「総パラメータ数は大きいが、1 回の推論で使う計算は抑える」という設計です。大雑把に言えば、巨大な専門家チームを用意しつつ、各入力では最も向いた担当者 1 人だけを呼ぶようなものです。
課題
この技術が解決しようとしている課題は、モデルを大きくすると性能は伸びやすい一方で、計算コスト、通信コスト、学習安定性が急速に厳しくなることです。
何が難しいのかというと、密な Transformer は全トークンが全層をフルに通るため、モデルを大きくするほど FLOPs もメモリもほぼ正面から増えていく点です。GPT 系でも T5 系でも、性能向上のためにモデルを太らせる戦略は有効ですが、学習時間や推論コストは無視できません。
既存の MoE 系手法は、この問題に対して「入力ごとに使う重みを変える」という方向を示していました。ただし、複数エキスパートへのルーティング、トークンの分配、デバイス間通信、負荷の偏り、学習不安定性といった実装上の難しさが大きく、広く使われるには障壁がありました。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは、単に精度が高いだけでは足りないからです。大量文書を扱う RAG、複数回 LLM を呼ぶ AI エージェント、多言語アシスタント、社内推論基盤では、モデルの賢さと同じくらい、学習効率、GPU 台数、推論あたり原価が重要です。Switch Transformer は、パラメータ数を増やす軸と、実際に使う計算量の軸を切り分けた点に価値があります。
用語解説
- Mixture of Experts（MoE）
- 複数の専門家ネットワークを用意し、入力ごとに使う専門家を切り替える考え方です。この論文の中心概念であり、「モデル全体は巨大だが、各入力で使う計算は限定する」ための基本設計になっています。
- Router
- 各トークンをどの expert に送るか決める仕組みです。Switch Transformer では、この router をできるだけ単純にし、各トークンを 1 つの expert にだけ送ることで、MoE の複雑さを大きく下げています。
- Top-1 Routing
- router が出した候補のうち、最も確率が高い expert 1 つだけを選ぶ方式です。従来の top-2 以上のルーティングより通信量と実装負荷を減らせるため、この論文の「Switch」という名前の核心になっています。
- Expert Capacity
- 各 expert が 1 バッチ内で処理できるトークン数の上限です。MoE は負荷が偏ると一部 expert にトークンが集中するので、capacity をどう決めるかが計算効率と品質の両方に効きます。
- Load Balancing Loss
- 特定の expert ばかりにトークンが偏らないようにする補助損失です。Switch Transformer は routing を単純化する一方で、この補助損失を使って各 expert がきちんと使われるようにしています。
技術の仕組み
Switch Transformer の面白さは、「MoE を使った」こと自体よりも、「MoE をかなり単純な形で回るようにした」ことにあります。ここでは基本アイデアから、実際に効いている実装上の工夫まで順に見ます。
基本アイデア
Transformer の各ブロックには、Self-Attention の後ろに FFN 層があります。通常はこの FFN が 1 つだけありますが、Switch Transformer ではこの FFN を複数の experts に置き換えます。
ただし、すべての expert を毎回使うわけではありません。各トークンについて router が expert を 1 つ選び、その expert の FFN だけを実行します。これにより、モデル全体のパラメータ数は expert 数に応じて大きく増やせますが、1 トークンが通る計算経路は限定されます。
重要なのは、dense な拡張では「モデルを大きくすると毎回の計算も増える」のに対し、Switch Transformer では「モデル全体の容量だけを増やしやすい」ことです。性能を決める要素として、FLOPs だけでなく「どれだけ多くの知識を格納できるか」という軸を強く押し出した設計だと言えます。
ルーティングの単純化
従来の MoE では、1 トークンを複数 expert に送る top-2 などの設計が使われていました。これは柔軟ですが、各トークンの出力を複数 expert から集約する必要があり、通信と実装が重くなります。
Switch Transformer はそこを割り切って、top-1 routing を採用しました。各トークンは最も確率の高い expert にだけ送られます。これにより、以下の利点が出ます。
モデル構造上の利点
- ルーティング計算が軽くなります。
- 1 トークンにつき 1 expert 分だけ計算すればよいので、expert あたりの capacity を抑えやすくなります。
- expert 間の通信や出力集約が単純になります。
この「1 トークン 1 expert」というルールが、Switch Transformer を実用寄りの MoE にした最大の工夫です。理論的に少し大胆な簡略化ですが、論文では品質を大きく落とさず、むしろ速度や安定性の面で有利だと示しています。
Expert Capacity とトークンあふれ対策
MoE で実務上すぐ問題になるのは、ある expert にトークンが集中してしまうことです。Switch Transformer では、各 expert が処理できるトークン数を tokens per batch / number of experts × capacity factor で決めます。
もしある expert に上限以上のトークンが来ると、その超過分はその層で処理されず、残差経路で次へ流されます。つまり、routing に失敗したトークンが一部「素通り」する可能性があります。capacity factor を大きくすればあふれは減りますが、空きスロットも増えて計算効率は悪くなります。
この設計は、単にメモリ節約の話ではありません。MoE では計算効率、通信量、品質が capacity の取り方でかなり変わるので、実装時には非常に重要です。論文では、負荷分散をうまく効かせることで、ドロップされるトークンを典型的には 1% 未満に抑えられたと報告しています。
Load Balancing Loss
top-1 routing は単純ですが、何も工夫しないと特定の expert ばかり使われやすくなります。そこで著者らは、router の出力確率と実際の割り当て頻度が偏りすぎないように、補助的な load balancing loss を加えています。
役割としては、「全 expert を均等に使え」と厳密に縛るというより、「一部だけが過剰に使われて capacity overflow を起こす状態を防ぐ」ことに近いです。Switch Transformer は routing を簡略化した代わりに、この補助損失で実運用上の偏りを抑えています。
学習安定化の工夫
MoE は hard routing が入るため、学習が不安定になりやすいです。特に大規模化すると、router 周辺の数値誤差や負荷偏りが問題になりやすくなります。
この論文では、router の局所計算だけを float32 にし、それ以外を bfloat16 のままにする selective precision を導入しています。全面 bfloat16 だと学習が発散した一方、この selective precision では float32 に近い品質を保ったまま、bfloat16 相当の速度をほぼ維持できたと報告されています。
また、expert 数を増やしても学習が崩れにくいように初期化や正則化も調整しています。つまり、Switch Transformer の貢献は routing のアイデアだけではなく、「大きな sparse model を実際に学習可能にしたレシピ」にもあります。
分散学習との相性
Switch Transformer は、expert を複数デバイスに分散しやすい構造です。各 core が自分の expert を持ち、router の結果に応じてトークンを all-to-all で送り合う形なので、データ並列、モデル並列、expert 並列を組み合わせてスケールできます。
この点は、単に巨大モデルを作れるというだけではありません。大規模学習基盤では、各 GPU や TPU に何を載せるかがボトルネックになります。Switch Transformer は「総パラメータ数は大きいが、各デバイスの常時アクティブな計算は限定する」ので、分散戦略と相性が良いです。
実験と結果
論文では、Switch Transformer が本当に有効かを、学習速度、下流タスク性能、多言語学習、極端な大規模化の 4 つの観点から検証しています。
何を検証したのか
主な検証ポイントは次の通りです。
- dense な T5 と比べて、同じ計算予算でより速く良い品質に到達できるか
- 下流タスクでも性能改善が残るか
- 多言語学習のようなマルチタスク環境でも効くか
- 1 兆パラメータ級まで拡張しても成立するか
評価対象は主に T5 系をベースにした比較で、Switch-Base、Switch-Large、Switch-XXL、Switch-C などのバリエーションが使われています。
データセットや評価指標
事前学習では C4、比較的大規模な多言語設定では mC4 を使っています。下流評価では GLUE、SuperGLUE、CNN/DailyMail、SQuAD などを利用しています。指標は、事前学習では negative log perplexity、下流では各タスクの一般的なスコア、要約では ROUGE 系、QA では exact match 系が中心です。
論文全体として、単一のベンチマークでたまたま勝ったというより、「事前学習効率が良いことが複数条件でも維持されるか」をかなり丁寧に見ています。
pre-training の速度と品質
最も象徴的なのは、Switch-Base と dense な T5-Base/T5-Large の比較です。論文では、同じ計算資源で Switch 系が最大 7 倍の pre-training speedup を示したと報告しています。つまり、同じ壁時計時間で見ると、dense モデルより速く良い perplexity に到達しやすいということです。
さらに大規模側では、T5-XXL と比べて、最大 1.6 兆パラメータの Switch-C が同じ計算予算で fixed perplexity に達するまで約 4 倍速いとされています。ここで重要なのは、「1.6 兆パラメータだから 1.6 兆ぶん毎回計算している」わけではなく、疎な活性化で実効計算量を抑えている点です。
下流タスクでの改善
事前学習だけ速くても、下流タスクで弱いと実用性は落ちます。この論文では、SuperGLUE でも改善が見られます。たとえば Table 8 では、223M パラメータの T5-Base が SuperGLUE 74.6 なのに対し、同 FLOPs 設定の 7.4B パラメータ Switch-Base は 81.3 を記録しています。
この差はかなり大きく、疎な expert を入れても下流性能が犠牲になっていないどころか、むしろ知識容量の増加が効いていることを示しています。単なるスピードハックではなく、モデル表現力の拡張として意味がある結果です。
多言語学習での結果
多言語設定では、mC4 上の 101 言語で mT5-Base と比較しています。論文では、Switch 系が 101 言語すべてで dense ベースラインを上回り、平均で約 5 倍の speedup、さらに 91% の言語で 4 倍以上の speedup を達成したと報告しています。
これは MoE が単一言語だけに効くわけではなく、多言語のように多様な分布をまとめて扱う場面でも役立つことを示しています。言い換えると、expert 分割は「入力の種類に応じて表現を分ける」方向と相性が良いわけです。
モデル圧縮の示唆
論文では、巨大な sparse モデルを小さな dense モデルへ distillation する実験も行っています。結果として、最大 99% 近い圧縮をしながらも、教師モデルによる改善分の約 30% を保持できるとしています。
これは、Switch Transformer をそのまま本番デプロイできない場合でも、強い teacher として使う余地があることを示します。研究用には巨大 sparse、配布用には軽量 dense という二段構えが考えやすくなります。
何に使える？
Switch Transformer は、単に「巨大モデルの歴史的マイルストーン」というだけではありません。今の開発現場でも、条件付き計算と専門化の考え方としてかなり応用余地があります。
大規模 RAG や社内検索
RAG では検索後に大量の文脈を読み、回答を生成するため、推論コストが積み上がりやすいです。Switch Transformer 系の発想は、毎回すべての重みを使わず、問い合わせや文脈に応じて一部 expert だけを使う構成と相性があります。特に多ドメイン検索では、法務、営業、製品仕様など入力タイプごとに暗黙の専門化が働く可能性があります。
AI エージェント
エージェントは、計画、検索、ツール実行後の再推論など、1 タスク中に LLM を何度も呼びます。このとき、dense モデルを単純に大きくするだけだと原価が厳しくなります。MoE 的な設計は、知識容量を増やしながら実効計算を抑える方向なので、長いワークフローを持つエージェント基盤と相性が良いです。
多言語アプリケーション
論文で 101 言語すべてに改善が出ている点は実務上かなり重要です。多言語チャット、翻訳補助、海外展開 SaaS では、全言語を 1 つの dense モデルで強化しようとするとコストが重くなります。MoE は、言語やスクリプトの違いを expert 側へ分担させやすいので、多言語最適化の設計として参考になります。
ドメイン特化モデル
金融、医療、法務、製造など、異なる知識分布を 1 モデルでまとめて扱いたい場面でも MoE の考え方は有効です。すべてを均一な dense 重みで抱え込むより、専門家を分けたほうが知識干渉を抑えやすい可能性があります。これは論文の直接検証範囲を少し超える推測ですが、実際に多ドメイン LLM ではかなり自然な応用です。
開発や事業へのヒント
この論文から得られるヒントは、「性能を上げるにはモデルを全部重くするしかない」という前提を崩せることです。特に、AI プロダクトを作る側にとっては、計算量とモデル容量を分けて考える発想が重要です。
すべてを dense にしない設計
自分でモデルを作る場合でも、すべての入力に同じ重みを毎回使う必要が本当にあるのかは見直す余地があります。RAG や社内アシスタントでも、質問の種類は偏っています。入力カテゴリごとに計算経路を分ける設計は、将来的な差別化ポイントになりえます。
推論原価の改善余地
事業として AI を載せると、最終的には 1 リクエストあたりの原価が効きます。Switch Transformer 自体をそのまま導入しなくても、「モデル容量は増やしたいが、常時計算は絞る」という発想は、speculative decoding、early exit、retrieval gating など他の最適化とも通じます。つまり、原価改善の打ち手を dense モデル一本に閉じないほうがよい、という示唆があります。
小規模チームでも活かせる考え方
小規模チームが 1 兆パラメータモデルを学習するのは現実的ではありませんが、考え方は転用できます。たとえば、小さな adapter 群を用途別に切り替える、query classifier で処理経路を分ける、複数 retrieval policy を入力ごとに選ぶ、といった設計は広い意味で条件付き計算です。Switch Transformer は、その発想を大規模モデルで極端に押し進めた例として参考になります。
今後注目すべき方向性
今後も、LLM は単に dense に巨大化するだけではなく、MoE、動的 routing、専門化、分散最適化の方向へ進む可能性が高いです。特に、同じ基盤モデルの中に専門家構造を埋め込む流れは、モデル品質だけでなくサービング設計やマルチテナント運用にも影響します。基礎論文として早めに押さえておく価値があります。
限界
Switch Transformer にも注意点はあります。
まず、計算量は抑えやすい一方で、総パラメータ数は大きくなるため、重み保存、チェックポイント管理、分散配置の複雑さは増えます。学習時には expert 並列や all-to-all 通信も必要になるので、単純な dense モデルよりインフラ要求が軽いわけではありません。
次に、routing がうまく負荷分散されないと、一部 expert にトークンが偏り、capacity overflow や品質低下が起きます。load balancing loss や capacity factor の調整が必要で、ハイパーパラメータ設計の難しさは残ります。
また、論文でも学習安定性は重要なテーマで、router 周辺だけ float32 にする selective precision のような工夫が必要でした。つまり、MoE は「理論上よさそう」でも、そのまま雑に大規模化すると不安定になりやすいです。
実運用面では、dense モデルよりデバッグが難しいこともあります。なぜこの入力でこの expert が選ばれたのか、expert の役割がどう分化したのか、偏りがどう生まれたのかを観測する必要があります。モデルの可観測性まで含めて設計しないと、保守負荷が高くなる可能性があります。
さらに、この論文は強い結果を示した一方で、すべてのタスクで dense モデルを完全に置き換えるとまでは言っていません。特に大規模 sparse model の upstream 改善が downstream にどこまで素直に効くかは、追加研究が必要だと論文中でも触れられています。
よくある質問
Q. Switch Transformer は普通の Transformer と何が違うのですか？
A. 一番大きい違いは、FFN 層を 1 つの dense 層として使うのではなく、複数 expert に分け、各トークンがそのうち 1 つだけを通る点です。これにより、総パラメータ数を増やしても、1 トークンあたりの実効計算量を大きく増やさずに済みます。
Q. なぜ expert を 1 つだけ選ぶのですか？
A. top-2 以上にすると柔軟性は上がりますが、通信、集約、実装が重くなります。Switch Transformer は top-1 routing に割り切ることで、MoE を現実的な速度と複雑さで回せるようにしました。
Q. MoE ならどれでも同じように速いのですか？
A. そうではありません。どの expert を何個置くか、capacity をどう決めるか、負荷分散をどう安定化するかで結果はかなり変わります。Switch Transformer が重要なのは、単に MoE を使ったからではなく、実装上のボトルネックを減らしたからです。
Q. この技術は推論にも効きますか？
A. はい。ただし論文の主眼は pre-training 効率です。推論でも、毎回全 expert を使わないため理屈上は効率化しやすいですが、実際の速度はルーティング実装、デバイス配置、通信コストに大きく依存します。
Q. 小規模なプロダクト開発でも学ぶ価値はありますか？
A. あります。1 兆パラメータ級をそのまま真似する必要はありませんが、「入力ごとに処理経路を変える」「全入力に同じ重みをフル適用しない」という考え方は、小規模な AI アプリや社内ツールでも十分応用可能です。
今日の学び
この論文は、モデルを大きくしたいほど計算量も重くなるという課題を扱いました。そこに対して、FFN を複数 expert に分け、各トークンを 1 つの expert にだけ送る top-1 routing を採用した Switch Transformer を提案しました。
学びとして大きいのは、モデル性能を伸ばす方法は dense な拡張だけではないことです。条件付き計算で「総容量」と「実際に使う計算」を切り分ける発想は、LLM、RAG、エージェント、推論基盤のどれを作る場合でも強いヒントになります。