ViTと競合するDeepMindは、スパースモデルからソフト混合エキスパートモデルに移行

大規模モデルが驚くべきパフォーマンスを示したため、モデルのサイズはモデルのパフォーマンスに影響を与える重要な要素の 1 つになりました。一般に、Transformer モデルの場合、モデルが大きくなるほどパフォーマンスは向上しますが、計算コストも増加します。最近の研究では、与えられたトレーニング コンピューティング バジェットを最適に使用するには、モデル サイズとトレーニング データを一緒に拡張する必要があることが示されています。

スパースな専門家の混合 (MoE) モデルは、より低い計算コストでモデル サイズを拡大できる有望な代替手段です。 Sparse MoE Transformer には、各入力トークンにどのモジュールを使用するかを決定するという重要な離散最適化問題があります。これらのモジュールは通常、エキスパートと呼ばれる MLP です。

トークンを専門家とうまくマッチングさせるために、線形計画法、強化学習、最適輸送など、さまざまな方法が設計されてきました。多くの場合、エキスパートの使用率のバランスを取り、割り当てられていないトークンを最小限に抑えるために、ヒューリスティックな補助損失が必要になります。これらの課題は、配布外のシナリオではさらに悪化する可能性があります。

現在、Google DeepMindの研究チームは、これらの問題の多くを解決する「Soft MoE」と呼ばれる新しい手法を提案しています。

論文アドレス: https://arxiv.org/pdf/2308.00951.pdf

トークンとエキスパート間のハード割り当てを実行するためにスパースおよび離散ルーターを使用する代わりに、Soft MoE はトークンを混合してソフト割り当てを実行します。この方法では、すべてのトークンの複数の加重平均を計算します。重みはトークンとエキスパートに依存し、各加重平均は対応するエキスパートによって処理されます。

一般的なスパース MoE アルゴリズムは通常、いくつかのルーター パラメータを学習しますが、これらのアルゴリズムのパフォーマンスはランダム固定ルーティングよりもさらに悪くなることがあります。 Soft MoE では、各ルーティング (またはミキシング) パラメータが単一の入力トークンに基づいて直接更新されるため、ルーターのトレーニング中に安定性を提供できます。研究チームはまた、トレーニング中に、入力トークンのほとんどがネットワーク内の個別のルートを同時に変更できることも観察しました。

さらに、エキスパート モジュールの数が多いとハード ルーティングが困難になる可能性があるため、ほとんどの研究では数十個のエキスパート モジュールのみをトレーニングします。対照的に、Soft MoE は数千のエキスパート モジュールに拡張され、構築によってバランスをとることができます。最後に、Soft MoE は推論中にバッチ効果の影響を受けません。

この研究では、ソフトMoE法の実際の効果を調査するために一連の実験を実施しました。実験結果によると、Soft MoE L/16 はアップストリームタスク、少数サンプルタスク、微調整において ViT H/14 に勝っています。さらに、Soft MoE L/16 はトレーニング時間が半分しかかからず、推論速度は ViT H/14 の 2 倍です。注目すべきは、Soft MoE B/16 のパラメータ数は ViT H/14 の 5.5 倍であるにもかかわらず、Soft MoE B/16 の推論速度は ViT H/14 の 5.7 倍であることです。

さらに、この研究では、ソフトルーティングを通じて学習された表現が画像とテキストの配置の利点を保持していることを実験的に示しています。

ソフト萌えモデル

アルゴリズムの説明

Soft MoE ルーティング アルゴリズムを以下の図 2 に示します。研究者は、入力トークンのシーケンスを表すために X ∈ R^m×d を使用します。ここで、m はトークンの数、d は次元です。各MoEレイヤーは、単一のトークンに適用されるn個のエキスパート関数のセット、つまり{f_i:R^d→R^d}_1:nを使用します。各エキスパートは p 個のスロットを処理します。各スロットには対応する d 次元のパラメータ ベクトルがあります。これらはこれらのパラメータを表すために使用されます。

研究者たちはスパース MoE の一般的な設計に従い、Transformer MLP ブロックのサブセットを Soft MoE ブロックに置き換えました。これは通常、MLP ブロックの後半部分を置き換えることによって行われます。時間の複雑さはエキスパートの数ではなくスロットの数に依存するため、スロットの合計数は Soft MoE レイヤーの重要なハイパーパラメータです。たとえば、スロットの数を入力シーケンスの長さと同じに設定して、同等の高密度トランスフォーマーの FLOP と一致させることができます。

ソフト萌えの特徴

まず、Soft MoE は完全に微分可能です。 Soft MoE のすべての操作は連続的かつ完全に微分可能です。ソフトマックススコアによる加重平均はソフト分類として解釈することができ、これが Soft MoE アルゴリズムの名前の由来です。対照的に、スパース MoE 法では通常、ハード分類が使用されます。

第二に、Soft MoE にはトークンのドロップやエキスパートの不均衡がありません。各スロットがすべてのトークンの加重平均で埋められるため、ソフト MoE は基本的にこれら 2 つのポイントの影響を受けません。また、ソフトマックスにより、すべての重みは厳密に正になります。

やはりソフトMoEは速いですね。その主な利点は、遅く、一般的にハードウェア アクセラレータに適していないソートやトップ k 操作を完全に回避できることです。したがって、下の図 6 に示すように、Soft MoE はほとんどのスパース MoE よりも大幅に高速です。

Soft MoE には、疎な特性と密な特性の両方が備わっています。スパース MoE のスパース性は、エキスパート パラメータが入力トークンのサブセットにのみ適用されるという事実から生じます。ただし、各スロットはすべての入力トークンの加重平均であるため、Soft MoE は技術的にはスパースではありません。そして、各入力トークンは、すべてのモデル パラメータのごく一部をアクティブにします。同様に、すべての出力トークンも、少数のすべてのスロット (およびエキスパート) に依存します。また、Soft MoE は、高密度 MoE (各エキスパートがすべての入力トークンを処理する) ではなく、各エキスパートがスロットのサブセットのみを処理することにも注意してください。

最後に、Soft MoE は順次実行されます。各入力シーケンスのすべてのトークンを結合するため、グループ サイズを 1 つの大きなシーケンスに設定するだけで済みます。各エキスパートは各入力トークンを処理するため、高レベルの特殊化の量がある程度制限される可能性があります。これはまた、Soft MoE がインスタンスごとに決定論的かつ高速にレンダリングすることを意味しますが、これは Sparse MoE の一般的なインスタンスには当てはまりません。

成し遂げる

時間の複雑さ。単一のエキスパート関数のトークンあたりのコストが O(k) であると仮定すると、ソフト MoE レイヤーの時間計算量は O(mnpd + npk) になります。各エキスパートに対して p = O (m/n) スロット、つまりトークンの数をエキスパートの数で割った値を選択すると、コストは O (m^2d + mk) に削減されます。

正規化。 Transformer では、通常、MoE レイヤーは各エンコーダー ブロックのフィードフォワード レイヤーを置き換えるために使用されます。したがって、ほとんどの最新の Transformer アーキテクチャのように事前正規化を使用する場合、MoE レイヤーへの入力は「レイヤー正規化」されます。

分散モデル。研究者たちは標準的な技術を使用して、モデルを多数のデバイスに配布しました。分散モデルは通常、モデルのコストを増加させますが、上記で導出した FLOP ベースの時間計算量分析では、これが考慮されていません。したがって、すべての実験において、研究者は FLOP だけでなく、TPUv3 チップ時間での壁時計時間も測定しました。

画像分類実験結果

研究者らは、画像分類に関する 3 種類の実験を発表しました。

• パレートフロンティアのトレーニング
• 推論時のモデルの最適化
• モデルアブレーション

研究者らは、最新バージョンに40億枚以上の画像と29,000のカテゴリが含まれる独自のデータセットであるJFT-4Bデータセットでモデルを事前トレーニングしました。

事前トレーニング中に、JFT-4B の上流検証精度 1 と ImageNet 10 ショット精度という 2 つのメトリックの評価結果が提供されます。また、ImageNet-1k (130 万枚の画像) を微調整した後の ImageNet-1k 検証セットの精度も提供されます。

研究者らは、2つの一般的なMoEルーティングアルゴリズム、Tokens ChoiceとExperts Choiceを比較しました。

パレート最適化モデルのトレーニング

研究者らは、VIT-S/8、VIT-S/16、VIT-S/32、VIT-B/16、VIT-B/32、VIT-L/16、VIT-L/32、および VIT-H/14 モデルと、それらのスパース モデルをトレーニングしました。

下の図 3a と 3b は、それぞれのトレーニング コスト/パフォーマンス パレート フロンティア上にある各カテゴリのモデルの結果を示しています。上記の両方の評価指標において、任意の FLOP または時間予算に対して、Soft MoE は密な方法やその他のスパースな方法よりも大幅に優れています。

長距離トレーニング

研究者らは、より短い実行とアブレーション実験に加えて、より長い実行（数百万ステップ）用にモデルをトレーニングし、より大規模な計算規模で Soft MoE のパフォーマンスをテストしました。

まず、研究者らは、Small から Huge までのさまざまなサイズの ViT モデルと Soft MoE モデルを 400 万ステップ実行してトレーニングしました。結果を下の図4と表2に示します。

図 4 は、JFT-4B の精度、ImageNet 10 ショットの精度、Soft MoE と ViT の ImageNet 微調整精度、および ExaFLOPS でのトレーニング コストを示しています。表 2 にすべての結果を示します。与えられた計算予算では、Soft MoE モデルは ViT モデルよりもはるかに優れたパフォーマンスを発揮します。

推論に最適化されたSoft MoEを見てみましょう。より小さなバックボーンを持つ Soft MoE がより大きな ViT モデルと競合できるという事実は、非常に低い推論コストでより高品質のモデルを取得するために、引き続き小さなバックボーンをトレーニングすることを奨励しています。

研究者らは、実行時間が長くなると、クールダウン期間が長くなる（学習率が直線的に 0 に減少する）ほど、Soft MoE のパフォーマンスが向上することを観察しました。そのため、クールダウンを 50,000 ステップから最大 500,000 ステップに増やしました。下の図 5 にこれらのモデルを示します。

結果から、1k TPUv3 日でトレーニングされた Soft MoE B/16 は、同様の時間予算でトレーニングされた ViT H/14 よりも優れており、FLOP 推論のコストは 10 倍、実時間では 5.7 倍削減されました。

ViT-H/14 のトレーニング予算を 2 倍にした場合でも (200 万ステップ、トレーニング日数 2039.8 日)、Soft MoE B/16 (1011.4 日) はほぼ同等です。

さらに、Soft MoE L/16 はすべてのモデルを大幅に上回り、推論速度は ViT H/14 のほぼ 2 倍です。

より詳しい技術的な詳細と実験結果については、原著論文を参照してください。