Switch Transformerとは？巨大モデルを効率化するMoE

Switch Transformerは、MoE（Mixture of Experts：複数のExpertから一部だけを使う構造）のroutingを1 tokenにつき1 Expertへ単純化したTransformerです。model全体では多数のparameterを持ちながら、各tokenで実行するExpertを1つに限定できます。

ただし、「1.6兆parameterなのに計算costは小さい」とだけ覚えると誤解します。Expertのweightを保持するmemory、tokenをdevice間で移す通信、routingの偏り、Expertの容量超過は残るからです。

3行で分かるSwitch Transformer

• Switch Transformerは、Transformerの一部のFFN（Feed-Forward Network：tokenごとにhidden stateを変換する層）を複数のExpertへ置き換え、Routerが最大確率の1 Expertだけを選びます。
• Top-1選択でExpert計算と通信を単純化しつつ、選択Expertの出力へ連続値のRouter確率を掛けることでRouterを学習します。
• 効率は自動的には得られません。Expert Capacity、負荷分散loss、数値精度、分散配置を一緒に設計する必要があります。

この記事では、Switch Transformers原論文を基に、1回のforward passを数式と具体例で追います。前提となるMoE一般の構造は、前回のMoEとは？複数の専門家を使うLLM構造で解説しています。

Switch Transformerが単純化したもの

一般的なSparse MoE（入力ごとに一部のExpertだけを動かすMoE）では、Routerが上位\(k\)個のExpertを選び、複数の出力を重み付きで合成します。Switch Transformerは\(k=1\)まで絞りました。

Top-1は「Routerをなくす」設計ではありません。全Expertのscoreを計算して確率を作り、その最大値を持つExpertだけを実行します。単純化したのは、選択後のExpert計算とtoken配送です。

原論文は、従来MoEの複雑なrouting、通信cost、学習の不安定さが普及を妨げていると問題設定しました。Top-1だけで品質を保てるなら、複数Expertへtokenを送るcostを削減できます。これが「Switch」という名前の中心です。

TransformerのどこへSwitch layerが入るのか

Switch Transformerの原論文は、T5を基盤とするencoder-decoder modelです。Self-Attention全体をExpertへ分けるのではなく、Transformer blockにあるDense FFNを、Routerと複数のFFN Expertから成るSwitch layerへ置き換えます。

一つのblockを単純化すると、処理順は次のようになります。

1. Self-Attentionがtoken間の情報を混ぜる
2. 残差接続と正規化を通る
3. Routerが各tokenのhidden stateを調べる
4. 選ばれたExpert FFNがtokenを変換する
5. 元のtoken順へ戻し、残差経路と合流する

Attentionはtoken同士を関係付け、Expert FFNは各token位置のhidden stateを個別に変換します。Switchが疎にするのは主に後者です。Transformerの基本構造を理解すると、置換位置を区別しやすくなります。

なお、論文中の代表的な設定では、すべてのFFNではなく一部の層をSwitch layerへ置き換えています。「Switch Transformerでは全層がMoEになる」とは限りません。

Top-1 Routingを数式で追う

Expert数を\(N\)、一つのtokenのhidden stateを\(x\)とします。Routerは学習可能な行列\(W_r\)を使って、Expertごとのlogitを計算します。

softmaxを通すと、Expert \(i\)を選ぶ確率\(p_i(x)\)が得られます。

Switchでは最大確率のExpertを一つ選びます。

出力は、選択したExpert \(E_{i^*}\)の出力へ、そのExpertのRouter確率を掛けたものです。

たとえば4 Expertの確率が次の値だったとします。

このtokenで実行するのはE3だけで、出力は\(0.63E_3(x)\)です。E1、E2、E4のFFNは実行しません。

argmaxは離散なのに、Routerはなぜ学習できるのか

argmaxで選ぶExpert IDは離散値です。選択境界をまたがない限り、argmax自体から通常のgradientは得られません。それでもSwitchの出力には、選択された確率\(p_{i^*}(x)\)が連続値として残っています。

lossを\(L\)とすると、選択Expertに対する経路では概念的に次の項が生じます。

このgradientがsoftmaxと\(W_r\)へ戻ります。つまり、選択はhardでも、選択Expertの寄与の大きさはsoftなgate値で学習できるという構造です。

ただし、main taskのlossだけでは一部Expertへtokenが集中する可能性があります。そこで後述する補助負荷分散lossが必要になります。

tokenをExpertへ送って戻すforward pass

数式だけでなく、batch全体の処理として見るとSwitch layerの制約が分かります。

Routerは各tokenに対してExpert IDとgate値を出します。Dispatcherは同じExpertを選んだtokenをまとめ、Expertが置かれたdeviceへ送ります。各Expertは割り当てられたtokenをまとめてmatrix multiplicationし、結果を元のdeviceとtoken順へ戻します。Combinerはgate値を掛けて残差経路へ合流させます。

この並べ替えは、単一deviceならmemory上のgather/scatterです。Expertを複数deviceへ分散するとAll-to-All（各deviceが相互にdataを交換する集合通信）になり得ます。Top-1でExpert計算を減らしても、通信が遅ければ実測throughputは伸びません。

また、tokenが均等にExpertへ割り当てられる保証はありません。人気Expertだけが大量のtokenを受け取ると、他のdeviceが待たされます。ここでExpert Capacityが登場します。

Expert Capacityはtokenの受付上限

原論文のTop-1設定では、1 Expertが1 batchで受け入れるtoken数の基準は次の式です。

\(T\)はbatch内のtoken数、\(N\)はExpert数、\(c\)はCapacity Factor（均等配分に対して何倍のslotを確保するか）です。実装では\(C\)を整数へ丸めます。

16 token、4 Expertで考える

\(T=16\)、\(N=4\)、\(c=1.0\)なら、各Expertのcapacityは4 tokenです。Routerの割当数が[7, 4, 3, 2]だった場合、E1の3 tokenがoverflow（受付上限超過）します。

capacityを増やせばdropは減りますが、確保したslotが空でもtensor shapeとmemoryを消費します。小さくすれば計算は締まる一方、人気Expertへ入れなかったtokenが増えます。

Switch Transformerでは、overflow tokenはそのExpertのFFN計算を飛ばし、残差接続を通って次へ進みます。token自体がsequenceから消えるわけではありません。しかし、その層で期待したExpert変換を受けられないため、drop率が高ければ学習と品質へ影響します。

論文では実験中のdropが概ね1%未満だったと報告しています。これは同論文の設定で得られた値であり、任意のbatch size、Expert数、data、分散構成で保証される上限ではありません。

補助負荷分散lossは何を均等にするのか

Routerが一部Expertだけを選ぶと、選ばれたExpertほど学習dataを受け取り、さらに選ばれやすくなる偏りが起こり得ます。Switch Transformerはmain taskのlossへ補助lossを加えます。

Expert数を\(N\)、batch内token数を\(T\)とすると、Expert \(i\)へ実際に割り当てられたtoken比率\(f_i\)は次の値です。

一方、Expert \(i\)に対するRouter確率のbatch平均\(P_i\)は次の値です。

補助lossは、この二つの内積として定義されます。

\(f_i\)はargmax後のhardな割当、\(P_i\)はsoftmax後の連続確率です。両方がuniformなら各値は\(1/N\)になり、補助lossは基準値\(\alpha\)になります。特定Expertの\(f_i\)と\(P_i\)が同時に大きいとpenaltyが増え、同じExpertへの集中を避ける方向にRouterを導きます。

ここで\(alpha\)を大きくしすぎると、main taskより均等配分を優先しかねません。小さすぎれば偏りを抑えられません。原論文で使われた値を、別modelの万能な既定値とはみなせません。

運用ではloss値だけでなく、次も一緒に監視します。

• Expertごとの割当token数
• capacityに対する使用率
• overflow / dropped token率
• Router確率のentropy
• deviceごとのExpert処理時間とAll-to-All時間

均等なtoken数でも、tokenの計算costや通信距離が違えば実行時間は均等にならない点にも注意が必要です。

Routerだけfloat32で計算するSelective Precision

大規模学習ではbfloat16などの低precisionを使うとthroughputを上げやすくなります。しかし、softmaxを含むRouterは小さな数値差で選択Expertが変わるため、低precisionの影響を受けやすい場所です。

原論文は、model全体をfloat32へ戻すのではなく、Routerの局所計算だけfloat32へcastするSelective Precision（不安定な部分だけ高precisionで計算する方法）を提案しました。

処理境界は次のとおりです。

1. hidden stateはbfloat16でRouterへ入る
2. Router入力とweightをfloat32で計算する
3. softmaxとTop-1選択を行う
4. dispatch / combine用tensorをbfloat16へ戻す
5. Expert間のAll-to-Allはbfloat16で行う

float32の通信を避けるため、精度を上げる範囲をRouter内部へ閉じています。論文の32 Expertによる初期学習実験では、Selective Precisionはfloat32に近いqualityを保ちつつ、bfloat16と同じthroughputを記録しました。ただし、これは論文中の特定model、hardware、学習区間の比較です。

小さい初期化とExpert Dropout

論文は、Transformer既定より小さいweight初期化scaleも学習安定化に有効だったと報告しています。多数のExpertを持つ大規模modelでは、初期activationやRouterの偏りが不安定さを増幅し得るためです。

fine-tuningでは、共有層よりExpert層へ強いdropoutを入れるExpert Dropoutも検討されました。原論文の小規模downstream taskにおける探索結果であり、すべてのtaskやmodel sizeへそのまま移す設定ではありません。

Expert Parallelismでparameterをdeviceへ分ける

総parameterが巨大でも、すべてのExpertを1 deviceへ載せる必要はありません。Expert Parallelism（Expertごとに異なるdeviceへ配置する並列化）を使い、Routerの割当結果に従ってtokenを移動できます。

実際の大規模学習では、これらを組み合わせます。Switch-Cのような巨大modelで重要なのは、総parameter数の大きさだけでなく、各tokenが通るactive pathと、そのpathを支えるdevice配置です。

Expert数を増やすとmodel capacityは増えますが、次のcostも増えます。

• 全Expert weightを保持するaccelerator memory
• checkpointの保存容量と読み込み時間
• Expertを跨ぐ通信と同期
• 少数tokenしか来ないExpertのmatrix multiplication効率
• serving時のbatchingとload balanceの難しさ

そのため、parameter数を増やしてもlatencyが不変になるとは限りません。

論文の「高速化」をどう読むか

arXivのabstractは、同じ計算資源でT5-BaseやT5-Largeと比較し、最大約7倍のpretraining speedupを報告しています。ここでいう効率は、あらゆるhardwareやserving条件で7倍速いという意味ではありません。

原論文の実験を読むときは、少なくとも次の指標を分けます。

論文のSwitch-Base 64 Expertは、T5-Baseが450,000 stepで達したqualityへ60,000 stepで到達しました。これは同じ学習設定でのsample / step efficiencyの強い結果です。一方、低level kernelやnetwork topologyが違えばwall-clockの倍率は変わります。

最大規模のSwitch-Cは約1.571兆parameter、2048 Expertを持ちました。原論文が示したのは、Top-1の条件付き計算と分散配置によって、この規模を学習可能な形へ持ち込めたことです。「1.6兆parameterを毎token実行した」わけではありません。

Switch Transformerの限界

Switch TransformerはMoEの複雑さを減らしましたが、Dense Transformerを常に置き換える設計ではありません。

原論文でも、pretraining lossの改善がすべてのdownstream taskへ同じように転移したわけではありません。特に大規模modelでは、upstreamの改善量から下流taskの改善量を単純には予測できないと報告されています。

また、Switchのsparsityはparameter / Expertの活性化が疎という意味です。Attention matrixの接続を疎にして長いsequenceの計算量を減らすSparse Attentionとは別の軸です。次回扱うLong Contextでは、sequence長に伴うAttention costの問題を切り分けます。

実装や論文を読むときの確認項目

Switch系modelを比較するときは、「MoEかどうか」だけでなく次を確認すると、実際のcostを判断しやすくなります。

1. Switch layerは何層に一度入るか
2. Expert数と各ExpertのFFN次元はいくつか
3. Routerはtoken単位か、Top-1か
4. Capacity Factorとoverflowの処理は何か
5. 補助loss、Router z-loss、noiseなど何を使うか
6. Routerをどのprecisionで計算するか
7. Expertをどのdeviceへ配置し、どの通信を使うか
8. 比較する総parameter、active parameter、FLOPs、throughputの条件はそろっているか

Hugging FaceのSwitchTransformers documentationでも、num_experts、expert_capacity、router_dtype、Routerのauxiliary lossなどが別々の設定・出力として現れます。実装を触るときも、これらを一つの「MoE設定」にまとめず追うのが安全です。

まとめ

Switch Transformerの核心は、MoEをただ巨大化したことではなく、routingをTop-1へ絞って学習と分散実装を成立させる条件を整理したことです。

• Routerは全Expertの確率を出すが、実行するFFNは最大確率の1個だけ
• 選択Expertの出力へgate確率を掛けるため、Routerへgradientを戻せる
• Capacity Factorはdropと空きslotのtrade-offを調整する
• 補助loss、Selective Precision、小さい初期化でroutingと学習を安定させる
• 総parameter数、active計算、通信、実測throughputは別々に評価する

Top-1は「1個だけ選べば簡単」という省略ではありません。capacity、負荷分散、precision、Expert配置まで含めたsystemとして設計したからこそ、trillion-scaleのSparse Transformerを現実的な学習対象へ近づけました。

参考文献

• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity（JMLR）
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity（arXiv）
• SwitchTransformers model documentation（Hugging Face）