HNSW 分布式构建实践

2024-12-11
浙江
本文字数：3740 字
阅读完需：约 12 分钟

作者：魏子敬

一、背景

随着大模型时代的到来，向量检索领域面临着前所未有的挑战。embedding 的维度和数量空前增长，这在工程上带来了极大的挑战。智能引擎事业部负责阿里巴巴搜推广及 AI 相关工程系统的设计和建设，我们在实际业务迭代与发展中遭遇了 embedding 维度和数量扩张带来的诸多问题，其中索引构建时间问题尤为突出。

图 1 HNSW

以 HNSW[1]算法为代表的近似图算法因其高性价比与高召回率特性，在向量召回领域已逐渐成为主流技术选择。尤其在淘宝天猫、拍立淘、闲鱼等平台的搜索召回场景中，近似图算法扮演了至关重要的角色，其应用范围极为广泛。但是近似图算法被诟病的一个主要问题就是索引构建时间过长，特别是在高维海量数据场景下，这一缺点被进一步放大。

在分布式场景下，基于分而治之的思想，原始数据会被划分成多个正交子集，然后每个节点只负责其中的一个子集，这样就可以用多个独立的计算(存储)节点来解决单机无法解决的海量数据问题。在向量召回场景中也不例外，原始的 embedding 数据集会被按照 pk 哈希被均匀划分到多个分列中，这样通过一次分片就可以极大缩小单个实例内需要处理的向量数量，然后查询时只需要分别在各列的索引中先做召回，最后再把各列的结果汇总就可以了。

那是不是在不考虑成本的情况下，不管有多少数据只要无线分片下去就能解决一切规模膨胀带来的问题？

1.1 分布式的可扩展性困境(Scalability)

在我们的搜索推荐场景，业务上一般只要求最终一致性，这使我们无需支付高昂的成本来支持复杂的事务处理或维护严格一致性。而宽松的一致性要求使得水平扩展更加容易。然而即使如此，分布式环境的不稳定性以及网络开销等因素仍然限制了这种扩展能力，更不用说成本问题。

在向量召回场景中，仅因离线索引构建过慢而扩列是不划算的。以 HNSW 为例，如果撇开其他索引和检索时间，同样是在召回 topk 情境下，不进行分列并将所有数据构建成一个图，在检索计算量上来看是必然最优的。因此一个自然的问题是：如何加速近似图的构建呢？利用现代处理器的多线程并发处理能力是一个显而易见的方向。

1.2 多线程并发构建的瓶颈

实际上在 HNSW 原始的论文[1]中就提到了只要在几个关键节点加上锁就能实现构图的并行，而且图的质量几乎不受影响，这得感谢建图过程中绝大部分时间各个执行流之间都是完全隔离，同步点很少，在主流的向量库中也大都支持了 HNSW 的多线程构建，在十并发以内几乎可以实现接近线性的加速比，但是在我们的实践中也很快发现当并发数量加到三十之后加速比差不多就到了极限，即便增加更多并发，构建时间也难以进一步缩短。而且在线上环境中，由于混布环境资源稳定性的限制，很难确保所有子图构建都能始终处于高速状态。

二、分治构图

这就引出了最终的方案: 分布式构图，和水平分列分片不同，我们这里的分布式构图指的是使用多个实例来构建一张图，即分治构图，通过一些调研工作，我们也发现业界和学术界其实早就有了一些较为成熟的方案[2][3][4]，这些方案中除了 Pyramid[4]的方案(详见附录)外其他都大同小异，都能抽象成下边三个分治步骤:

1. 按照某种方法拆分原数据集，将原始数据集合 X 拆分成多个 X', X'之间可能会有交集。

2. 分别对每个 X'进行构图 G'。

3. 将所有 G'合并回一个大图 G，这一步通常是简单的边合并，随后可能还会对大 G 进行进一步的 Fine-grained optimization 处理。

第一步拆分可以选择的方案有 k-means、随机拆分、pca 主从分析后的多重随机拆分，如果原始数据集合过大，可能会先有一个随机采样的过程。

第二步通常都是正常的子图构建。

第三步都会有边合并，合并后的图可能会需要进一步的处理，如简单邻居传播或者 nn-desent[5]。如果在第一步拆分的过程中各个集合没有交集，在第三部合图时还需要加入一些胶水节点来保证图的连通性。

图 2 两种合并图的方式

2.1 分布式 HNSW

由于近似图算法的理论基础不足，工程上通常很难提前预估方案的有效性, 经过大量的线下实验验证后，结合工程实现上的考量，我们最终采用了 DISKANN[2]的分布式构图方案：第一步首先从原始数据集中采样一些 doc 做 k-means 聚集出 k 个聚心，然后把每篇 doc 分配到距离它最近的α个聚心上，这样每个聚心都得到了一个集合，每个聚心对应的集合分别构建 HNSW，最后在 reduce 过程直接对相同点做边合并，对于合并后度数超限的点也会进行 prune，利用 k 个集合之间的相同点来合并原图。

图 3 分治构图

图 4 分治构图

这种方案的优势就是改动小，舍弃了合图后的 Fine-grained optimization 过程，工程实践上简单，缺点就是虽然大部分计算都被分散到了子节点中，但总计算量其实比较大，最终参与计算的向量数量会膨胀α - 1 倍。但是在实践中子图的构图 ef 参数可以适当调小，用来抵消膨胀的计算量。我们在多个公开数据集以及部分线上生产数据集中实验后论证了这种分治构图法对图的质量影响并不大，大部分情况下甚至有正收益。

剩下的问题就是如何在 BuildService[6]中支持上述的三部构图流。BuildService 是智能引擎团队推出的分布式索引构建系统，在阿里搜推场景有着广泛的应用。幸运的是，凭借 BuildService 强大的分布式 DAG 图执行能力，HNSW 的分布式构建流程图可以非常简单的使用 BuildService 的索引定制框架来表达。

目前分布式构图功能已经在拍立淘、闲鱼等多个场景上线，离线全量时间及稳定性都得到极大的改善，平均全量时间由原来的 7~10h 降低至平均 3h 左右。

三、其他相关工作

3.1 Filtered HNSW

过滤召回或者带类目(标签)召回也是向量召回的一个很重要的应用场景，因为 embedding 本身很难完全表征完整的类目信息，所以针对一些特殊场景需要将向量召回和普通的类目过滤相互结合。

针对这种问题平凡的解决方案(查询时通过 filter 过滤)只在查询阶段对过滤条件进行了处理，并没有关注索引构建环节，所以并没有很好的利用到问题的特性，而且最坏情况下可能需要把全图遍历完才能凑够需要召回的向量个数。

另一个解决方案是为每个离散类目单独构建向量索引，但是当类目数量过多或者每个向量关联到多个类目时这个方案的存储成本就会变得难以接受。

在这方面 DISKANN[7]的另一篇文章提出了针对离散类目(标签)场景的两种带类目元信息的向量索引构建方案: FilteredVamana 和 StitchedVamana ，一种是基于流式处理的算法，另一种是基于批量构建的算法。这两种算法的核心都是在索引构建时不仅要考虑向量距离度量，还要考虑到相关的类目标签。

我们参考 StitchedVamana 批量构建的方法利用上文提到的分布式 HNSW 实现了 Stitched HNSW 索引: 给每个类目都构建一个子图，这样因为每个点有多个标签，所以每个点可能会属于多个子图，最后利用子图间的点重叠做图归并。

合并的过程中，如果出现度数超过限制的情况就进行三角裁边，这里的三角裁边使用了论文中的带类目三角裁边 FilteredRobustPrune，在最终的合并图上会给每个类目都保留一个导航点，这样查询时每种类目都有自己的导航点，包含相同类目的点在一个连通块内，查询过程中也只会访问包含这些类目的点，因此可以解决后置方案过滤中计算放大的问题，同时因为最终还是合并到一个图上，可以解决为每个类目单独建索引带来的索引膨胀问题。我们在拍立淘业务上线此功能后，大幅提升了带类目查询的向量召回性能。