研究聚焦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 之所以有趣，是因为它建立在一个此前很大程度上被忽视的数学结构之上：将查询映射到与集合中任意键的最大内积的函数，是该集合的支撑函数（support function），这是凸分析中一个被深入研究的对象。该函数具有两个重要性质：第一，它是凸函数，因此可以用一类专为凸函数设计的神经网络——输入凸神经网络（ICNN）——来近似。第二，该支撑函数在任意查询处的梯度恰好返回最佳匹配的键本身。

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

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

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

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

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

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