Attenzione Sparse Rettificata

Abstract

La generazione efficiente di sequenze lunghe rappresenta una sfida cruciale per i Modelli Linguistici di Grande Dimensione. Sebbene i recenti metodi di decodifica sparsa migliorino l'efficienza, essi soffrono di disallineamento della cache KV, dove gli errori di approssimazione si accumulano e degradano la qualità della generazione. In questo lavoro, proponiamo l'Attenzione Sparsa Rettificata (ReSA), un metodo semplice ma efficace che combina l'attenzione sparsa a blocchi con una rettifica densa periodica. Aggiornando la cache KV a intervalli fissi mediante un passaggio in avanti denso, ReSA limita l'accumulo di errori e preserva l'allineamento con la distribuzione di pre-addestramento. Esperimenti su ragionamento matematico, modellazione del linguaggio e task di retrieval dimostrano che ReSA raggiunge una qualità di generazione quasi senza perdite con un'efficienza significativamente migliorata. In particolare, ReSA offre un'accelerazione end-to-end fino a 2,42 volte durante la decodifica con lunghezza di sequenza di 256K, rendendola una soluzione pratica per l'inferenza scalabile in contesti lunghi. Il codice è disponibile all'indirizzo 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.