Apple Researchers Introduce Amortized Maximum Inner Product Search to Accelerate AI Retrieval Systems

At the core of modern artificial intelligence infrastructure—spanning recommendation systems, large-scale search engines, and retrieval-augmented generation (RAG) frameworks—lies a fundamental computational operation known as maximum inner product search, or MIPS. Given a query vector, a MIPS system must rapidly scour a database consisting of millions of candidate key vectors to identify the single entry that best matches the query. In the context of RAG, which underpins modern question-answering systems, the document or passage associated with that best-matching key is typically retrieved and prefixed to the user’s prompt to provide essential context.

However, as databases scale to millions of items and vector dimensions expand into the thousands, the computational cost of performing an exhaustive, brute-force comparison between a query and every single key becomes prohibitively expensive. To address this challenge, researchers have historically relied on a variety of approximate methods designed to trade off search accuracy for speed, typically leveraging geometric techniques such as clustering and quantization.

While these baseline methods rely on clever indexing structures built from the properties of the key set, they often suffer from a significant limitation: their underlying geometric arguments loosely assume that the incoming query distribution will mirror the distribution of the keys, or they bypass considerations of query characteristics entirely. This mismatch can cause severe performance bottlenecks, particularly in real-world applications like RAG. In practice, user-generated queries are typically short, conversational questions, whereas database keys represent long, information-dense, and meticulously curated text passages.

Furthermore, the vast majority of production-level MIPS deployments serve a stable, well-characterized base of active users whose query patterns remain predictable over time. Traditional, query-agnostic indexing methods leave this valuable behavioral signal entirely untapped. In modern machine learning, aligning optimization techniques with the actual source of incoming inputs—rather than striving for generalized performance across every theoretically possible input—is a foundational principle for maximizing computational efficiency.

Addressing this gap, researchers from Apple presented a paper at the International Conference on Machine Learning (ICML) titled "Amortizing Maximum Inner Product Search with Learned Support Functions." The work introduces an innovative approach that trains small neural networks once, offline, on a fixed database of keys. These lightweight networks learn to anticipate how typical queries are matched, effectively routing incoming requests faster and more accurately at inference time by amortizing upfront computational costs.

Rethinking Geometric Routing and Convex Analysis

Traditional MIPS systems fundamentally rely on the geometry of the key dataset to accelerate search operations. A common approach involves partitioning the database into a predetermined number of clusters, denoted as $k$, where $k$ is typically much smaller than the total number of keys, $n$. The system maintains a representative centroid for each cluster. When a new query arrives, it is compared against these $k$ centroids, and the system routes the query toward the cluster whose centroid it lies closest to.

Once inside the selected cluster, the query is compared exhaustively against all keys stored within that specific partition. If a database contains $n$ keys distributed evenly across $k$ clusters, the effective search space for the secondary comparison shrinks to roughly $n$ over $k$, yielding a theoretical factor-of-$k$ speedup over a naive, exhaustive search. Unfortunately, this centroid-based routing mechanism is inherently prone to error. Centroid routing is an approximation, and a cluster frequently contains keys that sit far from their assigned centroid yet happen to lie very close to a specific query. When a query targets such an outlier, centroid routing misjudges the cluster’s relevance, assigns it a low score, and skips it entirely, resulting in a missed nearest neighbor.

What makes MIPS particularly compelling from a mathematical and machine learning perspective is its foundation in a structure that had been largely overlooked prior to this research: the function that maps any given query to the largest inner product with any key in a set is mathematically equivalent to the support function of that set, a well-established concept within convex analysis.

This support function possesses two critical properties. First, it is inherently convex, meaning it can be effectively approximated by a specialized class of neural networks designed specifically for convex functions, known as input convex neural networks (ICNNs). Second, computing the gradient of that support function at any given query location directly yields the optimal matching key vector itself.

Building on these mathematical insights, the Apple researchers proposed learning a neural network capable of mimicking the ground-truth support function—or directly computing its gradient—at a fraction of the computational cost. This objective is achieved through regression, a classical machine learning technique. By precomputing a vast dataset of training examples composed of sampled input vectors, such as user-generated queries, paired with their corresponding MIPS outputs—either the optimal inner product score or the optimal key vector itself—a suitably parameterized neural network can be trained to replicate the desired output given the input.

SupportNet and KeyNet Architectures

Although solving the regression problem perfectly via the mathematical principles outlined above would theoretically allow a single neural network to both evaluate the support function and output the optimal key through automatic differentiation, empirical observations indicate that training separate, specialized networks for each task yields superior performance in practice.

The first architecture, designated as SupportNet, learns the support function directly by employing a loosely constrained input convex neural network. This methodology is mathematically rigorous and harmonizes seamlessly with the underlying convex structure of the problem. Researchers found that SupportNet excels when deployed to compute individual support functions for each data cluster, thereby dramatically improving routing quality.

The second architecture, KeyNet, models are trained to directly output the best-matching key vector when presented with a query. To ensure these predictions remain geometrically grounded in the underlying support function—despite the absence of explicitly enforced convexity—the researchers incorporated a consistency loss derived from Euler’s theorem for homogeneous functions.

Both model types are trained on supervised datasets sampled from the relevant query distribution. For every sampled query, the true ground-truth best match is computed beforehand via exhaustive search and stored, capturing the complete MIPS solution landscape for the specific distribution of queries that matter to the production environment.

Optimizing Partition Routing with SupportNets

At larger scales, partitioning keys into clusters and evaluating queries against cluster centroids provides a fast, coarse-grained routing rule. However, as noted, centroids fail to capture the internal dispersion of their respective clusters. If a cluster holds a key located far from the centroid but proximate to the query, centroid routing will overlook it.

Access to the true support function of each cluster transforms this dynamic, enabling precise routing because each support function quantifies exactly how well any key within a cluster can match the query with a high score. Leveraging this theoretical breakthrough, the researchers proposed learning these $k$ support functions jointly using a trained SupportNet capable of outputting $k$ scores simultaneously. The system then utilizes the largest of these scores to route incoming queries.

Unlike simple centroid comparisons, SupportNets address a much more meaningful question: how well does the absolute best key inside a cluster match the incoming query? These networks are trained via standard regression techniques by generating sample queries drawn from the target distribution, computing the ground-truth support functions at those query points, and training the neural network to replicate the values.

Experimental evaluations were conducted using SupportNets of varying capacities—extra-small, small, medium, and large, representing roughly 1 percent, 5 percent, 10 percent, and 20 percent of the database size—evaluated on the Natural Questions dataset partitioned into 128 clusters. Natural Questions is a widely recognized retrieval benchmark pairing general-knowledge questions with passages drawn from a vast document database comprising approximately 2.6 million unique passages.

Performance metrics evaluated retrieval efficacy in terms of recall relative to computational cost. Routing recall measures the fraction of queries whose true partition—the one containing the exact best key—is successfully captured among the clusters selected by the router. Evaluations demonstrated that learned routers achieve target retrieval performance at significantly lower latency than traditional centroid routing. For instance, the small SupportNet model achieved 95 percent routing recall roughly 1.9 times faster than standard centroid routing, recording a latency of 1.74 milliseconds compared to 3.36 milliseconds, with even greater efficiency gains observed in larger model variants. When evaluated strictly on routing quality independent of neural network execution overhead, learned routers identified the correct partition significantly more often, recovering the true partition in up to 92.6 percent of cases when scanning just 5 percent of the database, compared to 87.5 percent for the centroid baseline.

Enhancing Approximate Search Indices with KeyNets

While SupportNet routing replaces conventional centroid-based routing found in inverted file index (IVF) systems, the KeyNet approach offers an alternative integration path. KeyNet allows existing IVF infrastructure to remain entirely unchanged by acting as a specialized query rewriter positioned directly in front of a standard approximate search index.

Instead of passing the raw query to the index, KeyNet predicts where the optimal match is likely to reside within the embedding space, transforming the query before submission. Because standard approximate search indices are heavily optimized, they can locate target results significantly faster when assisted by a pre-processed query that is already close to its intended target.

KeyNets were evaluated across the same capacity tiers, demonstrating profound recall-versus-cost trade-offs using a FAISS IVF index configured on the Natural Questions dataset. When measuring search budgets via wall-clock latency, reaching 99 percent recall required up to eight times less time using the largest KeyNet model compared to utilizing raw queries directly. When examining retrieval quality relative to the fraction of the database scanned by the index, queries mapped through KeyNets retrieved the correct key between 89 percent and 97 percent of the time while scanning a mere 0.001 percent of the database, compared to 86 percent for raw queries. Reaching 99 percent recall required up to 11 times fewer scanned keys when utilizing KeyNet pre-processing.

Visual comparisons of top retrieval hits across various Natural Questions queries—such as inquiries regarding musical groups, literature, and television shows—demonstrate that utilizing KeyNet-mapped queries consistently yields more accurate top-3 results than querying with raw user text alone across the entire search budget spectrum.

Ultimately, the choice between deploying SupportNet or KeyNet depends on practical engineering constraints within a production pipeline. KeyNets are ideally suited for scenarios where inference speed is the primary bottleneck and approximating the best key directly is paramount. SupportNets are preferred when calibrated scores are required for downstream tasks such as re-ranking or thresholding, or when computational budgets necessitate smaller, highly competitive model footprints.

Both architectures share compatible training methodologies, integrate cleanly into existing approximate search infrastructure, and exhibit robust resilience to mild distribution shifts in user query patterns. In massive, hyper-scale MIPS deployments where key databases are exceptionally vast, production systems may even combine both approaches by deploying a SupportNet router across all primary clusters, followed by individual KeyNets within each cluster to accelerate final index convergence.

This amortized approach to MIPS reflects a fundamental shift in how information retrieval can be engineered when query distributions are understood in advance. By encoding MIPS solution landscapes directly into neural network weights, this methodology trades a one-time offline training investment for ultra-fast, lightweight inference, establishing a practical blueprint for building learned retrieval systems explicitly tailored to the real-world queries they serve.