基于 GPU 的 ANN 检索
导读
近似最近邻（ANN）向量检索的 CPU 方案已被广泛地应用于在线检索等多种场景中并取得了不错的效果。GPU 相比 CPU 拥有更强大的并行计算能力，如何将 GPU 引入 ANN 检索获取更大收益，成为了业界重点研究的难题之一。百度与 NVIDIA 技术团队，基于 RAFT[1]开源代码库设计并实现了一种基于 GPU 的 ANN 在线检索方案，在一类高检索流量业务场景下获得了显著的成本收益。本文主要介绍整体方案并总结一些思考及展望。
01 ANN 简介
随着自然语言处理和深度学习的高速发展，向量检索为提升搜索用户体验发挥出重要作用，近似最近邻（ANN）搜索是在有限时空条件下实现向量搜索的主要方式。ANN 的核心思想是近似，所有 ANN 算法都致力于在低规模计算开销下保证最后选择的邻居是查询向量的最近邻概率很高，但不追求 100%的准确性。根据近似思路的不同，ANN 检索算法主要分为四类：基于树的 ANN 算法、基于 LSH 的 ANN 算法、基于量化的 ANN 算法和基于图的 ANN 算法。接下来对本文涉及的几种 ANN 算法原理进行简要介绍。
△IVF_FLAT
IVF_FLAT：如上图所示，基于量化[2]的方案是通过一组特定的点来表示全部的向量空间，每个特定点代表一个子空间，它常与倒排索引相结合[3]，首先对向量空间进行划分，每个数据点按照与其最近的子空间质心形成倒排拉链，这样每条倒排拉链对应一个子空间，拉链中元素为落在该子空间的向量 id。检索时首先通过与质心的距离计算选取若干个子空间，然后只计算所选子空间内数据点与查询向量的距离，排序后得到最后结果，该思路被称作 IVF_FLAT 方案。
△IVF_PQ
IVF_PQ：除了借助向量量化选取子空间减少计算以外，量化对于 ANN 检索的贡献还在于可以使用量化点计算量化距离从而代替精确距离的计算，这样减少了计算但会带来精度损失，这就要求量化的精度不能太低。乘积量化[3]（PQ）可以借助几组子向量器的质心，产生大量的质心，降低量化学习的复杂度。它与残差量化[4]相结合可以在有限的时空条件下提升整体量化的精度。基于乘积量化的 IVF 方案经典思路如上图所示，它是在构建时先对所有数据进行一次向量量化，根据这次向量量化的簇心形成倒排拉链，然后计算出所有数据与其质心的残差，然后再对残差进行乘积量化，记录对应其量化点的编码。检索时同样先与向量量化的质心计算决定候选拉链减少计算量，然后计算查询与质心的残差后对候选点进行距离计算，计算距离时是计算查询残差与候选点的量化距离，最后进行距离排序后返回检索结果，该思路被称作 IVF_PQ 方案。
△Cagra
Cagra[6]：基于图的 ANN 方案进行向量检索的过程都是通过图索引中点的邻接关系逼近到和查询向量点距离最近的若干个结果，图索引的构造和检索策略是基于图的 ANN 搜索方案研究的重点，需要考虑准召、搜索效率和空间开销等多方面因素。经过业界的不懈研究，以 HNSW[5]为代表的诸多基于图的方案凭借优异性能都得到了广泛应用，但大部分的图算法在构建流程和检索路由过程中很难充分利用到 GPU 的并行能力，NVIDIA 针对图算法这一挑战自研了名为 Cagra[6]的基于 GPU 的 ANN 图算法，如上图所示。该算法以 NN-descent[7]图为基础，然后借鉴 NGT[8]的路径调整的思路对边进行剪枝，最后对剪枝后的正向图与对应的反向图各取一半的边进行合并，生成最终的 Cagra 图索引。检索过程维护一个 top-M（M 大于等于返回结果数目 K）的列表和一个长度固定的候选列表，每轮迭代局部排序出距离最小的 top-M 并更新，然后将 top-M 的未被检索的邻居添加到候选列表中，直至 top-M 列表的节点都被遍历，然后选取 top-K 作为结果返回。Milvus[9]数据库就是借助 Cagra 算法实现了出色的性能加速。
02 我们选择了什么业务场景来使用 ANN on GPU
在百度的某一类搜索业务场景下，对单次用户查询，会进行扩展检索以召回更多样化的相关结果。具体的业务场景需求：
约束：ms 级低延迟、>90%的高召回率、亿级索引量
场景：超高检索吞吐
目标：更低的检索成本
在这样的超高检索吞吐场景下，基于 GPU 来做高并行度计算会比基于 CPU 的检索有更高的性价比。因此，我们选择在该场景下来做 ANN on GPU 的尝试。
03 我们是怎么做的
3.1 索引算法选择
RAFT[1]代码库是 NVIDIA 设计开发的一套开源代码库，包含了可广泛应用于机器学习和信息检索场景的基本算法和原语，这些算法经过 CUDA 加速优化并可以通过构建块使编程简易化。它提供了四种 GPU-ANN 算法：BruteForce、IVF_FLAT、IVF_PQ 和 Cagra，这些算法都通过相同名称的顶层接口来实现索引的构建（build）、存储（serialize）、加载（deserialize）、检索（search），为用户的使用提供了便利，以 IVF_FLAT 为例：
我们对 RAFT 支持的四种算法进行了对比测试以找到适合我们场景最佳的算法选型，主要结论如下：
BruteForce：
暴力全量计算，召回率 100%
单卡索引规模：目前仅支持 float 类型，A10 单卡最多可支持 2kw * 256 维数据
检索延时：约 70ms（batch_size<10）-> 无法满足我们场景下的延时要求 Cagra：
单卡索引规模：在百万量级数据集上相比 IVF 系列有优势，但在千万规模数据集上效果下降较严重
检索成本：从成本考虑，单卡数据规模越大越好 -> 无法满足我们场景下的成本约束
IVF_PQ：
单卡索引规模：支持 float 和 int8，会对任何输入类型都转成 FP32，从而利用 TF32 的 Tensor Core 加速能力，所以不同数据类型下的检索性能都差不多；核心优势是显存占用小（索引不需保存原始向量，只需要保存编码）适合更大索引规模。
召回率：在我们的数据集上，PQ 精度损失较大，量化距离排序不够准确，召回率损失较严重 -> 无法满足我们场景下的召回率要求
IVF_FLAT：
单卡索引规模：支持 float 和 int8 两种类型，float 类型索引在大规模数据量下显存开销高，索引建库时间长；int8 类型四分之一的压缩比可有效缓解显存受限问题
召回率：在合理选择 int8 量化方式后，IVF_INT8 精度损失很低，召回率不受限，加大检索倒排拉链数目后召回率提升很快
经过测试后，IVF_INT8 算法的查询延时和整体吞吐均表现优异
因此我们的场景下，最终选择采用 IVF_INT8 索引算法。
3.2 离线建库
单分片索引容量提升
目标：成本考虑，需要在显存受限的情况下，尽可能提升单分片索引量
思路：量化+超显存
量化：int8 量化
超显存：数据量达到一定规模后，有限的显存资源可能放得下索引，但却放不下输入数据无法构建出索引；为解决这种问题，我们在离线构建索引时使用了 GPU 的超显存技术（借助内存与显存的交换实现显存需求的超用），该技术可以在显存超用不严重的情况下依旧保持很高的计算性能，下方代码给出了使用示例：
索引调参
拉链长度控制：
IVF_INT8 在构建阶段需要训练簇心，训练数据从构建索引数据中按一定比例获取，由 kmeans_trainset_fraction 参数控制。而构建多少条倒排链由 n_lists 参数决定，该建库参数影响着达到目标召回率时的检索计算量，因此需要调参以获得更佳的检索性能。
我们实践的经验规律是控制平均链长在一千到两千之间时，有较好的综合性能表现。
均衡性控制：
由于各条链与查询计算时有天然的并行度，为了避免计算长尾构建索引时应尽量使各链链长均匀，这一点可以借助 raft 代码库中的 kmeans_balanced 组件实现：
3.3 在线检索
下图展示了我们所设计的基于 GPU 的 ANN 在线检索方案，通过 batch 检索来达到充分发挥 GPU 并行计算能力的目的。在上游查询到达时，查询将会放入查询队列中，检索线程会按照 batchsize 取出对应数量的 query 组成 BatchQuery，并将其作为输入进行基于 GPU 的 ANN 检索。在检索过程中需要先将 BatchQuery 拷贝到显存中，然后在 GPU 上进行 IVF_INT8 算法检索，返回对应 query 顺序的 BatchResult 结果并将其从显存拷贝取出，检索服务模块对 BatchResult 按 query 拆分后进行算分、过滤等后置操作后返回结果到上游。
△在线检索方案
向量距离计算并行化 & 显存数据复用
对选中的多条链与查询向量进行距离计算，是 IVF_INT8 在检索时计算量最大的部分，RAFT 代码库中使用 interleaved_scan_kernel 函数对这一阶段做高效并行计算，它是保障 IVF_INT8 检索性能的关键。
由于各链和查询之间的距离计算并无依赖，该接口将每条链的数据用一个 CUDA 块进行扫描，并将容量不大的 BatchQuery 向量加载到共享显存中，这样可以使得同时读取索引数据和查询时节省全局显存带宽。
每个 CUDA 块都会将一条链的数据按照合理粒度分别承载给多个 warp 并行计算获取距离结果，最后调用 matrix::detail::select_k 函数对多个不同块的结果 merge 后挑选最终 top-K 返回。
基于实时流量统计的动态 batchsize 控制
在合理的范围内增大 batchsize 可以增加实例的整体吞吐，但 batchsize 较大而流量较小时会出现为了组 batch 某些 query 等待延时过长的问题，从而导致查询超时；对应的，batchsize 较小而流量较大时会有大量 query 入队却无法及时处理的情况，因此根据流量选择适合的 batchsize 至关重要。
我们采用训练数据拟合的方式得到流量相关的 batchsize 预估函数。满足一定时间阈值时，会使用预估函数预估当前流量下适合的 batchsize，如果与之前设定的 batchsize 差距大于阈值，就会对 batchsize 进行调整。batchsize 没有实时更新而是选择定时预估的方式是因为在 GPU 上进行检索时，会提前根据 batchsize 预先分配 BatchQuery 和 BatchResult 的显存，这部分耗时较长，更新不能过于频繁，因此时间阈值不能过短，否则会影响检索效率；但时间阈值也不能过长，否则无法根据最近时间段流量情况找到最适合的 batchsize。
GPU 及 CPU 分数打平
我们对构建索引向量和查询向量的 float 数据进行 INT8 量化的思路是最通用的 min-max 量化，即通过训练统计最小值 min_val 和最大值 max_val 后进行量化。该方法会先计算两者差值为 diff_val = max_val - min_val，然后借助量化公式 y = (x - min_val) / diff_val * 255 - 128 将任何浮点数 x 量化为在 int8 数据表示范围的 y。而最后返回的量化距离可以借助量化公式的线性系数 255 / diff_val 转化为原数据的真实距离，这样可以跟其他 CPU 方案 ANN 索引的检索距离直接进行比较，以便上游模块进行排序等操作。
整体上，通过借助 RAFT 代码中 IVF_INT8 算法的高效实现、动态 batchsize 控制和显存数据复用逻辑充分发挥了 GPU 的并行计算优势，使得检索服务单实例可以达到更高的查询的吞吐能力，降低了高流量场景下副本数，大大降低了所需成本。
3.4 达到的效果
根据上述方案实现并测试后，对 2500w256 维数据集构建索引，索引所占显存<10G，建库时间<10 分钟，检索服务单实例在返回结果 top30 的条件下可以承载约 2w qps，平均查询延时在 ms 级。在我们的高检索流量落地场景中，总成本相比于 CPU 方案分别可以节省 30%~50%资源。
04 思考和展望
4.1 GPU 在 ANN 检索场景的讨论
什么场景下适合用 GPU，什么场景下不适合用
GPU 与 CPU 相比拥有更强大的并行处理能力，它可以借助成百上千个更小的、更高效的处理核心做到同时处理大量的数据，有益于提升计算吞吐。而是否选择 GPU 需要注意以下几点：
计算规模是否适合 GPU 发挥：
虽然 CPU 的计算能力相比 GPU 有不少差距，但是对于计算需求不大的场景 CPU 就足以胜任，这种情况下使用 GPU 无法充分发挥算力只会放大其价格昂贵的劣势。
ANN 的核心思路是通过近似方式来减少算量，虽然近似的思路各有不同，但普遍趋势上算量越大，召回率越高，所以一般总体计算量大、召回率要求非常高的场景适合使用 GPU；而对召回率要求较低、总计算任务不重的场景更适合 CPU。
计算逻辑是否能够并行：
CPU 拥有着可以支持复杂逻辑、更为强大的计算单元，而 GPU 虽然计算单元众多支持高并发，但控制逻辑相对简单，所以对于数据依赖严重、计算过程串行阶段较多的场景，可能很难设计并行算法发挥出 GPU 高带宽的计算能力，这类场景不适合 GPU。
在 ANN 算法设计上，CPU 方案可以进行步骤繁杂的流程，而 GPU 方案需要简化搜索逻辑，减少数据依赖和同步阶段，加强并行度以发挥 GPU 的高并发能力。所以对延时要求相对宽容、查询流量较高的场景更适合 GPU 方案。
是否受限于 GPU 显存容量：
由于 GPU 的计算任务需要将数据存放在显存中计算，而主存和显存之间的拷贝延时很长，所以基于 GPU 的 ANN 算法设计一般要避免过多的主存和显存的数据拷贝和交换，这就要求检索中需要用到的全部或关键数据需要提前存放在显存中。然而 GPU 显存又是有限的，因此 GPU 不仅身为计算单元有着其发挥计算能力的局限要求，还面临着类似“存储单元”的瓶颈——显存受限。
在 GPU 方案的 ANN 算法设计中需要想办法缓解这一问题，使显存中存放的数据更能发挥作用。所以单实例支持数据规模过大、数据冗余严重或需要数据备份等场景更适合 CPU 使用内存的方案；而单实例数据规模适度，数据不需频繁变动和过多冗余备份的场景更适合 GPU。
我们的场景下，GPU 比 CPU 节省成本的本质原因
对于大规模数据 ANN 在线检索场景，由于单实例资源有限，全部数据索引很难存放于同一个实例中，因此要对数据进行划分成若干个库层，查询会分发给每个库层进行检索得到局部结果后进行合并从而得到最终结果；另一方面，每个实例能够处理的流量也是有限的，因此根据上游下发的流量需要建立若干个副本共同承担查询流量。一个库种的检索成本与其库层数和每个库层的副本数高度相关。
对于低流量下的 ANN 在线检索场景，每个库层只需要少数副本，因为 CPU 资源足以承担计算，应用 GPU 只会浪费算力并放大 GPU 价格昂贵的缺点，导致总资源成本增加。而对于高流量下 ANN 在线检索场景，高流量所带来的巨额总计算量为 GPU 取代 CPU 创造了可能。因为 CPU 算力有限，基于 CPU 的 ANN 方案处理检索的能力也是有限的，因此需要对每个库层增添大量副本，这无疑会对成本造成极大的压力。虽然 GPU 的价格要比 CPU 昂贵，但如果能够有效发挥出 GPU 的并发处理计算能力，那么使用 GPU 替代 CPU 进行 ANN 检索势必能够实现高吞吐处理请求的效果，这样单副本可以承载更多的流量，从而减少每个库层的副本数节省资源。
假定一个库种承载的总流量为 x ，单个库层的总成本为 y，每个副本所使用的 GPU 总成本 p2 要高于使用 CPU 的总成本 p1，GPU 方案和 CPU 方案下每个副本能承载的流量分别为 q2 和 q1，那么 GPU 方案单个库层的总成本为 y2 = ceil (x / q2) * p2，CPU 方案单个库层的总成本为 y1 = ceil (x / q1) * p1。下图展示了 GPU 和 CPU 方案在不同流量下的成本趋势，可以看出在低流量时，即使充分发挥 GPU 作用（即 q2 远大于 q1 时），CPU 成本也比 GPU 要低；而在高流量场景下，GPU 的成本要比 CPU 低得多，而且随着流量的增加所获得的成本收益也越高。所以高流量下的 ANN 在线检索场景成为了本文 GPU 方案的应用场景。
△不同吞吐下的检索成本对比示意
4.2 未来展望
本文以 RAFT-LIB IVF_INT8 算法为核心算法，借助灵活 batch 检索充分发挥 GPU 的并行处理能力，完成了 ANN on GPU 在百度的一类实际业务场景的应用落地，在吞吐、延迟、召回效果等多个方面保证了很好的效果，并在成本方面取得了不小的收益。后续我们将会和 NVIDIA 技术团队一起，对 RAFT 代码库的 IVF_INT8 算法做进一步优化，并对其它 GPU 相关 ANN 前沿算法进行深入探究，以期待 GPU 在 ANN 领域获得更多应用。
参考文献：
[1] https://github.com/rapidsai/raft
[2] Robert M. Gray and David L. Neuhoff. Quantization. IEEE Trans. Inf. Theory, 1998, 44(6): 2325 ∼ 2383.
[3] Hervé Jégou, Matthijs Douze and Cordelia Schmid. Product Quantization for Nearest Neighbor Search. IEEE Trans. Pattern Anal. Mach. Intell. 2011, 33(1): 117 ∼ 128.
[4] Yongjian Chen, Tao Guan and Cheng Wang. Approximate Nearest Neighbor Search by Residual Vector Quantization. Sensors, 2010, 10(12): 11259 ∼ 11273.
[5] Yury A. Malkov and Dmitry A. Yashunin. Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs. CoRR, 2016, abs/1603.09320.
[6] H Ootomo, A Naruse, C Nolet, et al. 2023. CAGRA: Highly parallel graph construction and approximate nearest neighbor search for GPUs. arXiv preprint arXiv:2308.15136 (2023).
[7] Mengzhao Wang, Xiaoliang Xu, Qiang Yue, and Yuxiang Wang. A Comprehensive Survey and Experimental Comparison of Graph-Based Approximate Nearest Neighbor Search. arXiv:2101.12631 [cs], May 2021. arXiv: 2101.12631.
[8] Masajiro Iwasaki and Daisuke Miyazaki. Optimization of Indexing Based on k-Nearest Neighbor Graph for Proximity Search in Highdimensional Data, October 2018. arXiv:1810.07355 [cs].
[9]https://github.com/milvus-io/milvus
———— END ————
推荐阅读
百度官方技术账号 2021-01-22 加入
关注我们，带你了解更多百度技术干货。
评论