通过预测重要词元实现高效稀疏注意力，推理速度提升2.67倍。
STS: Efficient Sparse Attention with Speculative Token Sparsity
- 利用小模型预测大模型的关键词元，动态生成稀疏掩码。
- 在NarrativeQA上实现90%稀疏度，速度提升2.67倍，精度损失极小。
- 无需重训练，适合长序列推理与智能体应用。
注意力机制的二次复杂度对大语言模型推理造成严重内存与计算瓶颈，尤其在需处理千万级词元序列的新兴智能体应用中更为突出。本文提出STS，一种无需模型重训练的稀疏注意力机制。其核心思想是：小规模草稿模型识别的重要词元，高度预测大规模目标模型中的关键词元。通过集成至推测解码框架，STS复用草稿模型的注意力得分，动态构建词元与头级别的稀疏掩码，有效剪枝目标LLM中昂贵的注意力计算。评估表明，STS在代表性基准NarrativeQA上实现了约90%的稀疏度，推理速度提升2.67倍，与密集注意力相比精度下降可忽略。该方法在稀疏度-精度权衡上达到新SOTA，显著优于先前技术，在相同精度预算下支持更高稀疏度。
原文摘要 · Abstract (English)
The quadratic complexity of attention imposes severe memory and computational bottlenecks on Large Language Model (LLM) inference. This challenge is particularly acute for emerging agentic applications that require processing multi-million token sequences. We propose STS, a sparse attention mechanism that requires no model retraining. STS leverages the key insight that tokens identified as important by a smaller draft model are highly predictive of important tokens for a larger target model. By integrating into speculative decoding frameworks, STS repurposes the draft model's attention scores to dynamically construct a token-and-head-wise sparsity mask. This mask effectively prunes the expensive attention computation in the target LLM. Our evaluation shows that STS achieves a 2.67x speedup operating at approximately 90% sparsity on representative benchmark NarrativeQA, maintaining negligible accuracy degradation compared to dense attention. STS establishes a new state-of-the-art on the sparsity-accuracy trade-off, outperforming prior techniques by enabling higher sparsity levels for a given accuracy budget.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。