Google Brainの研究者は、自然言語処理 (NLP) AIモデルであるSwitch Transformerをオープンソース化した。このモデルは、最大1.6兆のパラメータにスケールアップし、T5 NLPモデルと比較して最大7倍のトレーニング時間を改善し、同等の精度を実現する。
チームは、arXivで公開された論文でモデルについて説明した。Switch Transformerは、混合エキスパート (mixture-of-experts: MoE) パラダイムを使用して、複数のTransformer Attention ブロックを組み合わせる。モデルのサブセットのみが特定の入力の処理に使用されるため、計算コストを一定に保ちながら、モデルパラメータの数を増やすことができる。Googleの最先端のT5 NLPモデルと比較して、Switch Transformerのベースラインバージョンは、トレーニング時間の 1/7 でターゲットとなるトレーニング前の Perplexity メトリックを達成できた。1.6兆のパラメータバージョンは、半分のデータでトレーニングを行ったにもかかわらず、Perplexity メトリックでT5-XXLを上回り、ダウンストリームNLPタスクで同等以上のパフォーマンスを発揮する。
Transformerアーキテクチャは、NLP研究に使用される主要な深層学習モデルになった。最近の取り組みは、パラメータの数で測定されるこれらのモデルのサイズを大きくすることに焦点を当てており、結果は人間のパフォーマンスを超える可能性がある。GPT-3モデルの作成者であるOpenAiのチームは、べき乗則の関係に従って、NLPのパフォーマンスが実際にパラメータの数に比例することを発見した。Switch Transformerの開発において、Google Brainチームは、トレーニング例ごとのFLOPSの数を一定に保ち、「比較的少量のデータ」でトレーニングしながら、パラメータ数を最大化しようとした。
これを達成するために、モデルは混合エキスパート（MoE）スキームを使用する。MoEは、深層学習のパイオニアでありSwitch Transformerの共同作成者であるGeoff Hinton氏を含む研究チームによって、1991年に開発された。その後、トロント大学で、そして現在はGoogle Brainで開発された。2017年、Hinton氏とGoogle Brainの同僚は、MoEを使用して、1370億パラメータのリカレント (再帰型) ニューラルネットワーク (RNN) に基づくNLPモデルを作成し、言語モデリングと機械翻訳ベンチマークで最先端の結果を達成した。
Switch Transformerは、Switch Routingと呼ばれる修正MoEアルゴリズムを使用する。複数のエキスパートをアクティブにしてそれらの出力を組み合わせる代わりに、Switch Routingは特定の入力を処理する単一のエキスパートを選択する。これにより、ルーティングの計算が簡素化され、個々のエキスパートモデルが異なるGPUデバイスでホストされるため、通信コストが削減される。ただし、このスキームの欠点の1つは、「ハード」な切り替えの決定により、特に精度の低い演算を使用する場合に、トレーニングが不安定になる可能性が高くなることだ。チームは、モデルパラメータを初期化するためのスケール係数を減らすことにより、これを軽減した。
チームは、Mesh-TensorFlow (MTF) を使用して、データとモデルの並列処理を利用してモデルをトレーニングした。さまざまなスケールでのアーキテクチャのパフォーマンスを調査するために、チームは2億2300万パラメータから1.6兆パラメータまで、さまざまなサイズのモデルをトレーニングし、「スケーリングの最も効率的なディメンション (次元)」はエキスパートの数であることがわかった。事前トレーニングおよびダウンストリームNLPタスクでのモデルのパフォーマンスを、サンプルごとに同様のFLOPを必要とするT5モデルと比較した。ベースラインサイズのSwitch Transformerモデルは、GLUE、SuperGLUE、およびSQuADベンチマークでT5を上回り、トレーニング前の時間で7倍のスピードアップを達成した。1.6兆パラメータと2048のエキスパートを備えた大規模なSwitch Transformerは、トレーニング前の Perplexity において130億パラメータのT5モデルを上回り、1/4の時間で終了した。
Redditのディスカッションで、コメント提供者は、Google BrainチームがモデルのパフォーマンスをGPT-3と比較しなかったことを指摘した。これは、OpenAIの公開結果に情報が不足しているためだと推測している。別のコメント投稿者は次のように述べている:
ハードウェア要件を犠牲にしてはいますが、精度が向上するまでの時間は驚くべきものです。これらすべてはGoogleにとって問題ではありませんが、少なくともこれまでのところ、OpenAIがこれらのモデルにあまり熱心ではない理由がわかります。
GoogleはSwitch Transformerの事前トレーニング済みモデルの重みをリリースしていないが、実装コードはGitHubで入手できる。