[論文レビュー] SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Hanrui Wang, Zhekai Zhang+1 more|arXiv (Cornell University)|Dec 17, 2020

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

ひとこと要約

SpAtten は、スパースアテンションアーキテクチャをハードウェア・ソフトウェア共同設計により提案し、段階的トークンおよびヘッドプリーニングに加え、段階的量子化を用いて、アテンション機構におけるメモリアクセスと計算を顕著に削減する。動的かつ入力依存のプリーニングにより、不要なトークンとヘッドをリアルタイムで削除し、信頼度に応じてビット幅を適応的に変更することで、TITAN Xp GPU に対して最大 4059× のエネルギー効率向上と 162× の高速化を達成し、精度に損失なしに実現した。

ABSTRACT

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, attention becomes the compution bottleneck because of its quadratic computational complexity to input length, complicated data movement and low arithmetic intensity. Moreover, existing NN accelerators mainly focus on optimizing convolutional or recurrent models, and cannot efficiently support attention. In this paper, we present SpAtten, an efficient algorithm-architecture co-design that leverages token sparsity, head sparsity, and quantization opportunities to reduce the attention computation and memory access. Inspired by the high redundancy of human languages, we propose the novel cascade token pruning to prune away unimportant tokens in the sentence. We also propose cascade head pruning to remove unessential heads. Cascade pruning is fundamentally different from weight pruning since there is no trainable weight in the attention mechanism, and the pruned tokens and heads are selected on the fly. To efficiently support them on hardware, we design a novel top-k engine to rank token and head importance scores with high throughput. Furthermore, we propose progressive quantization that first fetches MSBs only and performs the computation; if the confidence is low, it fetches LSBs and recomputes the attention outputs, trading computation for memory reduction. Extensive experiments on 30 benchmarks show that, on average, SpAtten reduces DRAM access by 10.0x with no accuracy loss, and achieves 1.6x, 3.0x, 162x, 347x speedup, and 1,4x, 3.2x, 1193x, 4059x energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

研究の動機と目的

• 一般用途プラットフォーム（CPU や GPU）におけるアテンション機構の非効率性、特に高いメモリ帯域幅と低い算術強度の問題に対処する。
• CNN や RNN に最適化された既存のアクセラレータが、アテンション機構には適していないという制限を克服する。
• 人間の言語における構造的再冗長性とアテンションヘッドの再冗長性を活用して、DRAM アクセスと計算量を削減する。
• 推論時に動的かつ入力依存のトークンおよびヘッドプリーニングを可能にするドメイン特化アクセラレータを設計し、リアルタイムでの効率向上を実現する。
• メモリとエネルギーのオーバーヘッドを最小限に抑えることで、資源制約のあるデバイスへの大規模 NLP モデル（例：BERT や GPT-2）の効率的デプロイを可能にする。

提案手法

• 複数層にわたる累積的アテンション確率に基づいて、不要なトークンを削除する段階的トークンプリーニングを提案し、文の長さに応じてプリーニング比を動的に調整する。
• 複数層にわたる累積的ヘッドマグニチュードスコアを用いて、冗長なアテンションヘッドを順位付け・削除する段階的ヘッドプリーニングを導入する。
• リアルタイムでのトークンおよびヘッドの重要度スコアの順位付けを可能にする、高スルーレートかつ並列処理可能な top-k エンジンを設計し、O(n) の時間計算量を実現する。
• まず最も有意義なビット（MSB）での計算を実施し、信頼度が低い場合にのみ、最も意味の薄いビット（LSB）を取得して再計算する段階的量子化を実装する。これにより、計算量を増加させながらもメモリ使用量を削減する。
• パイプライン化されたデータパス、最適化されたメモリ階層、動的スパarsity および可変ビット幅をサポートする専用ハードウェアアーキテクチャを構築し、理論的な節約効果を実際の実装に反映する。
• アルゴリズムとハードウェアを共同設計し、プリーニングおよび量子化を入力に依存して推論時に適用することで、柔軟性を確保するとともに、精度損失を最小限に抑える。

実験結果

リサーチクエスチョン

• RQ1 動的かつ入力依存のトークンおよびヘッドプリーニングは、モデル精度を損なわずにアテンション計算量とメモリアクセスを削減できるか？
• RQ2 低遅延な top-k エンジンを用いて、リアルタイムで重要なトークンおよびヘッドの順位付けと選択を効率的に行うハードウェアアクセラレータの実装方法は何か？
• RQ3 段階的量子化は、推論の信頼度を維持しながら、どれほどメモリ帯域幅とエネルギー消費を削減できるか？
• RQ4 A³ や MNNFast といった最先端アクセラレータと比較して、段階的プリーニングと段階的量子化は、性能およびエネルギー効率においてどの程度優れているか？
• RQ5 提案された共同設計アプローチは、判別的（例：BERT）および生成的（例：GPT-2）なトランスフォーマーモデルの両方に対して一般化可能か？

主な発見

• SpAtten は 30 のベンチマークで平均 10.0× の DRAM アクセス削減を達成し、精度に損失なしに実現した。
• A³ アクセラレータに対して 1.6×、MNNFast アクセラレータに対して 3.0× の高速化を達成し、それぞれ 1.4× および 3.2× のエネルギー効率向上を実現した。
• TITAN Xp GPU 上では、162× の高速化と 1193× のエネルギー効率向上を達成した。
• Xeon CPU 上では、347× の高速化と 4059× のエネルギー効率向上を達成した。
• top-k エンジンにより、O(n) の時間計算量で効率的かつ高スルーレートのトークンおよびヘッドの重要度順位付けが可能となり、リアルタイムプリーニングに不可欠である。
• 段階的量子化により、動的ビット幅の適応が可能となり、繰り返しの精錬により推論の信頼度を維持しながらメモリ帯域幅を削減できる。