Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jeffrey Willette, Heejun Lee, Sung Ju Hwang

TL;DR

Delta Attention addresses a critical bottleneck in transformer inference: the distributional shift introduced by sparse attention prefill, which harms query-key alignment and long-context performance. By adding a lightweight post-processing delta-correction every $\gamma$ rows, the method realigns sparse outputs toward the full quadratic attention distribution without sacrificing sparsity or incurring large latency. Across long-context benchmarks like 131K RULER and PG19 Long QA, it yields substantial gains (around $36$pp on average and about $88\%$ recovery of quadratic accuracy) with modest overhead, and can deliver large speedups (up to ~32x) over dense attention. The approach is kernel-agnostic and readily integrable into existing sparse attention pipelines, offering a practical path to faster, more accurate long-sequence Transformer inference.

Abstract

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.

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

TL;DR

Delta Attention addresses a critical bottleneck in transformer inference: the distributional shift introduced by sparse attention prefill, which harms query-key alignment and long-context performance. By adding a lightweight post-processing delta-correction every rows, the method realigns sparse outputs toward the full quadratic attention distribution without sacrificing sparsity or incurring large latency. Across long-context benchmarks like 131K RULER and PG19 Long QA, it yields substantial gains (around pp on average and about recovery of quadratic accuracy) with modest overhead, and can deliver large speedups (up to ~32x) over dense attention. The approach is kernel-agnostic and readily integrable into existing sparse attention pipelines, offering a practical path to faster, more accurate long-sequence Transformer inference.

Abstract

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.
Paper Structure (19 sections, 2 theorems, 21 equations, 17 figures, 9 tables, 2 algorithms)

This paper contains 19 sections, 2 theorems, 21 equations, 17 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background & Related Work
  3. Method
  4. Experiments
  5. Ablation & Analysis
  6. Discussion & Limitations
  7. Conclusion
  8. Acknowledgments
  9. Appendix Contents
  10. Broader Impact
  11. Implementation Details
  12. Compute Resources
  13. Latency of MInference with Delta Attention
  14. Approx Window Size Calculation
  15. Restatement and proof of
  16. ...and 4 more sections

Key Result

Lemma 1

w.l.o.g. Consider an arbitrary row in the attention matrix $\mathbf{a}$ and arbitrary column of the values $\mathbf{v}$, with both $\mathbf{a}$ and $\mathbf{v}$ being sorted according to rank of $\mathbf{a}$ such that $\mathbf{a} = (a_{r(1)} \leq a_{r(2)} \leq \dots \leq a_{r(N)})$. For a top-$k$ sp

Figures (17)

  • Figure 1: RULER 131K Subsets. At long context lengths, sparse attention can degrade performance by a large margin. Our simple $\mathbf{\Delta}$ correction improves performance and only requires an additional 1.5% of the full quadratic attention computation.
  • Figure 1: $\mathbf{\Delta}$ Attention Algorithm
  • Figure 2: Comparing RULER 131K prefill attention latency and accuracy for sparse attention methods.
  • Figure 2: $\alpha,\beta,\gamma$ Filter
  • Figure 3: Comparing sparse attention methods to quadratic attention. Our $\mathbf{\Delta}$ correction results in outputs that are more similar to quadratic attention.
  • ...and 12 more figures

Theorems & Definitions (4)

  • Lemma 1
  • proof
  • Lemma : name=Lemma \ref{['lem:delta']}, restated
  • proof