Выпрямленное разреженное внимание

Аннотация

Эффективная генерация длинных последовательностей представляет собой важную задачу для больших языковых моделей. Хотя современные методы разреженного декодирования повышают эффективность, они страдают от проблемы несоответствия кэша ключей и значений (KV cache misalignment), где ошибки аппроксимации накапливаются и ухудшают качество генерации. В данной работе мы предлагаем метод Rectified Sparse Attention (ReSA), простой, но эффективный подход, который сочетает блочно-разреженное внимание с периодической плотной коррекцией. Обновляя кэш ключей и значений через фиксированные интервалы с использованием плотного прямого прохода, ReSA ограничивает накопление ошибок и сохраняет соответствие распределению, полученному при предварительном обучении. Эксперименты в задачах математического рассуждения, языкового моделирования и поиска демонстрируют, что ReSA достигает почти без потерь качества генерации при значительном повышении эффективности. В частности, ReSA обеспечивает ускорение до 2,42 раза при декодировании последовательностей длиной 256K, что делает его практичным решением для масштабируемого вывода в длинных контекстах. Код доступен по адресу 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.