[論文レビュー] Faster Causal Attention Over Large Sequences Through Sparse Flash Attention

Matteo Pagliardini, Daniele Paliotta+2 more|arXiv (Cornell University)|Jun 1, 2023

Topic Modeling被引用数 5

ひとこと要約

本論文は、因果自己注意の動的かつ不規則なスパarsityパターンをサポートするFlashAttentionの拡張であるSparse Causal Flash Attention (SCFA) を紹介する。SCFAは、長文系列における効率的な計算を可能にし、モデルの perplexity を損なわずにFlashAttentionよりも最大3.3倍速い学習を達成する。これは、ハッシュベースおよびクエリ/キーのドロップによるスパarsityパターンを最適化されたスパースカーネル実行により効率的に処理することで実現される。

ABSTRACT

Transformer-based language models have found many diverse applications requiring them to process sequences of increasing length. For these applications, the causal self-attention -- which is the only component scaling quadratically w.r.t. the sequence length -- becomes a central concern. While many works have proposed schemes to sparsify the attention patterns and reduce the computational overhead of self-attention, those are often limited by implementations concerns and end up imposing a simple and static structure over the attention matrix. Conversely, implementing more dynamic sparse attentions often results in runtimes significantly slower than computing the full attention using the Flash implementation from Dao et al. (2022). We extend FlashAttention to accommodate a large class of attention sparsity patterns that, in particular, encompass key/query dropping and hashing-based attention. This leads to implementations with no computational complexity overhead and a multi-fold runtime speedup on top of FlashAttention. Even with relatively low degrees of sparsity, our method improves visibly upon FlashAttention as the sequence length increases. Without sacrificing perplexity, we increase the training speed of a transformer language model by $2.0 imes$ and $3.3 imes$ for sequences of respectively $8k$ and $16k$ tokens.

研究の動機と目的

• 長文系列Transformerモデルにおける因果自己注意の計算ボトル neck を解決すること。
• ハッシュやクエリ/キーのドロップによるものなど、動的かつ不規則な注意スパarsityパターンに対して、効率的かつ高性能な推論および学習を可能にすること。
• 現代のスパース注意メカニズムで一般的な非三角行列の因果マスクに対しても、FlashAttentionの効率性を拡張すること。
• 計算複雑性を増加させず、モデル品質を損なわずにFlashAttentionよりも実用的な高速化を達成すること。
• 幅広いスパarsityパターンを最小限の実装コストでサポートする柔軟でオープンソースのGPUカーネルを提供すること。

提案手法

• クエリごとのキー範囲として表現可能な任意のスパarsityパターンをサポートするようFlashAttentionを拡張し、不規則な因果マスクを可能にする。
• Tritonを活用した低レベル最適化により、計算複雑性のオーバーヘッドなしにスパース因果注意を計算するGPUカーネルを導入する。
• 幾何学的ハッシュ（例：Reformer風のLSH）を用いて動的スパarsityを適用するが、近似を一切行わず正確な計算を保証する。
• ヘッドごとの細かいクエリおよびキーのドロップを実装し、完全なヘッドプルーニングを行わずに比例的な計算削減を可能にする。
• ブロック単位のメモリアクセスパターンを採用することで、スパースアクセスでも高いメモリ帯域幅利用率を維持する。
• カスタムのカーネル融合戦略を採用し、カーネル起動のオーバーヘッドを最小限に抑え、現代のGPUにおけるオカペーシティを最大化する。

Figure 3: Comparing several hash-based sparse attention implementations with FlashAttention. Similarly to QK-dropping-based sparsity in Fig. 7 , due to the non-triangular causal mask resulting from re-ordering the tensors based on the hash buckets (see Fig. 1 ), a naive implementation would force th

実験結果

リサーチクエスチョン

• RQ1 FlashAttentionを非三角行列の因果マスクに拡張することは可能か？ これにより、不規則なスパarsityパターンに対する効率的な計算が可能になるか？
• RQ2 追加の計算複雑性を伴わず、スパース注意メカニズムにおいてFlashAttentionよりも実用的な高速化を達成できるか？
• RQ3 動的で細かいスパarsity（例：クエリ/キーのドロップ）は、モデル性能を損なわず、学習速度の向上を測定可能か？
• RQ4 SCFAによるハッシュベース注意は、元のReformer LSHと比較して、速度およびカバー範囲の点で優れているか？
• RQ5 シーケンス長が増加する際、特に8kトークンを超えた場合でも、SCFAは学習効率を維持または向上できるか？

主な発見

• SCFAは、8kおよび16kトークンのシーケンスにおいて、それぞれ最大2.0倍および3.3倍の高速化を達成し、 perplexity を損なわない。
• 8192トークンのシーケンスで8および16のバケットを使用した場合、SCFAによるハッシュスパース注意は、FlashAttentionベースラインと比較して学習時間を1.4倍および1.8倍短縮する。
• 訓練中にも一貫した高速化が得られ、H-LMモデルでは初期段階から加速され、高いスループットを維持する。
• SCFAのハッシュベース注意は正確な計算を実現し、ReformerのLSHよりも高速な実行時間となる。これにより、低カバー範囲や近似誤差の問題を回避できる。
• SCFAによるクエリおよびキーのドロップは、ドロップされたペアの割合に比例して計算量を削減でき、細かく制御可能な効率性を実現する。
• 実行時間の高速化は、シーケンス長およびスパarsityに比例して向上するが、バケット数が一定以上に達すると効果が薄れる。

(a) Forward attn. runtimes