乐趣区

关于数据库:论文赏析十亿级别单机向量检索方案DiskAnn

jiezi

摘要

”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。因为图索引内存占用比拟大，在单机内存受限的状况下，常驻内存的计划能解决点集的规模就非常无限。

许多利用须要疾速在十亿级别数据规模上做基于欧几里得间隔的近似查问，目前有两种支流的办法：

  1. 基于 倒排表 + 量化 的办法。毛病是召回率不高，因为量化会产生误差。尽管能够进步 topk 以改善召回率，然而相应的会升高 qps。
  2. 基于 分治 。将数据集分成若干个不相交的子集，每个子集建基于内存的索引，最初对后果做归并。这种办法比拟耗内存和机器。例如对 100M, 128 维的 float 数据集上建 NSG 索引，假如出度下限为 50，大概须要 75G 内存。因而解决十亿级别的数据就须要多台机器。在阿里巴巴淘宝的理论利用中，将 20 亿 128 维浮点数据分到 32 个分片中别离建 NSG 索引，recall@100 为 98% 的时延大概在 5ms。当数据规模减少到千亿级别时须要数千台机器。[2]

以上两种办法的局限性在于太依赖内存。所以这篇论文思考设计一种索引常驻 SSD 的计划。索引常驻 SSD 的计划次要面临的挑战是如何缩小随机拜访 SSD 的次数和缩小发动 SSD 拜访申请的数量。将传统的基于内存的 ANNS 算法放到 SSD 上的话均匀单条查问会产生数百个读磁盘操作，这会导致极高的时延。

论文奉献

这篇论文提出了可能无效反对大规模数据的常驻 SSD 的 ANNS 计划：DiskANN。该计划基于这篇论文中提出的另一个基于图的索引：Vamana，前面会具体介绍。这篇论文的次要奉献包含但不限于：

  1. DiskANN 能够在一台 64G 内存的机器上对十亿级别的维度大于 100 的数据集进行索引构建和提供查问服务，并在单条查问 recall@1 > 95% 的状况下均匀时延不超过 5ms。
  2. 提出了基于图的新索引 Vamana，该索引相比目前最先进的 NSG 和 HNSW 具备更小的搜寻半径，这个性质能够最小化 DiskANN 的磁盘拜访次数。
  3. Vamana 搜寻性能不慢于目前最好的图索引 NSG 和 HNSW。
  4. DiskANN 计划通过将大数据集分成若干个相交的分片，而后对每个分片建基于内存的图索引 Vamana，最初将所有分片的索引合并成一个大索引，解决了内存受限的状况下对大数据集建设索引的问题。
  5. Vamana 能够和现有的量化办法如 PQ 联合，量化数据能够缓存在内存中，索引数据和向量数据能够放在 SSD 上。

Vamana

这个算法和 NSG[2][4] 思路比拟像（不理解 NSG 的能够看参考文献 2，不想读 paper 的话能够看参考文献 4），次要区别在于裁边策略。精确的说是给 NSG 的裁边策略上加了一个开关 alpha。NSG 的裁边策略次要思路是：对于指标点街坊的抉择尽可能多样化，如果新街坊相比指标点，更凑近指标点的某个街坊，咱们能够不用将这个点退出街坊点集中。也就是说，对于指标点的每个街坊节点，四周方圆 dist(指标点，街坊点）范畴内不能有其余街坊点。这个裁边策略无效管制了图的出度，并且比拟激进，所以缩小了索引的内存占用，进步了搜寻速度，但同时也升高了搜寻精度。Vamana 的裁边策略其实就是通过参数 alpha 自在管制裁边的尺度。具体作用原理是给裁边条件中的 dist( 某个街坊点，候选点）乘上一个不小于 1 的参数 alpha，当 dist( 指标点，某个候选点）大于这个被放大了的参考间隔后才抉择裁边，减少了指标点的街坊点之间的互斥容忍度。

Vamana 的建索引过程比较简单：

  1. 初始化一张随机图；
  2. 计算终点，和 NSG 的导航点相似，先求全局质心，而后求全局离质心最近的点作为导航点。和 NSG 的区别在于：NSG 的输出曾经是一张近邻图了，所以间接在初始近邻图上对质心点做一次近似最近邻搜寻就能够了。然而 Vamana 初始化是一张随机近邻图，所以不能在随机图上间接做近似搜寻，须要全局比对，失去一个导航点，这个点作为后续迭代的起始点，目标是尽量减少均匀搜寻半径；
  3. 基于初始化的随机近邻图和步骤 2 中确定的搜寻终点对每个点做 ANN，将搜寻门路上所有的点作为候选街坊集，执行 alpha = 1 的裁边策略。这里和 NSG 一样，抉择从导航点登程的搜寻门路上的点集作为候选街坊集会减少一些长边，无效缩小搜寻半径。
  4. 调整 alpha > 1(论文举荐 1.2）反复步骤 3。因为 3 是基于随机近邻图做的，第一次迭代后图的品质不高，所以须要再迭代一次来晋升图的品质，这个对召回率很重要。

论文比拟了 Vamana、NSG、HNSW 3 种图索引，无论是建索引性能还是查问性能，Vamana 和 NSG 都比拟靠近，并且都稍好于 HNSW。具体数据能够看下文试验局部。

为了直观地体现建 Vamana 索引过程，论文中给出了这么一张图，用 200 个二维点模仿了两轮迭代过程。第一行是用 alpha = 1 来裁边，能够发现改裁边策略比拟激进，大量的边被裁剪。通过放大 alpha，裁边条件放松后，显著加回来了不少边，并且第二行最右这张图，即最终的图中，显著加了不少长边。这样能够无效缩小搜寻半径。

DiskAnn

一台只有 64G 内存的个人电脑连十亿原始数据都放不下，更别说建索引了。摆在咱们背后的有两个问题：1. 如何用这么小的内存对这么大规模的数据集建索引？2. 如果原始数据内存放不下如何在搜寻时计算间隔？

本文提出的办法：

  1. 对于第一个问题，先做全局 kmeans，将数据分成 k 个簇，而后将每个点分到间隔最近的 I 个簇中，个别 I 取 2 就够了。对每个簇建基于内存的 Vamana 索引，最初将 k 个 Vamana 索引合并成一个索引。
  2. 对于第二个问题，能够应用量化的办法，建索引时用原始向量，查问的时候用压缩向量。因为建索引应用原始向量保障图的品质，搜寻的时候应用内存能够 hold 住的压缩向量进行粗粒度搜寻，这时的压缩向量尽管有精度损失，然而只有图的品质足够高，大方向上是对的就能够了，最初的间隔后果还是用原始向量做计算的。

DiskANN 的索引布局和个别的图索引相似，每个点的街坊集和原始向量数据存在一起。这样做的益处是能够利用数据的局部性。

后面提到了，如果索引数据放在 SSD 上，为了保障搜寻时延，尽可能减少磁盘拜访次数和缩小磁盘读写申请。因而 DiskANN 提出两种优化策略：

  1. 缓存热点：将终点开始 C 跳内的点常驻内存，C 取 3~4 就比拟好。
  2. beam search：简略的说就是预加载，搜寻 p 点时，如果 p 的街坊点不在缓存中，须要从磁盘加载 p 点的街坊点。因为一次大量的 SSD 随机拜访操作和一次 SSD 单扇区拜访操作耗时差不多，所以咱们能够一次加载 W 个未拜访点的街坊信息，W 不能过大也不能过小，过大会节约计算和 SSD 带宽，太小了也不行，会减少搜寻时延。

试验

试验分三组：

一、基于内存的索引比拟：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 [3] 里列出的参数中抉择出性能比拟好的那组，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。

基于合并的索引步骤：

  1. 将数据集用 kmeans 训练出 40 个簇；
  2. 每个点分到最近的 2 个簇；
  3. 对每个簇建 L = 50，R = 64，alpha = 1.2 的 Vamana 索引；
  4. 合并每个簇的索引。

这个索引生成了一个 384G 的索引，均匀出度 92.1。这个索引在 64G DDR4 的机器上跑了 5 天。

比拟后果如下（Figure 2a）：

论断：

  1. 一次建成的索引显著优于基于合并的索引；
  2. 基于合并的索引也很优良；
  3. 基于合并的索引计划也实用于 DEEP1B 数据集（Figure 2b）。

三、基于磁盘的索引：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 的源码曾经开源：GitHub 地址

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 索引，最初合并。

    • 6.1：partition_with_ram_budget: 依据参数 M 确定分片数 k。对数据集采样做 kmeans，每个点分到最近的两个簇中，对数据集进行分片，每个分片产生两个文件：数据文件和 id 文件。id 文件和数据文件一一对应，id 文件中每个 id 对应数据文件中每条向量一一对应。这里的 id 能够认为是对原始数据的每条向量按 0 ～ n-1 编号。这个 id 比拟重要，跟前面的合并相干。

      • 6.1.1：用 1500000 / n 的采样率全局平均抽样出训练集

      • 6.1.2：初始化 num_parts = 3。从 3 开始迭代。

        • 6.1.2.1：对步骤 6.1.1 中的训练集做 num_parts-means++；
        • 6.1.2.2：用 0.01 的采样率全局平均采样出一个测试集，将测试集中的分到最近的 2 个簇中；
        • 6.1.2.3：统计每个簇中的点数，除以采样率估算每个簇的点数；
        • 6.1.2.4：按 Vamana 索引大小估算步骤 6.1.2.3 中最大的簇须要的内存，如果不超过参数 M，转步骤 6.1.3，否则 num_parts ++ 转步骤 6.1.2；
      • 6.1.3：将原始数据集分成 num_parts 组文件，每组文件包含分片数据文件和分片数据对应的 id 文件。
    • 6.2：对步骤 6.1 中的所有分片独自建设 Vamana 索引并存盘；

    • 6.3：merge_shards: 将 num_parts 个分片 Vamana 合并成一个全局索引。

      • 6.3.1：读取 num_parts 个分片的 id 文件到 idmap 中，这个 idmap 相当于建设了分片 ->id 的正向映射；
      • 6.3.2：依据 idmap 建设 id-> 分片的反向映射，晓得每条向量在哪两个分片中。
      • 6.3.3：用带 1G 缓存的 reader 关上 num_parts 个分片 Vamana 索引，用带 1G 缓存的 writer 关上输入文件，筹备合并；
      • 6.3.4：将 num_parts 个 Vamana 索引的导航点落盘到中心点文件中，搜寻的时候会用到；
      • 6.3.5：按 id 从小到大开始合并，依据反向映射顺次读取每条原始向量在各个分片的街坊点集，去重，shuffle，截断，写入输入文件。因为当初切片时是全局有序的，当初合并也按程序来，所以最终的落盘索引中的 id 和原始数据的 id 是一一对应的。
      • 6.3.6：删除临时文件，包含分片文件、分片索引、分片 id 文件；
  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. 开始搜寻：

    • 2.1：从搜寻队列中看不超过 beam_width + 2 个未拜访过的点，如果这些点在缓存中，退出缓存命中的队列，如果不命中，退出未命中的队列，保障未命中的队列大小不超过 beam_width；
    • 2.2：对于未命中队列中的点，发送异步磁盘拜访申请；
    • 2.3：对于缓存命中的点，用原始数据和查问数据算准确间隔退出后果队列，而后对这些点未拜访过的街坊点，用 pq 算间隔后退出搜寻队列，搜寻队列长度受参数限度；
    • 2.4：解决步骤 a 中缓存未命中的点，过程和步骤 c 一样；
    • 2.5：搜索队列为空时完结搜寻，返回后果队列 topk。

总结

这篇论文加代码花了一点工夫，总体来说还是很优良的。论文和代码思路都比拟清晰，通过 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.
  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. doi: 10.14778/3303753.3303754. URL http://www.vldb.org/pvldb/vol12/p461- fu.pdf.
  3. [NSG GitHub] https://github.com/ZJULearnin…
  4. Search Engine For AI：高维数据检索工业级解决方案
  5. Dmitry Baranchuk, Artem Babenko, and Yury Malkov. Revisiting the inverted indices for billion-scale approximate nearest neighbors.

笔者简介

李成明，Zilliz 研发工程师，东南大学计算机硕士。次要关注大规模高维向量数据的类似最近邻检索问题，包含但不限于基于图和基于量化等向量索引计划，目前专一于 Milvus 向量搜索引擎 knowhere 的研发。喜爱钻研高效算法，享受实现纯正的代码，热衷压迫机器的性能。他的 Github 账号为 https://github.com/op-hunter

退出移动版