Amortized MIPS: SupportNet and KeyNet improve IVF match rates

Amortizing Maximum Inner Product Search with Learned Support Functions, published July 2026 at ICML, proposes amortized MIPS: neural networks trained to predict MIPS solutions for queries drawn from a known distribution over a fixed key database. The paper is authored by Theo X. Olausson, João Monteiro, Michal Klein and Marco Cuturi; the authors note † MIT and ** work done while at Apple. The authors provide code at https://github.com/apple/ml-amips.

What does the paper introduce?

The paper introduces amortized MIPS, two learned models (SupportNet and KeyNet) and the observation that the MIPS value function is the support function of the set of keys, whose gradient yields the optimal key. SupportNet is an input-convex neural network trained to regress that support function; KeyNet is a vector-valued network that directly regresses the optimal key. SupportNet can act as a cluster router to steer queries toward relevant database partitions. KeyNet can be fed as a drop-in replacement for the original query to off-the-shelf indexing pipelines.

The authors position this as a regression-based approach that amortizes the repeated cost of solving MIPS when queries are drawn from a known distribution and the key database is fixed. The paper frames the approach around document embeddings and retrieval workloads.

How do SupportNet and KeyNet work?

SupportNet regresses the support function using an input-convex architecture; KeyNet directly outputs a vector intended to match the optimal key for a given query. SupportNet leverages convexity so that its gradient points to the maximizer of the inner product across the key set, while KeyNet skips that intermediate step and predicts the key vector itself.

Concretely, the paper treats the MIPS value function as a convex support function. Quoting the paper: "the MIPS value function is the support function of the set of keys," and the gradient of that function yields the optimal key. That insight motivates training SupportNet to fit the support function and using its gradient as a router, and training KeyNet to map queries straight to key vectors that can be returned or passed into existing index structures.

What were the experimental results?

On the BEIR benchmark for document embeddings the paper reports that learned SupportNets and KeyNets significantly improve IVF match rates when accounting for compute effort. The authors evaluate improvement under three compute-aware measures: floating-point operations (FLOPs), number of probes, and wall-clock time. The experiments show the learned models achieve higher IVF match rates relative to baselines once compute effort is considered.

The evaluation focuses on document-embedding retrieval tasks and positions the learned regressors either as routers that reduce search scope or as query replacements compatible with existing indexing pipelines. The paper’s code is available on GitHub at https://github.com/apple/ml-amips for others to reproduce or extend the experiments.

Why it matters

Amortizing MIPS shifts work from per-query search toward up-front learned prediction, which can lower repeated retrieval costs for workloads with predictable query distributions. For systems that rely on IVF-style indexes, a learned router or a predicted query vector could reduce the number of probes or FLOPs required to hit the same or better match rates, turning raw compute budgets into improved retrieval quality or latency. The approach also gives system designers two concrete integration paths: a router that narrows search and a direct query regressor that plugs into existing pipelines.

What to watch

Look for replication across other retrieval benchmarks and embedding types beyond BEIR, and for experiments that quantify absolute latency and end-to-end cost when the learned models run alongside standard indexes. Adoption signals will include forks and issues against the project at https://github.com/apple/ml-amips and follow-up work that compares routing versus direct-query replacement in production-style settings.