Atención Dispersa de Tokens: Inferencia Eficiente en Contextos Largos con Selección Intercalada de Tokens

Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

February 3, 2026

Autores: Dongwon Jo, Beomseok Kang, Jiwon Song, Jae-Joon Kim

cs.AI

Resumen

La complejidad cuadrática de la atención sigue siendo el principal cuello de botella en la inferencia de contextos largos para los modelos de lenguaje grandes. Los métodos de aceleración previos o bien esparsifican el mapa de atención con patrones estructurados o eliminan permanentemente tokens en capas específicas, lo que puede retener tokens irrelevantes o depender de decisiones tempranas irreversibles a pesar de la dinámica capa a capa y por cabezal de la importancia de los tokens. En este artículo, proponemos Token Sparse Attention, un mecanismo ligero y dinámico de esparsificación a nivel de token que comprime las matrices Q, K, V por cabezal a un conjunto reducido de tokens durante la atención y luego descomprime la salida de vuelta a la secuencia original, permitiendo que la información del token sea reconsiderada en capas subsiguientes. Además, Token Sparse Attention expone un nuevo punto de diseño en la intersección de la selección de tokens y la atención dispersa. Nuestro enfoque es totalmente compatible con las implementaciones de atención densa, incluyendo Flash Attention, y puede componerse sin problemas con kernels de atención dispersa existentes. Los resultados experimentales muestran que Token Sparse Attention mejora consistentemente la relación precisión-latencia, logrando una aceleración de la atención de hasta 3.2x en contextos de 128K tokens con una degradación de precisión inferior al 1%. Estos resultados demuestran que la esparsificación dinámica e intercalada a nivel de token es una estrategia complementaria y efectiva para la inferencia escalable de contextos largos.

English

The quadratic complexity of attention remains the central bottleneck in long-context inference for large language models. Prior acceleration methods either sparsify the attention map with structured patterns or permanently evict tokens at specific layers, which can retain irrelevant tokens or rely on irreversible early decisions despite the layer-/head-wise dynamics of token importance. In this paper, we propose Token Sparse Attention, a lightweight and dynamic token-level sparsification mechanism that compresses per-head Q, K, V to a reduced token set during attention and then decompresses the output back to the original sequence, enabling token information to be reconsidered in subsequent layers. Furthermore, Token Sparse Attention exposes a new design point at the intersection of token selection and sparse attention. Our approach is fully compatible with dense attention implementations, including Flash Attention, and can be seamlessly composed with existing sparse attention kernels. Experimental results show that Token Sparse Attention consistently improves accuracy-latency trade-off, achieving up to times3.23 attention speedup at 128K context with less than 1% accuracy degradation. These results demonstrate that dynamic and interleaved token-level sparsification is a complementary and effective strategy for scalable long-context inference.