更快的检索让 RAG 更高效

更快的检索，更好、更高效的 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 所示。

将 MIPS 视为支撑函数

从机器学习的角度来看，MIPS 之所以有趣，是因为它建立在一个在此项工作之前很大程度上被忽视的数学结构之上：将查询映射到与集合中任意键的最大内积的函数，是该集合的支撑函数，这是凸分析中一个被深入研究过的对象。该函数有两个重要性质：第一，它是凸函数，因此可以用一类专为凸函数设计的神经网络——输入凸神经网络（ICNN）——来近似。第二，该支撑函数在任意查询处的梯度会返回最佳匹配的键本身。

这项工作背后的总体思路是：学习一个神经网络，以低得多的计算成本来模拟真实支撑函数（或直接模拟其梯度）。这可以通过一种被称为回归的经典机器学习技术实现：只需预先计算大量训练样本，这些样本由采样的输入向量（用户查询）构成，并配以 MIPS 产生的相关输出信号（最优得分或最优键向量本身）。随后只需训练一个参数化适当的神经网络，使其在给定输入时复现该输出。

学习最大内积搜索的两种方式

如果按照上述数学方法完美地解决回归问题，那么用于模拟支撑函数的同一个神经网络也可以通过评估其梯度（通过自动微分）来输出最优键。在实践中，我们观察到，为这两项任务分别训练专门化的独立网络效果最佳：

SupportNet 神经网络架构可以使用（约束较宽松的）ICNN 直接学习支撑函数。该方法在数学上有原则性，并与问题的凸结构相一致。我们发现，当用于为每个数据簇计算一个支撑函数以提升路由质量时，该方法可以表现出色。

KeyNet 模型学习在输入查询时直接输出最匹配的键向量。由齐次函数的欧拉定理导出的 一致性损失 使其预测在几何上仍以底层支撑函数为基础，尽管并未显式强制凸性。

两个模型都在从相关查询分布中抽取的样本上训练：对于每个采样查询，通过穷举搜索计算一次真实最优匹配并存储，从而创建一个监督数据集，该数据集捕捉了对实际重要的查询分布下的 MIPS 解空间。

使用 SupportNet 的学习式路由

如图 1 所示，处理更大规模的最简单方法是将键划分为簇，将查询与每个簇的质心进行比较，并基于这 k 个得分将查询路由到最有希望的簇。这提供了一种廉价但粗糙的规则，容易出错，因为仅凭质心无法反映其簇的分布范围。事实上，一个簇可能包含一个键，它远离自身质心却靠近查询；当查询接近这样一个离群点时，质心路由会给该簇打低分并跳过它，从而错过真正的最优匹配。

如果能够获得每个簇的支撑函数，那么就可以进行精确路由，因为每个支撑函数都精确量化了每个簇内任意键能够以多大得分匹配查询。图 2 对此进行了说明，其中针对在嵌入空间中变化的查询，绘制了四个簇各自的支撑函数。

基于这一理论洞见，我们提出用一个训练好的 SupportNet 联合学习这 k 个支撑函数，该网络可以一次输出 k 个得分。然后我们可以使用这 k 个得分中最大的一个来路由任何新查询。SupportNet 模型可以回答一个比简单与质心比较更有用的问题：“对于这个查询，这个簇内最好的键有多好？”SupportNet 以经典机器学习方式训练，也称为回归：我们生成从分布中抽取的样本查询（通常是文本查询的嵌入），计算这些查询处的真实支撑函数，并训练神经网络来复现这些值。

我们在划分为 128 个簇的 Natural Questions 上训练了规模递增的 SupportNet（XS、S、M 和 L → 数据库大小的 1%、5%、10% 和 20%）。Natural Questions 是一个检索基准，将常识性问题与来自大型文档数据库的段落配对，在我们的设置中约有 260 万个唯一段落。

下面的图 4 展示了以 recall 衡量的检索性能随不同成本概念的变化。Routing recall 是指真实分区（即包含确切最佳键的分区）落在路由器选择探测的分区之中的查询比例。在两幅图中，都在与查询匹配的前 s 个簇中执行穷举搜索，其中 s 以 2 的幂次增长（1、2、4、8、…、128）。自然，当所有方法遍历全部 128 个簇时，recall 都达到 1。

第一个标签页显示，对于固定的检索性能，学习型路由器以比 centroid routing 更低的延迟达到该性能。小型 SupportNet 模型达到 95% routing recall 的速度比 centroid routing 快约 1.9 倍（1.74 ms 对 3.36 ms），而大型模型的提升更大，在 1.74 ms 内达到 97%。该图明确计入了运行我们的小型 NN 所带来的额外计算量，而非仅仅与 centroids 比较。

第二个标签页的图仅关注路由质量，忽略了运行 NN 带来的开销。在固定的探索簇预算下，它们更频繁地找到正确分区：仅扫描数据库的 5%，学习型路由器就有约 90–93% 的概率恢复真实分区——中型（M）最高达 92.6%——而 centroid 基线为 87.5%。

SupportNet Router Size Family

Plugging Into Existing Indices

SupportNet routing 取代了现有倒排文件索引（IVF）系统中基于 centroid 的简单路由。借助 KeyNet，我们也可以保持 IVF 系统不变，只是向它输入查询的修改版本。在这种情况下，KeyNet 自然地充当标准近似搜索索引前的“查询重写器”：预测最佳匹配在嵌入空间中可能的位置，然后要求你常用的 IVF 系统去那里查找，而不是从原始查询开始。由于 IVF 索引高度优化，在 KeyNet 的辅助下它能更快地找到结果。我们还依赖回归来训练 KeyNet，如图 5 所总结：样本查询与其 ground-truth、最佳匹配的键配对，训练 KeyNet 神经网络来重现这种关联。

我们评估了规模递增的 KeyNet（XS、S、M 和 L → 数据库大小的 1%、5%、10% 和 20%），并在下面的图 6 中展示了这些不同规模的 recall–cost 权衡。在第一个标签页中，搜索预算是 wall-clock 延迟：使用最大的模型达到 99% recall 所需时间最多减少约 8 倍（0.23 ms 对原始查询的 1.8 ms）。与之前一样，第二个标签页仅关注用 KeyNet 变换替换原始查询所带来的“质量”提升。x 轴是索引实际扫描的数据库比例：仅扫描数据库的 0.001%，KeyNet 映射的查询就有 89–97% 的概率检索到真实键——最大模型（L）最高达 96.7%——而原始查询为 86%，达到 99% recall 所需键数最多减少约 11 倍（0.009% 对原始 0.10%）。相同成本下更好的 recall，或以其中一小部分成本获得相同的 recall。

KeyNet Searcher Size Family

我们展示了从前面讨论过的同一 Natural Questions 数据集获得的示例，这些示例通过 FAISS IVF 索引提供。它是我们研究的若干数据集之一，且趋势在这些数据集上保持一致。在整个搜索预算扫描范围内，用 KeyNet 预测的 key 查询索引始终优于用原始问题查询。

即，下面的图 7 展示了五个 Natural Questions 示例：原始查询与 KeyNet 映射查询的 top-3 命中结果。

• 图库项 1

• 图库项 2

• 图库项 3

• 图库项 4

• 图库项 5

应该使用哪个网络？

两者之间的选择最终是一个实用性的问题。当推理速度是瓶颈、且你只需要近似最佳 key 时，使用 KeyNet。当你还想要一个经过校准的分数（用于重排序、阈值判断或下游使用），或者在较小模型规模下计算资源紧张时——在这种情形下 SupportNet 仍具竞争力——使用 SupportNet。两者的训练方式相似，都能与现有的近似搜索基础设施集成，并且对查询模式中轻微的分布偏移都相当稳健。最终，对于 key 集合非常庞大的超大规模 MIPS 系统，可能两者都需要：在所有簇上训练一个 SupportNet 路由器，然后每个簇训练一个 KeyNet，以帮助其常规索引更快地锁定最佳 key。

学习式检索的新方向

当查询分布事先已知时——这是生产 ML 系统中的常见情形——摊销式 MIPS 反映了检索方法上的一种转变。通过将 MIPS 解直接编码进网络权重，该方法以一次性的训练（即所谓的摊销）成本换取快速、轻量的推理，并为明确针对其所服务查询而定制的学习式检索系统开辟了一条实用路径。

相关阅读与更新。

用学习到的支撑函数摊销最大内积搜索

2026 年 7 月 2 日研究领域 知识库与搜索，研究领域 方法与算法会议 ICML

最大内积搜索（MIPS）是机器学习中的一个关键子程序，需要在数据库（key）中找出与给定查询最对齐的向量。我们提出摊销式 MIPS：一种基于回归的方法，训练神经网络直接预测 MIPS 解，从而为从固定 key 数据库上的已知分布中抽取的查询摊销反复求解 MIPS 的成本。我们的关键洞见是……

使用 Wally 的可扩展私有搜索

2024 年 10 月 16 日研究领域 隐私

本文提出 Wally，一个私有搜索系统，支持针对大型数据库的高效语义和关键词搜索查询。当有足够多的客户端发起查询时，Wally 的性能显著优于以往的系统。在以往的私有搜索系统中，对于每个客户端查询，服务器必须对每个数据库条目执行至少一次昂贵的密码学操作。因此，性能会下降……

在机器学习中发现机遇。

我们在机器学习领域的研究每天都在开辟新天地。

与我们合作