Attention via Black-Box Vector Search

arXiv:2610.10135 · cs.DS, cs.LG · Submitted 2026-10-07 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Attention via Black-Box Vector Search".

Jane: The gist The authors answer whether given blackbox access to a Maximum Inner Product Search (MIPS) oracle, how many keys must be retrieved to output an ε-accurate attention estimate.

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So we’re talking about this paper today titled "Attention via Black-Box Vector Search". Basically, the authors are tackling a really fundamental question in how we can estimate attention using just a black-box access to a Maximum Inner Product Search oracle.

Jane: That’s right, Tom. The central thesis is figuring out exactly how many keys you need to retrieve from that search oracle to get an attention estimate that's accurate within some epsilon level, epsilon.

Tom: They’re unifying several previous ideas by using a framework called priority sampling to look at the trade-off between how many MIPS indices you build and how many actual keys you have to pull.

Lu: It’s interesting because they show that with just a single MIPS index, the number of retrieved keys both sufficient and necessary is (sqrt n /epsilon), which is a square root term involving n and epsilon >

Meng: That sounds computationally heavy for inference. So, what’s the alternative they propose when you use more indices?

Tom: Well, if you use (n) indices instead of just one, they can give you an algorithm that only needs to retrieve O(n + one/epsilon two) keys and prove that this is near-optimal <ref:2610.10135#pg1,keys and prove that this is near-optimal>.

Jane: That makes sense because it’s a logarithmic term instead of a square root term, which is much better for scaling up as the sequence length n gets really big.

Lu: The framework involves defining a meta-algorithm through priority sampling where each key gets a random scale, you call it u i from Unif(zero one], and the priority pi i is defined as w i/u i, and then you just retrieve the highest K+one priorities pi i to act as a threshold for reweighting the first K items.

Tom: Okay, so that’s the mechanism behind their approach to characterizing that tradeoff between oracle budget and retrieval cost.

Jane: And this leads into some pretty interesting theoretical results about adaptive and non-adaptive algorithms achieving these specific tradeoffs between retrieval cost and oracle cost.

Meng: I see a theorem here, Theorem one which shows for any oracle budget T greater than or equal to one, and any epsilon less than one, there’s an epsilon-accurate adaptive MIPS attention algorithm where the oracle cost is at most T and the retrieval cost is bounded by something like O(T + n one/(T+one) epsilon-two + two/(T+one)) <ref:2610.10135#pg3>.

Paper summary: Tom: That bound on retrieval cost seems pretty tight for an adaptive algorithm, showing how much you pay in terms of keys even when you have some budget left for the oracle calls.

Lu: But there's also Theorem four which establishes a general lower bound for any T, showing that any epsilon-accurate MIPS attention algorithm must have a retrieval cost of at least (n n n + one/epsilon two) <ref:2610.10135#pg3,log n \log \log n>.

Jane: That lower bound is pretty strong because it shows that no matter what clever algorithm you design, you can’t do better than this minimum number of keys required to get that accuracy.

Tom: And then they have this concept called Softmax Lift, where a slight relaxation of the model allows them to implement priority sampling much more efficiently by lifting the keys and queries into one extra dimension.

Meng: Lifting them into an extra dimension sounds like it simplifies the search structure significantly, which is what we want when dealing with these high-dimensional vector searches.

Lu: Theorem five proves that there’s an epsilon-accurate attention algorithm Z where, after lifting the keys and queries into R(d+one), the oracle cost is one, and the retrieval cost drops to O(one/epsilon two) <ref:2610.10135#pg3>.

Jane: That result is significant because it effectively bypasses that lower bound of (sqrt n /epsilon) that they showed for the case where you only use a single MIPS index.

Tom: So, it means they found a way to get better performance by changing the underlying model assumptions, moving from the black-box search directly to this lifted version.

Meng: I’m curious about how this translates into something practical for running models, like Llama-three point two-1B-Instruct on long contexts <ref:2610.10135#pg3>.

Jane: The empirical testing integrates these algorithms into Llama three point two-1B Instruct long context inference using a custom harness, and the Softmax Lift algorithm consistently outperforms both TOPK and TOPK+SAMPLING retrieval strategies in terms of quality to efficiency tradeoff across various oracles <ref:2610.10135#pg3>.

Tom: It seems like the authors are suggesting that this Softmax Lift approach provides the best quality to efficiency tradeoff among all those tested retrieval strategies.

Lu: The empirical results suggest that Softmax Lift scales well even up to 1M tokens in their testing, and they found it was quite robust to approximate search oracles <ref:2610.10135#pg1>.

Paper summary: Jane: They also proved lower bounds by reducing the problem down to signed sum estimation, showing that any epsilon-accurate attention algorithm must have a retrieval cost of (n n n + one/epsilon two) <ref:2610.10135#pg1>.

Tom: So, while they have these theoretical limits on how many keys you need, the Softmax Lift algorithm gives us a concrete way to achieve good performance in practice.

Meng: It’s interesting how they show that combining Softmax Lift with an approximate oracle appears to be a promising overall candidate for real-world applications.

Jane: They leave open the theoretical problem of designing algorithms that are provably robust to approximate search oracles, which is something they haven't fully settled yet.

Tom: The paper "Attention via Black-Box Vector Search" shows us how to systematically characterize the tradeoff between using MIPS indices and retrieving keys for attention estimation.

Lu: It really highlights how agnostic you can be about the search implementation if you stick to this framework of priority sampling, even when dealing with a black-box oracle.

Meng: From an engineering standpoint, it tells us that we might not need to implement the most complex MIPS algorithms ourselves if we focus on these smarter retrieval strategies instead.

Jane: This paper suggests that for attention estimation in large models, focusing on how you structure your key and query augmentation can yield much better performance than just relying on simple top-k methods.

Tom: So, to summarize the main point of "Attention via Black-Box Vector Search," they give us a way to precisely map out the cost—both in terms of oracle calls and retrieval keys—for achieving a specific accuracy level in attention estimation.

Lu: It’s about taking the complexity away from just brute-forcing the search and instead using structured sampling ideas to manage that complexity.

Meng: The implication for us is that we can look at these lower bounds, like that (n n n + one/epsilon two) one, and use it to set realistic expectations for how much data we actually need to process <ref:2610.10135#pg1>.

Jane: And the authors’ Softmax Lift algorithm shows us a concrete method that gets close to those theoretical limits in a way that works well on actual models.

Tom: That's what they call the Softmax Lift algorithm, and it seems like it's making attention estimation much more scalable for large context tasks than previous methods.

Conclusion: Tom: So we've been diving into this paper, "Attention via Black-Box Vector Search."

Jane: Yeah, it’s about figuring out how many keys you actually need to pull from a black-box search to get an attention estimate that's accurate within some error level.

Tom: Exactly. The authors are looking at the trade-off between how much work you do in searching and how much data you have to retrieve.

Jane: It’s essentially mapping out that tradeoff using priority sampling as their main framework, showing how many keys are needed for a given number of search indices used.

Tom: And they prove that with a smart setup, you can get near-optimal results by only retrieving O(n + one/epsilon two) keys when you use logarithmic indices.

Jane: That’s pretty cool because it suggests we don't need to retrieve a huge amount of data just to get a decent attention score when we structure the search right.

Tom: The big picture here is that they provide these concrete limits, these lower bounds on how many keys are necessary, and then introduce this Softmax Lift method that actually performs well in practice.

Jane: Right. It seems like the authors are showing us a way to make attention estimation much more efficient for large models by changing how we look at the search process itself.

Tom: So, what does this mean for us right now? It's about understanding the fundamental costs involved when we try to use vector search for attention in these massive AI systems.

Jane: It really shifts our focus from just brute-forcing retrieval to using these more structured sampling ideas, which is a big step.

Tom: We'll be talking next on how this actually plays out in real-world inference with Llama models and what it means for scaling those things up.

Columbia University · University of Pennsylvania

cs.DS, cs.LG

Submitted: 2026-10-07

Updated: 2026-10-09

Comments: added funding acknowledgement

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 82/100

The gist: The gist The authors answer whether given blackbox access to a Maximum Inner Product Search (MIPS) oracle, how many keys must be retrieved to output an ε-accurate attention estimate.

Key concepts

Priority Sampling: This framework assigns a random scale to each key and calculates a priority based on its weight divided by this scale. The goal is then to retrieve the top priorities, which serve as a threshold to reweight the most important items, helping select relevant keys efficiently.
Softmax Lift: This technique relaxes model assumptions by lifting keys and queries into an extra dimension. This modification allows for a more efficient implementation of priority sampling, achieving high accuracy with only one oracle query while significantly reducing the required key retrieval cost.
MIPS Oracle: This is a black-box search mechanism where the algorithm can query for the maximum inner product between a given vector and keys stored in an index. The algorithm must estimate attention using only these limited, costly queries instead of full access to the underlying data.
Retrieval Cost vs. Oracle Cost Tradeoff: This refers to balancing two costs: how many keys the algorithm has to retrieve from the index (retrieval cost) versus how many times it must query the expensive black-box MIPS oracle (oracle cost). The paper shows algorithms that optimize this balance for a given accuracy requirement.

Summary

The gist The authors answer whether given blackbox access to a Maximum Inner Product Search (MIPS) oracle, how many keys must be retrieved to output an ε-accurate attention estimate.

How it works

The paper unifies prior approaches through the framework of priority sampling to characterize the tradeoff between the number of MIPS indices used by an algorithm and the number of keys it retrieves for an ε-accurate estimate. They show that with a single MIPS index, Θ(√n/ε) retrieved keys are both sufficient and necessary (Page 2). With Θ(log n) indices, they give an algorithm that retrieves only O(log n + 1/ε2) keys and prove that this is near-optimal (Page 2).

The framework involves defining a meta-algorithm via priority sampling where each key is assigned a random scale ui ∼ Unif(0, 1] and the priority pi:= wi/ui (Page 6). The goal is to retrieve the highest K + 1 priorities πi, which acts as a threshold to reweight the first K items (Page 6).

Theoretical Results

The authors provide adaptive and non-adaptive algorithms achieving specific tradeoffs between retrieval cost and oracle cost (Page 2). Theorem 1 shows that for any oracle budget T ≥ 1 and parameter ε ∈ (0, 1), there is an ε-accurate adaptive MIPS attention algorithm A with Coracle(A) ≤ T and Cretrieval(A) ≤ O(T + n1/(T +1)ε−2+2/(T +1)) (Page 3). Theorem 4 establishes a general lower bound for any T, showing that any ε-accurate MIPS attention algorithm must have Cretrieval(A) ≥ omegamin n, log n log log n + 1/ε2 (Page 3).

Softmax Lift

They show that a slight relaxation of the model allows priority sampling to be implemented much more efficiently by lifting keys and queries into one extra dimension, leading to the SOFTMAXLIFT algorithm (Page 2). Theorem 5 proves that there is an ε-accurate attention algorithm Z that, after lifting the keys and queries into R(d+1), has Coracle(Z) = 1 and Cretrieval(Z) = O(1/ε2) (Page 3). This bypasses the lower bound of omega√n/ε for T=1 (Page 6).

Empirical Results

The empirical testing integrates these algorithms into Llama-3.2-1B-Instruct long context inference using a custom harness (Page 4). SOFTMAXLIFT consistently outperforms TOPK and TOPK+SAMPLING retrieval strategies in terms of quality–efficiency tradeoff across various oracles (Page 8). The results suggest that SOFTMAXLIFT provides the best quality–efficiency tradeoff among tested retrieval strategies (Page 8).

Lower Bounds for MIPS Attention

The authors prove lower bounds by reducing the problem to signed sum estimation (Page 9). Theorem 2.2 proves that any ε-accurate signed sum estimation algorithm has retrieval cost omegamin n, log n log log n + 1/ε2 (Page 23). This result implies that any ε-accurate attention algorithm must retrieve at least this many keys regardless of its oracle cost (Page 23).

The paper concludes that the SOFTMAXLIFT algorithm, when integrated into LLM inference, substantially improves the accuracy–latency tradeoff over top-k and sampling approaches common in prior work (Page 8). The experiments provide preliminary evidence that SOFTMAXLIFT scales well to even larger contexts up to 1M tokens (Page 8). The findings suggest that combining SOFTMAXLIFT with an approximate oracle appears to be a promising overall candidate (Page 8). The results indicate that they are very robust to approximate search oracles (Page 9). The authors leave open the theoretical problem of designing algorithms provably robust to approximate search oracles (Page 9).

The paper uses generative AI tools to assist in the construction and analysis of all Algorithms and Theorems, including providing ingredients for proving the quality claim and assistance in writing the proofs, and in writing code for our implementations (Page 10). The authors take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI (Page 9). The paper is available as arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 (Page 4). The references section lists numerous works, including those by Mussmann and Ermon (Page 9). The paper is published in arXiv:2610.10135v1 [cs.DS] 7 Oct 2026 <ref:2610.

Improvements for AI systems

Bold key retrieval strategy selection: The improved system can dynamically choose between strategies like TOPK, TOPK+SAMPLING, or SOFTMAXLIFT based on real-time latency versus required accuracy, as demonstrated by SOFTMAXLIFT consistently outperforming prior approaches in the empirical results.
Adaptive oracle budget management: The system can implement an adaptive query strategy, where it queries oracles based on the frontier of bucket priorities to ensure that the (K + 1) keys with the largest priorities seen coincide exactly with the global K + 1 keys of highest priority, as described in Algorithm 3.
Context length scaling: The system can achieve sparse attention at context lengths up to 1M tokens by utilizing SOFTMAXLIFT, which substantially improves the accuracy–latency tradeoff over the top-k and sampling approaches common in prior work and shows preliminary evidence that SOFTMAXLIFT can enable efficient sparse attention at context lengths of up to 1M tokens.
Low-cost approximation for signed sums: The system can perform epsilon-accurate estimation of signed sums with an oracle cost of 1, achieving a retrieval cost of omega(min(n, √n/ε)) by leveraging the reduction shown in Theorem 14.
Efficient long-context attention: By integrating SOFTMAXLIFT with approximate oracles, the system can generate answer tokens at a roughly 3x larger rate than the exact attention done on the GPU while still succeeding on some tasks for 1M context examples.

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 epsilon-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 Θ(sqrt n /epsilon) retrieved keys are both sufficient and necessary. With Θ(n) indices, we give an algorithm that retrieves only O(n+1/epsilon 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/epsilon 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.