Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (2505.11254v1)

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.

Overview of $\mathbf{\Delta}$ Attention: Addressing Distributional Shift in Sparse Attention Mechanisms

The paper, "$\mathbf{\Delta}$ Attention: Fast and Accurate Sparse Attention Inference by Delta Correction," presents a novel approach targeting the inefficiencies observed in sparse attention mechanisms within transformer models. Specifically, the authors highlight the quadratic complexity involved in the computation of the attention matrix—an inherent challenge when dealing with long sequences. Sparse attention methods aim to mitigate this complexity, yet they are often accompanied by significant performance degradation due to distributional shifts in attention outputs.

Key Contributions

The authors make several vital contributions to understand and mitigate the effects of sparse attention inference:

1. Identification of Distributional Shift: Sparse attention calculations lead to a shift in attention outputs' distributions, impacting the query-key alignment during decoding. This misalignment results in decreased performance, particularly evident in long-context scenarios.
2. Development of $\mathbf{\Delta}$ Attention: A simple post-processing technique that corrects for the distributional shift, aligning sparse attention outputs closer to those observed in quadratic attention matrices. The $\mathbf{\Delta}$ Attention method is versatile and can be applied to existing sparse attention kernels, enhancing the performance without substantial computation overhead.
3. Significant Performance Improvement: When applied to various sparse attention methods, $\mathbf{\Delta}$ Attention achieves an average performance increase of 36 percentage points, recovering 88% of quadratic attention’s accuracy on long-context benchmarks like the 131K RULER, while maintaining high sparsity and reducing latency substantially compared to conventional methods.

Evaluation and Results

The evaluation was conducted using metrics such as perplexity (PPL), LongPPL, and accuracy on several challenging benchmarks including PG19 and RULER. The authors demonstrate that $\mathbf{\Delta}$ Attention consistently improves performance across these metrics:

• Perplexity Metrics: On the PG19 Long QA dataset, $\mathbf{\Delta}$ Attention effectively lowers both PPL and LongPPL compared to baseline sparse methods, indicating better language understanding and sequence prediction capabilities.
• RULER Benchmark: Across tasks involving long-context understanding, $\mathbf{\Delta}$ Attention shows superior accuracy, particularly at 131K context length. The method improves the sparse methods' capability to recall context, closing the gap significantly with full attention methods.

Implications and Future Work

The findings suggest that distributional shifts are a critical factor in the performance degradation seen in sparse attention methods. By correcting for these shifts, the paper positions $\mathbf{\Delta}$ Attention as a crucial improvement for inference efficiency and performance stability over long contexts. Looking ahead, the approach paves the way for:

• Enhanced Sparse Models: Incorporating distributional correction techniques in sparse attention frameworks could herald new architectures capable of handling extremely long sequences without sacrificing depth and accuracy.
• Broader Application: Beyond transformers, the principles outlined could inspire similar enhancements in other models reliant on attention mechanisms, broadening their efficacy in complex tasks.

Overall, the paper offers a significant advancement in efficiently handling sparse attention matrices, promising reduced computational costs and improved model reliability for extensive text sequences—valuable for those seeking scalability in AI applications.