通过学习支持函数实现最大内积搜索的 amortization
机器学习
2026-03-10 v1 机器学习
摘要
最大内积搜索（MIPS）是机器学习中的关键子程序，需识别与给定查询最匹配的关键向量。我们提出 amortized MIPS：一种基于学习的方法，通过训练神经网络直接预测 MIPS 解决方案，以摊销查询（来自固定分布）与固定密钥集合之间的匹配计算成本。我们的关键洞察在于，MIPS 价值函数，即查询与密钥之间的最大内积，也称为密钥集合的支持函数。支持函数是凸的、1齐次的，其对查询的梯度即为数据库中的最优密钥。我们采用两种互补方法来近似支持函数：（1）训练输入凸神经网络（SupportNet）直接建模支持函数；可通过（autodiff）梯度计算恢复最优密钥；（2）使用向量值网络（KeyNet）直接回归查询的最优密钥，在推理时完全跳过梯度计算。为学习 SupportNet，我们将得分回归与梯度匹配损失结合，并提出 homogenization 包装器以强制神经网络的正1齐次性，理论上将函数值与梯度关联。为训练 KeyNet，我们引入源自齐次函数欧拉定理的得分一致性损失。我们的实验表明，学习的 SupportNet 或 KeyNet 能实现高匹配率，并为以特定查询分布为中心压缩数据库打开了新方向。
引用
@article{arxiv.2603.08001,
title = {Amortizing Maximum Inner Product Search with Learned Support Functions},
author = {Theo X. Olausson and João Monteiro and Michal Klein and Marco Cuturi},
journal= {arXiv preprint arXiv:2603.08001},
year = {2026}
}