更快的检索以获得更好、更高效的 RAG

最大内积搜索 (MIPS) 是推荐系统、搜索和检索增强生成 (RAG) 的核心。给定一个查询向量，MIPS 系统在包含数百万个候选关键向量的数据库中搜索与查询最匹配的一个。然后，与该最佳匹配键关联的原始数据点用于处理查询：在 RAG 中，最适合回答问题的文档通常作为查询的前缀。随着键的数量增加到数百万，向量维度达到数千，对所有键进行详尽比较的查询的计算成本急剧增加。已经提出了许多近似方法来使用聚类和量化来权衡精度和速度。虽然这些基线方法侧重于键集的属性来构建巧妙的索引结构，但它们的几何参数松散地假设查询集本身将类似于键中找到的查询集，或者绕过对查询的外观的任何考虑。这可能在两个方面存在问题：在关键的 RAG 案例中，查询源自用户生成的简短问题，而键代表长的、内容丰富且经过精心策划的文本片段，这使得键和查询有很大不同。此外，在生产中部署的大多数 MIPS 用例都针对现有的用户群，这些用户群具有稳定、特征良好的查询分布，这表明通用的、与查询无关的方法尚未得到开发。制定最适合可能的输入源的方法（而不是针对所有可能的输入都取得良好的性能）是现代机器学习的核心原则，有助于更有效地重定向计算资源。在 ICML 2026 上发表的一篇论文《Amortizing Maximum Inner Product Search with Learned Support Functions》中，Apple 的研究人员针对固定的键数据库离线训练一次小型神经网络 (NN)，学习如何更快地将查询路由到最佳匹配。他们的方法权衡了前期摊销成本，以了解典型查询往往如何匹配，以换取通过这些小型神经网络在推理时更快、更准确的检索。 MIPS 的一个基本原则是利用密钥集的几何形状来加速搜索：典型的 MIPS 系统将密钥集分成 k 个簇 - k 通常远小于密钥总数 n - 跟踪每个簇的代表性质心。当收到新的查询时，系统将其与这 k 个质心进行比较。然后将查询路由到最接近的代表性质心的集群，并随后与该集群中存储的所有键进行详尽的比较。如果数据库有 n 个键和 k 个均匀分布的簇，则第二步中将考虑的键的有效数量约为 n / k — 与朴素穷举方法中进行的 n 次比较相比，加速了 k 倍。不幸的是，这种路由机制可能会导致错误，因为质心路由不是一种精确的方法，如下图 1 所示。图 1：几何搜索方法可能会失败并错过最近的邻居。这是因为他们依赖于根据到质心的距离来划分空间。在动画中，最近的邻居被错过，因为它与最近的质心无关。从机器学习的角度来看，MIPS 的有趣之处在于，它建立在一种数学结构之上，而在这项工作之前，这种数学结构在很大程度上被忽视了：将查询映射到集合中任何键的最大内积的函数是该集合的支持函数，这是凸分析中经过充分研究的对象。这个函数有两个重要的性质：首先，它是凸的，因此可以用一类专门为凸函数设计的神经网络来近似，即输入凸神经网络（ICNN）。其次，任何查询中该支持函数的梯度都会返回最佳匹配的键本身。这项工作背后的广泛思想是学习一个可以以低得多的计算成本模拟地面实况支持函数（或其直接梯度）的神经网络。这可以通过称为回归的经典 ML 技术来实现：只需预先计算通过采样输入向量（用户查询）形成的大量训练示例，并将它们与 MIPS 生成的相关输出信号（最佳分数或最佳密钥向量本身）配对。然后所需要做的就是训练一个适当参数化的神经网络来复制给定输入的输出。如果要完美解决回归问题，根据上述数学方法，经过训练以模仿支持函数的同一神经网络也可以通过评估其梯度（通过自动微分）来输出最佳密钥。在实践中，我们观察到为每个网络训练单独的、专门的网络