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 \emph{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 \SupportNet{}s and \KeyNet{}s 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.

翻译：最大内积搜索（MIPS）是机器学习中的关键子程序，需要从数据库（键）中找出与给定查询最匹配的向量。我们提出摊销MIPS：一种基于回归的方法，训练神经网络直接预测MIPS解，从而摊销对固定键数据库上已知分布查询反复求解MIPS的计算成本。我们的关键洞察是MIPS值函数实际上是键集合的\emph{支持}函数——一种经过充分研究的凸函数，其梯度能给出最优键。这激发了两种互补的摊销模型：SupportNet（一种输入凸神经网络，用于回归支持函数）和KeyNet（一种向量值网络，直接回归最优键）。SupportNet可作为聚类路由器，将查询引导至相关数据库分区，而KeyNet可作为原始查询的即插即用替代品，直接输入现成索引流水线。我们在BEIR基准上的实验表明，对于文档嵌入，无论以FLOPs、探测次数还是挂钟时间衡量计算开销，学习到的SupportNet和KeyNet均能显著提升IVF匹配率。我们的代码开源在：https://github.com/apple/ml-amips。