Google Brainの研究者たちは、より大きく、より良いものを追求し続ける中で、新たに提案したSwitch Transformer言語モデルを、計算コストを抑えながら1.6兆個のパラメータにまでスケールアップした。研究チームは、Mixture of Experts (MoE) ルーティングアルゴリズムを単純化し、データ、モデル、専門家の並列性を効率的に組み合わせ、この「法外な数のパラメータ」を可能にした。

最近の多くのシンプルなディープラーニングアーキテクチャは、より複雑なアルゴリズムよりも優れた性能を発揮していますが、このような性能向上には膨大な計算予算、膨大なデータセット、膨大なパラメータ数が必要となる。研究チームは、ディープラーニングモデルはすべての入力に対して同じパラメータを再利用する傾向があるのに対し、Mixture of Experts（MoE）モデルは異なるパラメータを使用することに注目している。研究チームは、各入力例に対するニューラルネットワークの重み（パラメータ）のサブセットのみを使用して言語モデルを大規模に訓練することに注目しており、このスパース度はMoEパラダイムを単純化するために新たに提案された技術に由来している。

ディープラーニングアーキテクチャの文脈では、MoEルーティングアルゴリズムは、モデルが複数のエキスパートネットワークの出力を組み合わせることを可能にする。このようにして、学習されたゲーティングネットワークは、エキスパートネットワークの出力を混合して最終的な出力を生成することが不可欠となる。「MoEは、言語モデリングや機械翻訳のベンチマークにおいて、最先端の結果をもたらした」と研究者は説明している。

この研究の主な貢献の1つは、単純化されたMoEパラダイムが通信と計算コストを削減したことだ。これまでに実績のあるMoE戦略では、複数のエキスパートネットワークにルーティングを行い、ルーティング関数の非自明な勾配を可能にしていましたが、提案されたモデルでは、1つのエキスパートのみを使用している。

同チームによると、提案された単純化された手法は、各デバイスのメモリと計算フットプリントを管理可能に維持しながら、デバイスの数に応じてモデルの重みが増加することを確実にしているとのことです。32個のTPUコアを使用してColossal Clean Crawled Corpus (C4)で事前学習したSwitch Transformerは、慎重に調整された密なモデルとMoEモデルの両方を上回る性能を発揮しながら、より少ない計算量を実現しました。

実験では、Switch Transformerは、Common Crawlデータセット（mC4）の多言語バリアントの101の言語で、多言語T5ベース（mT5）モデルを改良しました。Switch Transformerはまた、101言語のうち91%の言語で4倍のスピードアップが見られ、mT5ベースラインよりも平均的なプレトレーニングのスピードアップを達成しました。さらに、チームは、T5-XXLモデルの4分の1の時間で1.6兆個のパラメータを持つSwitch Transformerを事前学習することで、現在の言語モデルの規模を押し上げる可能性を実証した。