arXiv：cs.LG· Stepan Zharkov, Krish Singal, Ashwin Padaki, Alexandr Andoni·· 4 小时前AI 评分40
通过黑盒向量搜索实现注意力机制
Attention via Black-Box Vector Search
AI 导读
研究提出用 priority sampling 框架统一稀疏注意力方法，并分析 MIPS 黑盒检索键数量的理论边界：单个 MIPS 索引下需检索 Θ(√n/ε) 个键，使用 Θ(log n) 个索引时仅需 O(log n+1/ε²) 个键且近最优。若允许增广键和查询，单索引加 O(1/ε²) 个键即可绕过该下界。集成到 LLM 推理后，该方法优于现有 top-k 与采样方案，更适配长上下文。
正文
Abstract:Sparse attention mechanisms estimate attention over $n$ tokens using a small subset of keys. Many existing approaches use maximum inner product search (MIPS) to retrieve the heaviest keys, which motivates the following question: given black-box access to a MIPS oracle, how many keys must be retrieved to output an $\varepsilon$-accurate attention estimate?
We answer this question by unifying prior approaches through the framework of priority sampling. With a single MIPS index, we show that $\Theta(\sqrt{n}/\varepsilon)$ retrieved keys are both sufficient and necessary. With $\Theta(\log n)$ indices, we give an algorithm that retrieves only $O(\log n+1/\varepsilon^2)$ keys and prove that this is near-optimal. More generally, we design algorithms that establish a smooth tradeoff between the number of MIPS indices and number of retrieved keys. We then show that if we allow augmentation of keys and queries, we can bypass the above lower bounds: there exists a simple priority-sampling estimator using a single MIPS index and $O(1/\varepsilon^2)$ retrieved keys. When integrated into LLM inference, our algorithms outperform top-$k$ and sampling approaches used in prior work and yield attention approximation that scales favorably to long contexts.
| Subjects: | Data Structures and Algorithms (cs.DS); Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.10135 [cs.DS] |
| (or arXiv:2610.10135v1 [cs.DS] for this version) | |
| https://doi.org/10.48550/arXiv.2610.10135 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Ashwin Padaki [view email]
[v1]
Wed, 7 Oct 2026 14:14:58 UTC (66 KB)
来源：arXiv：cs.LG · arxiv.org