# Amortized Maximum Inner Product Search with Learned Support Functions

arXiv: https://arxiv.org/abs/2603.08001 · HTML: https://mlsanity.com/paper/2603.08001/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 38 references that are in the map.
- First posted on arXiv 9 Mar 2026.

## Builds on (most cited first, 27 of 27)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [The Case for Learned Index Structures](https://mlsanity.com/paper/1712.01208/) (2017)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)](https://mlsanity.com/paper/1405.5869/) (2014)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [Tutorial on amortized optimization](https://mlsanity.com/paper/2202.00665/) (2022)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Large-scale Multi-label Learning with Missing Labels](https://mlsanity.com/paper/1307.5101/) (2013)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Quantization based Fast Inner Product Search](https://mlsanity.com/paper/1509.01469/) (2015)
- [How to Scale Your EMA](https://mlsanity.com/paper/2307.13813/) (2023)
- [On Symmetric and Asymmetric LSHs for Inner Product Search](https://mlsanity.com/paper/1410.5518/) (2014)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Principled Weight Initialisation for Input-Convex Neural Networks](https://mlsanity.com/paper/2312.12474/) (2023)
- [SOAR: Improved Indexing for Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2404.00774/) (2024)
- [RoarGraph: A Projected Bipartite Graph for Efficient Cross-Modal Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2408.08933/) (2024)
- [Norm-Explicit Quantization: Improving Vector Quantization for Maximum Inner Product Search](https://mlsanity.com/paper/1911.04654/) (2019)
- [A Learning-to-Rank Formulation of Clustering-Based Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2404.11731/) (2024)
- [LeanVec: Searching vectors faster by making them fit](https://mlsanity.com/paper/2312.16335/) (2023)

## Built on by (most cited first, 1 of 1)

- [Nectar: Neural Estimation of Cached-Token Attention via Regression](https://mlsanity.com/paper/2605.09778/) (2026)

## Abstract (arXiv)

Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring the identification of a vector taken within a database (the keys) that best aligns with a given query. We propose amortized MIPS: a regression-based approach that trains neural networks to directly predict MIPS solutions, amortizing the cost of repeatedly solving MIPS for queries drawn from a known distribution over a fixed key database. Our key insight is that the MIPS value function is the support function of the set of keys, a well-studied convex function whose gradient yields the optimal key. This motivates two complementary amortized models: SupportNet, an input-convex neural network trained to regress the support function, and KeyNet, a vector-valued network that directly regresses the optimal key. SupportNet can serve as a cluster router, steering queries toward relevant database partitions, while KeyNet can be used as a drop-in replacement for the original query, fed directly to off-the-shelf indexing pipelines. Our experiments on the BEIR benchmark show that, for document embeddings, learned SupportNets and KeyNets significantly improve IVF match rates when accounting for compute effort, whether measured in FLOPs, number of probes, or wall-clock time. Our code is available at: https://github.com/apple/ml-amips.