Google Brainは、T5の最大7倍の事前トレーニング速度を備えた簡素化されたスパースアーキテクチャを提案しています。

先ほど、Google Brainのシニア研究科学者であるBarret Zoph氏が、言語モデルのパラメータ数を1.6兆個（GPT-3は1750億個）まで拡張できる「Switch Transformer」と呼ばれる簡素化されたスパースアーキテクチャを設計したと投稿しました。同じコンピューティング リソースを使用すると、Switch Transformer のトレーニング速度は T5 モデルの 4 ～ 7 倍になります。

ディープラーニングの分野では、モデルはすべての入力に対して同じパラメータを再利用することがよくあります。例外は、各入力例に対して異なるパラメータを選択するMixture of Experts (MoE) モデルです。その結果、スパース活性化モデルが生成されます。パラメータの数は膨大ですが、計算コストは​​一定です。

現在、MoEモデルは機械翻訳の分野で目覚ましい成果を上げていますが、モデルの複雑性の高さ、通信コストの高さ、トレーニングの不安定さなどにより、その広範な応用には一定の障害があります。

これらの問題を解決するために、Google Brain の研究者は Switch Transformer を提案しました。 Switch Transformerの設計では、MoEのルーティングアルゴリズムを簡素化し、直感的に操作できる改良モデルを設計しました。新しいモデルでは、通信コストと計算コストが大幅に削減されました。さらに、彼らが提案したトレーニング手法はトレーニングの安定性を向上させ、大規模なスパースモデルも低精度 (bfloat16) でトレーニングできることを初めて示しました。

論文リンク: https://arxiv.org/pdf/2101.03961.pdf

コードリンク:

https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py

研究者らは、新しいモデルを T5-Base および T5-Large と比較しました。その結果、同じコンピューティング リソースで、新しいモデルは事前トレーニングの速度が最大 7 倍向上したことが示されました。この改善は多言語設定にも拡張可能で、mT5-Base バージョンに対する新モデルのパフォーマンス向上は 101 言語すべてで測定されました。

最後に、研究者らは Colossal Clean Crawled Corpus で事前トレーニングを実行し、言語モデルのパラメータの数を数兆に増やし、T5-XXL モデルと比較して 4 倍の高速化を達成しました。

研究者らはまた、この研究は規模に焦点を当てているが、スイッチトランスフォーマーアーキテクチャはスーパーコンピューターのある環境で利点があるだけでなく、コンピューティングコアが少数しかないコンピューターでも効果的であることを示していると述べた。さらに、研究者が設計した大規模なスパース モデルは、スパース モデルの 30% の品質向上を維持しながら、小型で高密度のバージョンに凝縮できます。

スイッチトランスの設計原理

Switch Transformer の主な設計原則は、Transformer モデルのパラメータの数をシンプルかつ計算効率の高い方法で最大化することです。スケーリングの利点は Kaplan ら (2020) によって徹底的に研究されており、モデルとデータセットのサイズ、および計算予算に応じてべき乗則スケーリングが行われることが明らかになっています。重要なのは、この研究では、計算上最適なアプローチとして、比較的少量のデータで大規模なモデルをトレーニングすることを提唱していることです。

これらに基づいて、研究者は例ごとの FLOP を変更せずにパラメータの数を増やしました。パラメータの数は実行される計算の総量とは無関係であり、独立してスケーリングできる重要なコンポーネントであると想定しています。そこで研究者たちは、密行列乗算用に設計された GPU や TPU などのハードウェアを効率的に利用できる、スパースに活性化されたモデルを設計することで、この目標を達成しました。

分散トレーニング設定では、モデルのスパース アクティベーション レイヤーに、さまざまなデバイス間で一意の重みが割り当てられます。その結果、デバイスの数が増えるにつれてモデルの重みが調整され、各デバイスのメモリと計算スペースが管理可能な範囲に維持されます。

スイッチトランスフォーマーのエンコーダーブロックを図 2 に示します。

スパースルーティングの簡素化

Shazeerら（2017）は、トークン表現xを入力として受け取り、最も決定的な上位k人の専門家（N人の専門家^N_i=1から選択）に送信する自然言語MoEレイヤーを提案しました。彼らは、ルーティング関数の意味のある勾配を得るためには、トークン表現を k>1 の専門家に送信する必要があると仮定しています。彼らは、少なくとも 2 人の専門家の能力を比較しなければルーティングを学習することはできないと主張しています。

これらのアイデアとは異なり、この Google Brain の研究では簡略化された戦略を採用し、トークン表現を 1 人の専門家にのみ送信します。研究により、この簡素化戦略によりモデルの品質が維持され、ルーティング計算が削減され、パフォーマンスが向上することが示されています。研究者たちはこの k=1 戦略をスイッチ レイヤーと呼んでいます。

下の図 3 は、異なるエキスパート容量係数を使用したルーティングの例を示しています。

効率的なスパースルーティング

研究者らは、TensorFlow に類似したセマンティクスと API を備え、効率的な分散データとモデル並列アーキテクチャを促進する Mesh-Tensorflow ライブラリ (MTF) を使用しました。研究者たちは、静的なサイズ設定を必要とする TPU を念頭に置いてモデルを設計しました。

分散スイッチトランスフォーマーの実装: すべてのテンソルの形状はコンパイル時に静的に決定されますが、トレーニングおよび推論中のルーティングの決定により計算は動的になります。これを考慮すると、専門家の能力をどのように設定するかという重要な技術的問題が発生します。

エキスパート容量 (各エキスパートによって計算されるトークンの数) は、バッチあたりのトークン数をエキスパートの数で割り、容量係数を掛けて計算されます。式（３）に示すように：

エキスパートに送信されるトークンが多すぎる場合（以下、「破棄されたトークン」と呼びます）、計算はスキップされ、トークン表現は残余接続を介して次のレイヤーに直接渡されます。ただし、エキスパート容量を増やすことには欠点がないわけではありません。値が高すぎると、計算とメモリが無駄になります。トレードオフは上の図 3 に示されています。

経験的研究により、破棄されるトークンの割合を低く保つことがスパース エキスパート モデルのスケーラビリティにとって重要であることがわかっています。設計上の決定がモデルの品質と速度に与える影響を以下の表 1 に示します。

スイッチトランス

研究者らはまず、マスク言語モデリングタスクを使用して、Colossal Clean Crawled Corpus (C4) データセットで Switch Transformer を事前トレーニングし、テストしました。事前トレーニング設定では、Raffel ら (2019) によって特定された最適なスキームに従い、トークンの 15% を削除し、マスキング シーケンスを単一のセンチネル トークンに置き換えます。モデルのパフォーマンスを比較するために、研究者は負の対数パープレキシティの結果を示しています。

スイッチトランスフォーマーとMoEトランスフォーマーの比較結果を以下の表1に示します。結果は、Switch Transformer が、速度品質の面で慎重に調整された高密度モデルや MoE Transformer よりも優れており、固定量の計算と実時間で最高の結果を達成することを示しています。Switch Transformer の計算フットプリントは MoE Transformer よりも小さく、Switch Transformer は低い容量係数 (1.0、1.25) でより優れたパフォーマンスを発揮します。

トレーニングと微調整を改善するためのヒント

通常の Transformer モデルと比較すると、Sparse Expert モデルのトレーニングはより困難になる可能性があります。これらのレイヤーでのハードスイッチング (ルーティング) の決定はすべて、モデルの不安定性につながる可能性があります。さらに、bfloat16 のような低精度の形式では、ルーターのソフトマックス計算の問題が悪化する可能性があります。研究者たちは、トレーニングの困難を克服し、安定したスケーラブルなトレーニングを実現するために、いくつかの手法を採用しました。

大規模スパースモデルによる選択的精度

安定性のためのパラメータ初期化の縮小

大規模スパースモデルの正規化

事前トレーニングのスケーラビリティ

事前トレーニング中に、研究者は Switch Transformer のスケーラビリティを研究しました。その過程で、彼らは計算能力とデータの両方が無制限となるメカニズムを検討しました。データ制限を回避するために、研究者は 1800 億のターゲット トークンを含む大規模な C4 データベースを使用しました。彼らは、収穫逓減が観察されるまで訓練を続けました。

専門家の数は、モデルをスケーリングするための最も効果的な次元です。モデルはエキスパートの総数に関係なく、トークンごとに 1 人のエキスパートのみを選択するため、エキスパートの数を増やしても計算コストは​​ほとんど変わりません。ルーターはより多くのエキスパートに基づいて確率分布を計算する必要がありますが、これは軽量な計算コスト O(d_model × num expert) です。このうち、d_model はレイヤー間で渡されるトークンの埋め込み次元です。このセクションでは、固定の計算コストでステップ数と時間に基づいてスケーラビリティを検討します。

ステップ数に基づくスケーラビリティ

下の図 4 は、トレーニング ステップの数が一定で、エキスパートの数が増加した場合の複数のモデルのスケーラビリティの向上を示しています。傾向が見られます: パラメーター (エキスパート) を増やすと、トークンあたりの FLOPS を一定に保ちながらトレーニング速度を上げることができます。

時間ベースのスケーラビリティ

上記の図 4 に示すように、専門家の数が増えるにつれて、モデルのパフォーマンスは向上し続けます。このモデルはトークンあたりの FLOPS がベースラインとほぼ同じですが、スイッチ トランスフォーマーによってデバイス間の通信コストが増加し、ルーティング メカニズムによって計算オーバーヘッドも増加します。したがって、ステップベースの設定で観察されるサンプル効率の向上は、必ずしも時間の経過とともにモデル品質の向上につながるとは限りません。こうなると、次のような疑問が湧いてきます。

トレーニング時間と計算コストが固定されている場合、密なモデルをトレーニングするべきでしょうか、それとも疎なモデルをトレーニングするべきでしょうか?

以下の図 5 と 6 はこの問題を説明しています。図 5 は、トレーニング時間が長くなるにつれて、事前トレーニング済みモデルの品質がどのように変化するかを示しています。トレーニング時間と計算コストの両方が固定されている場合、Switch Transformer の速度の利点は非常に明白です。この設定では、Switch-Base 64 エキスパート モデルのトレーニング時間は、T5-Base モデルの 1/7 のみで、同様の難解さを達成できます。

スケーラビリティと大規模高密度モデル

上記の実験は、無制限の計算を備えた密なモデルが Switch によって上回られたことを示しています。図 6 は別のケースを考えています。計算リソースをより大きな密なモデルに分割するとどうなるでしょうか。

この問題を検証するために、研究者は Switch-Base とより強力なベースラインである T5-Large を比較しました。実験結果によると、T5-Large は Switch-Base よりもトークンあたり 3.5 倍の FLOP を使用しますが、後者は依然としてサンプル効率が高く、前者よりも 2.5 倍高速です。さらに、Switch モデル (Switch-Large) が T5-Large と同じ FLOP で設計されている場合、上記の改善はさらに顕著になります。

下流タスクの結果

微調整

ここで使用されるベースライン メソッドは、223M のパラメータを持つ高度に調整された T5-Base と 739M のパラメータを持つ T5-Large です。これら 2 つのモデルに対して、本研究の著者らは、より多くのパラメータを持つ FLOP マッチング スイッチ トランスフォーマーを設計しました。

Switch Transformer は、複数の自然言語タスクで大幅なパフォーマンスの向上をもたらします。最も顕著なのは SuperGLUE で、FLOP マッチング Switch Transformer により、T5-Base と T5-Large のパフォーマンスがそれぞれ 4.4% と 2% 向上しています。Winogrande、closed book Trivia QA、XSum でも同様の結果が得られています。パフォーマンスの向上が見られなかった唯一のベンチマークは、AI2 推論チャレンジ (ARC) データセットでした。ARC チャレンジ データセットでは、T5-Base が Switch-Base を上回り、ARC イージー データセットでは、T5-Large が Switch-Large を上回りました。

全体として、Switch Transformer モデルは、複数の推論および知識タスクで大幅なパフォーマンスの向上をもたらします。これは、モデル アーキテクチャが事前トレーニングに役立つだけでなく、微調整を通じて下流のタスクに品質の改善を転送できることを示しています。

蒸留

数十億または数兆のパラメータを持つ大規模なニューラル ネットワークを展開するのは簡単ではありません。この目的のために、本論文では、大規模なスパース モデルを小規模な密なモデルに抽出する方法を研究します。以下の表 7 は、この研究で使用された蒸留技術を示しています。

研究者らは、表 7 の最適な蒸留手法を使用して、複数のスパース モデルを密なモデルに蒸留しました。彼らは、専門家の数の違いによりパラメータサイズが 11 億から 147 億に及ぶ Switch-Base モデルを抽出しました。この研究では、11億のパラメータを持つモデルを82%圧縮しながら、パフォーマンスを37%向上させることができました。最も極端なケースでは、モデルは 99% 圧縮され、パフォーマンスは 28% 向上しました。

最後に、研究者たちは微調整されたスパースモデルを高密度モデルに精製しました。以下の表 9 は、74 億のパラメータを持つ Switch-Base モデル (SuperGLUE タスク用に微調整済み) - 223M T5-Base の蒸留結果を示しています。事前トレーニングの結果と同様に、蒸留モデルでも 30% のパフォーマンス向上が維持されています。これにより、微調整タスクのための特定の専門家を識別して抽出し、モデルの圧縮を向上させることができます。

多言語学習

下流タスク実験では、研究者らはモデルの品質と速度のトレードオフを測定し、モデルは 101 種類の言語で事前トレーニングされました。下の図 7 は、すべての言語における mT5-Base と比較した Switch T5 Base モデルの品質向上 (負の対数パープレキシティ) を示しています。両モデルに対して 100 万ステップの事前トレーニングを行った後、Switch Transformer の最終的な負の対数パープレキシティはベースラインと比較して改善されました。

下の図 8 は、mT5-Base と比較した Switch Transformer の各ステップの高速化を示しています。前者は平均 5 倍の高速化を達成し、言語の 91% で少なくとも 4 倍の高速化を達成しています。これは、Switch Transformer が効率的なマルチタスクおよび多言語学習者であることを示しています。

データ、モデル、エキスパートの並列処理を使用してモデルを設計する

専門家の数を恣意的に増やすと、収穫逓減の問題が発生します (上の図 4 を参照)。この研究では、データ、モデル、専門家の並列処理を組み合わせるトレードオフを含む補完的な拡張戦略をいくつか紹介します。

データ、モデル、専門家を並行して組み合わせて、1兆パラメータのモデルを構築する

Switch Transformer の設計中、研究者はトークンあたりの FLOP とパラメータの数のバランスを取ろうとしました。エキスパートの数が増えると、パラメータの数も増えますが、トークンあたりの FLOP は変わりません。 FLOP を増やすには、d_ff 次元を増やす必要があります (これにより、パラメータの数も増えますが、比較的少なくなります)。これはトレードオフです。d_ff 次元を増やすと各コアのメモリが不足するため、m を増やす必要があります。しかし、コア数 N は N = n × m に固定されているため、n を減らす必要があり、それはより小さなバッチ サイズを使用することを意味します。

モデル並列処理とエキスパート並列処理を組み合わせた後、適切なエキスパートにトークンを送信し、モデル並列処理によって発生する内部の全削減通信によって、全対全通信コストが発生します。これら 3 つのアプローチを組み合わせると、コアあたりの FLOP、通信コスト、メモリのバランスをとることが非常に複雑になります。

この研究では、データ、モデル、エキスパートの並列処理を組み合わせて、それぞれ3,950 億のパラメータと1.6 兆のパラメータを持つ 2 つの大規模な Switch Transformer モデルを設計し、上流の事前トレーニング済み言語モデルと下流の微調整タスクにおけるこれらのモデルのパフォーマンスを調査します。さまざまなモデルのパラメータ数、FLOP、ハイパーパラメータを以下の表 10 に示します。

スイッチトランスフォーマーについてはまだ多くの疑問が残っています

論文の最後で、Google Brain の研究者は、Switch Transformer とスパース エキスパート モデル (ここでのスパースは重みを指し、注意パターンを指しません) に関するいくつかの問題を調査しました。

質問 1: スイッチトランスフォーマーの方がパフォーマンスは優れていますか?理由はパラメータの数が膨大だからでしょうか?

パフォーマンスは確かに向上していますが、その理由はパラメータの数ではなく、設計にあります。パラメータはニューラル言語モデルのスケーリングに役立ち、モデルが大きいほどパフォーマンスが向上します。しかし、この研究で提案されたモデルは、同じコンピューティングリソースを使用しながら、より高いサンプル効率を備えています。

質問 2: スーパーコンピューターがなくてもこの方法を使用できますか?

この論文は非常に大規模なモデルに焦点を当てていますが、研究者たちは、パフォーマンスを向上させながら、一般的な GPU または TPU のメモリ制限内に簡単に収まる 2 人のエキスパートを備えたモデルを発見しました。したがって、研究者たちは、この技術は小規模な設定で使用できると考えています。

質問 3: 速度と精度のパレート曲線では、スパース モデルの方が密なモデルよりもパフォーマンスが優れていますか?

はい。さまざまなモデル サイズにわたって、スパース モデルは、ステップあたりおよび実時間の両方の点で密なモデルよりも優れています。制御された実験では、一定の計算量と時間では、スパース モデルが密なモデルよりも優れていることが示されています。

質問 4: 1 兆パラメータのモデルを展開できません。モデルを縮小できますか?

1兆パラメータモデルの品質を完全に維持することはできませんが、スパースモデルを密なモデルに蒸留することで、10〜100倍の圧縮率を達成し、エキスパートモデルよりも約30％の品質向上が得られます。

質問 5: 並列高密度モデルの代わりにスイッチトランスフォーマーを使用するのはなぜですか?

時間の観点から見ると、Switch Transformer は、シャードされたパラメータを使用する密なモデルよりもはるかに効率的です。同時に、この選択は相互に排他的ではなく、モデルの並列処理は Switch Transformer でも使用できます。これにより、トークンあたりの FLOP を増やすことができますが、従来のモデルの並列処理の速度低下にもつながります。

質問 6: スパース モデルが広く使用されていないのはなぜですか?

スパースモデルを試すというアイデアは、密なモデルの圧倒的な成功によって阻止されました。さらに、スパース モデルは、モデルの複雑さ、トレーニングの難しさ、通信コストなど、さまざまな問題に直面します。スイッチトランスフォーマーはこれらの問題を軽減します。