Attention parcimonieuse par jetons : Inférence efficace en contexte long avec sélection entrelacée de jetons
Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection
February 3, 2026
papers.authors: Dongwon Jo, Beomseok Kang, Jiwon Song, Jae-Joon Kim
cs.AI
papers.abstract
La complexité quadratique de l'attention reste le principal goulot d'étranglement dans l'inférence à contexte long pour les grands modèles de langage. Les méthodes d'accélération antérieures soit éparsifient la carte d'attention avec des motifs structurés, soit suppriment définitivement des tokens à des couches spécifiques, ce qui peut conserver des tokens non pertinents ou reposer sur des décisions précoces irréversibles, malgré la dynamique variable par couche et par tête de l'importance des tokens. Dans cet article, nous proposons Token Sparse Attention, un mécanisme d'éparsification léger et dynamique au niveau des tokens qui compresse les Q, K, V par tête vers un ensemble réduit de tokens pendant l'attention, puis décompresse la sortie vers la séquence originale, permettant de reconsidérer l'information des tokens dans les couches suivantes. De plus, Token Sparse Attention introduit un nouveau point de conception à l'intersection de la sélection de tokens et de l'attention éparse. Notre approche est entièrement compatible avec les implémentations d'attention dense, y compris Flash Attention, et peut être composée de manière transparente avec les noyaux d'attention éparse existants. Les résultats expérimentaux montrent que Token Sparse Attention améliore constamment le compromis précision-latence, atteignant jusqu'à 3,23 fois d'accélération de l'attention pour un contexte de 128K avec une dégradation de précision inférieure à 1%. Ces résultats démontrent que l'éparsification dynamique et entrelacée au niveau des tokens est une stratégie complémentaire et efficace pour une inférence à contexte long évolutive.
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.