更快的检索，带来更好、更高效的 RAG
最大内积搜索（MIPS）是推荐系统、搜索和检索增强生成（RAG）的核心。给定一个查询向量，MIPS 系统会在包含数百万个候选键向量的数据库中搜索与查询最匹配的那一个。随后，与该最佳匹配键相关联的原始数据点被用于处理该查询：在 RAG 中，最适合回答该问题的文档通常会被前置到查询之前。
随着键的数量增长到数百万、向量维度达到数千，将查询与所有键逐一进行穷举比较的计算成本会急剧上升。人们已经提出了许多近似方法，利用聚类和量化来以精度换取速度。这些基线方法关注的是键集合的性质，以构建巧妙的索引结构，但它们的几何论证粗略地假设查询集合本身会与键集合相似，或者完全回避了对查询形态的任何考量。这在两个方面可能带来问题：在关键的 RAG 场景中，查询来自用户生成的简短问题，而键则代表冗长、信息丰富且经过精心整理的文本片段，这使得键与查询存在显著差异。 此外，生产环境中部署的大多数 MIPS 用例面向的是已有的用户群体，其查询分布稳定且特征明确——这一信号表明，通用的、与查询无关的方法尚有未被利用的空间。针对可能的输入来源设计最有效的方法——而不是追求在所有可能输入上都表现良好——是现代机器学习的核心原则之一，有助于更高效地重新分配计算资源。
在 《用学习到的支撑函数摊销最大内积搜索》这篇论文发表于 ICML 2026，来自 Apple 的研究人员离线一次性训练小型神经网络（NN），针对固定的键数据库，让它们学会如何更快地将查询路由到最佳匹配。该方法以一次性的摊销成本来理解典型查询通常如何被匹配，从而在推理阶段通过这些小神经网络实现更快、更准确的检索。
最大内积搜索基础
MIPS 的一个基本原则是利用键集合的几何结构来加速搜索：典型的 MIPS 系统将键集合划分为 k 个簇——k 通常远小于键的总数 n——并为每个簇保留一个代表性质心。当收到新查询时，系统会将该查询与这 k 个质心进行比较。随后，查询会被路由到与其最接近的代表性质心所在的簇，接着再与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键和 k 个均匀分布的簇，那么在此过程中实际需要考虑的键的有效数量 个簇—— 通常远小于键的总数 n——并为每个簇保留一个代表性质心。当收到新查询时，系统会将该查询与这 k 个质心进行比较。随后，查询会被路由到与其最接近的代表性质心所在的簇，接着再与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键和 k 个均匀分布的簇，那么在此过程中实际需要考虑的键的有效数量 通常远小于键的总数 ——为每个簇保留一个代表性质心。当收到新查询时，系统会将该查询与这 k 个质心进行比较。随后，查询会被路由到与其最接近的代表性质心所在的簇，接着再与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键和 k 个均匀分布的簇，那么在此过程中实际需要考虑的键的有效数量约为 n/k 量级——相比 ——为每个簇保留一个代表性质心。当收到新查询时，系统会将该查询与这些 个质心进行比较。随后，查询会被路由到与其最接近的代表性质心所在的簇，接着再与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键和 k 个均匀分布的簇，那么在此过程中实际需要考虑的键的有效数量约为 n/k 量级——相比朴素穷举方法中的 n 次比较，实现了 k 倍的加速。不幸的是，这种路由机制可能导致错误， 质心。随后，查询会被路由到与其最接近的代表性质心所在的簇，接着再与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键和 k 约为 n/k 量级——相比朴素穷举方法中的 n 次比较，实现了 k 倍的加速。 n/k——相比朴素穷举方法中的 n 次比较，实现了 k 倍的加速。/——相比朴素穷举方法中的 n 次比较，实现了 k 倍的加速。 ——相比 倍的加速。 倍的加速。 n 不幸的是，这种路由机制可能导致错误，因为质心路由并非一种精确的方法，如图 图 1 所示。
MIPS 作为一种支撑函数
从机器学习的角度来看，MIPS 之所以有趣，是因为它建立在一个在此项工作之前基本被忽视的数学结构之上：将查询映射到与集合中任意键的最大内积的那个函数，正是该集合的 支撑函数 ，这是凸分析中一个被充分研究的对象。该函数具有两个重要性质：第一，它是凸函数，因此可以用一类专门为凸函数设计的神经网络来近似， 输入凸神经网络（ICNN）其次，该支撑函数在任意查询点处的梯度返回的就是最佳匹配的键本身。
这项工作背后的总体思路是， 学习 一个能够以低得多的计算成本模拟真实支撑函数（或其梯度）的神经网络。这可以通过一种被称为回归的经典机器学习技术来实现：只需预先计算大量训练样本，这些样本通过采样输入向量（用户查询）并将其与 MIPS 产生的相关输出信号（要么是最优分数，要么是最优键向量本身）配对而成。之后只需训练一个参数化适当的神经网络，使其在给定输入时复现该输出即可。
学习最大内积搜索的两种方法
如果按照上述数学方法完美地解决回归问题，那么同一个被训练来模拟支撑函数的神经网络，也可以通过评估其梯度（通过自动微分）来输出最优键。在实践中，我们观察到，为这两项任务分别训练专门的网络效果最佳：
SupportNet 神经网络架构可以直接学习支撑函数，使用（约束较宽松的）ICNN。该方法在数学上有原则，并与问题的凸结构相一致。我们发现，当用于为每个数据簇计算一个支撑函数以改进路由质量时，这种方法可以表现出色。KeyNet 模型学习在输入查询时直接输出最佳匹配的键向量。由齐次函数的欧拉定理推导出的一致性损失可保持 SupportNet 神经网络架构可以直接学习支撑函数，使用（约束较宽松的）ICNN。该方法在数学上有原则，并与问题的凸结构相一致。我们发现，当用于为每个数据簇计算一个支撑函数以改进路由质量时，这种方法可以表现出色。
KeyNet 模型学习 直接输出最佳匹配的键向量。 在输入查询时输出最佳匹配的键向量。由齐次函数的欧拉定理推导出的一致性损失，使其预测在几何上始终锚定于底层的支撑函数，即使并未显式施加凸性约束。
两个模型都使用从相关查询分布中抽取的样本进行训练：对于每个采样查询，通过穷举搜索计算一次真实的最佳匹配并将其存储，从而创建一个监督数据集，该数据集捕捉了真正重要的查询分布下的 MIPS 解空间。
使用 SupportNet 的学习式路由
如图 1 所示，处理更大规模的最简单方法是将键划分为多个簇，将查询与每个簇的质心进行比较，并根据这些 k 个分数将查询路由到最有希望的簇。这提供了一种廉价但粗略的规则，容易出错，因为仅凭质心无法反映其所在簇的分布范围。事实上，一个簇中可能包含某个键，它远离自身质心却靠近查询；当查询接近这样的离群点时，质心路由 图 1分数将查询路由到最有希望的簇。 分数。这提供了一种廉价但粗略的规则，容易出错，因为仅凭质心无法反映其所在簇的分布范围。事实上，一个簇中可能包含某个键，它远离自身质心却靠近查询；当查询接近这样的离群点时，质心路由会给该簇打低分并跳过它，从而错过真正的最佳匹配。如果能够获得每个簇的支撑函数，那么就可以实现精确路由，因为每个支撑 分数。这提供了一种廉价但粗略的规则，容易出错，因为仅凭质心无法反映其所在簇的分布范围。事实上，一个簇中可能包含某个键，它远离自身质心却靠近查询；当查询接近这样的离群点时，质心路由会给该簇打低分并跳过它，从而错过真正的最佳匹配。
如果能够获得每个簇的支撑函数，那么反而可以实现精确路由，因为每个支撑函数都精确量化了每个簇内的任意键能够在多大程度上以高分匹配查询。这一点在 图 2中，针对一个在嵌入空间中变化的查询，绘制了四个簇各自的支撑函数。
依托这一理论洞见，我们提出学习这些 支撑函数，并与一个可一次性输出 k 个分数的训练好的 SupportNet 联合学习。随后，我们可以改用这 k 个分数中最大的一个来路由任何新查询。与简单地与质心进行比较相比，SupportNet 模型能够回答一个更有用的问题：“对于这个查询，该簇内最佳键有多好？”SupportNet 以经典的机器学习方式（也称为回归）进行训练：我们生成从分布中抽取的样本查询（通常是嵌入 支撑函数，并与一个经过训练的 SupportNet 联合使用，该网络可以一次性输出 k 个分数。随后，我们可以改用这 k 个分数中最大的一个来路由任何新查询。与简单地与质心进行比较相比，SupportNet 模型能够回答一个更有用的问题：“对于这个查询，该簇内最佳键有多好？”SupportNet 以经典的机器学习方式（也称为回归）进行训练：我们生成从分布中抽取的样本查询（通常是嵌入 分数。随后，我们可以改用这些分数中最大的一个来 分数来路由任何新查询。与简单地与质心进行比较相比，SupportNet 模型能够回答一个更有用的问题：“对于这个查询，该簇内最佳键有多好？”SupportNet 以经典的机器学习方式（也称为回归）进行训练：我们生成从分布中抽取的样本查询（通常是嵌入 分数来路由任何新查询。与简单地与质心进行比较相比，SupportNet 模型能够回答一个更有用的问题：“对于这个查询，该簇内最佳键有多好？” 最佳 这个查询在该簇内的最佳键有多好？”SupportNet 以经典的机器学习方式（也称为回归）进行训练：我们生成从分布中抽取的样本查询（通常是文本查询的嵌入），计算这些查询处的真实支撑函数值，并训练一个神经网络来复现这些值。
我们在 Natural Questions 上训练了规模递增的 SupportNet（XS、S、M 和 L → 分别为数据库大小的 1%、5%、10% 和 20%），该数据集被划分为 128 个簇。Natural Questions 是一个检索基准，将常识问题与大型文档数据库中的段落配对，在我们的设置中约有 260 万个唯一段落。下面的图 4 展示了以召回率衡量的检索性能随不同成本概念的变化。路由召回率是指真实分区（即包含确切最佳 key 的分区）落在路由器选择探测的分区之中的查询比例。对于这两幅图，一个穷举 Natural Questions 划分为 128 个簇。Natural Questions 是一个检索基准，它将常识性问题与来自大型文档数据库的段落配对，在我们的设置中约有 260 万个唯一段落。
图 4 下图展示了以召回率衡量的检索性能随不同成本概念的变化。路由召回率是指真实分区（即包含确切最佳 key 的分区）落在路由器选择探测的分区之中的查询比例。对于这两幅图，均在排名前个与查询匹配的簇中进行穷举搜索，其中 s 按 2 的幂次增长（1、2、4、8、…、128）。自然，当所有方法遍历全部 128 个簇时，召回率都达到 1。第一个标签页显示，在固定检索性能下，学习型路由器以比质心路由更低的延迟达到该性能。小型 SupportNet 模型达到 95% 路由召回率的速度比质心路由快约 1.9 倍（1.74 ms 对 3.36 ms），而大型模型的提升更大，在 1.74 ms 内达到 97%。该图明确 个与查询匹配的簇，其中 按 2 的幂次增长（1、2、4、8、…、128）。 按 2 的幂次增长 (1, 2, 4, 8, …, 128)自然，当所有方法遍历全部 128 个簇时，召回率都达到 1。
第一个标签页显示，在固定检索性能下，学习型路由器以比质心路由更低的延迟达到该性能。小型 SupportNet 模型达到 95% 路由召回率的速度比质心路由快约 1.9 倍（1.74 ms 对 3.36 ms），而大型模型的提升更大，在 1.74 ms 内达到 97%。该图明确计入了运行我们的小型神经网络所带来的额外计算量，而不是简单地与质心进行比较。
第二个标签页的图表仅关注路由质量，忽略了运行神经网络所带来的开销。在探索簇数量固定的预算下，它们更频繁地找到正确的分区：仅扫描数据库的 5%，学习型路由器就有约 90–93% 的概率恢复真实分区——中型（M）模型最高可达 92.6%——而质心基线为 87.5%。
SupportNet 路由器规模系列
接入现有索引
SupportNet 路由取代了现有倒排文件索引（IVF）系统中基于质心的简单路由。借助 KeyNet，我们也可以保持 IVF 系统不变，只向其输入修改后的查询。在这种情况下，KeyNet 自然充当一个 “查询重写器” 在标准近似搜索索引之前：预测最佳匹配在嵌入空间中可能出现的位置，然后让你常用的 IVF 系统从那里开始查找，而不是从原始查询开始。由于 IVF 索引经过高度优化，在 KeyNet 的辅助下它能更快地找到结果。我们还依赖回归来训练 KeyNet，概述如下： 图 5：样本查询与其真实标签、最佳匹配键进行匹配，并训练 KeyNet 神经网络来复现这种关联。
我们评估了规模递增的 KeyNet（XS、S、M 和 L → 分别为数据库大小的 1%、5%、10% 和 20%），并在 图 6 下面，我们展示这些不同规模下的召回率–成本权衡。在第一个标签页中，搜索预算是实际耗时：使用最大模型达到 99% 召回率所需时间最多可减少约 8 倍（0.23 毫秒，而原始查询为 1.8 毫秒）。与之前一样，第二个标签页仅关注用 KeyNet 变换替换原始查询所带来的“质量”提升。横轴是索引实际扫描的数据库比例：仅扫描数据库的 0.001% 时，经 KeyNet 映射的查询有 89–97% 的概率检索到真实键——最大模型（L）最高可达 96.7%——而原始查询为 86%；达到 99% 召回率所需键数最多可减少约 11 倍（0.009% 对原始查询的 0.10%）。相同成本下获得更高召回率，或以一小部分成本获得相同召回率。
KeyNet 搜索器规模系列
我们展示了从同一数据集获得的示例。 Natural Questions 数据集，通过一个 FAISS IVF 索引。它是我们研究过的几个数据集之一，而且这一趋势在这些数据集上是一致的。在整个搜索预算扫描范围内，用 KeyNet 预测的键来查询索引，始终优于用原始问题来查询。
即， 图 7 下图展示了五个 Natural Questions 示例：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果。
图 7：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果：“What is Matchbox Twenty?”
图 7：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果：“Who wrote The Eye of the World?”
图 7：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果：“What is a wind controller used for?”
图 7：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果：“When will the Summertime Ball be on?”
图 7：原始查询与经 KeyNet 映射的查询各自的前 3 条命中结果：“When was the show Victorious on Nickelodeon?”你应该使用哪个网络？
当推理速度是瓶颈、且你只需要近似最佳 key 时，使用 KeyNet。当你还需要经过校准的分数（用于重排序、阈值判断或下游使用），或在较小模型规模下计算资源紧张时——在这种情形下 SupportNet 仍具竞争力——则使用 SupportNet。两者的训练方式相似，都能与现有的近似搜索基础设施集成，并且对查询模式中轻微的分布偏移具有相当的鲁棒性。最终，对于一个超大规模 MIPS 系统， KeyNet 当推理速度是瓶颈、且你只需要近似最佳 key 时，使用 KeyNet。 SupportNet 当你还想要经过校准的分数（用于重排序、阈值判断或下游使用），或在较小模型规模下计算资源紧张时——在这种情形下 SupportNet 仍具竞争力——则使用 SupportNet。两者的训练方式相似，都能与现有的近似搜索基础设施集成，并且对查询模式中轻微的分布偏移具有相当的鲁棒性。最终，对于一个 key 集合极其庞大的超大规模 MIPS 系统，可能两者都需要：先在所有簇上训练一个 SupportNet 路由器，然后为每个簇训练一个 KeyNet，以帮助其常规索引更快地锁定最佳 key。
学习式检索的新方向
摊销式 MIPS 反映了当查询分布可预先获知时——这是生产级 ML 系统中的常见情形——检索方式所发生的一种转变。通过将 MIPS 解直接编码进网络权重，该方法以一次性的训练（即所谓的“ 摊销成本，实现快速、轻量的推理，并开辟了一条切实可行的路径，通向明确针对其所服务查询而量身定制的学习式检索系统。
相关阅读与更新。
用学习到的支撑函数摊销最大内积搜索
2026年7月2日研究领域：知识库与搜索, 研究领域：方法与算法ICML 会议
最大内积搜索（MIPS）是机器学习中的一个关键子程序，它需要在数据库（键）中找出与给定查询最对齐的向量。我们提出摊销式 MIPS：一种基于回归的方法，训练神经网络直接预测 MIPS 解，从而将针对固定键数据库、从已知分布中抽取的查询反复求解 MIPS 的成本摊销。我们的关键洞见是……
阅读更多使用 Wally 实现可扩展的隐私搜索
2024年10月16日研究领域：隐私
本文提出了 Wally，一种私有搜索系统，它支持针对大型数据库的高效语义和关键词搜索查询。当有足够多的客户端发起查询时，Wally 的性能显著优于以往的系统。在以往的私有搜索系统中，对于每个客户端查询，服务器必须对每个数据库条目至少执行一次昂贵的密码学操作。因此，性能会下降……
阅读更多
探索机器学习领域的机遇。
我们在机器学习领域的研究每天都在开辟新的天地。