论文
排名前回想一下：相似性引导的 Top-$K$ 重用可实现高效的长上下文注意力
Recall Before You Rank: Similarity-Guided Top-$K$ Reuse for Efficient Long-Context Attention
摘要
Top-$K$ 稀疏注意力通过仅关注键值 (KV) 条目的一小部分子集来降低 Softmax 和值聚合的成本。然而，识别这个子集仍然需要根据完整的 KV 缓存对当前查询进行评分并执行全局 Top-$K$ 选择，从而使选择器成本与上下文长度呈线性关系，并限制了稀疏注意力对于长上下文解码的实际效率。在本文中，我们介绍了 ReTopK，这是一种 无需训练 方法，它通过重用历史检索决策来加速动态 Top-$K$ 注意力。 ReTopK 建立在这样的观察之上：类似的查询经常关注重叠的支持，并且部分重叠的支持仍然可以保留大部分 Exact Top-$K$ 注意力质量。对于每个注意力头，它维护历史查询支持对的有界缓存，为每个新查询检索最相似的缓存查询，将其存储的支持与最近的窗口联合起来，并使用精确的当前查询分数仅对结果紧凑候选集进行重新排序。当重用不可靠时，基于相似性的回退会调用完整历史 Exact Top-$K$，而定期精确刷新会限制缓存漂移。 ReTopK 保留完整的 KV 缓存，并且仅重用选定的索引，而不是历史分数、注意力权重或输出。在 16K--128K 的上下文中，ReTopK 在评估的近似方法中实现了最低的 PG19 困惑度以及最高的 NIAH 和 LongBench 分数。在 128K 且 $K=512$ 时，ReTopK 只比 Exact Top-$K$ 增加了 0.50\% 的困惑度，同时将注意力计算加速了 $3.07\times$。