通过黑盒向量搜索实现注意力

arXiv:2610.10135v1 公告类型：cross 摘要：稀疏注意力机制使用一小部分键来估计 $n$ 个 token 上的注意力。许多现有方法使用最大内积搜索（MIPS）来检索权重最大的键，这引出了以下问题：在以黑盒方式访问 MIPS oracle 的情况下，为了输出一个 $\varepsilon$-准确的注意力估计，需要检索多少个键？我们通过优先级采样框架统一以往方法，并回答了这一问题。使用单个 MIPS 索引时，我们证明检索 $\Theta(\sqrt{n}/\varepsilon)$ 个键既是充分的，也是必要的。使用 $\Theta(\log n)$ 个索引时，我们给出了一种仅检索 $O(\log n+1/\varepsilon^2)$ 个键的算法，并证明该结果已接近最优。更一般地，我们设计了多种算法，在 MIPS 索引数量与检索键数量之间建立了平滑的权衡关系。随后我们证明，如果允许对键和查询进行增广，就可以绕过上述下界：存在一种简单的优先级采样估计器，仅使用单个 MIPS 索引和 $O(1/\varepsilon^2)$ 个检索键。在集成到 LLM 推理中后，我们的算法优于以往工作中使用的 top-$k$ 和采样方法，并能以有利于长上下文的方式扩展注意力近似。