Amortizing Maximum Inner Product Search with Learned Support Functions

Theo X. Olausson, Joao Monteiro, Michal Klein, Marco Cuturi
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:94606-94639, 2026.

Abstract

Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring identification of database vectors that align most strongly with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of search across queries drawn from a known distribution. Our key insight is that the MIPS value function - the maximum inner product as a function of the query - is convex (as the pointwise maximum of linear functions), and its gradient at each query equals the optimal database vector. We explore two complementary architectures: (1) Input Convex Neural Networks (ICNNs) that learn the convex value function and recover the optimal match via gradient computation, and (2) VectorICNNs that directly regress the argmax, bypassing gradient computation entirely at inference time. For ICNNs, we combine score regression with gradient matching losses; for VectorICNNs, we introduce a score consistency loss derived from Euler’s theorem for homogeneous functions. We further propose homogenization wrappers that enforce positive 1-homogeneity, theoretically linking function values to gradients. Our experiments on retrieval benchmarks demonstrate that convexity provides an effective inductive bias, with learned potentials achieving high match rates while requiring only a single forward pass at inference. Our code is available at: https://github.com/apple/ml-amips.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-olausson26a, title = {Amortizing Maximum Inner Product Search with Learned Support Functions}, author = {Olausson, Theo X. and Monteiro, Joao and Klein, Michal and Cuturi, Marco}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {94606--94639}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/olausson26a/olausson26a.pdf}, url = {https://proceedings.mlr.press/v306/olausson26a.html}, abstract = {Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring identification of database vectors that align most strongly with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of search across queries drawn from a known distribution. Our key insight is that the MIPS value function - the maximum inner product as a function of the query - is convex (as the pointwise maximum of linear functions), and its gradient at each query equals the optimal database vector. We explore two complementary architectures: (1) Input Convex Neural Networks (ICNNs) that learn the convex value function and recover the optimal match via gradient computation, and (2) VectorICNNs that directly regress the argmax, bypassing gradient computation entirely at inference time. For ICNNs, we combine score regression with gradient matching losses; for VectorICNNs, we introduce a score consistency loss derived from Euler’s theorem for homogeneous functions. We further propose homogenization wrappers that enforce positive 1-homogeneity, theoretically linking function values to gradients. Our experiments on retrieval benchmarks demonstrate that convexity provides an effective inductive bias, with learned potentials achieving high match rates while requiring only a single forward pass at inference. Our code is available at: https://github.com/apple/ml-amips.} }
Endnote
%0 Conference Paper %T Amortizing Maximum Inner Product Search with Learned Support Functions %A Theo X. Olausson %A Joao Monteiro %A Michal Klein %A Marco Cuturi %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-olausson26a %I PMLR %P 94606--94639 %U https://proceedings.mlr.press/v306/olausson26a.html %V 306 %X Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring identification of database vectors that align most strongly with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of search across queries drawn from a known distribution. Our key insight is that the MIPS value function - the maximum inner product as a function of the query - is convex (as the pointwise maximum of linear functions), and its gradient at each query equals the optimal database vector. We explore two complementary architectures: (1) Input Convex Neural Networks (ICNNs) that learn the convex value function and recover the optimal match via gradient computation, and (2) VectorICNNs that directly regress the argmax, bypassing gradient computation entirely at inference time. For ICNNs, we combine score regression with gradient matching losses; for VectorICNNs, we introduce a score consistency loss derived from Euler’s theorem for homogeneous functions. We further propose homogenization wrappers that enforce positive 1-homogeneity, theoretically linking function values to gradients. Our experiments on retrieval benchmarks demonstrate that convexity provides an effective inductive bias, with learned potentials achieving high match rates while requiring only a single forward pass at inference. Our code is available at: https://github.com/apple/ml-amips.
APA
Olausson, T.X., Monteiro, J., Klein, M. & Cuturi, M.. (2026). Amortizing Maximum Inner Product Search with Learned Support Functions. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:94606-94639 Available from https://proceedings.mlr.press/v306/olausson26a.html.

Related Material