[論文レビュー] 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メモリに収まらない10億スケールのデータセットに対し、GPU上で圧縮データを用いて高速な距離計算を実現するとともに、グラフと元のベクトルをCPU上に保持することで、1枚のGPUで100億スケールの近似最近傍探索（ANNS）を実現する新規なGPUベースの近似最近傍探索手法を提案する。高リコール（0.9）条件下で、最新のGPUベース手法と比較して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メモリ容量を超える100億スケールのデータセットにおけるGPUベースANNSのスケーラビリティのボトルネックを解消すること。
• GPUメモリに完全なグラフとデータを保持する必要がある、または複数GPU構成に依存する、既存のGPUベースANNS手法の制限を克服すること。
• GPU上で圧縮ベクトルを用いて計算をオフロードすることで、CPU上にグラフを保持しつつ、1枚のGPUで高スループットかつ高リコールのANNSを実現すること。
• 通信と計算のオーバーモードと知的なプリフェッチを活用して、CPUとGPU間のPCIeデータ転送のオーバーヘッドを最小限に抑えること。
• 単一GPUのみを用いて、最新のCPUおよび複数GPUベースANNSシステムと同等またはそれ以上の性能を達成すること。

提案手法

• BANGは、高次元ベクトルの圧縮にProduct Quantization（PQ）を用い、GPUメモリにフル精度ベクトルを格納せずに効率的な距離計算を可能にする。
• ディスクに保存されたVamanaグラフ構造（DiskANNから採用）をインデックスとして用い、CPU上で保持・管理する。一方、近傍の走査と距離計算は、圧縮ベクトルを用いてGPU上で実行する。
• GPUカーネルは、距離計算、ソーティング、ワークリスト更新、最近傍選定といったコア演算を高度に最適化しており、GPUのスレッドスケジューリングと並列性を最大限に活用している。
• CPUとGPUは、パイプライン化され、並列的に検索ステージを実行する：CPUは近傍情報の準備とデータのプリフェッチを実行する一方、GPUは距離計算を実行する。
• ブールフィルタを用いて、非有望なノードを早期に除外することで、グラフ走査を高速化し、不要なGPUワークロードとメモリアクセスを削減する。
• 非同期データ転送とデュアルバッファリングを活用して、通信と計算をオーバーラップさせることで、CPU-GPU間のデータ転送のオーバーヘッドを最小限に抑え、アイドルタイムとPCIeのボトルネックを低減する。

実験結果

リサーチクエスチョン

• RQ1 100億スケールのデータセットで、GPUメモリ容量を超える状況下でも、単一GPUのANNSシステムが高スループットと高リコールを達成できるか？
• RQ2 GPU-CPUワークロードの最適な分散をどのように実現すれば、グラフベースANNSにおけるデータ転送を最小限に抑え、並列性を最大化できるか？
• RQ3 PQによるデータ圧縮を用いることで、検索精度を損なわずにGPU上で高速な距離計算を実現できる程度はどの程度か？
• RQ4 パイプライン処理と通信と計算のオーバーラップにより、グラフデータをCPUにオフロードした場合の性能ペナルティを顕著に低減できるか？
• RQ5 標準ベンチマークにおいて、提案手法は最新のGPUおよび複数GPUベースANNSシステムと比較して、スループットとリコールの両面で優位性を示せるか？

主な発見

• 100億スケールのデータセットにおいて、BANGはリコール0.9の条件下で、競合するGPUベース手法と比較して40×～200×の高いスループットを達成し、顕著な性能優位性を示した。
• 100万スケールのデータセットにおいて、BANGは多数のベンチマークケースで最先端の性能を上回り、最も優れた競合手法と比較して算術平均のスループットが2倍にのぼった。
• SIFT1B や GIST1M といった大規模ベンチマークを含む、すべての評価対象データセットで、BANGは高リコール（≥0.9）を維持した。これは、従来のGPUベース手法がリコールとスループットの両立に苦慮する状況でも同様の結果を達成したことを示している。
• GPU上で圧縮データを用いることで、メモリ帯域幅の圧迫を軽減し、全データセットとグラフをCPU上に保持している状況でも、効率的な計算が可能になった。
• パイプライン処理と通信と計算のオーバーラップにより、CPU-GPU間のデータ転送オーバーヘッドが低減され、BANGはPCIe帯域幅の制限に強く対応できるようになった。
• BANGは、複数GPU構成を必要とせず、GPU内に完全なインデックスを格納する必要がないという点で、単一GPU上で100億スケールのデータセットを効率的に処理できる最初のGPUベースANNSシステムである。