基于磁盘量身定制，十亿规模高效向量检索方案

编者按：

“SPANN: Highly-efficient Billion-scale Approximate Nearest Neighbor Search”［1］是发表于 2021 年 人工智能领域全球顶会 NeurIPS 的一篇关于近似最近邻检索（ANNS）的论文，提出了高效的内存-磁盘混合索引和搜索方案 SPANN。

目前，大部分的 ANNS 方案主要基于向量数据能够完全放入内存中的假设，然而在实际业务中，有限的内存往往无法容纳日益增长的数据和索引文件。因此，向量检索引擎越来越青睐于使用外部存储设备与内存混合存储，从而满足容量、访问速度与召回率的平衡。
在设计上，SPANN 以较大粒度进行磁盘访问，充分利用了外部存储设备按块传输的访问特性，有效提高了外部存储设备的有效数据带宽。在十亿级向量数据检索场景中，SPANN 与现有的混合存储索引方案 DiskANN 相比，在内存开销相同的情况下，仅使用 1/3 到 1/2 的检索时间，即可达到相同的召回率。

本文将提炼这篇论文中的前沿思想、关键技术、论文实验，并针对这篇论文进行总结和思考。

论文作者信息

Qi Chen，HaiDong Wang，Mingqin Li 等人来自于微软亚洲研究院和必应团队，主要研究方向是云计算、云架构和大规模数据的深度学习。在大规模向量数据近似检索领域中，他们团队在 2018 年发表了“SPTAG: A library for fast approximate nearest neighbor search”，2021 年在 NeurIPS 发表了 “SPANN: Highly-efficient Billion-scale Approximate Nearest Neighbor Search”，以及在 GitHub 上开源了向量检索库（https://github.com/microsoft/...）。

论文动机

高维向量是将信息进行特征化表示的一种常用方式。ANNS 对特征化的向量快速地检索，避免了对向量检索引擎中的数据一一比对。例如，基于倒排表的检索方法对所有高维空间中的向量采用 K-Means 方法聚类到多个聚类集合中，对与查询向量靠近的少数的聚类集合进行查找来避免向量检索引擎中全量数据查找。

纯内存的索引方案已经逐渐难以容纳日益增长的数据量，近年来相关技术的探索热点逐渐开始向内存和非易失存储设备共同支持的方案转移。然而，混合存储的索引构建和查找方案在索引构建和查询这两个过程都会频繁地发生内存和磁盘之间的数据交换，不可避免地带来高昂的磁盘访问开销而导致性能降低。

现有在混合设备上构建的 ANNS 索引主要是基于图实现的，如 HM-ANN[2] 和 DiskANN[3]，数据以邻接表的形式存储在磁盘中。这类型的索引的数据访问是大量的细粒度的对存储在磁盘上的邻接表的随机数据访问，这种访问模式产生了大量的磁盘访问开销。GRIP[5] 结合了多种 ANNS 方法，内存中存储以聚类中心点集合构建的图索引和经过乘积量化压缩的向量，磁盘中存储无损的原始数据。查询向量首先经过图索引快速定位相近的中心点与其所在的聚类，然后比较这些聚类中的所有向量在内存中的量化表示来筛选需要磁盘访问获得的需要进一步比较的原始向量。GRIP 的磁盘访问过程也是对 SSD 的大量的细粒度的单条数据访问。总而言之，上述的单个磁盘访问粒度在几十到几百 KB 不等，且不具备访问的局部性，因此不能有效地利用外部存储器件的预读机制和操作系统的缓存机制，同时产生读放大。

论文提出了 SPANN，针对大规模向量近似搜索场景，采用小内存和大硬盘混合存储的策略。SPANN 基于倒排文件设计，能够有效地将相似的向量以小规模聚类集合的方式连续地存储在磁盘上，通过加载有限个数的聚类集合来减少磁盘访问。

基本思想

SPANN 实现为倒排文件结构，倒排文件通过 K-Means 算法将向量检索引擎中的相互靠近的数据聚类到一起，并使用中心点来表示该聚类集合。查询向量在倒排文件的检索时，首先会找到与查询向量靠近的多个聚类集合，然后在多个聚类集合中进行进一步的搜索，因此能够避免对整个向量检索引擎中的数据进行搜索。

每个聚类的中心点能够近似代表整个聚类集合中的向量，即如果查询向量离聚类的中心点很近，则可认为该查询向量离该聚类集合中所包含的所有向量都很近。因此，在倒排文件结构搜索过程中，查询点首先会和中心点进行一一比对，找到有限个相近的中心点向量，进一步对中心点向量所在的聚类进行搜索。

相比于向量检索引擎的数据规模，中心点集合的规模要小得多。SPANN 把中心点向量集合常驻在内存中提供快速的聚类候选集合的检索，定位存储在磁盘中的大量的小规模聚类集合。进一步， SPANN 通过将存储在磁盘中的多个聚类集合加载到内存中进行搜索。

磁盘访问优化方案

为了提供高效的磁盘访问，SPANN 在倒排文件的建索引和检索过程进行了优化。

对于倒排文件索引，增加聚类的个数，能够有效减少获得高召回所需要搜索的向量数。SPANN 利用了这一现象，采用了多层级的负载均衡的聚类算法，生成大量的聚类集合来减少总的向量检索引擎的数据搜索数目。并且限制了每个聚类的规模，使每个聚类大小尽可能地均匀，这是为了最小化不同的查询向量的访问开销的方差。

在先前的工作中，GRIP[5] 通过实验分析验证了这一现象。GRIP 进行了在相同的数据集下采用不同的聚类集合数目构建索引的召回比较实验。下图为实验结果，其中 NC 为聚类集合的个数。从图中可以看出，在相同的搜素召回下，搜索聚类的数目并不随着索引生成的聚类数目线性相关。尤其在高召回的情况下，不同聚类数目的索引所搜索的聚类数目相近。同时，在相同搜索聚类个数的情况下，聚类数目大的索引的向量搜索总数比聚类数目小的索引要小得多。

另一方面，SPANN 在对查询向量检索时实现动态剪枝，减少了磁盘访问的次数。因为查询向量具有“难易”之分，“难”的查询向量需要搜索大量的聚类才能获得高召回，“易”的查询向量需要搜索少量的聚类就能获得高召回。“易”的查询向量在搜索一定的聚类后，再增加搜索聚类的数目所获的召回提高并不显著。SPANN 通过动态剪枝可以避免召回收益不高的聚类的查询。

需要解决的问题

由于较均匀地对向量检索引擎中的数据进行聚类会产生大量的小规模的聚类集合，使聚类间的边缘点增多，进而影响召回率。下图可以解释这个现象：图中存在两个聚类集合，在两个聚类集合之外的黄色点为查询向量，聚类内灰色点为中心点。当需要检索的聚类数有限，查询向量因为和绿色集合中的中心点更近，而选择绿色的聚类集合进行进一步的检索，实际上蓝色聚类集合中的红色边缘点里离查询向量更近。针对磁盘访问优化策略所带来的召回率的损失，利用了磁盘的大容量，SPANN 对边缘点采用分配到多个聚类集合中的策略。

在了解 SPANN 的磁盘访问优化方案和弥补召回下降策略的大致思想后，产生了几点疑问：（1）如何在高昂的大量聚类集合训练过程支持向量的均衡划分？（2）向量冗余放置在多个聚类集合中所带来的额外的磁盘访问如何处理？（3）不同的查询向量需要查找的聚类集合个数不同，如何调整查询策略来应对“难”和“不难”的查询向量？SPANN 通过以下的三种关键技术来解决这些问题。

关键技术

关键技术 1：均衡多层聚类算法

为了限制聚类集合的规模，SPANN 采用了多约束平衡聚类算法[4] 将数据集均匀地划分到大量的聚类集合中。损失函数表示为公式(1)：

其中，X 表示向量检索引擎中的向量矩阵，|X| 表示向量检索引擎中的向量个数，C 表示聚类中心点的矩阵，N 表示聚类中心点的个数，H 是表示向量划分的$$|X| N 的位矩阵，λ* 是控制聚类和条件约束比重的超参数。多约束聚类算法的训练过程和 K-means 的大致相同，所用的损失函数是在 K-means 的损失函数的基础上多增加了一个均衡划分的条件约束，采用的条件约束条件可以理解为最小化划分后的聚类大小的方差。

SPANN 为了缓解训练大量聚类集合所带来的高昂的时间开销，采用了层次聚类的方法将时间复杂度从 O(|X|dimN) 降为 O(|X|dimk*log_k(N)) ，如下图所示。SPANN 采用自顶向下的层次聚类方式。层次聚类中的根节点会根据数据集中的所有向量采用较小的 K 值进行聚类，并把所有的向量划分到 K 个聚类中去。根节点和内部节点所生成的 K 个聚类作为它们的子节点。这 K 个子节点会根据其数据规模来区分是叶节点或是内部节点。如果节点中的数据集合大小小于特定的限制值，则作为落盘的叶子节点；否则作为下一层级的内部节点进行下一层聚类。反复迭代这个过程，直至所有的聚类集合都小于特定的限制值并落盘。 同时，SPANN 为了快速地检索聚类的候选集，采用 SPTAG 索引结构维护存储在内存中的中心点。

关键技术 2：边缘点的冗余放置

为了缓解上述的边缘问题而导致的召回下降，SPANN 选择将聚类中的边缘点冗余放置到多个相近的聚类集合中，来增加这些向量的可见性。SPANN 将向量检索引擎中的向量分配到距离差不多的最近的几个聚类集合中，如下图所示，“距离差不多”这一条件能够筛选出边缘点而避免了向量检索引擎中所有向量的大量冗余所带来的高时延的磁盘访问开销。聚类中心点根据与向量检索引擎中的向量 x 的距离进行排序，向量 x 会被分配到中心点满足公式(2) 的聚类集合中。公式(2) 中使用了一个参数 \epsilon_1 来限制边缘点重复放置的规模，SPANN 对边缘点根据与各个中心点的距离来筛选出在最小距离 Dist(x, c_{i1}) 和 (1 + \epsilon_1) 倍最小距离 (1 + \epsilon_1) * Dist(x, c_{i1}) 之间的中心点。

SPANN 采用 RNG 规则进一步减少边缘点的冗余放置。边缘点可以不重复放置在相近的聚类内，因为相近的聚类中心点相似度很高，被同时召回的概率也会提高，从而相同的数据会多次分配到磁盘中。对于聚类集合的中心点 c_{ij}，如果 x 和 c_{ij} 的相对邻域内出现中心点 c_{il}，中心点 c_{il} 和向量 x 的距离小于中心点 c_{ij} 和向量 x 的距离，则认为中心点 c_{ij} 和中心点 c_{il} 相近，向量 x 不会分配到 c_{ij} 的聚类中。

关键技术 3：查询感知的动态剪枝

SPANN 采用查询感知的剪枝方法来调整不同查询向量需要加载的候选聚类的数目。通常，在倒排文件的检索过程中，首先会找到相同个数个最近的聚类中心点，然后对固定个聚类进行全量的搜索。然后，由于查询向量具有差异，有的“容易”查询向量只需要检索少数的聚类就能够获得高召回，而有的“难”查询向量需要检索更多的聚类。倒排文件使用相同的查找聚类个数会导致“容易”的查询向量需要检索很多召回收益不高的聚类。

SPANN 通过查询向量与聚类中心点的距离远近程度来确定对该聚类的搜索是否是高召回收益的。SPANN 利用查询向量到与其最近的聚类中心点 c_{i1} 的距离 Dist(q, c_{i1}) 作为尺度，使用搜索参数 \epsilon_2 来控制动态剪枝的程度。当查询向量和某聚类中心点的距离大于 (1 + \epsilon_2) * Dist(q, c_{i1})，则认为是查询向量和中心点距离较远，对这一聚类进行进一步搜索的收益不高，可以进行剪枝，不对其进行搜索。

论文实验

论文的实验在 2600MHz 主频的 52 核 CPU，128GB 的内存和 2.6TB 的 SSD 机器上进行。采用召回、时延和 VQ 作为评价指标。其中，VQ 是索引处理的数据量规模和 QPS 的乘积，同时反应查询延迟和内存开销，VQ 越大，索引所消耗的资源越少。

论文使用十亿级规模的数据集（SIFT1B、DEEP1B 和 SPACEV1B）与 DiskANN 进行比较，百万级规模的数据集（SIFT1M）与多个比较出名的内存索引方案（HNSW、NSG、NGT 等）进行比较。

相比于 DiskANN，SPANN 能够更快地获得召回的峰值。如下图所示，SPANN 仅用 DiskANN 一半的搜索延迟在 SIFT1B 数据集达到 95% 的召回率。

SPANN 和现有的内存索引方案的比较如下图所示。在不同的召回下，SPANN 的 VQ 值最高，说明SPANN 能够处理更大数据规模的场景。

总结和思考

总体而言，SPANN 基于倒排文件表索引进行了针对磁盘访问的数据布局和访问的优化方案设计。相比于DiskANN，避免了向量量化所带来的数据精度损失和细粒度的磁盘随机访问，因此能够在较短的搜索时间内提供高召回。

个人认为，SPANN 还可以在以下方面有改进：

1. SPANN 的系统参数过多而导致难以获得最优的性能。在 SPANN 的构建和检索过程中，针对不同的硬件配置和场景，需要考虑到如何设置每个聚类的数据规模、K-Means 的损失函数中控制聚类均衡的超参数、边缘点冗余放置参数和动态查询剪枝参数。
2. 层次聚类和小规模的聚类集合这两种方法都会使聚类之间的边缘点增加而导致召回率的下降。SPANN 是在牺牲召回的情况下，减少了索引的训练时间和磁盘访问时间。从实验结果看出，SPANN 搜索延迟随着搜索聚类个数增加，同时召回率也增加。但是，过多的边缘点带来了召回率上升的瓶颈，当搜索的聚类数超过一定阈值，召回率不再随着搜索代价的增加而增加。
3. 大部分的服务器环境具有比较充裕的内存空间，因此 SPANN 不仅可以存储中心点的索引结构，还可以作为磁盘设备的缓存，利用查询向量集合的局部性存储频繁访问的聚类集合，进一步减少磁盘的访问次数。

本文作者

程倩雅， Zilliz 研发工程师，华中科技大学计算机系统结构硕士。目前在 Zilliz 的研究团队，主要工作是 ANNS 领域算法的研究和开发以及数据库系统方向预研。

参考文献

[1] Chen Q, Zhao B, Liu H W M L C, et al. SPANN: Highly-efficient Billion-scale Approximate Nearest Neighbor Search[J]. arXiv preprint arXiv:2111.08566, 2021.

[2] Ren J, Zhang M, Li D. Hm-ann: Efficient billion-point nearest neighbor search on heterogeneous memory[J]. Advances in Neural Information Processing Systems, 2020, 33.

[3] Subramanya S J, Kadekodi R, Krishaswamy R, et al. DiskANN: Fast accurate billion-point nearest neighbor search on a single node[C]//Proceedings of the 33rd International Conference on Neural Information Processing Systems. 2019: 13766-13776.

[4] Liu H, Huang Z, Chen Q, et al. Fast Clustering with Flexible Balance Constraints[C]//2018 IEEE International Conference on Big Data (Big Data). IEEE, 2018: 743-750.

[5] Zhang M, He Y. Grip: Multi-store capacity-optimized high-performance nearest neighbor search for vector search engine[C]//Proceedings of the 28th ACM International Conference on Information and Knowledge Management. 2019: 1673-1682.

---

Zilliz 以重新定义数据科学为愿景，致力于打造一家全球领先的开源技术创新公司，并通过开源和云原生解决方案为企业解锁非结构化数据的隐藏价值。
Zilliz 构建了 Milvus 向量数据库，以加快下一代数据平台的发展。Milvus 数据库是 LF AI & Data 基金会的毕业项目，能够管理大量非结构化数据集，在新药发现、推荐系统、聊天机器人等方面具有广泛的应用。