Large-Scale Approximate k-NN Graph Construction on GPU

TL;DR

This paper presents a GPU-optimized redesign of NN-Descent for large-scale approximate k-NN graph construction, significantly improving speed and scalability for datasets that exceed GPU memory.

Contribution

Redesigns NN-Descent for GPU architecture, reducing memory accesses and enabling efficient merging of graphs for out-of-memory datasets.

Findings

100-250x faster than single-thread NN-Descent

2.5-5x faster than existing GPU approaches

Enables construction of high-quality k-NN graphs for large datasets

Abstract

k-nearest neighbor graph is a key data structure in many disciplines such as manifold learning, machine learning and information retrieval, etc. NN-Descent was proposed as an effective solution for the graph construction problem. However, it cannot be directly transplanted to GPU due to the intensive memory accesses required in the approach. In this paper, NN-Descent has been redesigned to adapt to the GPU architecture. In particular, the number of memory accesses has been reduced significantly. The redesign fully exploits the parallelism of the GPU hardware. In the meantime, the genericness as well as the simplicity of NN-Descent are well-preserved. In addition, a simple but effective k-NN graph merge approach is presented. It allows two graphs to be merged efficiently on GPUs. More importantly, it makes the construction of high-quality k-NN graphs for out-of-GPU-memory datasets tractable. The results show that our approach is 100-250x faster than single-thread NN-Descent and is 2.5-5x faster than existing GPU-based approaches.