Rectified Sparse Attention

Zusammenfassung

Effiziente Langsequenzgenerierung stellt eine entscheidende Herausforderung für große Sprachmodelle dar. Während aktuelle sparsame Dekodierungsmethoden die Effizienz verbessern, leiden sie unter einer Fehlausrichtung des KV-Caches, bei der sich Approximationsfehler ansammeln und die Generierungsqualität beeinträchtigen. In dieser Arbeit schlagen wir Rectified Sparse Attention (ReSA) vor, eine einfache, aber effektive Methode, die block-sparse Attention mit periodischer dichter Rektifizierung kombiniert. Durch das Aktualisieren des KV-Caches in festen Intervallen mittels eines dichten Vorwärtsdurchlaufs begrenzt ReSA die Fehlerakkumulation und bewahrt die Ausrichtung mit der Vortrainingsverteilung. Experimente in den Bereichen mathematisches Denken, Sprachmodellierung und Retrieval-Aufgaben zeigen, dass ReSA eine nahezu verlustfreie Generierungsqualität bei deutlich verbesserter Effizienz erreicht. Insbesondere ermöglicht ReSA eine bis zu 2,42-fache End-to-End-Beschleunigung bei der Dekodierung von Sequenzen mit einer Länge von 256K, was es zu einer praktischen Lösung für skalierbare Langkontext-Inferenz macht. Der Code ist verfügbar unter 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.