STS: Efficient Sparse Attention with Speculative Token Sparsity

📄 arXiv: 2605.15508v1 📥 PDF

作者: Ceyu Xu, Jiangnan Yu, Yongji Wu, Yuan Xie

分类: cs.LG, cs.CL

发布日期: 2026-05-15

备注: 14 pages, 12 figures

💡 一句话要点

提出STS：一种基于推测Token稀疏性的高效稀疏注意力机制，加速LLM推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 稀疏注意力 大型语言模型 推测解码 模型加速 长序列处理

📋 核心要点

  1. LLM推理中注意力机制的二次复杂度导致严重的内存和计算瓶颈，限制了其在长序列任务中的应用。
  2. STS利用小draft模型预测大target模型的重要token，构建稀疏掩码，剪枝不重要的注意力计算，无需重新训练。
  3. 实验表明，STS在保持精度基本不变的情况下，实现了显著的加速，并在稀疏度-精度权衡方面优于现有技术。

📝 摘要（中文）

本文提出了一种名为STS的稀疏注意力机制，旨在解决大型语言模型（LLM）推理过程中因注意力机制的二次复杂度而导致的内存和计算瓶颈问题，尤其是在需要处理数百万token序列的新兴Agent应用中。STS无需模型重新训练，其核心思想是：较小的draft模型识别出的重要token能够高度预测较大target模型的重要token。通过集成到推测解码框架中，STS利用draft模型的注意力分数动态构建token和head级别的稀疏掩码，从而有效地剪枝target LLM中昂贵的注意力计算。评估结果表明，在代表性的NarrativeQA基准测试中，STS在约90%的稀疏度下实现了2.67倍的加速，与密集注意力相比，精度损失可忽略不计。STS在稀疏度-精度权衡方面建立了新的state-of-the-art，通过在给定的精度预算下实现更高的稀疏度水平，优于以往的技术。

🔬 方法详解

问题定义：大型语言模型（LLM）在处理长序列时，注意力机制的计算复杂度呈二次方增长，导致内存和计算瓶颈。这严重限制了LLM在需要处理数百万token序列的新兴Agent应用中的应用。现有的稀疏注意力方法通常需要重新训练模型，或者在精度上有所损失。

核心思路：STS的核心思想是利用一个较小的“draft”模型来预测较大“target”模型中哪些token是重要的，从而构建一个稀疏掩码，避免对所有token进行注意力计算。这种方法基于一个关键观察：draft模型识别出的重要token，对于target模型来说，也很有可能是重要的。通过这种方式，可以在target模型中只关注重要的token，从而降低计算复杂度。

技术框架：STS集成到推测解码框架中。首先，使用draft模型计算注意力分数。然后，基于这些分数，动态构建一个token和head级别的稀疏掩码。这个掩码用于剪枝target LLM中的注意力计算，只保留重要的token和head。最后，使用剪枝后的注意力计算结果进行后续的推理。整体流程包括draft模型注意力计算、稀疏掩码构建和target模型剪枝注意力计算三个主要阶段。

关键创新：STS的关键创新在于利用draft模型的注意力分数来动态构建稀疏掩码，从而避免了对target模型进行重新训练。此外，STS实现了token和head级别的细粒度稀疏性，能够更有效地剪枝不重要的注意力计算。与现有方法相比，STS能够在更高的稀疏度下保持较高的精度。

关键设计：STS的关键设计包括：1) 如何选择合适的draft模型；2) 如何根据draft模型的注意力分数来确定token和head的重要性；3) 如何有效地将稀疏掩码应用到target模型的注意力计算中。具体来说，可以使用注意力分数的阈值来确定token和head的重要性。阈值的选择需要根据具体的任务和模型进行调整，以达到最佳的稀疏度-精度权衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

STS在NarrativeQA基准测试中实现了显著的性能提升。在约90%的稀疏度下，STS实现了2.67倍的加速，与密集注意力相比，精度损失可忽略不计。此外，STS在稀疏度-精度权衡方面优于以往的技术，能够在给定的精度预算下实现更高的稀疏度水平。这些结果表明，STS是一种高效且有效的稀疏注意力机制。

🎯 应用场景

STS具有广泛的应用前景，尤其是在需要处理长序列的Agent应用中，例如对话系统、文档摘要、代码生成等。通过降低LLM推理的计算成本，STS可以使得这些应用能够更高效地运行，并能够处理更长的上下文信息。此外，STS还可以应用于资源受限的设备上，例如移动设备和嵌入式系统，从而使得LLM能够在这些设备上运行。

📄 摘要（原文）

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.