ReLU がビジュアル Transformer のソフトマックスに取って代わり、DeepMind の新しい手法でコストが急速に削減される

Transformer アーキテクチャは、現代の機械学習で広く使用されています。 Attention はトランスフォーマーのコア コンポーネントであり、トークンの確率分布を生成するソフトマックス関数が含まれています。 Softmax は、指数計算とシーケンスの長さにわたる合計を実行するためコストが高く、並列化が困難になります。

Google DeepMind は、ソフトマックス演算を、必ずしも確率分布を出力しない新しい方法に置き換えるという新しいアイデアを思いつきました。また、シーケンスの長さで割った ReLU を使用した注意は、ビジュアル Transformer で使用すると従来のソフトマックス注意に近づくか、それに匹敵する可能性があることも確認しました。

論文: https://arxiv.org/abs/2309.08586

この結果により、ReLU アテンションはシーケンス長次元で並列化できるため、従来のアテンションよりも少ない収集操作で済むため、並列化のための新しいソリューションがもたらされます。

方法

注意

アテンションは、d 次元のクエリ、キー、値 {q_i、k_i、vi_i} を 2 段階のプロセスで変換することによって機能します。

最初のステップでは、注目度の重みは次の式で得られます。

ここでϕは通常ソフトマックスです。

次に、このアテンション重みを使用して出力を計算します。この論文では、ϕ を置き換えるためにポイントごとの計算ソリューションを使用する方法を検討します。

ReLUアテンション

DeepMind は、式 1 の ϕ = softmaxが適切な代替手段であることを観察しました。彼らが使用するアテンションは、ReLU アテンションと呼ばれます。

詳細なポイントごとの注意

また、 α∈[0, 1]かつh∈{relu,relu², gelu,softplus, identity,relu6,sigmoid}となるより広い範囲の選択肢を実験的に探索します。

シーケンス長の拡張

また、シーケンスの長さ L を含む項でスケーリングすると、高い精度が達成されることも観察されました。ソフトマックスを削除しようとした以前の研究では、この拡張機能は使用されていませんでした。

ソフトマックスアテンションを使用した現在のTransformer設計では、これは必要条件ではないものの、初期化時の複雑さが この条件を維持することで、ソフトマックスを置き換えるときに他のハイパーパラメータを変更する必要性が軽減される可能性があります。

初期化時にqとkの要素はO(1)なので、それらもO(1)になります。 ReLUのような活性化関数はO(1)を維持するので、nをnの複雑度にするにはnの係数が必要です。

実験と結果

主な結果

図 1 は、ImageNet-21k トレーニングで ReLU アテンションとソフトマックス アテンションが同様にスケーリングされることを示しています。 x 軸は、実験に必要なカーネル計算時間の合計 (時間) を示します。 ReLU アテンションの主な利点は、シーケンス長次元で並列化できるため、ソフトマックス アテンションよりも必要な収集操作が少なくて済むことです。

配列長の延長の効果

図 2 は、シーケンス長拡張法の結果と、ソフトマックスの他のさまざまなポイントごとの代替方法の結果を比較しています。具体的には、softmax を relu、relu²、gelu、softplus、identity などの方法で置き換えることです。 x軸はαです。 Y 軸は、S/32、S/16、および S/8 ビジュアル Transformer モデルの精度です。通常、α が 1 に近い場合に最良の結果が得られます。明確な最適な非線形性は存在しないため、主な実験ではより高速な ReLU を使用しました。

qk-layernormの効果

主な実験では qk-layernorm を使用します。この実験では、クエリとキーがアテンションの重みを計算する前に LayerNorm に渡されます。 DeepMindは、デフォルトでqk-layernormを使用する理由として、モデルサイズを拡大縮小する際に不安定性を防ぐ必要があるためだと述べています。図 3 は、qk-layernorm を削除した場合の影響を示しています。この結果は、qk-layernorm がこれらのモデルにほとんど影響を与えないことを示唆していますが、モデルのサイズが大きくなった場合にはそうなる可能性もあります。

ドア効果の追加

ソフトマックスの除去に関するこれまでの研究では、ゲーティング ユニットを追加するアプローチが採用されていましたが、このアプローチはシーケンスの長さに応じて拡張されません。具体的には、ゲート アテンション ユニットには、出力を生成するための追加の投影があり、出力投影の前に要素ごとの乗算によって結合されます。図 4 は、ゲートの存在によってシーケンス長の拡張の必要性がなくなるかどうかを示しています。全体的に、DeepMind は、ゲートの有無にかかわらず、シーケンス長の拡張によって最高の精度が達成されたことを確認しました。また、ReLU を使用する S/8 モデルの場合、このゲーティング メカニズムにより、実験に必要なコア時間が約 9.3% 増加することにも注意してください。