Atenção Esparsa por Tokens: Inferência Eficiente em Contextos Longos com Seleção 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

Resumo

A complexidade quadrática da atenção permanece o principal gargalo na inferência de contexto longo para grandes modelos de linguagem. Métodos de aceleração anteriores ou esparsificam o mapa de atenção com padrões estruturados ou removem permanentemente tokens em camadas específicas, o que pode reter tokens irrelevantes ou depender de decisões precoces irreversíveis, apesar da dinâmica camada-a-camada e cabeça-a-cabeça da importância dos tokens. Neste artigo, propomos a Token Sparse Attention, um mecanismo leve e dinâmico de esparsificação em nível de token que comprime Q, K, V por cabeça para um conjunto reduzido de tokens durante a atenção e, em seguida, descomprime a saída de volta para a sequência original, permitindo que a informação do token seja reconsiderada em camadas subsequentes. Além disso, a Token Sparse Attention expõe um novo ponto de projeto na interseção entre seleção de tokens e atenção esparsa. Nossa abordagem é totalmente compatível com implementações de atenção densa, incluindo o Flash Attention, e pode ser composta de forma transparente com *kernels* de atenção esparsa existentes. Resultados experimentais mostram que a Token Sparse Attention melhora consistentemente o *trade-off* entre precisão e latência, alcançando uma aceleração de atenção de até 3,23 vezes em contexto de 128K com menos de 1% de degradação de precisão. Estes resultados demonstram que a esparsificação dinâmica e intercalada em nível de token é uma estratégia complementar e eficaz para inferência escalável de contexto longo.

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.