万億パラメータモデルをよりシンプルで効率的に——Switch Transformerの深い解読

論文のアドレス：https://arxiv.org/abs/2101.03961
著者：Google チーム

はじめに#

より強力でよりインテリジェントな人工知能モデルを追求する中で、「大きいほど良い」という法則が広く認識されるようになりました。しかし、モデルの「大きさ」は驚異的な計算コストとトレーニングの難易度を伴います。従来の深層学習モデルは、各入力を処理する際に全てのパラメータを使用するため、モデルのスケールを拡張することが非常に高価になります。

「専門家混合（Mixture of Experts, MoE）」モデルは、有望な解決策と見なされていました。これは、モデルが異なる入力に対して異なるパラメータのサブセット（つまり「専門家」）を活性化することを可能にし、膨大なパラメータを持ちながら計算コストを一定に保つことができます。しかし、MoE モデルの複雑さ、高い通信コスト、そして厄介なトレーニングの不安定性は、広範な応用を難しくしています。

Google の研究チームは、論文『Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity』でブレークスルーをもたらしました。彼らが提案した Switch Transformer アーキテクチャは、一連の優雅で効率的な設計を通じて、超大規模な言語モデルの構築とトレーニングの方法を根本的に変えました。

この記事では、Switch Transformer の核心的なアイデアを深く掘り下げます：

専門家ルーティングメカニズムをどのように簡素化したのか？
トレーニングの不安定性を克服するためにどのような革新的な技術を採用し、初めて低精度（bfloat16）で大規模なスパースモデルのトレーニングを実現したのか？
このアーキテクチャはどれほどの性能向上をもたらしたのか？

スケーリングのジレンマ：密なモデル vs スパース MoE パラダイム#

従来の密なモデル、例えば Transformer アーキテクチャでは、自己注意モジュールの計算複雑度とストレージ使用量は O (n²) であり、入力シーケンスが増えるにつれて計算操作は平方級に増加します。長文タスクにとってはコストが高く、LLM の顕著なボトルネックとなります。さらに重要なのは、密なモデル（Dense Model）のパラメータが各サンプルの計算コスト（FLOPs）と密接に結びついていることです。モデルが大きくなるほど、計算は高価になります。混合専門家（MoE）モデルは、各入力トークンに対してネットワーク内の小さなパラメータサブセット（つまり「専門家」）を動的に選択して活性化する「スパースアクティベーション」戦略を提案しました。この設計の核心的な魅力は、モデルパラメータのスケールを非常に大きく（例えば数兆レベル）拡張できる一方で、各トークンの実際の計算量を相対的に一定に保つことができるため、「モデルが大きいほど計算が高価」という呪縛を理論的に打破することができる点です。しかし、初期の MoE モデルは、モデル設計の複雑さ、専門家間の通信コストの高さ、トレーニングプロセスの不安定性など、一連の厳しい課題に直面しました。このような背景の中で、Switch Transformer が登場し、これらの課題に真正面から取り組み、MoE の巨大な潜在能力を真に解放することを目指しています。

アーキテクチャの深い解析#

Switch の革新：簡素化されたスパースルーティング#

従来の MoE では、1 つのトークンが Top k の専門家を活性化しますが、これはハイパーパラメータ k によって決まります。Switch Transformer はスパースルーティングを簡素化し、1 つのトークンが 1 つの専門家を活性化するだけです。まず、入力トークン $x\in \mathbb{R}^{1 \times d}$ に対して、ルーティングパラメータ $W_r\in \mathbb{R}^{d \times N}$ （N は専門家の数）があります：

$h(x) =W_r · x$

これにより、 $x$ に対応する各専門家のスコアが得られます。その後、softmax 関数を通じて正規化され、専門家 i へのルーティング確率が得られます：

$p_i(x) = \frac{e^{h(x)_i}}{\sum_{j=1}^{N} e^{h(x)_j}}$

最終的に MoE 層の出力は次のようになります：

$y = \sum_{i \in \mathcal{T}} p_i(x) E_i(x).$

ここで T はルーティングされた専門家を示し、E は対応する専門家によるトークンの演算操作です。

Switch Transformer の単一専門家ルーティングは、従来の Top k 専門家ルーティングに比べて計算量を削減し、専門家のバッチ処理サイズの要求を減少させ、通信コストを低下させます。

専門家の負荷管理：容量とトークン処理#

専門家容量（Expert Capacity）は、1 つの専門家が処理できるトークンの数であり、その評価基準は次の式で計算されます：

$\text{expert capacity} = \left(\frac{\text{tokens per batch}}{\text{num of experts}}\right) \times \text{capacity factor}.$

ここで $capacity factor$ は容量因子であり、トークンのオーバーフローを解決するのに役立ちます。デフォルトは 1 で、値が大きいほどバッファ能力が良くなり、より多くのトークン演算を収容できます。専門家は過剰なトークンを処理できず、残差接続（Residual Connection）を介して次の層に直接渡されます（これをトークンのドロップと呼びます）。また、より大きな容量因子はドロップされるトークンを減少させますが、計算とメモリのオーバーヘッドを増加させることに注意が必要です。

公平性の確保：微分可能な負荷均衡損失#

MoE のルーティングメカニズム（すなわちスパースアクティベーション）は、Dense モデルに比べて多くの利点がありますが、負荷の不均衡が存在します。ルーティングプロセス中に、特定の専門家が何度も選ばれる一方で、一部の専門家が無参加となる場合があり、これは N 個の専門家を設計した目的に反します。したがって、各専門家がスパースアクティベーションの演算に均等に参加できるようにするための負荷均衡戦略が必要です。Switch Transformer は、この問題を解決するために微分可能な補助負荷均衡損失関数（Auxiliary Load Balancing Loss）を提案しました。

$\text{loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i$ 、ここで $\alpha$ は損失を制御するハイパーパラメータ、N は専門家の総数、 $f_i$ は実際に専門家 i に割り当てられたトークンの割合、 $P_i$ は専門家 i がルーティングで選ばれる確率を示します。

$f_i = \frac{1}{T} \sum_{x \in \mathcal{B}} \mathbb{1}\{\text{argmax } p(x) = i\}$ ； $P_i = \frac{1}{T} \sum_{x \in \mathcal{B}} p_i(x).$

損失関数からは、損失が小さいほど負荷が均衡していることが直感的にわかります。また、 $f_i$ と $P_i$ が等しいとき、損失は $\alpha$ となります。論文では、 $\alpha$ を 0.01 に設定したときに最も良い結果が得られました。

安定で効率的な大規模トレーニングの実現#

混合精度トレーニング戦略#

著者は、bfloat16 精度でのスパースモデルのトレーニングが非常に不安定であることを指摘し、混合精度トレーニングの使用を検討しました。以前の研究では通常 float32 精度でトレーニングされている状況を参考にし、Switch Transformer は float32 と bfloat16 の混合トレーニング戦略を選択しました。ルーティング関数内で入力を float32 精度に強制変換して重要な計算（softmax など）を行いますが、高価な all-to-all 通信の前にテンソルを再度 bfloat16 に変換します。実験では、float32 精度の安定性を維持しつつ、bfloat16 に近いトレーニング速度を得られることがわかりました。

安定したトレーニングのさらなる探求：小さなパラメータ初期化#

著者は、モデルパラメータの初期化が Switch Transformer のトレーニングの安定性に重要な役割を果たすことを発見しました。段階的正規分布を使用し、平均 $\mu=0$ 、標準偏差 $\sigma = \sqrt{\frac{s}{n}}$ を設定します。ここで s はスケーリングハイパーパラメータで、値は 0.1、n はその層に接続される入力特徴の数です（重みテンソルが適用されるニューロン）。

微調整における過学習の抑制：ドロップアウト#

著者は、下流タスクのトレーニングデータが少ない場合、Switch Transformer が過学習のリスクが高いことを特に強調しています。したがって、深層学習で最も広く使用されているドロップアウトを採用しました。具体的には、Expert Layer でより高いドロップアウト値を適用し、Non Expert Layer では比較的小さなドロップアウト値を使用しました。また、T5-Base モデルと Switch-Base モデルが GLUE、CNNDM、SQuAD、SuperGLUE データセットで異なるドロップアウト値を使用して微調整した結果を比較し、最終的な結論を得ました。

モデル能力の飛躍：蒸留の応用#

著者は、大規模スパースモデルの能力を小さなパラメータ量のモデルに蒸留することを提案しました。この取り組みは、大型スパースモデルの実際のシナリオでの展開可能性を高めるために重要です。数千億から万億パラメータの巨大なモデルは、その推論コストとハードウェア要件が多くのアプリケーションにとって耐え難いものです。Switch Transformer の論文で探討された蒸留方法には以下が含まれます：

重み初期化のテクニック：研究者たちは、蒸留時にスパース教師モデルの「非専門家層」重みを密な学生モデルの対応する層に初期化することで、一定の性能向上が得られることを発見しました。これは、全体のモデルアーキテクチャが異なっていても（スパース vs 密）、多くの基本的な非専門家層（注意メカニズムの一部、埋め込み層など）が機能的に類似しており、そのトレーニング済みの重みが学生モデルにとって貴重な先行知識であるためです。
損失関数の組み合わせ：蒸留プロセスは、学生モデルの標準データセットでのパフォーマンス（ハードラベル）だけでなく、教師モデルの出力確率分布（ソフトラベル）も学習します。論文では、教師モデルの出力確率と真のラベル（ground truth）を一定の比率で混合（例えば、0.25 の教師確率 + 0.75 の真のラベル損失）することで、学生モデルの最適化目標がより良い蒸留効果を得られることが示されています。

これらの蒸留技術を通じて、Switch Transformer は大型スパースモデルの約 30% の性能向上を、パラメータ量を数十倍から近百倍に圧縮した小型密なモデルに移行することに成功しました。例えば、数十億パラメータを持つ Switch-Base モデルは、蒸留後にパラメータ量を 95% 以上削減しても、学生モデルは元の教師モデルの性能向上の近三成を保持することができます。この成果は、事前トレーニング段階の言語モデルに適用されるだけでなく、特定の下流タスク（例えば SuperGLUE）で微調整されたモデルの蒸留にも検証され、これらの先進的なスパースアーキテクチャの実用価値を大いに高めました。

卓越した成果：SOTA 性能の概要#

その革新的なアーキテクチャとトレーニング戦略により、Switch Transformer はモデルの規模とトレーニング効率において突破を達成しただけでなく、複数の自然言語処理ベンチマークテストで当時の SOTA（State-of-the-Art、最先端）または非常に競争力のある成績を収めました：

前例のない規模と効率：論文では、パラメータ量が 1.6 兆に達する Switch-C モデルと、T5-XXL（110 億パラメータ）と同等の計算量を持ちながら 3950 億パラメータを持つ Switch-XXL モデルを成功裏にトレーニングしました。さらに重要なのは、これらの巨大なスパースモデルが非常に高いトレーニング効率を示していることです。例えば、Switch-XXL は T5-XXL に対して 4 倍の事前トレーニング速度向上を実現し、T5-Base などのモデルに対しては、Switch Transformer は 7 倍の事前トレーニング加速を達成し、これらはすべて同じ計算リソースで行われました。
下流タスクでの強力なパフォーマンス：事前トレーニングされた Switch Transformer モデルは、微調整後に多くの下流タスクで強力な実力を示しました。GLUE、SuperGLUE などの包括的な言語理解ベンチマークや、SQuAD などの読解タスク、XSum などのテキスト要約タスクにおいて、Switch モデルは FLOPs にマッチした T5 密なベースラインモデルを大幅に上回りました。特に、知識集約型のクローズドブック質問応答タスク（TriviaQA、Natural Questions、WebQuestions など）において、Switch Transformer は SOTA またはそれに近いパフォーマンスを達成しました。
多言語能力の普遍的な向上：Switch Transformer アーキテクチャの利点は単一言語に限られません。101 言語を含む mC4 データセットでトレーニングされた mSwitch-Base モデルは、評価に参加したすべての言語で強力な mT5-Base ベースラインを上回り、多タスクおよび多言語学習におけるそのアーキテクチャの普遍性と効率性を示しました。
注意すべきは、論文が率直に指摘しているように、事前トレーニングの困惑度などの指標で大きな進歩を遂げたにもかかわらず、当時、これらの利点を特定の下流タスク（特に複雑な推論タスク）の SOTA 性能に完全に転換することは、最大規模のモデルにとって依然として継続的な探求の領域であるということです。

まとめと深遠な影響#

Switch Transformer の論文は、深層学習、特に大規模言語モデルの発展史において重要なマイルストーンであることは間違いありません。以下のいくつかの核心的な貢献を通じて、後続の研究方向に深く影響を与えました：

MoE アーキテクチャの簡素化と普及：単一専門家ルーティング（k=1）メカニズムを導入することで、従来の MoE の複雑さを大幅に簡素化し、理解、実装、トレーニングが容易になりました。
大規模スパースモデルのトレーニングにおける重要な課題の克服：選択的精度トレーニング、小規模パラメータ初期化、負荷均衡損失、専門家ドロップアウトなどの一連の革新的な技術を提案し、トレーニングの安定性と効率を大幅に向上させました。
スパース化の巨大な潜在能力の証明：数兆パラメータ規模に達しながら計算効率を維持するスパースアクティベーションモデルの巨大な潜在能力を、揺るぎない実験結果で示し、「大きな力で奇跡を起こす」ための新しい、より経済的な道を提供しました。
オープンソースエコシステムと後続の革新の推進：Mesh TensorFlow およびその後の T5X/JAX に基づく関連コードのオープン化は、学術界と産業界における MoE および関連するスパース技術の探求と応用を大いに促進しました。

後続の研究への影響

MoE が主流技術路線の重要な推進力に：Switch Transformer の成功により、MoE は比較的小規模な探求方向から、トップクラスの大型言語モデル（Mixtral シリーズやその他のインスパイアされたモデル）の構築における重要な技術の 1 つへと進化しました。
ルーティングメカニズムの継続的な革新を刺激：単一専門家ルーティングはシンプルで効率的ですが、より動的でインテリジェントなルーティングアルゴリズムの研究を刺激しました。例えば、モデルが専門家の容量を自適応的に学習する方法や、タスクやデータ特性に基づいてトークン配分を最適化する方法などです。
スパースモデルの微調整と展開に対する理解の深化：論文における微調整と蒸留に関する初期の探求は、これらの巨大なスパースモデルを効率的に利用し、展開する方法を指し示し、専門家特化、プルーニング、量子化などの研究を生み出しました。
大規模トレーニングの経験の普及：トレーニングの不安定性の解決、通信の最適化、混合精度の利用に関する経験は、大規模 AI モデルのトレーニング分野全体に貴重な参考を提供しました。

要するに、Switch Transformer は単なるモデルの名称ではなく、一つの設計哲学と一連の実行可能な方法論を代表しています。精巧なエンジニアリング設計とアルゴリズムの革新を通じて、膨大な知識を持つスパースな巨獣を制御し、よりシンプルで効率的な方法でますます複雑なインテリジェントアプリケーションにサービスを提供できる能力を証明しました。また、後続の DeepSeek による MoE の負荷均衡戦略の改善などの研究の基盤を築きました。