轻量级稀疏路由块让大模型长文本推理快3.9倍且不丢精度
SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
- 通过可插拔的轻量选择器，动态选出每查询关注的前K个关键记忆
- 在128K上下文长度下保持密集模型准确率，解码速度比FlashAttention快3.9倍
- 支持量化到INT4/FP4微尺度，缓存缩小3.5倍无精度损失，适合部署
长序列已成为预训练大模型的标准，但其计算成本高昂：预填充阶段计算量随序列长度平方增长，解码时每步需重读线性增长的键值缓存。稀疏注意力通过仅关注过去相关标记来降低开销，但选择子集本身也代价不菲。我们提出SpotAttention，一种可插入冻结预训练变换器的轻量选择器，通过KL散度蒸馏学习估计注意力分布。该选择器为每个查询选取前K个键，因其输出为校准分布，可直接使用双层top-p规则读取每查询、每层的预算。在Qwen3（稠密型，4B-32B）和Qwen3.5（混合线性/全注意力，4B-9B）上，SpotAttention在长达128K标记的上下文中达到与稠密模型相当的精度，是训练长度的八倍。在L=128K时，解码速度比FlashAttention快3.9倍，比最强的无训练基线Twilight快1.8倍。将选择器的K缓存量化至INT4或FP4微尺度，可缩小3.5倍且无精度损失。
原文摘要 · Abstract (English)
Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。