arXiv:2602.03216cs.CLcs.LG2026-02被引 1

动态选择关键token,让大模型长文本推理更快更准。

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

  • 每层动态选关键token,可反复调整,不丢重要信息。
  • 128K上下文下速度提升3.23倍,准确率损失小于1%。
  • 兼容现有加速技术,适合需要长文本处理的场景。

注意力计算的二次复杂度仍是大语言模型长上下文推理的核心瓶颈。以往加速方法或采用固定模式稀疏化注意力矩阵,或在特定层永久丢弃token,可能导致保留无关内容,或依赖不可逆的早期决策,无法适应各层/头间token重要性的动态变化。本文提出Token Sparse Attention,一种轻量级、动态的粒度级稀疏化机制:在注意力计算中将每头的$Q$、$K$、$V$压缩至精简token集合,再将输出解压回原始序列,使token信息可在后续层被重新评估。该方法揭示了token选择与稀疏注意力结合的新设计空间。其完全兼容密集注意力实现(如Flash Attention),可无缝集成现有稀疏注意力核。实验表明,Token Sparse Attention持续改善准确率-延迟权衡,在128K上下文下实现最高3.23倍的注意力加速,且准确率损失低于1%。结果证明,动态、交错式粒度级稀疏化是可扩展长上下文推理的有效补充策略。

原文摘要 · Abstract (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 $\times$3.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.

长文本推理注意力机制稀疏化高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。