文章摘要
AI大规模落地后,训练数据构建、大模型与多模态数据管理等场景涌现大量大TopK向量检索需求,传统图索引处理该场景性能显著下降。文章提出三层优化方案:算法层选用IVF索引,执行层用Reservoir结构降低结果维护成本,分布式场景用AutoIndex动态分配检索预算,可有效提升性能、降低成本。

更低成本、更高效的向量检索,是当前AI大规模落地的核心需求之一。在传统的推荐、语义搜索或知识库场景中,向量检索系统通常默认用户仅需要少量最相似的结果,TopK数值大多控制在个位数到数百之间。但进入AI时代后,这一默认假设已经不再适用,各类大TopK查询需求开始大量涌现。

一、AI workload催生的大TopK检索需求

典型的大TopK场景主要集中在各类AI工作负载中,具体包括以下几类:

1. 训练数据检索与数据集构建

假设一个多模态训练平台维护着数十亿甚至上百亿条图片向量 embedding,用于文本生成图像或视觉模型的开发。当开发者输入一个文本提示词或一组参考图片后,往往需要一次检索获取数万到数十万张相关图片,用来构建特定领域的微调数据集或者寻找长尾类别中的相关数据。此时仅返回几十个结果显然无法支撑后续的数据处理流程。

2. 训练数据质量与分布分析

大TopK检索也常用于分析训练数据的质量和分布。例如团队可能希望了解与某个文本概念相关的数据总量、某个视觉类别在训练集中的覆盖充分程度、检索结果是否集中在少数来源或风格中,以及不同数据集之间是否存在大量重复或高度相似的样本。如果仅观察很小的TopK,分析结果很容易被最相似的一小部分样本主导,无法反映更完整的数据分布。

3. LLM / 多模态数据管理

在大语言模型数据管理、多模态数据分析和训练数据资产管理中,向量检索通常只是处理流程的第一步。一个生产集群可能管理约100TB数据和数十亿条图片向量embedding,并支持十万级别的TopK检索。返回的结果通常不会直接展示给终端用户,而是继续进入后续的数据处理流程,比如SQL过滤与聚合、元数据分析、数据质量检测、样本发现、数据集导出与构建等。这与在线问答或普通图片搜索有明显区别,后者只需要返回少量结果,而数据管理系统需要把一大批相关数据高效地交给下游任务。

4. 探索式相似性分析

大TopK的另一个典型使用方式是先检索一个较大的相似邻域,再从中寻找规律。研究人员可能先获取数十万个相关向量,然后继续执行分组和聚类、属性分布分析、条件过滤、Join操作、异常值检测、人工抽样与审核等任务。这类查询通常更关注候选集规模、吞吐量和整体处理效率,而不只是单次查询能否在几毫秒内完成。

二、传统图索引的局限性与IVF索引的优势

图索引(如HNSW等)在绝大多数常规向量检索场景下表现出极高的搜索效率,主要得益于其巧妙利用了“小世界”网络特性和贪心路由机制。这种索引构建了具有“小世界”拓扑结构的网络,同时包含用于局部精细搜索的短跳边和用于跨区域快速跳转的长跳边,使得算法能够在极其庞大的数据集中通过极少的步数快速逼近目标向量所在的大致区域,类似人类社会中的“六度空间理论”。

但在面对大TopK查询时,图索引的性能会显著变差,主要有以下几个原因:

  • 贪心搜索退化:图索引依赖贪心路由快速逼近局部最优解,适合寻找极其相近的少数点。当K值很大时,搜索范围被迫急剧扩大,算法需要遍历大量的节点,图索引依靠少量跳转完成搜索的优势会明显下降。

  • 优先队列维护成本剧增:对于每个新访问的节点,系统需要计算其距离,并根据距离更新优先队列。在较小的K值下,这些队列规模有限,插入、弹出和调整堆结构的成本通常可以接受。但在大TopK查询中,结果集合可能包含数万甚至数十万个元素,随着K增大,优先队列会占用更多内存,每个候选结果都可能触发堆调整,大量插入和弹出操作会消耗更多CPU,数据结构维护开始占据显著的查询时间。

  • 随机访存的劣势放大:图索引中的节点及其邻接关系通常分布在不同内存位置,搜索过程需要沿着图边不断跳转,因此会产生大量随机内存访问。当访问节点数量较少时,这个问题并不突出,但大TopK需要遍历更多节点,缓存未命中和内存访问延迟会被持续放大,最终查询性能可能不再受计算能力限制,而是受到内存带宽和缓存效率限制。

因此在处理大TopK查询时,基于倒排的IVF索引是一个更好的选择,它的优势主要体现在以下几个方面:

  • 内存连续存储,局部性好:IVF通过K-Means等算法将整个高维向量空间划分为多个聚类簇(Voronoi单元),同一个簇内的向量数据在物理内存中是紧凑且连续存储的。在进行大范围搜索时,IVF执行的是顺序内存读取,这极大地提高了缓存的命中率,避免图索引因离散节点遍历引发的内存带宽瓶颈。

  • 硬件指令集加速更友好:大TopK查询本质上需要对比海量的候选向量,IVF连续存储的特性使其完美契合现代CPU的SIMD指令集(如AVX-512)或者GPU的大规模矩阵运算核心。

  • 将导航和扫描分开:IVF先通过聚类中心缩小搜索范围,再对选中的列表进行批量扫描。这种结构把查询分成两个阶段:先判断哪些区域值得搜索,再高效扫描这些区域中的候选向量。当K较大时,第二阶段的扫描成本会成为主要开销,IVF可以在这一阶段充分利用连续存储和向量化计算,因此比图遍历更容易获得稳定的吞吐量。

三、用Reservoir结构优化结果维护

前面提到在图索引中,优先队列在处理大TopK查询时会成为性能瓶颈,这一点在IVF索引中也无法避免,因为无论使用什么索引,只要系统需要从大量候选向量中返回距离最近的K个结果,就需要维护一个候选结果集合。

传统方法通常使用大小为K的优先队列,对于每个新候选的处理流程如下:

  1. 如果结果数量不足K,则直接插入队列;

  2. 如果新候选优于当前最差结果,则插入新结果;

  3. 弹出当前最差结果;

  4. 重新调整堆结构。

当K很大时,这一流程可能被执行数百万次,即使单次操作复杂度不高,累计成本仍然非常可观。

但大TopK查询并不要求系统在扫描过程中始终维护一个严格有序的TopK集合,只要扫描结束时能够准确选出最优的K个结果,中间状态可以更加宽松。因此可以通过Reservoir的方式,把存储结果的数据结构设置成比TopK更大的一个容量。

在搜索过程中,比当前第K大元素(注意这不是严格的)小的结果可以直接追加到结果中,而不需要像操作优先队列一样需要做队列结构的调整,只有当数据规模达到Reservoir的容量以后,再选出当前状态的第K大元素即可。这样就可以把大量细粒度的堆操作,转换成次数更少的批量筛选操作,而且K越大,这种批量维护方式相对于传统优先队列的优势通常越明显。

四、分布式查询的智能优化:AutoIndex

大TopK带来的挑战不仅在于索引的选型上,还在于查询模式的重构。主流的分布式向量检索架构通常采用scatter-gather模式:Proxy会把查询分发到多个查询节点,每个节点搜索自己负责的分片,执行本地TopK检索,并返回本地候选结果,系统随后对所有节点返回的候选进行合并,得到最终TopK。

对于较小的K,这种流程效果不错,但在大TopK场景中,如果每个分片都机械地执行一次完整的大TopK查询,成本会被显著放大。更关键的是,不同分片对最终结果的贡献通常并不相同,如果对所有分片分配相同的检索预算,就会在低贡献分片上浪费大量计算。

为了解决这类问题,分布式检索平台推出了AutoIndex,这是一个基于机器学习和运行时统计信息的优化器,用于在召回率和查询性能之间选择更合适的执行参数。当处理大TopK查询时,AutoIndex可以通过统计信息分析每个分片对结果的贡献程度,为它们分配不同的检索预算。

这意味着系统不必要求每个分片都返回同样规模的本地TopK,高贡献分片可以返回更多候选,低贡献分片则可以减少搜索范围和结果数量。在维持目标召回率的前提下,这种方式可以同时减少分片内部的索引搜索开销、查询节点返回的中间结果数量、节点之间的数据传输,以及全局结果归并的计算压力。

总结

总体来说,优化大TopK向量检索的成本可以分为三步:

  1. 算法层:使用更适合批量扫描的IVF索引;

  2. 执行层:使用Reservoir结构降低结果维护成本;

  3. 分布式层:使用AutoIndex控制分片查询预算。

将这些优化汇总后,可以显著提升大TopK查询的整体性能,降低计算和存储成本。

大TopK查询的实际配置方式

当前主流的向量检索平台默认的TopK上限通常为16384,对于需要返回更多结果的业务场景,可以开启大TopK查询模式,开启后单次查询最高可支持返回百万级别的相似结果。

在创建集合时,可以通过属性设置查询模式:

client.create_collection(
    collection_name="test",
    schema=schema,
    num_partitions=1,
    properties={"query_mode": "large_topk"}
)

对于已经创建的集合,也可以通过修改集合属性开启该模式:

client.alter_collection_properties(
    collection_name="scenarios_corpus",
    properties={"query_mode": "large_topk"}
)
以上内容不代表本平台立场,仅供读者参考