Rectified Sparse Attention

2025-06-04

AI Summary

Large language models (LLMs) suffer from degraded generation quality in long-sequence decoding due to KV cache misalignment and error accumulation in sparse decoding. To address this, we propose an efficient decoding framework integrating block-sparse attention with periodic dense correction. Our method introduces, for the first time, a dense correction mechanism that resets the KV cache at fixed decoding steps—rigorously bounding error propagation and enabling high-fidelity alignment with the pretraining distribution under sparse decoding. Additionally, we incorporate dynamic KV cache refreshing and long-context inference optimization. Evaluated on 256K-length decoding tasks, our approach achieves an end-to-end speedup of 2.42× while preserving near-zero quality degradation across mathematical reasoning, language modeling, and retrieval benchmarks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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.42$ imes$ 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.

Problem

• Addresses KV cache misalignment in sparse decoding methods
• Improves efficiency in long-sequence generation for LLMs
• Reduces error accumulation while maintaining generation quality

Innovation

• Combines block-sparse attention with dense rectification
• Refreshes KV cache at fixed intervals
• Achieves near-lossless quality with high efficiency