Amortized Maximum Inner Product Search with Learned Support Functions

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.

Modern AI applications, from chatbots that retrieve relevant facts to recommendation systems, must constantly search through millions or billions of stored items to find the best match for a user's query. The cost of these ballooning search operation can be mitigated through e.g. geometric tools (such as clustering and inverted file indexing) or compression. We explore in this work a radically different approach: if the distribution of user queries is known (i.e. queries follow predictable patterns), a learning approach can be used instead: we can train a neural network to learn how to optimize this search problem in advance, and predict the answer directly at inference time, rather than spending time on search. We developed two complementary tools based on this idea. The first routes each query to the right part of the database. The second produces a transformed query that plugs directly into off-the-shelf search engines. No changes to the search infrastructure required, but search becomes significantly faster as the query is transformed into an approximation of its most similar key in the database before search is ran. On standard document retrieval benchmarks, our approach consistently improves the accuracy-versus-cost tradeoff, offering a practical path toward faster and more efficient retrieval at scale.