万億パラメータモデルをよりシンプルで効率的に——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 など）を行いますが、高価な全対全通信の前にテンソルを再び 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% 以上削減した場合でも、学生モデルは元の教師モデルの性能向上の約 30% を保持することができます。この成果は、事前トレーニング段階の言語モデルに適用されるだけでなく、特定の下流タスク（例えば 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 の負荷均衡戦略の改善などの研究の基盤を築きました。