Token Sparse Attention: Efficiënte Inferentie in Lange Contexten met Intergeplaatste Tokenselectie
Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection
February 3, 2026
Auteurs: Dongwon Jo, Beomseok Kang, Jiwon Song, Jae-Joon Kim
cs.AI
Samenvatting
De kwadratische complexiteit van aandacht blijft het centrale knelpunt bij inferentie met lange contexten voor grote-taalmodellen. Bestaande versnellingsmethoden sparsifiëren de aandachtmatrix met gestructureerde patronen of verwijderen tokens permanent op specifieke lagen, wat kan resulteren in het behoud van irrelevante tokens of onomkeerbare vroege beslissingen ondanks de laag- en hoofdafhankelijke dynamiek van tokenbelang. In dit artikel stellen we Token Sparse Attention voor, een lichtgewicht en dynamisch sparsificatiemechanisme op tokenniveau dat per hoofd de Q, K, V comprimeert naar een gereduceerde tokenset tijdens de aandachtswerking, en vervolgens de output terug decomprimeert naar de oorspronkelijke sequentie, waardoor tokeninformatie in volgende lagen opnieuw kan worden overwogen. Bovendien biedt Token Sparse Attention een nieuw ontwerppunt op het snijvlak van tokenselectie en sparse attention. Onze aanpak is volledig compatibel met dichte aandacht-implementaties, zoals Flash Attention, en kan naadloos worden gecombineerd met bestaande sparse attention-kernels. Experimentele resultaten tonen aan dat Token Sparse Attention consistent de nauwkeurigheid-latency trade-off verbetert, met een snelheidswinst van tot 3,23× voor aandacht bij een contextlengte van 128K en minder dan 1% nauwkeurigheidsverlies. Deze resultaten demonstreren dat dynamische en geïnterleefde sparsificatie op tokenniveau een complementaire en effectieve strategie is voor schaalbare inferentie met lange contexten.
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.