“DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node” [1]是 2019 年发表在 NeurIPS 上的论文。
该文提出了一种基于磁盘的 ANN 方案,该方案可以在单个 64 G 内存和足够 SSD 的机器上对十亿级别的数据进行索引、存储和查询, 并且能够满足大规模数据 ANNS 的三个需求: 高召回、低查询时延和高密度(单节点能索引的点的数量)。
该文提出的方法做到了在 16 核 64G 内存的机器上对十亿级别的数据集 SIFT1B 建基于磁盘的图索引,并且 recall@1 > 95% 的情况下 qps 达到了 5000, 平均时延不到 3ms。
Suhas Jayaram Subramanya:前微软印度研究院员工,CMU 在读博士。主要研究方向有高性能计算和面向大规模数据的机器学习算法。
Devvrit:Graduate Research Assistant at The University of Texas at Austin。研究方向是理论计算机科学、机器学习和深度学习。
Rohan Kadekodi:德克萨斯大学的博士研究生。研究方向是系统和存储,主要包括持久化存储、文件系统和 kv 存储领域。
Ravishankar Krishaswamy:微软印度研究院 Principal Researcher。CMU 博士学位。研究方向是基于图和聚类的近似算法。
Harsha Vardhan Simhadri:微软印度研究员 Principal Researcher。CMU 博士学位。以前研究并行算法和运行时系统,现在主要工作是开发新算法,编写编程模型。
目前有很多向量检索的 ANN 算法,这些算法在建索引性能、查询性能和查询召回率方面各有取舍。当前在查询时间和召回率上表现较好的是基于图的索引如 HNSW 和 NSG。由于图索引内存占用比较大,在单机内存受限的情况下,常驻内存的方案能处理点集的规模就十分有限。
许多应用需要快速在十亿级别数据规模上做基于欧几里得距离的近似查询,目前有两种主流的方法:
以上两种方法的局限性在于太依赖内存。所以这篇论文考虑设计一种索引常驻 SSD 的方案。索引常驻 SSD 的方案主要面临的挑战是如何减少随机访问 SSD 的次数和减少发起 SSD 访问请求的数量。将传统的基于内存的 ANNS 算法放到 SSD 上的话平均单条查询会产生数百个读磁盘操作,这会导致极高的时延。
这篇论文提出了能够有效支持大规模数据的常驻 SSD 的 ANNS 方案:DiskANN。该方案基于这篇论文中提出的另一个基于图的索引:Vamana,后面会详细介绍。这篇论文的主要贡献包括但不限于:
这个算法和 NSG[2] [4]思路比较像(不了解 NSG 的可以看参考文献 2,不想读 paper 的话可以看参考文献 4),主要区别在于裁边策略。准确的说是给 NSG 的裁边策略上加了一个开关 alpha。NSG 的裁边策略主要思路是:对于目标点邻居的选择尽可能多样化,如果新邻居相比目标点,更靠近目标点的某个邻居,我们可以不必将这个点加入邻居点集中。
也就是说,对于目标点的每个邻居节点,周围方圆 dist(目标点,邻居点)范围内不能有其他邻居点。这个裁边策略有效控制了图的出度,并且比较激进,所以减少了索引的内存占用,提高了搜索速度,但同时也降低了搜索精度。
Vamana 的裁边策略其实就是通过参数 alpha 自由控制裁边的尺度。具体作用原理是给裁边条件中的 dist(某个邻居点,候选点) 乘上一个不小于 1 的参数 alpha,当 dist(目标点,某个候选点)大于这个被放大了的参考距离后才选择裁边,增加了目标点的邻居点之间的互斥容忍度。
Vamana 的建索引过程比较简单:
论文比较了 Vamana、NSG、HNSW 3种图索引,无论是建索引性能还是查询性能, Vamana 和 NSG 都比较接近,并且都稍好于 HNSW。具体数据可以看下文实验部分。
为了直观地表现建 Vamana 索引过程,论文中给出了这么一张图,用 200 个二维点模拟了两轮迭代过程。第一行是用 alpha = 1 来裁边,可以发现改裁边策略比较激进,大量的边被裁剪。经过放大 alpha,裁边条件放松后,明显加回来了不少边,并且第二行最右这张图,即最终的图中,明显加了不少长边。这样可以有效减少搜索半径。
一台只有 64G 内存的个人电脑连十亿原始数据都放不下,更别说建索引了。摆在我们面前的有两个问题:1. 如何用这么小的内存对这么大规模的数据集建索引?2. 如果原始数据内存放不下如何在搜索时计算距离?
本文提出的方法:
DiskANN 的索引布局和一般的图索引类似,每个点的邻居集和原始向量数据存在一起。这样做的好处是可以利用数据的局部性。
前面提到了,如果索引数据放在 SSD 上,为了保证搜索时延,尽可能减少磁盘访问次数和减少磁盘读写请求。因此 DiskANN 提出两种优化策略:
实验分三组:
一、 基于内存的索引比较:Vamana VS. NSG VS. HNSW
数据集:SIFT1M(128 维), GIST1M(960 维), DEEP1M(96 维) 以及从 DEEP1B 中随机采样了 1M 的数据集。
索引参数(所有数据集都采用同一组参数):
HNSW:M = 128, efc = 512.
Vamana: R = 70, L = 75, alpha = 1.2.
NSG: R = 60, L = 70, C= 500.
论文里没有给搜索参数,可能和建索引参数一致。对于这个参数选择,文中提到 NSG 的参数是根据 NSG 的github repository 里列出的参数中选择出性能比较好的那组,Vamana 和 NSG 比较接近,因此参数也比较接近,但是没有给出 HNSW 的参数选择理由。在笔者看来,HNSW 的参数 M 选择偏大,同为图索引,出度应该也要在同一水平才能更好做对比。
在上述建索引参数下,Vamana、HNSW 和 NSG 建索引时间分别为 129s、219s 和 480s。NSG 建索引时间包括了用 EFANN [3] 构建初始化近邻图的时间。
召回率-qps 曲线:
从 Figure 3 可以看出,Vamana 在三个数据集上都有着优秀的表现,和 NSG 比较接近,比 HNSW 稍好。
比较搜索半径:
这个结果可以看 Figure 2.c,从图中可以看出 Vamana 相比 NSG 和 HNSW,在相同召回率下平均搜索路径最短。
二、 比较一次性建成的索引和多个小索引合并成一个大索引的区别
数据集:SIFT1B
一次建成的索引参数:L = 50, R = 128, alpha = 1.2. 在 1800G DDR3 的机器上跑了 2 天, 内存峰值大约 1100 G,平均出度 113.9。
基于合并的索引步骤:
这个索引生成了一个 384G 的索引,平均出度 92.1。这个索引在 64G DDR4 的机器上跑了 5 天。
比较结果如下(Figure 2a):
结论:
三、 基于磁盘的索引:DiskANN VS. FAISS VS. IVF-OADC+G+P
IVFOADC+G+P 是参考文献 [5] 提出的一种算法。
这篇论文只和 IVFOADC+G+P 比较,因为参考文献 [5] 中已经证明比 FAISS 要更优秀,另外 FAISS 要用 GPU,并不是所有平台都支持。
IVF-OADC+G+P 好像是一个 hnsw + ivfpq。用 hnsw 确定 cluster,然后在目标 cluster 上加上一些剪枝策略进行搜索。
结果在 Figure 2a里。图里的 16 和 32 是码本大小。数据集是 SIFT1B,用的 OPQ 量化。
DiskANN 的源码已经开源:https://github.com/microsoft/DiskANN
2021 年 1 月开源了磁盘方案的源码,稍微看了下,大概介绍一下磁盘方案的实现细节。
这个代码质量比较高,可读性比较强,建议有能力的可以读一下。
下面主要介绍建索引过程和搜索过程。
建索引
建索引参数有 8 个:
data_type: float/int8/uint8 三选一
data_file.bin: 原始数据二进制文件,文件前 2 个 int 分别表示数据集向量总数 n 和向量维度 dim。后 n * dim * sizeof(data_type) 字节就是连续的向量数据。
index_prefix_path: 输出文件的路径前缀,索引建完后会生成若干个索引相关文件,这个参数是公共前缀。
R: 全局索引的最大出度。
L: Vamana 索引的参数 L,候选集大小上界。
B: 查询时的内存阈值,单位 GB,控制 pq 码本大小。
M: 建索引时的内存阈值,决定分片大小,单位 GB。
T: 线程数。
建索引流程(入口函数:aux_utils.cpp::build_disk_index):
1. 根据 index_prefix_path 生成各种产出文件名。
2. 参数检查。
3. 读 data_file.bin 的 meta 获取 n 和 dim。根据 B 和 n 确定pq 的码本子空间数 m。
4. generate_pq_pivots: 用 p = 1500000/n 的采样率全局均匀采样出 pq 训练集训练 pq 中心点。
5. generate_pq_data_from_pivots: 生成全局 pq 码本,中心点和码本分别保存。
6. build_merged_vamana_index: 对原始数据集切片,分段建 Vamana 索引,最后合并。
7. create_disk_layout: 步骤 6 中生成的全局索引只有邻接表,而且还是紧凑的邻接表,这一步就是把索引对齐,邻接表和原始数据存在一起,搜索的时候加载邻接表顺便把原始向量一起读上去做精确距离计算。这里还有一个 SECTOR 的概念,默认大小是 4096,每个 SECTOR 只放 4096 / node_size 条向量信息,node_size = 单条向量大小+单节点邻接表大小。
8. 最后再做一个 150000 / n 的全局均匀采样,存盘,搜索时用来做 warmup。
搜索
搜索参数主要有 10 个:
index_type: float/int8/uint8 三选 1,和建索引第一个参数 data_type 是一样的;
index_prefix_path: 同建索引参数 index_prefix_path;
num_nodes_to_cache: 缓存热点数;
num_threads: 搜索线程数;
beamwidth: 预加载点数上限,如果为 0 由程序自己确定;
query_file.bin: 查询集文件;
truthset.bin: 结果集文件,“null” 表示不提供结果集,程序自己算;
K: topk;
result_output_prefix: 搜索结果保存路径;
L*: 搜索参数列表,可以写多个值,对于每个 L 都会做搜索并输出不同参数 L 下的统计信息。
搜索流程:
1. 加载相关数据:查询集、pq 中心点数据、码本数据、搜索起点等数据,读取索引 meta;
2. 用建索引时采样的数据集做 cached_beam_search,统计每个点的访问次数,将 num_nodes_to_cache 个访问频次最高的点加载到缓存;
3. 默认有一个 WARMUP 操作,和步骤 2 一样,也是用这个采样数据集做一次 cached_beam_search,不太明白这一步有什么用,因为步骤 2 其实已经起到了 warm up 的作用了,难道是再刷一遍系统缓存?
4. 根据给的参数 L 个数,每个参数 L 都会用查询集做一遍 cached_beam_search,输出召回率、qps 等统计信息。warm up 和 统计热点数据的过程不计入查询时间。
关于 cached_beam_search:
1. 从候选起点中找离查询点最近的,这里用的 pq 距离,起点加入搜索队列;
2. 开始搜索:
这篇论文加代码花了一点时间,总体来说还是很优秀的。论文和代码思路都比较清晰,通过 kmeans 分若干 overlap 的桶,然后分桶建图索引,最后合并的思路还是比较新颖的。至于基于内存的图索引 Vamana,本质上是一个随机初始化版本的可以控制裁边粒度的 NSG。查询的时候充分利用了缓存 + pipline,掩盖了部分 io 时间,提高了 qps。
不过按论文中说的,就算机器配置不高,训练时间长达 5 天,可用性也比较低,后续可以考虑对训练部分做一些优化。从代码来看,质量比较高,可以直接上生产环境那种。感觉建索引那段代码在算法和实现方面还是有加速空间的。
参考资料:
1. Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, Rohan Kadekodi. DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node. NeurIPS 2019. (https://www.microsoft.com/en-us/research/publication/diskann-fast-accurate-billion-point-nearest-neighbor-search-on-a-single-node)
2. Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai. Fast approximate nearest neighbor search with the navigating spreading-out graphs. PVLDB, 12(5):461 – 474, 2019. (http://www.vldb.org/pvldb/vol12/p461- fu.pdf)
3. Cong Fu and Deng Cai. ( https://github.com/ZJULearning/efanna)
4. Search Engine For AI:高维数据检索工业级解决方案 (https://zhuanlan.zhihu.com/p/50143204)
5. Dmitry Baranchuk, Artem Babenko, and Yury Malkov. Revisiting the inverted indices for billion-scale approximate nearest neighbors.(https://arxiv.org/abs/1802.02422)
✏️ 作者:李成明,Zilliz 研发工程师,东南大学计算机硕士。主要关注大规模高维向量数据的相似最近邻检索问题,包括但不限于基于图和基于量化等向量索引方案,目前专注于 Milvus 向量搜索引擎 knowhere 的研发。喜欢研究高效算法,享受实现纯粹的代码,热衷压榨机器的性能。Github 账号为 https://github.com/op-hunter
Github @Milvus-io|CSDN @Zilliz Planet|Bilibili @Zilliz-Planet
Zilliz 以重新定义数据科学为愿景,致力于打造一家全球领先的开源技术创新公司,并通过开源和云原生解决方案为企业解锁非结构化数据的隐藏价值。
Zilliz 构建了 Milvus 向量数据库,以加快下一代数据平台的发展。Milvus 目前是 LF AI & Data 基金会的毕业项目,能够管理大量非结构化数据集。我们的技术在新药发现、计算机视觉、推荐引擎、聊天机器人等方面具有广泛的应用。