Attention Rectifiée et Sparse

Résumé

La génération efficace de longues séquences constitue un défi majeur pour les grands modèles de langage. Bien que les méthodes récentes de décodage parcimonieux améliorent l'efficacité, elles souffrent d'un désalignement du cache KV, où les erreurs d'approximation s'accumulent et dégradent la qualité de la génération. Dans ce travail, nous proposons l'Attention Parcimonieuse Rectifiée (ReSA), une méthode simple mais efficace qui combine l'attention parcimonieuse par blocs avec une rectification dense périodique. En rafraîchissant le cache KV à intervalles fixes via un passage dense, ReSA limite l'accumulation d'erreurs et préserve l'alignement avec la distribution de pré-entraînement. Les expériences menées sur des tâches de raisonnement mathématique, de modélisation du langage et de recherche démontrent que ReSA atteint une qualité de génération quasi sans perte avec une efficacité significativement améliorée. En particulier, ReSA offre une accélération de bout en bout allant jusqu'à 2,42 fois lors du décodage de séquences de 256K, en faisant une solution pratique pour l'inférence contextuelle longue et évolutive. Le code est disponible à l'adresse https://aka.ms/ReSA-LM.

English

Efficient long-sequence generation is a critical challenge for Large Language Models. While recent sparse decoding methods improve efficiency, they suffer from KV cache misalignment, where approximation errors accumulate and degrade generation quality. In this work, we propose Rectified Sparse Attention (ReSA), a simple yet effective method that combines block-sparse attention with periodic dense rectification. By refreshing the KV cache at fixed intervals using a dense forward pass, ReSA bounds error accumulation and preserves alignment with the pretraining distribution. Experiments across math reasoning, language modeling, and retrieval tasks demonstrate that ReSA achieves near-lossless generation quality with significantly improved efficiency. Notably, ReSA delivers up to 2.42times end-to-end speedup under decoding at 256K sequence length, making it a practical solution for scalable long-context inference. Code is available at https://aka.ms/ReSA-LM.