[論文レビュー] Transformer Acceleration with Dynamic Sparse Attention

Liu Liu, Zheng Qu+3 more|arXiv (Cornell University)|Oct 21, 2021

Advanced Neural Network Applications参考文献 26被引用数 12

ひとこと要約

本稿では、計算コストを削減しつつモデルの精度を損なわずに、入力に依存するスパースなアテンションパターンを動的に特定・活用する、Dynamic Sparse Attention (DSA) を提案する。近似アテンションスコアを用いた軽量な予測パスを用いることで、DSAは最大95%のスパーシティと4.35倍の理論的計算削減を達成し、わずか1.33%のオーバーヘッドで、V100 GPU上でのアテンション出力で1.94倍、ソフトマックス計算で14.6倍の高速化を実現した。90%スパーシティ条件下で、

ABSTRACT

Transformers are the mainstream of NLP applications and are becoming increasingly popular in other domains such as Computer Vision. Despite the improvements in model quality, the enormous computation costs make Transformers difficult at deployment, especially when the sequence length is large in emerging applications. Processing attention mechanism as the essential component of Transformer is the bottleneck of execution due to the quadratic complexity. Prior art explores sparse patterns in attention to support long sequence modeling, but those pieces of work are on static or fixed patterns. We demonstrate that the sparse patterns are dynamic, depending on input sequences. Thus, we propose the Dynamic Sparse Attention (DSA) that can efficiently exploit the dynamic sparsity in the attention of Transformers. Compared with other methods, our approach can achieve better trade-offs between accuracy and model complexity. Moving forward, we identify challenges and provide solutions to implement DSA on existing hardware (GPUs) and specialized hardware in order to achieve practical speedup and efficiency improvements for Transformer execution.

研究の動機と目的

• 自己アテンションの高い計算コスト、特に長文系列長における問題に対処すること。
• 入力依存の動的スパーシティを捉えられず、固定されたスパースパターンに限界をもつ静的スパースアテンションの限界を克服すること。
• 完全なアテンションの表現力を維持しながら計算を削減する、実用的でハードウェアに配慮した手法を開発すること。
• 動的スパーシティとデータフロー最適化を通じて、既存および専用ハードウェア上での大規模Transformerの効率的デプロイを可能にすること。

提案手法

• 標準的な自己アテンションに、動的スパースパターンを同定するための軽量な予測パスを追加する。
• 計算オーバーヘッドを低減するため、近似アテンションスコアをスパーシティ予測子の入力として使用する。
• 予測されたスパーシティを学習可能なバイナリマスクとして表現し、ソフトマックスおよび出力計算における不要な計算をスキップする。
• GPUや専用アクセラレータのハードウェア最適化を目的として、細粒度の動的スパーシティをブロック単位やベクトル単位の構造的パターンへ拡張する。
• カスタムカーネルの設計とデータ再利用の機会の活用により、GPUおよび専用アクセラレータ上のメモリアクセス効率を向上させる。
• スパーシティ予測の低精度推論と、計算再順序化によるトークンレベル並列化を統合する。

実験結果

リサーチクエスチョン

• RQ1 自己アテンションにおける動的スパースパターンを、モデル精度を損なわずに効果的に同定・活用できるか？
• RQ2 精度と効率のトレードオフにおいて、動的スパーシティは静的または固定スパースパターンと比べてどのように優れるか？
• RQ3 DSAは、既存のGPUアーキテクチャおよび専用ハードウェアアクセラレータ上で実用的にどの程度の性能向上を達成できるか？
• RQ4 アルゴリズム的スパーシティをハードウェアと共同設計することで、最大の高速化とエネルギー効率を実現できるか？

主な発見

• DSAはLRAベンチマークにおいて、モデル精度に顕著な損失を生じさせることなく、最大95%のスパーシティを達成した。
• スパーシティ予測による追加計算オーバーヘッドは1.17%〜1.33%にとどまり、理論的計算量は4.35倍削減された。
• NVIDIA V100 GPU上では、90%スパーシティ条件下で、アテンションスコア計算で1.15倍の高速化、ソフトマックス計算で14.6倍の高速化、アテンション出力計算で1.94倍の高速化を達成した。
• ベクトル単位のスパーシティと最適化されたカーネル設計により、専用ハードウェア上での合計メモリアクセスを最大2.54倍削減した。
• 低ランク近似や過剰なプルーニングとは異なり、重要なアテンション重みを保持することで、完全なアテンションの表現力が維持された。
• ハードウェア特化により、低精度予測、細粒度スパース計算、および改善されたデータ局所性を通じてさらなる性能向上が実現された。