Attention Delta : Inférence Rapide et Précise de l'Attention Sparse par Correction Delta

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

May 16, 2025

Auteurs: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

Résumé

Le mécanisme d'attention d'un transformateur présente une complexité quadratique, entraînant des coûts d'inférence élevés et une latence importante pour les longues séquences. Cependant, les matrices d'attention sont majoritairement creuses, ce qui implique que de nombreuses entrées peuvent être omises du calcul pour une inférence efficace. Les méthodes d'inférence par attention creuse visent à réduire cette charge de calcul, mais elles s'accompagnent également d'une dégradation problématique des performances. Nous découvrons qu'une raison de cette dégradation est que le calcul creux induit un décalage distributionnel dans les sorties de l'attention. Ce décalage distributionnel entraîne un mauvais alignement des requêtes au moment du décodage avec les clés appropriées de l'étape de préremplissage, ce qui se traduit par une baisse de performance. Nous proposons une procédure simple, novatrice et efficace pour corriger ce décalage distributionnel, rapprochant ainsi la distribution des sorties de l'attention creuse de celle de l'attention quadratique. Notre méthode peut être appliquée par-dessus toute méthode d'attention creuse et entraîne une augmentation moyenne de 36 points de pourcentage des performances, récupérant 88 % de la précision de l'attention quadratique sur le benchmark RULER de 131K lorsqu'elle est appliquée par-dessus l'attention par fenêtre glissante avec des tokens de puits, tout en n'ajoutant qu'une faible surcharge. Notre méthode peut maintenir environ 98,5 % de creux par rapport à l'attention quadratique complète, rendant notre modèle 32 fois plus rapide que Flash Attention 2 lors du traitement de préremplissages de 1 million de tokens.

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.