最大内积搜索（MIPS）是机器学习中的一项基础操作，其目标是在给定的键数据库中找到与查询向量内积最大的键向量。许多应用场景中，查询往往从已知分布中反复采样，而键数据库保持不变。针对这一特点，来自MIT和Apple的研究者提出了一种名为“摊销MIPS”（Amortized MIPS）的回归方法，旨在通过训练神经网络直接预测MIPS的解，从而避免对每个查询都执行昂贵的精确搜索，将计算成本分摊到训练阶段。
论文的核心理论洞见在于：MIPS值函数实际上是键集合的支持函数（support function），这是一个经过深入研究的凸函数，其梯度恰好对应于最优键。这一性质为设计摊销模型提供了坚实的数学基础。基于此，作者提出了两种互补的神经网络架构。第一种是SupportNet，它是一种输入凸神经网络（input-convex neural network），专门训练来回归支持函数。SupportNet的输出可以作为集群路由器，将查询引导至相关的数据库分区，从而缩小搜索范围。第二种是KeyNet，它是一个向量值网络，直接回归最优键。KeyNet生成的键可以充当原始查询的替代品，直接输入到现成的索引管线中，无需修改现有系统。
为了验证方法的有效性，研究者在BEIR基准上进行了实验，重点关注文档嵌入场景。实验结果表明，当综合考虑计算量（分别以FLOPs、探测数和实际运行时间衡量）时，学习得到的SupportNet和KeyNet能够显著提升IVF（倒排文件）索引的匹配率。这意味着在相同的计算预算下，摊销方法能够找到更相关的键，从而提升检索质量。论文还提供了开源代码，地址为https://github.com/apple/ml-amips，方便社区复现和扩展。
该工作将凸分析中的支持函数概念与神经网络回归相结合，为MIPS问题提供了一种新颖的摊销视角。与传统的精确搜索或近似搜索不同，摊销MIPS通过利用查询分布的先验知识，在训练阶段学习查询到解的映射，从而在推理阶段实现快速预测。这种方法特别适用于查询模式稳定、数据库固定的场景，例如推荐系统或信息检索中的嵌入检索。未来，这一思路可能扩展到其他具有凸结构的搜索问题，或与更复杂的索引结构相结合。