Token Sparse Attention: Effiziente Inferenz mit langem Kontext durch verschachtelte Token-Auswahl
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
Die quadratische Komplexität der Aufmerksamkeit bleibt das zentrale Hindernis bei der Inferenz mit langen Kontexten in großen Sprachmodellen. Bisherige Beschleunigungsmethoden verdünnen entweder die Aufmerksamkeitsmatrix mit strukturierten Mustern oder entfernen Token dauerhaft in bestimmten Schichten, was irrelevante Token beibehalten oder auf irreversible Frühentscheidungen angewiesen sein kann – trotz der schicht- und kopfabhängigen Dynamik der Token-Bedeutung. In dieser Arbeit schlagen wir Token Sparse Attention vor, einen leichtgewichtigen und dynamischen Sparsifikationsmechanismus auf Token-Ebene, der pro Kopf Q, K, V auf einen reduzierten Token-Satz während der Aufmerksamkeitsberechnung komprimiert und die Ausgabe dann zurück auf die ursprüngliche Sequenz dekomprimiert, wodurch Token-Informationen in nachfolgenden Schichten neu bewertet werden können. Darüber hinaus eröffnet Token Sparse Attention einen neuen Gestaltungsansatz an der Schnittstelle von Token-Auswahl und sparser Aufmerksamkeit. Unser Ansatz ist vollständig kompatibel mit Implementierungen dichter Aufmerksamkeit, einschließlich Flash Attention, und kann nahtlos mit existierenden sparsen Aufmerksamkeitskernen kombiniert werden. Experimentelle Ergebnisse zeigen, dass Token Sparse Attention durchgängig die Genauigkeits-Latenz-Abwägung verbessert und eine bis zu 3,23-fache Beschleunigung der Aufmerksamkeitsberechnung bei 128K Kontextlänge mit weniger als 1% Genauigkeitsverlust erreicht. Diese Ergebnisse demonstrieren, dass dynamische und verschachtelte Sparsifikation auf Token-Ebene eine komplementäre und effektive Strategie für skalierbare Inferenz mit langem Kontext darstellt.
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.