Delta-Attention: Schnelle und präzise Inferenz von Sparse Attention durch Delta-Korrektur

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

May 16, 2025

Autoren: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

Zusammenfassung

Der Aufmerksamkeitsmechanismus eines Transformers weist eine quadratische Komplexität auf, was zu hohen Inferenzkosten und Latenz bei langen Sequenzen führt. Allerdings sind die Aufmerksamkeitsmatrizen größtenteils spärlich besetzt, was bedeutet, dass viele Einträge für eine effiziente Inferenz aus der Berechnung ausgelassen werden können. Methoden zur spärlichen Aufmerksamkeitsinferenz zielen darauf ab, diese Rechenlast zu verringern; sie gehen jedoch auch mit einer problematischen Leistungsverschlechterung einher. Wir stellen fest, dass ein Grund für diese Verschlechterung darin liegt, dass die spärliche Berechnung eine Verteilungsverschiebung in den Aufmerksamkeitsausgaben verursacht. Diese Verteilungsverschiebung führt dazu, dass Abfragen zur Dekodierungszeit nicht gut mit den entsprechenden Schlüsseln aus der Prefill-Phase übereinstimmen, was zu einem Leistungsabfall führt. Wir schlagen ein einfaches, neuartiges und effektives Verfahren vor, um diese Verteilungsverschiebung zu korrigieren und die Verteilung der spärlichen Aufmerksamkeitsausgaben näher an die der quadratischen Aufmerksamkeit heranzuführen. Unsere Methode kann auf jede spärliche Aufmerksamkeitsmethode angewendet werden und führt zu einer durchschnittlichen Leistungssteigerung von 36 Prozentpunkten, wobei 88 % der Genauigkeit der quadratischen Aufmerksamkeit auf dem 131K RULER-Benchmark wiederhergestellt werden, wenn sie auf die Sliding-Window-Aufmerksamkeit mit Sink-Tokens angewendet wird, während nur ein geringer Overhead hinzugefügt wird. Unsere Methode kann eine Spärlichkeit von etwa 98,5 % im Vergleich zur vollständigen quadratischen Aufmerksamkeit aufrechterhalten, wodurch unser Modell 32-mal schneller als Flash Attention 2 ist, wenn es 1M Token-Prefills verarbeitet.

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.