[论文解读] BANG: Billion-Scale Approximate Nearest Neighbor Search using a Single GPU

Karthik Viswanathan, Saim Khan+3 more|arXiv (Cornell University)|Jan 20, 2024

Data Management and Algorithms被引用 4

一句话总结

BANG 提出了一种新型的基于 GPU 的近似最近邻搜索方法，通过压缩数据实现快速距离计算，同时将图结构和原始向量保留在 CPU 上，从而在单个 GPU 上实现了十亿规模的 ANNS。在高召回率（0.9）下，其吞吐量比当前最先进的方法高出 40×–200×，得益于 GPU-CPU 流水线处理、优化的内核以及减少的 PCIe 数据传输。

ABSTRACT

Approximate Nearest Neighbour Search (ANNS) is a subroutine in algorithms routinely employed in information retrieval, pattern recognition, data mining, image processing, and beyond. Recent works have established that graph-based ANNS algorithms are practically more efficient than the other methods proposed in the literature. The growing volume and dimensionality of data necessitates designing scalable techniques for ANNS. To this end, the prior art has explored parallelising graph-based ANNS on GPU, leveraging its massive parallelism. The current state-of-the-art GPU-based ANNS algorithms either (i) require both the dataset and the generated graph index to reside entirely in the GPU memory, or (ii) they partition the dataset into small independent shards, each of which can fit in GPU memory, and perform the search on these shards on the GPU. While the first approach fails to handle large datasets due to the limited memory available on the GPU, the latter delivers poor performance on large datasets due to high data traffic over the low-bandwidth PCIe interconnect. We introduce BANG, a first-of-its-kind technique for graph-based ANNS on GPU for billion-scale datasets, that cannot entirely fit in the GPU memory. BANG stands out by harnessing a compressed form of the dataset on a single GPU to perform distance computations while efficiently accessing the graph index kept on the host memory, enabling efficient ANNS on large graphs within the limited GPU memory. BANG incorporates highly optimised GPU kernels and proceeds in phases that run concurrently on the GPU and CPU, taking advantage of their architectural specificities. Using a single NVIDIA Ampere A100 GPU, BANG achieves throughputs 50x-400x higher than competing methods for a recall of 0.9 on three popular billion-scale datasets.

研究动机与目标

• 解决基于 GPU 的 ANNS 在超过 GPU 内存容量的十亿规模数据集上面临的可扩展性瓶颈。
• 克服现有基于 GPU 的 ANNS 方法的局限性，这些方法要求将完整的图结构和数据存储在 GPU 内存中，或依赖多 GPU 配置。
• 通过将计算卸载到 GPU 上的压缩数据，同时将图结构保留在 CPU 上，实现在单个 GPU 上的高吞吐量、高召回率 ANNS。
• 通过重叠通信与计算以及智能预取，最小化 CPU 与 GPU 之间的 PCIe 数据传输开销。
• 仅使用单个 GPU 即可实现与当前最先进的 CPU 和多 GPU ANNS 系统相当或更优的性能。

提出的方法

• BANG 使用乘积量化（PQ）对高维向量进行压缩，使 GPU 上能够高效执行距离计算，而无需将完整精度的向量存储在 GPU 内存中。
• 采用 DiskANN 中的 Vamana 图结构作为底层索引，图结构存储并由 CPU 管理，而邻居遍历和距离计算则在 GPU 上使用压缩向量执行。
• GPU 内核针对核心操作（包括距离计算、排序、工作列表更新和最近邻选择）进行了高度优化，以最大化 GPU 占用率和并行性。
• CPU 和 GPU 以流水线、并发的方式执行搜索阶段：CPU 同时准备邻居信息并预取数据，而 GPU 执行距离计算。
• 使用布隆过滤器通过早期剪枝非有希望的节点来加速图遍历，从而减少不必要的 GPU 计算和内存访问。
• 通过异步数据传输和双缓冲技术，将 CPU 与 GPU 之间的通信与计算重叠，最大限度减少空闲时间和 PCIe 瓶颈。

实验结果

研究问题

• RQ1 单 GPU 的 ANNS 系统是否能在超过 GPU 内存容量的十亿规模数据集上实现高吞吐量和高召回率？
• RQ2 如何优化 GPU-CPU 工作负载分配，以最小化图基 ANNS 中的数据传输并最大化并行性？
• RQ3 通过 PQ 压缩在 GPU 上实现快速距离计算，在不牺牲搜索准确率的前提下，其潜力有多大？
• RQ4 流水线处理和通信与计算的重叠是否能显著降低将图数据卸载到 CPU 所带来的性能损失？
• RQ5 在标准基准测试中，所提出的方法与当前最先进的 GPU 和多 GPU ANNS 系统相比，在吞吐量和召回率方面表现如何？

主要发现

• 在十亿规模数据集上，BANG 在召回率为 0.9 的条件下，吞吐量比竞争的基于 GPU 的方法高出 40× 至 200×，展现出显著的性能优势。
• 在百万规模数据集上，BANG 在大多数基准测试中表现优于最先进水平，其算术平均吞吐量是最佳竞争方法的 2 倍。
• BANG 在所有评估的数据集中均保持了高召回率（≥0.9），包括 SIFT1B 和 GIST1M 等大规模基准，而以往的 GPU 方法在召回率与吞吐量之间难以平衡。
• 在 GPU 上使用压缩数据可减轻内存带宽压力，并在完整数据集和图结构均存储于 CPU 的情况下仍能实现高效计算。
• 流水线处理和通信与计算的重叠显著降低了 CPU-GPU 数据传输开销，使 BANG 对 PCIe 带宽限制具有更强的鲁棒性。
• BANG 是首个在单个 GPU 上高效处理十亿规模数据集的基于 GPU 的 ANNS 系统，无需多 GPU 配置或将完整索引存储在 GPU 中。