デルタアテンション：デルタ補正による高速かつ正確なスパースアテンション推論

Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction

May 16, 2025

著者: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

要旨

Transformerのアテンションメカニズムは二次複雑度を持ち、長いシーケンスに対して高い推論コストとレイテンシを引き起こします。しかし、アテンションマトリックスはほとんどがスパースであり、効率的な推論のために多くのエントリを計算から省略できることを意味します。スパースアテンション推論手法はこの計算負荷を軽減することを目指していますが、同時に厄介な性能低下も伴います。私たちは、この性能低下の一因が、スパース計算がアテンション出力の分布シフトを引き起こすことにあることを発見しました。この分布シフトにより、デコード時のクエリがプレフィル段階の適切なキーとうまく整合しなくなり、性能の低下を招きます。私たちは、この分布シフトを補正するためのシンプルで新規かつ効果的な手順を提案します。これにより、スパースアテンション出力の分布を二次アテンションの分布に近づけることができます。私たちの手法は、任意のスパースアテンション手法の上に適用可能であり、スライディングウィンドウアテンションとシンクトークンを組み合わせた場合、131K RULERベンチマークにおいて二次アテンションの精度の88%を回復し、平均36%ポイントの性能向上をもたらします。また、わずかなオーバーヘッドを追加するだけで、完全な二次アテンションに対して約98.5%のスパース性を維持できるため、1Mトークンのプレフィル処理においてFlash Attention 2よりも32倍高速なモデルを実現します。

English

The attention mechanism of a transformer has a quadratic complexity, leading to high inference costs and latency for long sequences. However, attention matrices are mostly sparse, which implies that many entries may be omitted from computation for efficient inference. Sparse attention inference methods aim to reduce this computational burden; however, they also come with a troublesome performance degradation. We discover that one reason for this degradation is that the sparse calculation induces a distributional shift in the attention outputs. The distributional shift causes decoding-time queries to fail to align well with the appropriate keys from the prefill stage, leading to a drop in performance. We propose a simple, novel, and effective procedure for correcting this distributional shift, bringing the distribution of sparse attention outputs closer to that of quadratic attention. Our method can be applied on top of any sparse attention method, and results in an average 36%pt performance increase, recovering 88% of quadratic attention accuracy on the 131K RULER benchmark when applied on top of sliding window attention with sink tokens while only adding a small overhead. Our method can maintain approximately 98.5% sparsity over full quadratic attention, making our model 32 times faster than Flash Attention 2 when processing 1M token prefills.