研究关注RAG检索效率：加速最大内积搜索以提升RAG性能

更快的检索，更高效、更出色的 RAG

最大内积搜索（Maximum inner product search，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 倍的加速。然而，这种路由机制可能带来误差，因为质心路由并非精确方法，如下图所示。

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

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

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

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

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

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

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）