通过复用历史查询相似性，加速长文本注意力计算。
Recall Before You Rank: Similarity-Guided Top-$K$ Reuse for Efficient Long-Context Attention
- 基于查询相似性复用过往选中的键值对索引，减少重复计算。
- 在128K上下文长度下，仅增加0.50%困惑度，速度提升3.07倍。
- 无需训练，适合长文本生成与高效推理场景。
Top-K稀疏注意力通过只关注少量键值对来降低Softmax和值聚合的开销。然而，确定这些键值对仍需对完整键值缓存进行评分并执行全局Top-K选择，导致选择器成本随上下文长度线性增长，限制了稀疏注意力在长上下文解码中的实际效率。本文提出ReTopK，一种无需训练的方法，通过复用历史检索决策加速动态Top-K注意力。ReTopK基于相似查询常指向重叠支持区域的观察，并发现部分重叠的支持仍可保留绝大部分精确Top-K注意力质量。每个注意力头维护一个有限的历史查询-支持对缓存，对新查询检索最相似的缓存查询，将其存储的支持与近期窗口合并，仅对所得紧凑候选集使用当前查询得分重新排序。基于相似性的回退机制在复用不可靠时触发全历史精确Top-K，周期性精确刷新则限制缓存漂移。ReTopK保留完整的键值缓存，仅复用选定的索引，而非历史得分、注意力权重或输出。在16K–128K上下文范围内，ReTopK在所评估的近似方法中实现了最低的PG19困惑度，以及最高的NIAH和LongBench得分。在128K上下文且K=512时，其困惑度仅比精确Top-K高出0.50%，同时将注意力计算速度提升3.07倍。
原文摘要 · Abstract (English)
Top-$K$ sparse attention reduces the cost of Softmax and value aggregation by attending to only a small subset of key--value (KV) entries. However, identifying this subset still requires scoring the current query against the full KV cache and performing global Top-$K$ selection, leaving selector cost linear in context length and limiting the practical efficiency of sparse attention for long-context decoding. In this paper, we introduce ReTopK, a training-free method that accelerates dynamic Top-$K$ attention by reusing historical retrieval decisions. ReTopK builds on the observation that similar queries often attend to overlapping supports and that partially overlapping supports can still preserve most of the Exact Top-$K$ attention mass. For each attention head, it maintains a bounded cache of historical query--support pairs, retrieves the most similar cached queries for each new query, unions their stored supports with a recent window, and reranks only the resulting compact candidate set using exact current-query scores. A similarity-based fallback invokes full-history Exact Top-$K$ when reuse is unreliable, while periodic exact refreshes limit cache drift. ReTopK retains the complete KV cache and reuses only selected indices, rather than historical scores, attention weights, or outputs. Across 16K--128K contexts, ReTopK achieves the lowest PG19 perplexity and the highest NIAH and LongBench scores among the evaluated approximate methods. At 128K with $K=512$, ReTopK incurs only a 0.50\% perplexity increase over Exact Top-$K$ while accelerating attention computation by $3.07\times$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。