トークン疎注意：インタリーブ型トークン選択による効率的な長文脈推論

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

February 3, 2026

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

cs.AI

要旨

大規模言語モデルにおける長文脈推論において、注意機構の二次計算複雑性は依然として中核的なボトルネックである。従来の高速化手法は、構造化パターンによる注意マップのスパース化、または特定層でのトークンの恒久的除去のいずれかを採用してきたが、これらは無関係なトークンを保持するか、トークン重要性が層・ヘッド単位で動的に変化するにも関わらず、不可逆的な早期決定に依存する問題があった。本論文では、Token Sparse Attentionを提案する。これは軽量かつ動的なトークンレベルのスパース化機構であり、注意機構処理においてヘッド単位のQ、K、Vを削減されたトークン集合に圧縮し、出力を元の系列に逆変換することで、後続層においてトークン情報を再評価可能にする。さらに、本手法はトークン選択とスパース注意の交差点に新たな設計ポイントを提供する。我々のアプローチはFlash Attentionを含む密な注意機構の実装と完全に互換性があり、既存のスパース注意カーネルとシームレスに組み合わせることが可能である。実験結果では、Token Sparse Attentionが精度と遅延のトレードオフを一貫して改善し、128Kトークンの文脈長において精度劣化1%未満で最大3.23倍の注意機構高速化を達成した。これらの結果は、動的かつ層間で交互に行われるトークンレベルのスパース化が、スケーラブルな長文脈推論における相補的かつ効果的な戦略であることを示唆する。

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.