Delta Attention: Snelle en Nauwkeurige Inferentie van Sparse Attention door Delta Correctie

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

May 16, 2025

Auteurs: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

Samenvatting

Het aandachtmechanisme van een transformer heeft een kwadratische complexiteit, wat leidt tot hoge inferentiekosten en latentie voor lange sequenties. Echter zijn aandachtmatrices meestal spaarzaam, wat impliceert dat veel elementen kunnen worden weggelaten bij de berekening voor efficiënte inferentie. Methoden voor spaarzame aandachtinferentie hebben als doel deze rekenlast te verminderen; ze gaan echter ook gepaard met een vervelende prestatievermindering. Wij ontdekken dat een reden voor deze vermindering is dat de spaarzame berekening een distributieverschuiving induceert in de aandachtuitvoer. Deze distributieverschuiving zorgt ervoor dat queries tijdens het decoderen niet goed aansluiten bij de juiste sleutels uit de prefill-fase, wat leidt tot een prestatieverlies. Wij stellen een eenvoudige, nieuwe en effectieve procedure voor om deze distributieverschuiving te corrigeren, waardoor de distributie van spaarzame aandachtuitvoer dichter bij die van kwadratische aandacht komt. Onze methode kan worden toegepast bovenop elke spaarzame aandachtmethode en resulteert in een gemiddelde prestatieverbetering van 36 procentpunt, waarbij 88% van de nauwkeurigheid van kwadratische aandacht wordt hersteld op de 131K RULER-benchmark wanneer deze wordt toegepast bovenop sliding window attention met sink tokens, terwijl slechts een kleine overhead wordt toegevoegd. Onze methode kan ongeveer 98,5% spaarzaamheid behouden ten opzichte van volledige kwadratische aandacht, waardoor ons model 32 keer sneller is dan Flash Attention 2 bij het verwerken van 1M token prefills.

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.