cs.DSOct 7, 2026

Attention via Black-Box Vector Search

Authors: Stepan Zharkov, Krish Singal, Ashwin Padaki, Alexandr Andoni

Organizations: University of Pennsylvania

Abstract

Sparse attention mechanisms estimate attention over nn 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 Θ(n/ε)Θ(\sqrt{n}/\varepsilon) retrieved keys are both sufficient and necessary. With Θ(log⁡n)Θ(\log n) indices, we give an algorithm that retrieves only O(log⁡n+1/ε2)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/ε2)O(1/\varepsilon^2) retrieved keys. When integrated into LLM inference, our algorithms outperform top-kk and sampling approaches used in prior work and yield attention approximation that scales favorably to long contexts.

Figures & tables

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Feb 9, 2026cs.LG

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.
May 7, 2026cs.LG

Sparse Attention as a Range Searching Problem: Towards an Inference-Efficient Index for KV Cache

Sparse attention improves LLM inference efficiency by selecting a subset of key-value entries, but at the cost of potential accuracy degradation. In particular, omitting critical KV entries can induce substantial errors in model outputs. Existing methods typically operate under fixed or adaptive token budgets and provide empirical robustness or partial theoretical guarantees, yet they do not ensure zero false negatives in decoding steps, particularly since the set of relevant tokens is both query- and step-dependent. Our empirical observations confirm that missing even one critical key can lead to sharp error spikes, especially in long reasoning tasks where the set of important tokens varies throughout decoding. This observation motivates the need for indexing methods that dynamically adapt to these variations across decoding steps while guaranteeing a full recall of the relevant keys above a certain threshold. We address this challenge by reformulating sparse attention as the halfspace range searching problem. However, existing range searching indices are not suitable for modern LLM inference due to their computational and implementation overheads. To overcome this, we introduce Louver, a novel index structure tailored for efficient KV cache retrieval. Louver (i) guarantees zero false negatives with respect to a specified threshold in both theory and practice, (ii) is lightweight to integrate into existing LLM pipelines, and (iii) incorporates hardware-aware optimizations for both CPU and GPU executions. Our experiments demonstrate that Louver outperforms prior sparse attention methods in both accuracy and runtime, and is faster than highly optimized dense attentions such as FlashAttention. These results highlight that recall guarantees are a critical and overlooked dimension of sparse attention, and open a new direction for building theoretically grounded, efficient KV cache indices.
May 7, 2026cs.DS

Nearly Optimal Attention Coresets

We consider the problem of estimating the Attention mechanism in small space, and prove the existence of coresets for it of nearly optimal size. Specifically, we show that for any set of unit-norm keys and values (K,V)(K,V) in Rd\mathbb{R}^d, there exists a subset (K′,V′)(K',V') of size at most O(deρ+o(ρ)/ε)O({\sqrt{d} e^{ρ+o(ρ)}/\varepsilon}) such that ∥Attn⁡(q,K,V)−Attn⁡(q,K′,V′)∥≤ε\left\| \operatorname{Attn}(q,K,V)- \operatorname{Attn}(q,K',V') \right\| \le \varepsilon simultaneously for all queries whose norm is bounded by ρρ. This outperforms the best known results for this problem. We also offer an improved lower bound showing that ε\varepsilon-coresets must have size Ω(deρ/ε)Ω({\sqrt{d} e^ρ/ε}).