Google の研究者が、無駄な言葉に脳の力を無駄にしない新しい AI 技術を開発

リンク一覧

1. 導入
2. 背景
3. 深さ混合トランスフォーマーの実装

   • 3.1. コンピューティング予算の定義

   • 3.2. トランスブロックの周りのルーティング

   • 3.3. ルーティングスキーム

   • 3.4. ルーティングの実装

   • 3.5. サンプリングと3.6. トレーニング方法

4. 結果
   • 4.1. トレーニング、等フロップス比較
   • 4.2. 自己回帰評価と 4.3. 深さと専門家の混合 (MoDE)
5. 議論と参考文献

トランスフォーマーベースの言語モデルは、入力シーケンス全体に FLOP を均一に分散します。この研究では、トランスフォーマーが代わりにシーケンス内の特定の位置に FLOP (または計算) を動的に割り当てることを学習し、モデルの深さ全体にわたってさまざまなレイヤーのシーケンスに沿った割り当てを最適化できることを示します。私たちの方法は、特定のレイヤーで自己注意と MLP 計算に参加できるトークンの数 (𝑘) を制限することで、合計計算予算を強制します。処理されるトークンは、トップ𝑘ルーティングメカニズムを使用してネットワークによって決定されます。𝑘 は事前に定義されているため、この単純な手順では、他の条件付き計算手法とは異なり、既知のテンソルサイズを持つ静的計算グラフが使用されます。ただし、𝑘 トークンの ID は流動的であるため、この方法では、時間とモデルの深さの次元全体で FLOP を不均一に消費する可能性があります。したがって、計算消費は合計で完全に予測可能ですが、トークンレベルでは動的でコンテキストに依存します。この方法でトレーニングされたモデルは、コンピューティングを動的に割り当てることを学習するだけでなく、効率的に割り当てます。これらのモデルは、同等の FLOPS とトレーニング時間のベースライン パフォーマンスと一致しますが、フォワード パスあたりの FLOPS はごくわずかで、トレーニング後のサンプリング中のステップが 50% 以上高速化されます。

1. はじめに

すべての問題を解決するのに、同じ時間や労力がかかるわけではありません。同様に、言語モデリングでは、すべてのトークンとシーケンスが、正確に予測を行うのに同じ時間や労力がかかるわけではありません。それでも、トランスフォーマー モデルは、フォワード パスでトークンごとに同じ量の計算を消費します。理想的には、トランスフォーマーは、不必要に計算を消費しないことで、合計の計算予算をより少なく使用します。

条件付き計算は、必要なときだけ計算を消費することで、全体の計算量を削減しようとする手法です (Bengio et al., 2016; Bengio, 2013; Bengio et al., 2013)。さまざまなアルゴリズムが、いつ、どれだけの計算を使用すべきかという解決策を提供しています (Ainslie et al., 2023; Bapna et al., 2020; Fedus et al., 2022)。ただし、この困難な問題の一般的な定式化は、動的な計算グラフを導入する傾向があるため、既存のハードウェア制約ではうまく機能しない可能性があります (Dehghani et al., 2018; Graves, 2016)。代わりに、最も有望な条件付き計算方法は、静的計算グラフを優先する現在のハードウェア スタックと調和し、ハードウェア使用率を最大化するように選択された既知のテンソル サイズである可能性があります。

ここでは、バニラ トランスフォーマーが使用するものよりも少なくできる静的な計算予算を使用して言語モデリングの問題を検討します。ネットワークは、各レイヤーでトークンごとに利用可能な予算から計算をどこに費やすかを決定することで、利用可能な計算を動的に割り当てる方法を学習する必要があります。私たちの実装では、合計計算はユーザーが定義し、トレーニング前に変更されず、ネットワークのオンザフライの決定の関数ではありません。したがって、メモリ フットプリントの削減やフォワード パスあたりの FLOP の削減など、ハードウェア効率の向上を事前に予測して活用できます。後で説明するように、これらの向上は全体的なパフォーマンスを犠牲にすることなく実現できます。

私たちは、ネットワークの深さ全体にわたって動的なトークン レベルのルーティング決定を行う、Mixture of Experts (MoE) トランスフォーマーに似たアプローチを活用します。MoE とは異なり、トークンに計算を適用するか (標準のトランスフォーマーの場合)、残余接続を通過させるか (変更されずに計算を節約) を選択します。また、MoE とは対照的に、このルーティングをフォワード MLP とマルチヘッド アテンションの両方に適用します。これは処理するキーとクエリにも影響するため、ルーティングでは、更新するトークンだけでなく、アテンションに使用できるトークンも決定します。この戦略を Mixture-of-Depths (MoD) と呼び、個々のトークンがトランスフォーマーの深さを通じて異なる数のレイヤーまたはブロックを通過する方法を強調しています (図 1 を参照)。

MoD 技術では、パフォーマンスと速度のトレードオフも可能です。一方では、同等のトレーニング FLOP (isoFLOP) で、同等のウォールクロック時間でトレーニングしながら、バニラ トランスフォーマーより最終ログ確率トレーニング目標を 1.5% も向上させる MoD トランスフォーマーをトレーニングできます。他方では、isoFLOP が最適なバニラ トランスフォーマーとトレーニング損失パリティを達成しながら、フォワード パスごとに FLOP の一部 (50% 以上) を使用するため、ステップが高速になる MoD トランスフォーマーをトレーニングできます。これらの結果を合わせると、MoD トランスフォーマーは、フォワード パスあたりの FLOP フットプリントが小さいにもかかわらず、シーケンスごとに同等以上のログ確率を達成できるため、インテリジェントにルーティングする (つまり、不要な計算をスキップする) ことを示唆しています。