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

Maximum inner product search, commonly known as MIPS, occupies a foundational position in the modern machine learning ecosystem. It serves as the computational engine powering recommendation systems, web search engines, and retrieval-augmented generation, or RAG, architectures. When an AI or software system receives a query vector, the MIPS process must search through a vast database containing millions of candidate key vectors to identify the single entry that achieves the highest mathematical alignment with the query. In practical applications such as RAG—where a foundational language model relies on external knowledge—the data point associated with that best-matching key, typically an informative text document, is retrieved and prefixed to the user’s prompt to formulate an accurate, context-aware response.

As digital databases scale to hold tens of millions of keys and vector dimensions expand into the thousands, executing an exhaustive mathematical comparison between a single query and every individual key becomes computationally prohibitive. To address this bottleneck, researchers and engineers have historically relied on various approximate methods. These baseline techniques typically trade off a degree of search accuracy in exchange for computational speed by leveraging clever indexing structures built around clustering and vector quantization.

However, these traditional geometric methods focus exclusively on analyzing the intrinsic properties of the set of database keys while largely ignoring the predictable nature of the queries themselves. They either assume that incoming user queries will naturally mirror the distribution of the database keys, or they abandon query-aware considerations altogether. This oversight creates significant friction in real-world deployments.

In the critical context of RAG, user-generated questions are typically short, conversational, and direct, whereas database keys represent long, information-dense, and heavily curated passages of text. This structural mismatch means keys and queries occupy fundamentally different spaces. Furthermore, almost all production-level MIPS deployments serve a stable, well-characterized base of users whose query distribution remains steady over time. Standard, query-agnostic indexing methods leave this valuable behavioral signal entirely untapped. Modern machine learning has long established that designing systems tailored to a likely source of inputs—rather than optimizing for every conceivable input—directs computational resources far more efficiently.

Shifting Paradigms with Amortized MIPS

In a research paper presented at the International Conference on Machine Learning (ICML) 2026, titled Amortizing Maximum Inner Product Search with Learned Support Functions, a team of researchers from Apple details a novel approach to this challenge. Instead of treating every query as an isolated, unpredictable event, the team trains small neural networks—once, entirely offline, for a fixed database of keys—that learn how to route incoming queries toward their optimal matches with significantly greater speed.

This methodology trades an upfront amortization cost to understand how typical queries map to their corresponding matches in exchange for substantially faster, more accurate retrieval performance at inference time. The core principle underlying traditional MIPS is the exploitation of the geometric arrangement of the database keys. Standard indexing systems typically partition the set of keys into a specific number of clusters, denoted as $k$, where $k$ is kept much smaller than the total number of keys $n$. The system maintains a representative centroid for each cluster. When a new query arrives, the system compares it against these $k$ centroids, routes the query toward the cluster whose centroid it sits closest to, and then conducts an exhaustive comparison solely against the keys stored within that specific cluster.

If a database contains $n$ keys distributed evenly across $k$ clusters, the effective number of keys evaluated during the secondary phase drops to approximately $n/k$. This yields a theoretical factor-of-$k$ speedup over a naive exhaustive search across all $n$ items. Unfortunately, this geometric routing mechanism frequently introduces errors. Centroid routing is an approximation, and a nearest neighbor can easily be missed if it happens to reside closer to a different cluster’s boundaries rather than the nearest geometric centroid.

Unlocking Convex Analysis and Neural Networks

From a machine learning perspective, MIPS possesses a mathematical structure that had been largely overlooked prior to this research: the function mapping 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 thoroughly studied concept in convex analysis. This support function exhibits 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, or ICNNs. Second, calculating the gradient of this support function at any given query directly yields the best-matching key itself.

The overarching philosophy of Apple’s research is to learn a neural network capable of mimicking this ground-truth support function, or directly predicting its gradient, at a fraction of the computational cost. This objective is achieved through classical machine learning regression techniques. Researchers precompute a massive dataset of training examples by sampling input vectors from actual user queries and pairing them with the corresponding output signals produced by MIPS, such as optimal scores or optimal key vectors. A suitably parameterized neural network is then trained to replicate these exact outputs given the corresponding inputs.

While solving this regression problem perfectly would allow a single neural network to output the optimal key simply by evaluating its gradient via automatic differentiation, empirical observations indicate that training separate, specialized networks for distinct tasks yields superior results. This realization led to the development of two distinct architectural models.

The first architecture, designated as SupportNet, learns the support function directly using a loosely constrained input convex neural network. This mathematically principled approach aligns naturally with the convex structure of the problem and excels when utilized to compute individual support functions for each data cluster, thereby dramatically improving routing quality.

The second architecture, KeyNet, is designed to directly output the best-matching key vector when presented with a query. To maintain geometric stability without enforcing strict mathematical convexity, KeyNet models incorporate a specialized consistency loss derived from Euler’s theorem for homogeneous functions. Both network types are trained on supervised datasets sampled from relevant query distributions, capturing the exact solution landscape for the queries that matter most in production environments.

Enhancing Cluster Routing with SupportNets

In traditional inverted file index systems, scaling up search operations involves partitioning keys into clusters, comparing incoming queries against each cluster centroid, and routing queries toward the most promising partitions. While efficient, this coarse rule remains vulnerable to error because a single centroid cannot account for the internal spread of its cluster. A cluster may contain a key located far from its centroid yet remarkably close to a given query. When a query approaches such an outlier, centroid routing miscalculates the cluster score, bypasses the partition, and misses the true best match entirely

.

Providing the system with access to the exact support function of each cluster transforms this dynamic, enabling precise routing. Each support function mathematically quantifies the degree to which any key within a given cluster can match a query with a high inner product score. Building upon this theoretical framework, the researchers trained SupportNets capable of evaluating multiple support functions simultaneously to generate scores for $k$ clusters at once, using the highest score to guide the routing process.

When evaluated on the Natural Questions retrieval benchmark—comprising approximately 2.6 million unique passages partitioned into 128 clusters—SupportNets of varying sizes demonstrated substantial performance advantages over baseline centroid routing. Across rigorous evaluations measuring retrieval recall against latency and database scanning budgets, learned routers consistently outperformed traditional methods.

For instance, small SupportNet configurations achieved a 95% routing recall roughly 1.9 times faster than standard centroid routing, reducing latency from 3.36 milliseconds down to 1.74 milliseconds. Larger model variants pushed performance even higher, achieving 97% recall at identical latency thresholds. When isolating routing quality from neural network execution overhead, learned routers successfully identified the correct database partition between 90% and 93% of the time while scanning merely 5% of the database, compared to 87.5% for the centroid baseline.

Query Rewriting via KeyNets

While SupportNet routing replaces traditional centroid-based approaches in inverted file index architectures, KeyNet offers an alternative integration path by functioning as an intelligent query rewriter positioned in front of standard approximate search indices. Instead of modifying the underlying index structure, KeyNet predicts where the best-matching key is likely to reside within the embedding space and transforms the original query accordingly. Because standard approximate search indices are heavily optimized, they locate results much faster when assisted by a pre-processed query that already aligns closely with its intended target.

Evaluations using FAISS inverted file indices across the Natural Questions dataset demonstrated that KeyNet-mapped queries consistently outperformed raw queries across the entire search budget spectrum. Reaching a 99% recall rate required up to eight times less total time when utilizing the largest KeyNet model compared to executing searches with raw queries. Furthermore, when analyzing the exact fraction of the database scanned, KeyNet-mapped queries achieved retrieval accuracies between 89% and 97% while scanning a mere 0.001% of the database, compared to 86% for raw queries. Reaching peak recall thresholds required up to eleven times fewer scanned keys overall.

Practical deployment considerations dictate the choice between these two complementary architectures. Engineering teams can select KeyNet when inference latency represents the primary operational bottleneck and the objective is to rapidly approximate the best-matching key. Conversely, SupportNet is ideal when calibrated scoring is required for downstream re-ranking and thresholding, or when compute budgets at smaller model scales demand high efficiency. For massive, large-scale MIPS deployments managing extensive key databases, system architects may ultimately combine both approaches, deploying a SupportNet router across all overarching clusters while pairing individual KeyNets with specific cluster indices to accelerate final retrieval.

Ultimately, amortized MIPS represents a fundamental shift in how search and retrieval systems can be engineered when query distributions remain stable and predictable in advance. By encoding known search solutions directly into the weights of neural networks, this methodology effectively trades a one-time offline training expense for lightweight, rapid inference, paving the way for next-generation retrieval systems explicitly tuned to the specific needs of the users they serve.