Atenção Esparsa Retificada

Resumo

A geração eficiente de sequências longas é um desafio crítico para Modelos de Linguagem de Grande Escala. Embora métodos recentes de decodificação esparsa melhorem a eficiência, eles sofrem com o desalinhamento do cache KV, onde erros de aproximação se acumulam e degradam a qualidade da geração. Neste trabalho, propomos a Atenção Esparsa Retificada (ReSA), um método simples, porém eficaz, que combina atenção esparsa em blocos com retificação densa periódica. Ao atualizar o cache KV em intervalos fixos usando uma passagem direta densa, a ReSA limita o acúmulo de erros e preserva o alinhamento com a distribuição de pré-treinamento. Experimentos em tarefas de raciocínio matemático, modelagem de linguagem e recuperação demonstram que a ReSA alcança qualidade de geração quase sem perdas com eficiência significativamente melhorada. Notavelmente, a ReSA oferece um aumento de velocidade de até 2,42 vezes de ponta a ponta na decodificação de sequências de 256K de comprimento, tornando-a uma solução prática para inferência escalável de contexto longo. O código está disponível em 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.