SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

要約

大規模なモデルには、二次時間の複雑さがあるため、効率的な注意実装が不可欠です。
幸いなことに、注意は一般的にスパース性を示します。つまり、注意マップの多くの値はゼロに近く、対応する計算を省略します。
多くの研究では、注意を促すためにまばらなパターンを利用しています。
ただし、ほとんどの既存の作品は、注意マップの特定のまばらなパターンを活用することにより、特定のモデル内で注意を最適化することに焦点を当てています。
多様なモデルのスピードアップとエンドツーエンドのパフォーマンスの両方を保証する普遍的なまばらな注意は、とらえどころのないままです。
この論文では、あらゆるモデルの普遍的なまばらで量子化された注意であるSpargeattnを提案します。
私たちの方法では、2段階のオンラインフィルターを使用しています。最初の段階では、注意マップを迅速かつ正確に予測し、いくつかのマトリックス乗算のスキップが注意を払うことを可能にします。
第2段階では、オーバーヘッドが追加されず、マトリックスの乗算をさらにスキップするオンラインソフトマックス認識フィルターを設計します。
実験では、この方法は、エンドツーエンドのメトリックを犠牲にすることなく、言語、画像、ビデオ生成を含む多様なモデルを大幅に加速することを示しています。
コードはhttps://github.com/thu-ml/spargeattnで入手できます。

要約(オリジナル)

An efficient attention implementation is essential for large models due to its quadratic time complexity. Fortunately, attention commonly exhibits sparsity, i.e., many values in the attention map are near zero, allowing for the omission of corresponding computations. Many studies have utilized the sparse pattern to accelerate attention. However, most existing works focus on optimizing attention within specific models by exploiting certain sparse patterns of the attention map. A universal sparse attention that guarantees both the speedup and end-to-end performance of diverse models remains elusive. In this paper, we propose SpargeAttn, a universal sparse and quantized attention for any model. Our method uses a two-stage online filter: in the first stage, we rapidly and accurately predict the attention map, enabling the skip of some matrix multiplications in attention. In the second stage, we design an online softmax-aware filter that incurs no extra overhead and further skips some matrix multiplications. Experiments show that our method significantly accelerates diverse models, including language, image, and video generation, without sacrificing end-to-end metrics. The codes are available at https://github.com/thu-ml/SpargeAttn.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google