更高效检索助力RAG性能提升

Faster Retrieval for Better, More Efficient RAG

Maximum inner product search (MIPS) sits at the core of recommendation systems, search, and retrieval-augmented generation (RAG). Given a query vector, MIPS systems search within a database of millions of candidate key vectors for the one that best matches the query. The original data point associated with that best-matching key is then used to handle the query: in RAG, the document that is best suited to answer the question is typically prefixed to the query.

As the number of keys grows to millions and vector dimensions reach thousands, the computational cost of running an exhaustive comparison of the query to all keys grows dramatically. Numerous approximate methods have been proposed to trade off accuracy for speed using clustering and quantization. While these baseline methods focus on the properties of the set of keys to build up clever indexing structures, their geometric arguments loosely assume that the set of queries itself will resemble that found in keys, or bypass any considerations about what the queries will look like. This can be problematic on two fronts: In the crucial RAG case, queries arise from short user-generated questions, whereas keys represent long, informative, and curated pieces of text, making keys and queries substantially different. Additionally, most MIPS use cases deployed in production target an existing base of users with a stable, well-characterized query distribution — a signal that generic, query-agnostic methods leave untapped. Crafting methods that work best for a likely source of inputs — rather than aiming for good performance on all possible inputs — is a central tenet of modern machine learning that helps redirect compute resources more efficiently.

In Amortizing Maximum Inner Product Search with Learned Support Functions, a paper presented at ICML 2026, researchers from Apple train small neural networks (NNs) — once, offline, for a fixed database of keys — that learn how to route queries toward their best match faster. Their approach trades off an upfront amortization cost to understand how typical queries tend to be matched in exchange for faster, more accurate retrieval at inference time through these small NNs.

The Basics of Maximum Inner Product Search

MIPS 的一个基本原理是利用键集合的几何结构来加速搜索：典型的 MIPS 系统将键集合划分为 k 个簇 —— k 通常远小于键的总数 n —— 并为每个簇跟踪一个代表性的质心。当接收到一个新的查询时，系统会将其与这 k 个质心进行比较。

随后，查询会被路由到与其最接近的代表性质心所在的簇，并随后与该簇中存储的所有键进行穷举比较。如果数据库有 n 个键且 k 个簇分布均匀，那么在第二步中需要考虑的有效键数量大约为 n/k —— 与 naive 穷举方法中进行的 n 次比较相比，这是一个 k 倍的加速。不幸的是，这种路由机制可能会导致错误，因为质心路由不是一种精确的方法，如图 1 所示。

MIPS 作为一种支持函数

从机器学习的角度来看，MIPS 的有趣之处在于它是建立在一种在此项工作之前很大程度上被忽视的数学结构之上的：将查询映射到集合中任何键的最大内积的函数是该集合的支持函数，这是凸分析中一个研究得很好的对象。

该函数有两个重要性质：首先，它是凸函数，因此可以用一类专门为凸函数设计的神经网络来近似，即输入凸神经网络（ICNN）。其次，该支持函数在任何查询处的梯度都会返回最佳匹配的键本身。

这项工作的核心思想是学习一个神经网络，以低得多的计算成本来模拟 ground-truth 支持函数（或其梯度）。这可以通过一种经典的 ML 技术，即回归来实现：只需预先计算大量训练样本，这些样本由采样输入向量（用户查询）组成，并将它们与 MIPS 生成的相关输出信号（即最优分数或最优键向量本身）配对。然后，只需要训练一个参数合适的神经网络，使其在给定输入的情况下复制输出。

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

如果回归问题能够被完美解决，根据上述数学方法，训练用于模拟支持函数的同一个神经网络也可以通过计算其梯度（通过自动微分）来输出最优键。在实践中，我们发现为这两个任务中的每一个分别训练专门的、专用的网络效果最好：

SupportNet 神经网络架构可以使用（约束较宽松的）ICNN 直接学习支持函数。该方法在数学上是合理的，与问题的凸结构相一致。我们发现，当用于计算每个数据簇的一个支持函数以改善路由质量时，这种方法可以发挥出色。

KeyNet 模型在接收到查询时，学会直接输出最佳匹配的键向量。一种源自欧拉齐次函数定理的一致性损失，确保了它们的预测在几何上植根于底层支撑函数，尽管并未显式强制执行凸性。

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

使用 SupportNet 进行学习路由

如图 1 所示，处理更大规模的最简单方法是先将键分成簇，将查询与每个簇的质心进行比较，并根据这 k 个分数将查询路由到最有希望的簇。这提供了一种廉价但粗略的规则，由于仅靠质心无法解释其簇的分布，因此容易出错。确实，一个簇可能包含一个距离自身质心很远但靠近查询的键，当查询靠近这样一个离群点时，质心路由会给该簇打低分并跳过它，从而错过了真正的最佳匹配。

如果能够访问每个簇的支撑函数，那么就可以进行精确路由，因为每个支撑函数精确量化了该簇内任何键以高分匹配查询的程度。这在图 2 中有所说明，图中绘制了在嵌入空间中变化的查询下四个簇中每个簇的支撑函数。

基于这一理论洞察，我们提出联合学习这 k 个支撑函数，以及一个可以同时输出 k 个分数的已训练 SupportNet。然后我们可以使用这 k 个分数中的最大值来路由任何新查询。SupportNet 模型可以回答一个比简单与质心比较更有用的问题：“该簇内最好的键对于这个查询有多好？

”SupportNets 以经典的机器学习方式（也称为回归）进行训练：我们生成从分布中抽取的样本查询（通常是文本查询的嵌入），计算这些查询处的真实支撑函数，并训练一个神经网络来复制这些值。

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

下图 4 显示了检索性能