DeepMind がトランスフォーマーをアップグレードし、フォワードパスの FLOP を最大半分に削減可能

ハイブリッドの深さを導入した DeepMind の新しい設計により、トランスの効率が大幅に向上します。

言うまでもなく、Transformer の重要性は現在、多くの研究チームがこの革新的なテクノロジーの改善に取り組んでいます。 Transformer のパフォーマンスを向上させます。これにより、不必要な計算を節約できるようになります。

Transformer アーキテクチャの提案者の 1 人であり、NEAR プロトコルの共同創設者である Illiya Polosukhin 氏は、少し前の Jen-Hsun Huang 氏との会話で次のように述べています。次のステップはコンピューティングです。特定の問題にどれだけのコンピューティング リソースが費やされるかに注意を払う必要があります。」

##実際、人間はその能力を持って生まれています。コンピューティングに適応するために - 人々は問題を解決しています。問題が異なれば、当然、割り当てられる時間とエネルギーの量も異なります。

#同様のことが言語モデリングにも当てはまります。正確な予測結果を得るために、すべてのトークンとシーケンスに同じ時間やリソースを投資する必要はありません。ただし、Transformer モデルは、フォワード パスの各トークンに同じ量の計算を費やします。このことを人々は嘆くでしょう。計算のほとんどが無駄になっているのです。

#不必要な計算の実行を回避できれば、Transformer の計算予算を削減できるのが理想的です。

#条件付き計算は、必要な場合にのみ計算を実行することで、総計算量を削減する手法です。これまでに多くの研究者が、いつ計算が実行されるか、どの程度の計算が使用されるかを評価できるさまざまなアルゴリズムを提案してきました。

ただし、この困難な問題に対して、一般的に使用される解決形式は、動的な計算グラフを導入する傾向があるため、既存のハードウェアの制限にうまく対処できない可能性があります。代わりに、最も有望な条件付き計算方法は、現在のハードウェア スタックを一貫して使用し、静的計算グラフとハードウェアの最大使用率に基づいて選択された既知のテンソル サイズの使用を優先するものである可能性があります。

最近、Google DeepMind はこの問題を研究し、より少ないコンピューティング予算を使用して、Transformer で使用される計算量を削減したいと考えています。

#論文のタイトル: 深さの混合: トランスフォーマーベースの言語モデルでのコンピューティングの動的割り当て
# #論文アドレス: https://arxiv.org/pdf/2404.02258.pdf

彼らは次のように構想しました: 各層で、ネットワークは意思決定を学習する必要がある利用可能なコンピューティング バジェットを動的に割り当てるために、トークンごとに作成されます。特定の実装では、総計算量は、ネットワークが動作する際の実行決定の関数ではなく、トレーニング前にユーザーによって設定され、変更されることはありません。これにより、ハードウェア効率の向上 (メモリ フットプリントの削減や順方向パスごとの FLOP の削減など) を事前に予測して活用することができます。チームの実験では、ネットワーク全体のパフォーマンスを損なうことなく、これらの利点を達成できることが示されています。

#DeepMind のチームは、Mixed Expert (MoE) Transformer と同様のアプローチを採用しており、動的なトークンレベルのルーティング決定がネットワーク深度全体にわたって実行されます。

MoE とは異なり、ここでの選択は、トークンに計算を適用するか (標準の Transformer と同じ)、残りの接続を介してトークンをラップするかです (変更せずにそのままにして、計算を保存します)。 MoE とのもう 1 つの違いは、このルーティング メカニズムが MLP とマルチヘッド アテンションの両方に使用されることです。したがって、これはネットワークによって処理されるキーとクエリにも影響するため、ルートはどのトークンが更新されるかだけでなく、どのトークンが注目に値するかを決定します。

DeepMind は、各トークンが Transformer の深さで異なる数のレイヤーまたはモジュールを通過するという事実を強調するために、この戦略を Mixture-of-Depths (MoD) と名付けました。ここではこれを「混合の深さ」と訳します。図 1 を参照してください。

#MoD は、ユーザーがパフォーマンスと速度を比較検討できるようにサポートします。一方で、ユーザーは通常の Transformer と同じトレーニング FLOP で MoD Transformer をトレーニングでき、最終的な対数確率トレーニング ターゲットを最大 1.5% 向上させることができます。一方、MoD Transformer は、通常の Transformer と同じトレーニング損失を達成するために使用する計算量が少なく、フォワード パスあたりの FLOP が最大 50% 少なくなります。

これらの結果は、MoD Transformer がインテリジェントなルーティング (つまり、不必要な計算のスキップ) を学習できることを示しています。

Mixed Depth (MoD) Transformer の実装

要約すると、戦略は次のとおりです。

• 同等の通常の Transformer が必要とする計算量よりも少ない静的計算バジェットを設定します。これは、量を制限することによって行われます。シーケンス内の計算の数 モジュール計算 (つまり、セルフアテンション モジュールと後続の MLP) に参加できるトークンの数。たとえば、通常の Transformer ではシーケンス内のすべてのトークンがセルフ アテンション計算に参加することを許可できますが、MoD Transformer ではシーケンス内のトークンの 50% のみの使用を制限できます。
• 各トークンについて、スカラー重みを与えるルーティング アルゴリズムが各モジュールにあります。この重みは、各トークンのルーティング設定 (モジュールの計算に参加するかどうか) を表します。過去を回避するために。
• 各モジュールで、上位 k 個の最大スカラー重みを見つけます。それに対応するトークンがモジュールの計算に参加します。 k 個のトークンのみがこのモジュールの計算に参加する必要があるため、その計算グラフとテンソル サイズはトレーニング プロセス中は静的です。これらのトークンは、ルーティング アルゴリズムによって認識される動的でコンテキスト関連のトークンです。

##ルーティング オプション

チームは 2 つの学習済みルーティングを検討しましたスキーム (図 2 を参照): トークン選択タイプとエキスパート選択タイプ。

トークン選択型ルーティング スキームでは、ルーティング アルゴリズムは、計算パス全体 (MoE Transformer のエキスパート ID 全体など) で各トークンの確率分布を生成します。その後、トークンは優先パス (つまり、最も確率が高いパス) に送信され、補助損失によりすべてのトークンが同じパスに収束しないことが保証されます。トークンが可能なパス間で適切に分割されることが保証されていないため、トークン選択ルーティングでは負荷分散の問題が発生する可能性があります。

エキスパート選択ルーティングは、上記のスキームを逆にします。トークンに優先パスを選択させる代わりに、各パスはトークンの優先順位 (top-k) に基づいて上位 k 個のトークンを選択します。これにより、各パスに常に k 個のトークンが保証されるため、完全な負荷分散が保証されます。ただし、一部のトークンが複数のパスの上位 k にあり、他のトークンには対応するパスがない可能性があるため、これにより一部のトークンが過剰または過小処理される可能性があります。

DeepMind が専門家による選択的ルーティングを使用することを選択した理由は 3 つあります。

#第一に、補助的なバランスを失う必要がありません。

第 2 に、上位 k を選択する操作は配線の重みの大きさに依存するため、この配線スキームでは相対的な配線の重みを使用できるため、現在の配線の重みを決定するのに役立ちます。モジュール どのトークンが最も必要であるかを計算します。ルーティング アルゴリズムは、重みを適切に設定することで、最も重要なトークンが上位 k に含まれるように試みることができます。これは、トークン選択ルーティング スキームでは実行できないことです。特定の使用例では、本質的に null 操作である計算パスがあるため、重要なトークンを null にルーティングすることは避けるべきです。

第三に、ルーティングは 2 つのパスのみを通過するため、単一の top-k 操作でトークンを効率的に 2 つの相互排他的なセットに分割できます (それぞれがパスのセットを計算します)。 、上記の過剰または処理不足の問題に対処できます。

このルーティング スキームの具体的な実装については、元の論文を参照してください。

サンプリング

エキスパートによる選択ルーティングには多くの利点がありますが、明らかな問題もあります。 -k 操作は非因果的です。つまり、特定のトークンのルーティング重みが上位 k に入るかどうかは、その後のルーティング重みの値に依存しますが、自己回帰サンプリングを実行する場合、これらの重みを取得することはできません。

この問題を解決するために、チームは 2 つの方法をテストしました。

最初の方法は、単純な補助損失を導入することです。言語モデリングの主な目標に対するその影響は 0.2% ～ 0.3% であることが実践的に証明されていますが、これにより、モデルは自己回帰的にサンプリングします。彼らは、ルーティング アルゴリズムの出力がロジットを提供するバイナリ クロス エントロピー ロスを使用し、これらのロジットの上位 k を選択することでターゲットを提供できます (つまり、トークンが上位 k にある場合は 1、それ以外の場合は 1)。 0)。

2 番目の方法は、入力がルーティング アルゴリズム (ストップ グラディエントを含む) と同じであるが、出力がルーティング アルゴリズムと同じである小さな補助 MLP 予測器 (別のルーティング アルゴリズムのような) を導入することです。は、トークンがシーケンスの上位 k にあるかどうかの予測結果です。このアプローチは言語モデリングの目標には影響せず、実験ではこのステップの速度に大きな影響を与えないことが示されています。

これらの新しいメソッドを使用すると、ルーティング先のトークンを選択して自己回帰サンプリングを実行したり、ルーティング アルゴリズムの出力に基づいてモジュールをバイパスしたりすることができます。依存関係なし 将来のトークンに関する情報。実験結果によると、これは比較的単純な補助タスクであり、すぐに 99% の精度を達成できることがわかりました。

結果

##トレーニング、isoFLOP 比較

まず、チームは比較的小さい FLOP バジェット (6e18) でいくつかのモデルをトレーニングし、最適なハイパーパラメーターを決定しました (以下の図 3 を参照)。

全体として、MoD Transformer がベースラインの isoFLOP 曲線を右下にドラッグしていることがわかります。言い換えれば、最適な MoD Transformer は、最適なベースライン モデルよりも損失が低く、より多くのパラメーターを持ちます。この効果は幸運な結果をもたらします。MoD モデル自体は、ハイパーパラメーター設定で isoFLOP 最適ではないにもかかわらず、最適なベースライン モデルと同等またはそれよりも優れたパフォーマンスを発揮するモデルがいくつかあります (段階的には高速ですが)。たとえば、2 億 2,000 万個のパラメーターを持つ MoD バリアント (図 3 のモデル番号 3) は、isoFLOP 最適ベースライン モデル (同じく 2 億 2,000 万個のパラメーター、図 3 のモデル番号 1) よりわずかに優れていますが、この MoD バリアントはトレーニング中のステップが60% 以上高速化されます。

以下の図 4 は、合計 FLOP が 6e18、2e19、および 1e20 である場合の isoFLOP 分析を示しています。見てわかるように、FLOP バジェットが大きくなると、この傾向が続きます。

# 以下の図 5 は、インターリーブ ルーティング モジュールを使用してトレーニングされた MoD Transformer のルーティング決定を示しています。モジュールのバイパスが多数あるにもかかわらず、この MoD Transformer は通常の Transformer よりも優れたパフォーマンスを実現します。

自己回帰評価

MoD バリアントの自己回帰サンプリングも評価しましたパフォーマンス、結果以下の図 6 に示します。これらの結果は、MoD Transformer によって達成される計算量の節約がトレーニング設定に限定されないことを示しています。

Mixed Depth with Expertise (MoDE)

MoD テクノロジーは MoE と自然に連携します。モデルはいわゆる MoDE モデルに統合されます。以下の図 7 は、MoDE とそれがもたらす改善点を示しています。

MoDE には、段階的 MoDE と統合 MoDE の 2 つのバリエーションがあります。

段階的 MoDE はセルフアテンション ステップの前にルーティング バイパスまたはリーチ トークン操作を実行しますが、統合 MoDE は通常の MLP エキスパート間で「操作なし」を統合します。 MoD ルーティングを実装する専門家。前者の利点は、トークンがセルフアテンション ステップをスキップできることですが、後者の利点は、ルーティング メカニズムが単純であることです。

チームは、統合された方法で MoDE を実装する方が、専門家の能力を直接低下させ、残留ルーティングを実現するためにトークンの破棄に依存する設計よりも大幅に優れていることに気づきました。

以上がDeepMind がトランスフォーマーをアップグレードし、フォワードパスの FLOP を最大半分に削減可能の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。