[논문 리뷰] BANG: Billion-Scale Approximate Nearest Neighbor Search using a Single GPU

Karthik Viswanathan, Saim Khan외 3명|arXiv (Cornell University)|2024. 01. 20.

Data Management and Algorithms인용 수 4

한 줄 요약

BANG은 고차원 벡터를 압축하여 GPU에서 빠른 거리 계산을 가능하게 하면서도 그래프와 원본 벡터는 CPU에 유지함으로써 단일 GPU에서 빌리언 스케일 ANNS를 구현하는 새로운 GPU 기반 근사 최근접 이웃 검색 방법을 제안한다. 이는 고성능 GPU-CPU 파ip라인, 최적화된 커널, 그리고 감소된 PCIe 데이터 전송을 통해 고정밀도(0.9)에서 기존 최고 수준의 방법보다 40×–200× 높은 처리량을 달성한다.

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 메모리 용량을 초과하는 빌리언 스케일 데이터셋을 위한 GPU 기반 ANNS의 확장성 문제를 해결한다.
• 기존 GPU ANNS 방법이 전체 그래프와 데이터를 GPU 메모리에 모두 요구하거나 다중 GPU 설정에 의존하는 한계를 극복한다.
• GPU에서 압축된 데이터를 사용하여 계산을 오프로드하면서도 그래프를 CPU에 유지함으로써 단일 GPU에서 고처리량, 고정밀도 ANNS를 가능하게 한다.
• 통신과 계산을 겹치고 지능적인 프리패칭을 통해 CPU와 GPU 간의 PCIe 데이터 전송 오버헤드를 최소화한다.
• 단일 GPU만을 사용하여 최신 CPU 및 다중 GPU ANNS 시스템과 성능가치를 동등하거나 초월한다.

제안 방법

• BANG은 고차원 벡터를 압축하기 위해 제품 양자화(PQ)를 사용하여 GPU에서 전체 정밀도 벡터를 GPU 메모리에 저장하지 않고도 효율적인 거리 계산을 가능하게 한다.
• 디스크ANN의 Vamana 그래프 구조를 기반 인덱스로 사용하며, 이는 CPU에 저장되고 관리되며, 이웃 탐색과 거리 계산은 압축된 벡터를 사용해 GPU에서 수행된다.
• GPU 커널은 거리 계산, 정렬, 워크리스트 갱신, 최근접 이웃 선택 등의 핵심 연산을 최적화하여 GPU 점유율과 병렬성을 극대화한다.
• CPU와 GPU는 파이프라인 방식으로 동시에 작업을 수행한다: CPU는 이웃 정보를 준비하고 데이터를 프리패칭하는 동안 GPU는 거리 계산을 수행한다.
• 비트맵 필터를 사용하여 비순위가 높은 노드를 조기에 제거함으로써 그래프 탐색을 가속화하여 불필요한 GPU 작업과 메모리 접근을 줄인다.
• 비동기 데이터 전송과 더블 버퍼링을 통해 통신과 계산을 겹쳐 CPU와 GPU 간의 통신 오버헤드를 최소화하고, 유휴 시간과 PCIe 병목 현상을 방지한다.

실험 결과

연구 질문

• RQ1 GPU 메모리 용량을 초과하는 빌리언 스케일 데이터셋에서 단일 GPU ANNS 시스템이 고처리량과 고정밀도를 동시에 달성할 수 있는가?
• RQ2 그래프 기반 ANNS에서 데이터 전송을 최소화하고 병렬성을 극대화하기 위해 GPU-CPU 작업 분배를 어떻게 최적화할 수 있는가?
• RQ3 PQ를 통한 데이터 압축이 검색 정확도를 훼손하지 않으면서도 GPU에서 빠른 거리 계산을 가능하게 할 수 있는가?
• RQ4 파이프라인 처리와 계산과의 겹침을 통해 그래프 데이터를 CPU로 오프로드할 경우 성능 저하를 상당부분 줄일 수 있는가?
• RQ5 표준 벤치마크에서 제안된 방법은 최신 GPU 및 다중 GPU ANNS 시스템과 비교해 처리량과 정밀도 측면에서 어떻게 성능을 내는가?

주요 결과

• 빌리언 스케일 데이터셋에서 BANG은 정밀도 0.9에서 경쟁하는 GPU 기반 방법보다 40×에서 200×까지 높은 처리량을 달성하여 뚜렷한 성능 우위를 보였다.
• 미리언 스케일 데이터셋에서는 대부분의 벤치마크 케이스에서 최신 기술을 초월하며, 평균 처리량이 최고의 경쟁 방법보다 2배 높았다.
• 모든 평가된 데이터셋에서 BANG은 높은 정밀도(≥0.9)를 유지했으며, SIFT1B와 GIST1M와 같은 대규모 벤치마크에서도 이전 GPU 방법이 정밀도와 처리량을 균형 있게 확보하기 어려운 상황에서 성능을 유지를 하였다.
• GPU에서 압축된 데이터를 사용함으로써 메모리 대역폭 부담을 줄이고, 전체 데이터셋과 그래프가 CPU에 저장되어 있어도 효율적인 계산이 가능해졌다.
• 파이프라인 처리와 계산과의 겹침을 통해 CPU-GPU 데이터 전송 오버헤드를 감소시켜 BANG이 PCIe 대역폭 제약에 대해 강건함을 입증했다.
• BANG은 단일 GPU에서 다중 GPU 설정 없이도 빌리언 스케일 데이터셋을 효율적으로 처리할 수 있는 최초의 GPU 기반 ANNS 시스템이다.