Faster Retrieval for Better, More Efficient RAG

content type highlight published November 5, 2026 Faster Retrieval for Better, More Efficient RAG Maximum inner product search (MIPS) sits at the core of recommendation systems, search, and retrieval-augmented generation (RAG). Given a query vector, MIPS systems search within a database of millions of candidate key vectors for the one that best matches the query. The original data point associated with that best-matching key is then used to handle the query: in RAG, the document that is best suited to answer the question is typically prefixed to the query. As the number of keys grows to millions and vector dimensions reach thousands, the computational cost of running an exhaustive comparison of the query to all keys grows dramatically. Numerous approximate methods have been proposed to trade off accuracy for speed using clustering and quantization. While these baseline methods focus on the properties of the set of keys to build up clever indexing structures, their geometric arguments loosely assume that the set of queries itself will resemble that found in keys, or bypass any considerations about what the queries will look like. This can be problematic on two fronts: In the crucial RAG case, queries arise from short user-generated questions, whereas keys represent long, informative, and curated pieces of text, making keys and queries substantially different. Additionally, most MIPS use cases deployed in production target an existing base of users with a stable, well-characterized query distribution — a signal that generic, query-agnostic methods leave untapped. Crafting methods that work best for a likely source of inputs — rather than aiming for good performance on all possible inputs — is a central tenet of modern machine learning that helps redirect compute resources more efficiently. In Amortizing Maximum Inner Product Search with Learned Support Functions , a paper presented at ICML 2026, researchers from Apple train small neural networks (NNs) — once, offline, for a fixed database of keys — that learn how to route queries toward their best match faster. Their approach trades off an upfront amortization cost to understand how typical queries tend to be matched in exchange for faster, more accurate retrieval at inference time through these small NNs. The Basics of Maximum Inner Product Search A fundamental principle in MIPS is to leverage the geometry of the set of keys to speed up search: Typical MIPS systems split the set of keys into k clusters — k is typically much smaller than the total number of keys n — keeping track of a representative centroid for each cluster. When receiving a new query, the system compares it with these k centroids. The query is then routed toward the cluster of the representative centroid it is closest to, and subsequently compared exhaustively with all keys stored in that cluster. If the database has n keys and k evenly distributed clusters, the effective number of keys that will be considered in that second step is on the order of n / k — a factor-of- k speedup when compared to the n comparisons done in the naive exhaustive approach. Unfortunately, this routing mechanism can result in errors, since centroid routing is not an exact approach, as illustrated in Figure 1 below. MIPS as a Support Function What makes MIPS interesting from a machine learning perspective is that it is built on a mathematical structure that had been largely overlooked before this work: The function that maps a query to the largest inner product with any key in a set is the support function of that set, a well-studied object in convex analysis. This function has two important properties: First, it is convex, so it can be approximated by a class of neural networks specifically designed for convex functions, input convex neural networks (ICNNs) . Second, the gradient of that support function at any query returns the best-matching key itself. The broad idea behind this work is to learn a neural network that can mimic the ground-truth support function (or its gradient directly) for a much lower computational cost. This can be achieved through a classical ML technique known as regression: one simply needs to precompute a large number of training examples formed by sampling input vectors (user queries) and pair them with the relevant output signal produced by MIPS (either the optimal score or the optimal key vector itself). All that is then needed is to train a suitably parameterized neural network to replicate the output given the input. Two Ways to Learn Maximum Inner Product Search If the regression problem were to be solved perfectly, according to the mathematical approach outlined above, the same neural network that is trained to mimic the support function could also output the optimal key by evaluating its gradient (through automatic differentiation). In practice, we observe that training separate, specialized networks for each of these two tasks works best: The SupportNet neural network architecture can learn the support function directly, using a (loosely constrained) ICNN. The approach is mathematically principled and aligns with the convex structure of the problem. We find that this approach can shine when used to compute one support function per data cluster to improve routing quality. KeyNet models learn to directly output the best-matching key vector when fed a query. A consistency loss derived from Euler’s theorem for homogeneous functions keeps their predictions geometrically grounded in the underlying support function, even though no convexity is explicitly enforced. Both models are trained on examples drawn from the relevant query distribution: for each sampled query, the ground-truth best match is computed once via exhaustive search and stored, creating a supervised dataset that captures the MIPS solution landscape for the distribution of queries that actually matter. Learned Routing with SupportNet As shown in Figure 1 , the simplest way to handle larger scales is to split the keys into clusters, compare the query to each cluster’s centroid, and route the query to the most promising clusters based on these k scores. This provides a cheap yet coarse rule that can be prone to mistakes because a centroid alone cannot account for the spread of its cluster. Indeed, a cluster can hold a key that sits far from its own centroid yet close to the query, and when the query is near such an outlier, centroid routing scores the cluster low and skips it, missing the true best match. If one were given access to the support function of each cluster, then one would instead be able to do exact routing, since each support function quantifies exactly to what extent any key within each cluster can match the query with a large score. This is illustrated in Figure 2 , in which the support function of each of the four clusters is plotted for a query that varies in the space of embeddings. Leaning on this theoretical insight, we propose to learn these k support functions jointly with a trained SupportNet that can output k scores at once. We can then use the largest of these k scores to route any new query instead. The SupportNet model can answer a more useful question than a simple comparison to the centroid can: “How good is the best key inside this cluster for this query?” SupportNets are trained in a classic machine learning fashion, also known as regression: we generate sample queries that are drawn from the distribution (typically embeddings for text queries), compute the ground-truth support functions at those queries, and train a neural network to replicate these values. We trained SupportNets of increasing sizes (XS, S, M, and L → 1%, 5%, 10%, and 20% of the database size) on Natural Questions partitioned into 128 clusters. Natural Questions is a retrieval benchmark that pairs general-knowledge questions with passages from a large document database, roughly 2.6M unique pass