1. 引言
最近邻搜索(Nearest Neighbor Search, NNS)是信息检索、模式识别、数据挖掘和推荐系统等领域的基础算法 [1] 。然而随着数据规模的持续增长和以及维度增长带来的维度灾难，现有的NNS算法已经不适用于大规模数据集与低延迟的场景。因此，大多数研究者提出了近似最近邻搜索(ANNS)以平衡查询精度和效率。现有的ANNS方法可以分为两个过程：离线索引构建和在线查询。在大部分算法的实现中，为了高效的索引构建和极低的查询延迟(通常在5 ms以内)，索引构建和查询过程都需要在内存中进行。
随着向量数据的增多，数据规模可以达到数亿甚至十亿，例如BigANN基准测试中使用的数据集(https://big-ann-benchmarks.com/)，向量数据的规模接近十亿级别。对于128维的uint8数据，仅将所有向量数据加载到内存中就需要119 GB的物理内存，这使得纯内存方案变得十分昂贵，尤其是当下流行的CPU也仅仅支持128 G。
为了解决这个问题，一种思路是将数据分组划分为许多切片。每个切片可以部署在一个单独的物理机上，从而构建一个分布式系统。当查询请求到来时，相同的查询应用于每台机器上，最后汇总获得全局的结果。这种方法在解决大规模数据问题很受欢迎。然而，它有一个致命缺点，每一个查询必须应用于集群中所有机器，这意味着增加机器数量不会提高查询速度。并且，这种方案也很难做到负载均衡，因为每一个查询必须应用到所有机器上。而另一个处理大规模向量数据检索的方案，则是使用高性能的固态硬盘(SSD)代替内存(DRAM)降低成本。SSD的价格比DRAM低许多，而且SSD可以存储TB级别的数据，这比DRAM的容量大数千倍。更重要的是，随着PCIe技术的进步，SSD在过去的十年中已经变得越来越强大。最广泛使用的PCIe 3.0 SSD甚至可以实现3000 MB/s的带宽，这意味着使用SSD代替DRAM已经成为了潜在的可能。已经有DiskANN [2] 、SPANN [3] 、Zoom [4] 、GRIP [5] 等多个过往的研究表明，使用SSD代替DRAM是完全可行的。
在本文中，我们介绍了一种新型基于固态硬盘(SSD)的近似最近邻搜索(ANNS)算法。该算法利用异步I/O框架io_uring来优化搜索延迟。通常，ANNS算法包括两个阶段：索引构建和在线搜索。对于索引构建阶段，我们的算法AIO-ANNS主要采用DiskANN的索引构建方法。在DiskANN的论文中指出，其图索引构建方有助于查询算法减少SSD的访问次数。在搜索阶段，我们提出了一个基于异步I/O的优化方法来降低搜索延迟。因为算法依赖的索引数据全部存放在SSD上，在搜索过程时，当需要访问一个索引结点时，必须要访问SSD。访问一个向量索引节点对应着一个单独的I/O请求。虽然SSD具有内在并行性，但显然等待一批I/O操作同时完成的时间比等待一个单独的IO操作要更长。以往提出的基于SSD的算法，通常在搜索迭代阶段提交一批I/O请求，然后等待所有请求完成。相反地，我们使用io_uring将I/O请求拆分为两个阶段：提交阶段和完成阶段。通过异步IO方式，算法消除了等待所有I/O请求完成的不必要时间，从而极大缓解了高延迟I/O请求带来的不利影响。
2. 相关工作
随着数据规模的持续增长，向量数据集扩展到数十亿甚至上千亿的规模，向量维度达到数百甚至上千，传统的基于内存的算法常常不能满足索引构建和查询处理的要求。为了应对这一挑战，出现了两种主流解决方案：对于一个向量数据集合
，其包含n个维度为m的向量数据。给定一个查询向量
，搜索的目标是需要从X中找到与查询向量最近的向量p*，即
。这样的向量p*被称作最近邻。类似的，可以定义k-最近邻问题。受制于暴力搜索的计算开销大和高查询延迟的缺陷，近似最近邻搜索被提出用于从底层数据库从检索出k个近似最近邻的向量，以满足低延迟和高吞吐的要求。对于给定查询q，ANNS算法会返回k个向量构成的集合S，也被称为候选集。我们可以定义：
ANNS算法的目标是在尽可能低的延迟下取得最大的召回率。并且，算法的索引构建时间和索引大小也是必须要考虑的因素，因此算法不仅需要在online的查询上取得很好的效果，也需要满足offline快速构建的需求。
过去的几十年里，研究者提出过诸多数据结构和算法用于解决近似最近邻搜索问题。这些算法可以划分为四类：基于树的算法；基于哈希的算法；基于量化的算法；基于临近图的算法。近些年来，基于量化的算法和图的算法成为了ANNS算法的主流。特别是各种临近图的算法，在大部分数据集上都取了了最优的性能，但却没有带来更多的索引构建时间。
k-d树是最早提出用于向量检索的算法，其思想类似于二叉搜索树。通过比较数据，把大于当前节点的数据放在树的右侧，小于节点的数据放于左侧，但是与二叉树不同的是，k-d树大小的比较是在向量的维度上进行的。ball-tree修改了kd-tree的划分方式，将划分的子空间从超长方体修改为球体。vp-tree的划分不发生在向量的具体维度上，其使用距离度量对于整个向量空间进行考虑。Muja等人 [6] 整合了一系列基于树的查询算法，推出了flann库。
基于哈希的算法目标是将高维向量编码为二进制向量，并且尽可能保留原本向量空间下向量之间的距离关系，主要可以分为局部敏感哈希 [7] 和哈希学习 [8] 。局部敏感哈希采用与数据无关的哈希函数，整个学习处理过程不依赖任何的数据内容。LSH通过一个局部敏感的哈希函数将相似的数据点以更高的概率映射到相同的哈希编码上。查询时，只需要在遍历查询向量落入的哈希桶中所有的向量，就可以找到最近邻。哈希学习通过机器学习机制将数据映射为二进制串的形式，能够显著减少数据的存储和通信开销，从而提高学习系统的效率。
向量量化(Vector Quantization, VQ)是信息论中进行数据压缩的传统算法 [9] ，其目的是降低数据维度但尽可能保持向量数据原始分布。乘积量化(Product Quantization, PQ)在向量量化的基础上发展而来，jegou等人 [10] 首次将PQ用于解决ANN问题。PQ算法的主要思想是对原始向量进行维度切分，对切分后得到的子向量进行向量量化后，再使用笛卡尔积进行组合。通常使用PQ算法有很多优点：1) PQ对原始向量进行压缩，使得向量数据可以全部存放在内存中；2) 查询向量和压缩向量之间的距离可以通过查表方式快速计算得到。3) PQ算法的实现和数据节点都很简单，这使得其可以与其他算法组成混合算法。因为PQ对子空间的切分没有考虑到向量数据的分布，划分完全与数据无关，但是现实中大部分数据集不同的维度分布是不均匀的，Ge等人提出了Optimized Product Quantization也就是OPQ来保留向量的原始分布，其做法是简单描述是通过一个正定矩阵R对于原始的向量空间进行了旋转 [11] 。Ge等人同时还提出了Locally Optimized Product Quantization (LOPQ)来进一步优化OPQ算法 [12] 。量化通常和倒排索引结合使用，例如IVFADC [13] 、FAISS [14] 、IVFOADC + G + P [15] 以及SCANN [16] 。通常这些算法会使用kmeans或者其他聚类算法，将整个向量数据集切分为多个分区，然后对于每个分区内的向量进行PQ操作。在查询时，仅仅需要对少量最近的分区进行访问，变量分区内的所有向量，通常每次查询访问的分区数远小于总的划分分区数。通常只需要少于64 GB的内存就可以存放十亿的压缩向量，并且fassi可以做到小于5 ms查询延迟。但是根据ANN benchmark的测试结果，这些算法在1-recall@1取得相当糟糕的性能，并且在高Recall (通常指99%以上)其查询速度急剧衰减。
最近十年内，取得了最优性能的ANNS算法通常都是基于临近图这一数据结构 [17] 。基于图的ANNS算法会使用原本的数据集(a)构建一个图索引(b)，图上的节点对应原始数据集中的向量。其查询算法通常使用贪心算法(Best-First-Search)，每一步选择最优节点从入口点逐渐逼近查询点。Malkov最先提出了基于插入增长式的构建索引算法(Navigable Small World Graphs, NSW) [18] ，基于NSW算法，Malkov进一步提出了更高效的算法——Hierarchical Navigable Small World Graphs (HNSW) [19] 。HNSW在选择节点邻居时，采用relative neighbour graph中的启发式算法，并且构建分层的图索引，使得算法可以高效地进行贪心查询。直到今天，HNSW仍然是大部分ANNS场景下的第一选择，算法健壮性得到广泛认可。由浙大和阿里团队提出的NSG (Navigating Spreading-Out Graph)算法同样取得了优异的性能表现 [20] ，不同于HNSW增长插入的构建方式，其使用KGraph构建初始图，并在初始图上进行优化。Wang等人 [1] 对基于图的ANNS算法进行了深入且全面的综述，涵盖了大部分主流图算法。
这些算法在该领域取得了显著的进展，为处理越来越大和复杂的数据集提供了新的策略和解决方案。但是随着数据规模的不断增长，向量数据集的尺度可以扩展到亿级甚至十亿级别，并且向量的维度也达到上百维甚至上千维。传统的纯内存算法已经不能满足构建和查询需求，甚至算法无法工作。为了解决这一问题，产生了两种主流的处理方案。1) NSG论文中提出的分布式解决方案。类似于分布式数据库的处理办法，将数据进行分片，在一台服务器上只对一个分片进行索引构建和查询执行。2) 以DiskANN算法为代表的一系列使用内存和磁盘混合的算法。
通常，向量数据分片的主流处理办法是随机划分，包括当下最成熟的分布式向量数据库milvus也是采样类似的处理。对于单个query而言，需要在集群中所有的机器都执行一遍查询。这无疑与分布式的思想相违背——对集群的扩展并没有带来速度的提升。因此使用磁盘和内存的混合算法得到了更多企业和研究人员的关注。
这里的提到的磁盘通常是数据库中广义的非易失性存储介质，不仅仅代表HDD，还包括SSD或者傲腾非易失性内存。微软也许在业务场景中对于混合算法有着迫切的需求，提出了一系列基于磁盘的ANNS算法。在2018年，微软提出了Zoom算法 [4] ，其思想很简单，将数据划分为需要很小的分区，每个分区的中心点组成的集合，使用HNSW进行索引构建，并全部放入内存中，而每个分区中的向量经过PQ量化后存放到SSD上。这样查询就变成了一个两阶段过程。次年，微软再一次发布了改进后的算法GRIP [5] ，和Zoom算法类似，但是GRIP将PQ量化后的压缩向量也存放到内存中，仅仅在SSD上存放全精度向量。在查询时，先通过内存导航图找到对应的partition，然后使用PQ向量获得PQ距离，再使用PQ距离对候选向量进行排序，从SSD上拉取PQ距离最近的全精度向量后，再使用全精度向量与查询向量的距离进行重排。这类算法往往在recall@1取得不错的效果，但是在大topK的情况下，算法几乎丧失优势。
2020年，微软提出了DiskANN算法，与Zoom和GRIP算法完全不同，DiskANN使用数据集中全部向量构建Vanama索引(Vanama和NSG只有轻微区别，它的构建步骤与NSG类似，不同于NSG以K最近图为初始图，在其论文中直接使用了随机图作为初始图，并使用一些启发式算法构建最终图索引)，并且构建好的图索引邻接表和全精度向量共同存放到SSD上。并且DiskANN还会同时训练PQ，来对所有向量进行量化，并将压缩后的向量全部存放在内存中。搜索时，其先用PQ压缩向量计算，索引节点向量与query向量的PQ距离。根据PQ距离大小从小到大排序，取前beam个离query向量最近的节点。从SSD上读取其对应的entry，这样我们就获得了该节点对应的全精度向量数据，以及其邻接表。我们使用全精度向量重新计算距离，作为最终查询结果的候选。而邻接表中记录的邻居，则计算这些邻居的PQ向量与query向量的距离，并重新排序，从SSD继续拉取新的entry。直到满足终止条件。
Chen等人在2021年提出了一种基于kmeans-tree的新型算法SPANN，在内存中存放导航图，而在SSD上存放经过kmeans划分后的partition，每个partition中含有多个全精度向量。Zilliz的Search团队同样研发了一种混合算法BBAnn [21] ，其与SPANN类似也是使用kmeans进行聚类，然后在内存中构建导航图，全精度向量放在SSD上。
3. 研究动机
基于图的ANNS算法可以分为构建索引阶段和搜索阶段，并且这些算法大部分在搜索阶段都使用了一种greedy search的搜索策略，即贪心搜索。greedy search会维护一个固定长度的侯选池来存储需要遍历的向量，并且根据查询向量和这些向量之前的距离进行排序。算法迭代式地从侯选池中检索获得最近的未访问结点，并且将其标记为已访问。接着，算法会检查此结点在图索引上的所有邻居节点，计算邻居结点所代表的向量与查询向量之间的距离。并使用(邻居向量，距离)对来更新侯选池。如果侯选池所有的结点都得到了访问，算法就终止。这种搜索模式会产生大量访问SSD的往返操作(round-trip)，其中一个往返操作涉及OS提交一个I/O请求到SSD，然后SSD将数据返回到OS。因此，我们使用beam search来尽可能地访问SSD的round-trip次数。如图1所示，beam search与greedy search的不同点在于，beam search每一次从侯选池中获取多个需要访问的结点，将原本多次小的IO请求合并一个大的IO请求，从而减少了round-trip的数量。虽然beam search最终需要读取的数据比greedy search要多一些，但是由于SSD本身具有内在并行性，在SSD上执行时，beam search的速度会显著超过greedy search。
Figure 1. Beam search compares with greedy search
图1. Beam search与greedy search的对比
在许多场景中，beam search已经可以取得非常优秀的性能，但是在图1遍历过程中，揭示了beam search一个显著的缺点。即beam search每一次需要提交一批读请求，并且等待这一批请求全部完成。这使得一次round-trip的延迟会由最慢的那个IO请求所决定。如果延迟出现不一致情况，则会显著伤害系统的性能。遗憾的是，这种情况在固态硬盘(SSD)、硬盘驱动器(HDD)和网络环境中经常出现。大多数IO操作都执行得很快，但少数几个可能需要异常长的时间，导致所谓的拖尾现象。
我们设计了一个针对SSD实验，在实验中我们制造了大量的随机读请求，并且记录了每个读请求完成的时间。结果完全符合我们的猜测：虽然98.5%的读请求的完成延迟低于100微秒，但最慢的请求却需要长达10,000微秒，结果如图2所示。因此，为了提高基于SSD的近似最近邻(ANNs)算法的性能，减少这些高延迟请求的影响变得至关重要。受到流水线概念的启发，我们提出了一种策略来改善这种效应。算法不再等待一个整批的IO请求完成，而是在每个请求完成时处理它——计算距离、更新候选池，然后继续处理下一个请求，直到批中的所有请求都完成。
Figure 2. CDF of I/O completion time
图2. I/O完成时间的分布曲线图
此外，大多数ANNs算法通常涉及两个池：一个用于导航的候选池，按照Product Quantization (PQ)压缩距离排序，以及一个按照全精度距离排序并作为最终结果呈现给用户的结果池。对于导航至关重要的候选池必须在下一次beam search迭代之前更新。另一方面，结果池不对搜索路径做出贡献，因此可以在搜索的任何时候进行更新。通过解耦这两个组件，允许在IO请求等待时间内更新结果池，可以提高算法的整体效率。
4. 算法细节
在之前的部分已经提到，AIO-ANN算法很大程度收到了DiskANN算法的启发，因此在索引结构上AIO-ANN与DiskANN算法保持一致。我们使用两轮构建的Vanama图作为我们的图索引，并且整个图索引数据全部放置在SSD上，而不是DRAM中。在beam search的过程中，我们在图上进行游走，迭代式地从SSD中读取图索引数据以及向量原始数据。这个过程也是整个算法主要耗时部分。此外，我们还采用了BBAnn中提到的分层kmeans技术来优化我们的索引在SSD上的数据布局。BBAnn是一种主要设计用于范围检索的算法，它使用分层均衡kmeans来将数据切分为多个蔟。我们将这一技术与DiskANN相结合，有助于优化数据布局。在内存中，我们使用PQ量化后的压缩向量来代替原始向量，这可以为我们的算法节省4到16倍的内存成本。算法具体的架构如图3所示。
在研究部分的延迟实验表明，SSD的读请求延迟并不稳定，并且受到拖尾效应的影响，这对算法的性能是有害的。此外，实现还暗示了一个有前景的策略——通过解耦候选池和结果池的更新，可以在IO请求完成期间隐藏更新结果池的时间。在图4的算法1中详细描了这一过程。
AIO-ANN将每一个IO请求过程分解为两个部分，分别为提交过程与完成过程。在提交过程，算法会一次性提交多个读请求，并且不会阻塞住，而是立刻返回提交的结果。在完成过程中则恰恰相反，算法会阻塞住，直到有IO请求完成并返回。然而，在linux的传统的文件操作中是不存在这样的系统调用的，例如read或者write都必须阻塞直到完成。因此我们使用io_uring这一工具来实现我们的目的。io_uring是linux kernel提供的最新文件操作API，它可以支持各种异步操作。Io_uring使用了一个共享的环形缓冲队作为内核与用户空间通信的方式，极大地降低了系统调用的开销。
在算法的实现代码中使用了IORING_SETUP_SQPOLL标志用于启用iouring的内核线程模式。在这种模式下，提交阶段与完成阶段都不涉及任何系统调用。在提交一批IO请求与等待任意一个IO请求完成之间，我们执行结果池的更新。这一过程包括了计算向量全精度距离，并根据这个结果重排最终的topk结果。通过这一优化策略，可以显著消除由长尾效应带来的负面影响。
在数据库领域还有一种常用的减少IO的策略，那就是内存缓存。在ANNs领域中同样可以实现类似的缓存机制，与数据库的缓存策略不同的是，数据库通常缓存的最小单位是内存页，而ANNs的缓存通常是缓存一个索引结点，包括这个索引结点的邻接表与原始向量。在实际搜索过程中，如果缓存命中，那么就可以减少IO请求。考虑到每个查询的IO请求数量通常不会超过100，如果缓存次数命中足够多，会带来巨大的性能提升。
我们设计的算法具体架构如图3所示，在最上层我们使用了DiskANN提出的Vanama索引结构图，在Vanama图上从起始点出发上使用beam search，迭代方式寻找最终结果。在每一步迭代过程，算法使用压缩PQ向量来计算近似距离，用于指导生成访问SSD的请求。在SSD上层，算法添加了Cache层，用于降低I/O。如果cache不命中，才会真正去访问SSD。
最后算法还采用了result-based技术来优化ANNS中初始选点的步骤。在实际的应用场景中，向量查询通常表现出局部性，这意味着一批查询向量在短时间内可能具有相似的特征。因此，我们使用了前一个查询的结果来初始化当前的侯选池，可以加速算法路由过程。
5. 实验结果
在本节中，我们将AIO-ANN算法与目前主流算法DiskANN进行对比基准测试。所有实验都在相同规格的主机上完成，配置如下：
操作系统：Ubuntu 22.04
处理器：AMD Ryzen 7 5700G CPU
内存：32GB 3200MHzDDR4 DRAM
存储：三星980 Pro 1TB
实验使用到了三个不同的数据集：
BigANN 10M：该数据集由大型图像数据集中提取的SIFT描述符组成。
SpaceV 10M：这个与网络搜索相关的数据集是微软Bing为竞赛发布的。它由Microsoft SpaceV Superior模型编码的文档和查询向量组成。
deep 10M：该数据集由Imagenet分类任务上预训练的GoogLeNet模型的最后一个全连接层的投影和归一化输出组成。
使用“Recall-with-QPS”度量来评估所提算法的性能，该度量包括两个组成部分：
Recall：在像ANNS这样的信息检索系统的上下文中，召回率量化了系统检索所有相关文档的能力。从数学上定义，它是检索到的相关结果数与现有总数的比率。越高的召回率表示系统检索能力越强。
QPS：表示每秒查询数，QPS是信息检索系统的常见性能度量。它测量系统每秒可以处理的查询数。高QPS表示一个能够在短时间内处理大量查询请求的高效系统。
我们在三个数据集上进行了实验，搜索线程数设置为4。结果如图5所示。在recall@1和recall@10的所有三个数据集上，我们提出的AIO-ANN算法的性能对比DiskANN提高了20%至30%。这一优势在从较低的90%到较高的95%的召回范围内始终可观察到。此外，我们的算法显示出更强大的稳定性，更少的性能波动。
Figure 5. AIO-ANN vs. DiskANN on (a) BigANN 10M, (b) deep 10M, and (c) SpaceV 10M.
图5. AIO-ANN与DiskANN对比：(a) BigANN 10M；(b) deep 10M；(c) SpaceV 10M
6. 结论
在本文，我们提出了一种新型的磁盘向量检索算法AIO-ANN。该算法主要使用异步I/O技术，尽可能降低了少量高延迟读请求对整体搜索性能的不利影响。在每一步迭代过程，算法并不等待所有I/O请求同时完成，而是在处理已有数据的同时等待其他I/O请求的完成。同时我们还讨论了两个可能优化，查询缓存与结果初始化，并应用在了算法的实现中。实验的结果表明，AIO-ANN在性能上超越了主流的DiskANN算法，并且具有更强的稳定性。AIO-ANN有很大潜力能提升已有的ANNS系统的性能，特别是在高I/O请求延迟的情况下。未来的工作将寻求探索额外的优化，并测试该算法在更大规模数据集上的性能。