新しい近似注意メカニズム HyperAttention: 長いコンテキストに適しており、LLM 推論が 50% 高速化します

トランスフォーマーは、自然言語処理、コンピューター ビジョン、時系列予測などの分野におけるさまざまな学習タスクに効果的に適用されてきました。これらのモデルは成功しているものの、アテンション レイヤーの正確な計算の実行時間とメモリの複雑さが 2 次 (シーケンス長) となるため、スケーラビリティの厳しい制限に直面しています。これは、Transformer モデルをより長いコンテキスト長に拡張するための根本的な課題となります。

二次時間的注意層の問題に対処するためにさまざまなアプローチが検討されてきましたが、注目すべき方向性の 1 つは、注意層の中間行列を近似することです。これを実現する方法には、スパース行列、低ランク行列、またはその両方の組み合わせによる近似が含まれます。

ただし、これらの方法では、注意出力行列の近似に対してエンドツーエンドの保証を提供することはできません。これらの方法は、注意の個々のコンポーネントをより速く近似することを目的としていますが、完全なドット積注意のエンドツーエンドの近似を提供するものはありません。これらのメソッドは、最新の Transformer アーキテクチャの重要なコンポーネントである因果マスクの使用をまだサポートしていません。最近の理論的限界によれば、一般的なケースでは、注目行列をサブ二次時間で項ごとに近似することは不可能であることが示されています。

しかし、KDEFormer と呼ばれる最近の研究では、注目行列のエントリが制限されているという仮定の下で、サブ 2 次時間で証明可能な近似値を提供できることが示されました。理論的には、KDEFormer の実行時間はおよそ です。KDEFormer はカーネル密度推定 (KDE) を使用して列ノルムを近似し、アテンション マトリックスの列をサンプリングする確率を計算できます。しかし、現在の KDE アルゴリズムは実用的な効率性に欠けており、理論上でも、KDEFormer の実行時間と理論的に実現可能な O(n) 時間アルゴリズムの間にはギャップがあります。この論文では、著者らは、同じ制限付きエントリ仮定の下で、ほぼ線形時間のアルゴリズムが可能であることを示しています。しかし、彼らのアルゴリズムでは、ソフトマックスを近似するために多項式法も使用されており、これはおそらく実用的ではありません。

この記事では、イェール大学、Google Research、その他の機関の研究者が、両方の長所を兼ね備えたアルゴリズムを紹介します。このアルゴリズムは実用的かつ効率的で、ほぼ線形の時間保証を実現できます。さらに、私たちの方法は、以前の研究では不可能だった因果マスキングをサポートしています。

論文アドレス: https://arxiv.org/abs/2310.05869

この論文では、大規模言語モデルで使用される長いコンテキストの複雑さが増すことによってもたらされる計算上の課題に対処するために、「HyperAttention」と呼ばれる近似的な注意メカニズムを提案しています。最近の研究では、注意行列のエントリが制限されているか、行列の安定ランクが低い場合を除き、最悪の場合、2 次時間が必要になることが示されています。

研究者らは、（1）正規化された注意行列の最大列ノルム、および（2）大きなエントリを検出して削除した後の非正規化された注意行列の行ノルムの比率を測定するための2つのパラメータを導入しました。問題の難易度を反映するために、これらのきめ細かいパラメータを使用します。上記のパラメータが小さい限り、行列に無制限のエントリがある場合や大きな安定ランクがある場合でも、線形時間サンプリング アルゴリズムを実装できます。

HyperAttention はモジュール設計を特徴としており、他の高速な基盤実装、特に FlashAttention を簡単に統合できます。経験的に、HyperAttention は、大きなアイテムを識別するために LSH アルゴリズムを使用する既存の方法よりも優れており、FlashAttention などの SOTA ソリューションと比較して大幅な高速化を実現します。研究者らは、さまざまな長いコンテキスト長のデータセットで HyperAttention のパフォーマンスを検証しました。

たとえば、HyperAttention により、32k のコンテキスト長で ChatGLM2 の推論時間が 50% 高速化され、パープレキシティは 5.6 から 6.3 に増加します。より大きなコンテキスト長 (例: 131k) と因果マスクを使用すると、HyperAttention は単一の注意層で 5 倍高速になります。

方法の概要

ドット積アテンションには、Q (クエリ)、K (キー)、V (値) の 3 つの入力行列の処理が含まれます。これらの行列はすべて nxd のサイズで、n は入力シーケンス内のトークンの数、d は潜在表現の次元です。このプロセスの出力は次のようになります。

ここで、行列 A := exp (QK^T) は、QK^T の要素指数として定義されます。 D は、A の行の合計から導出された n×n 対角行列です。この場合、行列 A は「注意行列」と呼ばれ、(D^-1)A は「ソフトマックス行列」と呼ばれます。注意行列 A を直接計算するには Θ(n²d) 回の演算が必要であり、それを保存すると Θ(n²) のメモリが消費されることに注意してください。したがって、Attを直接計算するには、Ω(n²d)の実行時間とΩ(n²)のメモリが必要です。

私たちの目標は、スペクトル特性を維持しながら、出力行列 Att を効率的に近似することです。彼らの戦略は、ほぼ線形時間で対角スケーリング行列 D の効率的な推定器を設計することです。さらに、サブサンプリングによってソフトマックス行列 D^-1A の行列積を素早く近似します。より具体的には、彼らの目標は、エラーの演算子の指定に関する次の制約を満たす、有限数の行と対角行列を持つサンプリング行列を見つけることです。

研究者らは、Vの行ノルムに基づいてサンプリング行列Sを定義することで、式（1）の注意近似問題の行列乗算部分を効率的に解くことができることを示している。さらに難しい問題は、対角行列 D の信頼できる近似値をどのように得るかということです。最近の研究では、Zandieh らは高速 KDE ソルバーを効果的に活用して、D の高品質な近似値を取得しました。 KDEformer 手順を簡素化し、カーネル密度ベースの重要度サンプリングを必要とせずに、均一サンプリングで目的のスペクトル保証を達成できることを示します。この大幅な簡素化により、実用的で証明可能な線形時間アルゴリズムを開発することができました。

これまでの研究とは異なり、私たちの方法では、境界付きエントリや境界付き安定ランクは必要ありません。さらに、注意行列内のエントリまたは安定ランクが大きい場合でも、時間的複雑さを分析するために導入された細粒度のパラメータは依然として小さい可能性があります。

その結果、HyperAttention は大幅に高速化され、シーケンス長 n = 131k の場合、前方伝播と後方伝播の両方が 50 倍以上高速化されました。因果マスクを処理する場合、この方法では依然として 5 倍の大幅な高速化が達成されます。さらに、この手法を事前トレーニング済みの LLM (例: chatqlm2-6b-32k) に適用し、ロングコンテキスト ベンチマーク データセット LongBench で評価すると、微調整を行わなくても元のモデルに近いパフォーマンス レベルが維持されます。研究者らは特定のタスクも評価し、要約とコード補完のタスクが質問応答タスクよりも近似注意層に大きな影響を与えることを発見しました。

アルゴリズム

Att を近似する際にスペクトルの保証を得るために、この論文の最初のステップは、行列 D の対角要素に対して 1 ± ε 近似を実行することです。 (D^-1)AとVの間の行列積は、Vの2乗行ℓ₂ノルムに従ってサンプリングすることによって近似されます。

D を近似するプロセスは 2 つのステップで構成されます。まず、定義 1 に示すように、ハミング ソート LSH に基づくアルゴリズムを使用して、注目行列の主なエントリを識別します。 2 番目のステップは、小さな数値 K をランダムに選択することです。この論文では、行列 A と D についていくつかの軽度の仮定の下で、この単純なアプローチによって推定行列のスペクトル境界を確立できることを示します。研究者の目標は、次の条件を満たす十分に正確な近似行列 D を見つけることです。

この論文の仮定は、ソフトマックス行列の列ノルムが比較的均一な分布を示すというものです。より正確には、任意のi∈[n]tに対して、となるようなものが存在すると仮定します。

アルゴリズムの最初のステップは、ハミングソートされた LSH (sortLSH) を使用してキーとクエリを均一なサイズのバケットにハッシュすることにより、アテンション マトリックス A 内の大きなエントリを識別することです。アルゴリズム 1 ではこのプロセスを詳細に説明し、図 1 ではこのプロセスを直感的に示しています。

アルゴリズム 1 は、注目行列の主なエントリを分離することを目的としたスパース マスクを返します。このマスクを与えられた研究者は、式（2）のスペクトル保証を満たすアルゴリズム2の行列Dの近似値を計算する。このアルゴリズムは、マスクに対応する注意値と、注意マトリックス内のランダムに選択された列のサブセットを組み合わせることによって機能します。私たちのアルゴリズムは汎用性が高く、注目マトリックス内の主要なエントリの位置を指定する定義済みマスクを効果的に使用できます。このアルゴリズムによって提供される主な保証は定理 1 に示されています。

近似値と値行列 V の間の近似対角線と行列積を積分するサブルーチン。そこで研究者らは、式(1)のスペクトル保証付き注意機構をほぼ線形時間で近似できる効率的なアルゴリズムであるHyperAttentionを導入した。アルゴリズム 3 は、注目行列内の主要なエントリの位置を定義するマスク MH を入力として受け取ります。このマスクはsortLSHアルゴリズム（アルゴリズム1）を使用して生成することも、[7]の方法と同様に事前定義されたマスクにすることもできます。大規模エントリマスク M^H は設計上スパースであり、非ゼロエントリの数は制限されていると仮定します。

図 2 に示すように、私たちのアプローチは重要な観察に基づいています。マスクされた注意 M^C⊙A は、それぞれが元の注意行列の半分のサイズである 3 つのゼロ以外の行列に分解できます。対角線より完全に下にあるブロックA_21は、マスクされていない注目点です。したがって、アルゴリズム 2 を使用して行の合計を近似的に計算できます。

図 2 に示されている 2 つの対角ブロックは因果的注意であり、元のブロックの半分のサイズしかありません。これらの因果関係を処理するために、研究者は再帰的な方法を使用して、それらをさらに小さなチャンクに分割し、プロセスを繰り返しました。アルゴリズム 4 はこのプロセスの疑似コードを示します。

実験と結果

研究者らは、既存の大規模言語モデルを拡張して長距離シーケンスを処理できるようにすることで、アルゴリズムのベンチマークを実施しました。すべての実験は、正確な注意計算のために FlashAttention 2 を使用して、単一の 40 GB A100 GPU 上で実行されます。

モンキーパッチング自己注意

研究者らはまず、事前トレーニング済みの 2 つの LLM で HyperAttention を評価し、実際のアプリケーションで広く使用されている、異なるアーキテクチャを持つ 2 つのモデル (chatglm2-6b-32k と phi-1.5) を選択しました。

動作中、最終的な ℓ 注意層を HyperAttention に置き換えることでパッチを適用します。ℓ の数は、0 から各 LLM 内のすべての注意層の合計数まで変化します。どちらのモデルでも注意には因果マスクが必要であり、入力シーケンスの長さ n が 4,096 未満になるまでアルゴリズム 4 が再帰的に適用されることに注意してください。すべてのシーケンス長に対して、バケット サイズ b とサンプリングされた列の数 m を 256 に設定します。彼らは、このタイプのモンキーパッチモデルのパフォーマンスを、困惑度と加速の観点から評価しました。

同時に、研究者らは、単一/複数文書の質問回答、要約、小サンプル学習、合成タスク、コード補完という 6 つの異なるタスクを含む、ロングコンテキスト ベンチマーク データセットのコレクションである LongBench を使用しました。彼らは、エンコードシーケンスの長さが 32,768 を超えるデータセットのサブセットを選択し、長さが 32,768 を超えた場合はプルーニングを実行しました。次に、各モデルの困惑度、つまり次のトークン予測の損失を計算します。長いシーケンスへのスケーラビリティを強調するために、HyperAttention または FlashAttention のどちらで実行されたかに関係なく、すべてのアテンション レイヤーの合計スピードアップも計算します。

結果は下の図 3 に示されています。HyperAttention のモンキー パッチ後でも、chatglm2-6b-32k は依然として妥当な困惑を示しています。たとえば、20 層を置き換えた後、困惑度は約 1 増加し、24 層目に到達するまでゆっくりと増加し続けます。注意層の実行時間が約 50% 向上します。すべてのレイヤーを置き換えると、困惑度は 12 に上昇し、実行速度は 2.3 増加します。 phi-1.5 モデルも同様の状況を示していますが、HyperAttention の数の増加に伴って、困惑度は直線的に増加します。

さらに、研究者らは、LongBenchデータセット上でモンキーパッチを適用したchatglm2-6b-32kのパフォーマンスを評価し、単一/複数文書の質問回答、要約、小サンプル学習、合成タスク、コード補完などの各タスクの評価スコアを計算しました。結果は下の表1に示されています。

HyperAttention を置き換えると一般的にパフォーマンスが低下しますが、その影響は手元のタスクによって異なることが観察されました。たとえば、要約とコード補完は、他のタスクに比べて最も堅牢性があります。

注目すべきことに、アテンション レイヤーの半分 (つまり 14 レイヤー) にパッチを適用した場合、ほとんどのタスクのパフォーマンスが 13% 以上低下しないことが研究者によって確認されています。特に、要約タスクのパフォーマンスはほとんど変化せず、このタスクが注意メカニズムの部分的な変更に対して最も堅牢であることを示しています。 n=32kの場合、アテンション層の計算速度は1.5倍に増加します。

単一の自己注意層

研究者らは、シーケンスの長さが 4,096 から 131,072 まで変化した場合の HyperAttention の加速をさらに調査しました。彼らは、FlashAttention を使用して計算した場合、または HyperAttention によって加速された場合の、前方および前方 + 後方操作のウォールクロック時間を測定しました。さらに、因果マスキングの有無による壁時計時間が測定されました。入力 Q、K、V はすべて同じ長さで、次元は d = 64 に固定され、アテンション ヘッドの数は 12 です。

HyperAttention では以前と同じパラメータを選択します。下の図 4 に示すように、因果マスクを適用しない場合は HyperAttention は 54 倍高速になり、因果マスクを適用すると 5.4 倍高速になります。因果マスキングと非マスキングの時間的な複雑性は同じですが、因果マスキングの実際のアルゴリズム (アルゴリズム 1) では、Q、K、V の分割やアテンション出力のマージなどの追加の操作が必要になり、実際の実行時間が増加します。シーケンスの長さ n が増加すると、加速度は高くなります。

これらの結果は、推論だけでなく、より長いシーケンスに対する LLM のトレーニングや微調整にも自己注意を拡張する道を開くものだと考えています。