用神经网络学习向量数据库的最优匹配,大幅加速查询。
Amortized Maximum Inner Product Search with Learned Support Functions
- 构建支持函数网络(SupportNet)和最优键预测网络(KeyNet),直接回归最佳匹配结果。
- 在BEIR数据集上,同等计算开销下,索引命中率提升显著,支持函数网络可作聚类路由。
- 适合需要快速检索的场景,如文档搜索、推荐系统等应用部署。
最大内积搜索(MIPS)是机器学习中的关键子程序,需从数据库(键)中找出与给定查询最对齐的向量。本文提出缓存式MIPS:一种基于回归的神经网络方法,通过训练模型直接预测MIPS解,从而分摊针对固定键库、来自已知分布的重复查询的计算成本。核心洞察是:MIPS值函数即键集的支持函数,一个经典的凸函数,其梯度可导出最优键。据此设计两种互补模型:SupportNet——输入凸神经网络,用于回归支持函数;KeyNet——向量值网络,直接回归最优键。SupportNet可作为聚类路由,引导查询至相关数据库分区;KeyNet则可作为原查询的直接替代品,接入现有索引管道。在BEIR基准测试中,对于文档嵌入,所学的SupportNet与KeyNet在考虑计算投入(以FLOPs、探测次数或实际时间衡量)时,显著提升了IVF匹配率。代码已公开于:https://github.com/apple/ml-amips。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。