토큰 희소 주의: 인터리브 토큰 선택을 통한 효율적인 장문 컨텍스트 추론

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

초록

어텐션의 2차 복잡도는 대규모 언어 모델의 장문 컨텍스트 추론에서 여전히 핵심 병목 현상으로 남아 있습니다. 기존 가속화 방법들은 구조화된 패턴으로 어텐션 맵을 희소화하거나 특정 계층에서 토큰을 영구적으로 제거하는 방식으로, 토큰 중요도의 계층/헤드별 역동성에도 불구하고 관련성 없는 토큰을 유지하거나 되돌릴 수 없는 초기 결정에 의존할 수 있습니다. 본 논문에서는 경량이면서 동적인 토큰 수준 희소화 메커니즘인 Token Sparse Attention을 제안합니다. 이는 어텐션 연산 중 헤드별 Q, K, V를 축소된 토큰 집합으로 압축한 후 출력을 원래 시퀀스로 다시 복원하여, 이후 계층에서 토큰 정보를 재고찰할 수 있도록 합니다. 더 나아가, Token Sparse Attention은 토큰 선택과 희소 어텐션의 교차점에 새로운 설계 지점을 제시합니다. 우리의 접근 방식은 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.