Atención Dispersa Rectificada

Resumen

La generación eficiente de secuencias largas es un desafío crítico para los Modelos de Lenguaje a Gran Escala (LLM). Aunque los métodos recientes de decodificación dispersa mejoran la eficiencia, sufren de desalineación en la caché KV, donde los errores de aproximación se acumulan y degradan la calidad de la generación. En este trabajo, proponemos Atención Dispersa Rectificada (ReSA), un método simple pero efectivo que combina la atención dispersa en bloques con una rectificación densa periódica. Al actualizar la caché KV a intervalos fijos mediante un pase hacia adelante denso, ReSA limita la acumulación de errores y preserva la alineación con la distribución de preentrenamiento. Los experimentos en tareas de razonamiento matemático, modelado de lenguaje y recuperación demuestran que ReSA logra una calidad de generación casi sin pérdidas con una eficiencia significativamente mejorada. Notablemente, ReSA ofrece una aceleración de hasta 2.42 veces en el proceso de decodificación con secuencias de 256K de longitud, convirtiéndolo en una solución práctica para la inferencia escalable en contextos largos. El código está disponible en 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.