Parallel Algorithms for Nearest Neighbor Search Problems in High Dimensions

The nearest neighbor search problem in general dimensions finds application in computational geometry, computational statistics, pattern recognition, and machine learning. Although there is a significant body of work on theory and algorithms, surprisingly little work has been done on algorithms for high-end computing platforms and no open source library exists that can scale efficiently to thousands of cores. In this paper, we present algorithms and a library built on top of Message Passing Interface (MPI) and OpenMP that enable nearest neighbor searches to hundreds of thousands of cores for arbitrary dimensional datasets. The library supports both exact and approximate nearest neighbor searches. The latter is based on iterative, randomized, and greedy KD-tree searches. We describe novel algorithms for the construction of the KD-tree, give complexity analysis, and provide experimental evidence for the scalability of the method. In our largest runs, we were able to perform an all-neighbors query search on a 13 TB synthetic dataset of 0.8 billion points in 2,048 dimensions on the 131K cores on Oak Ridge’s XK6 “Jaguar” system. These results represent several orders of magnitude improvement over current state-of-the-art methods. Also, we apply our method to non-synthetic data from machine learning data repositories. For example, we perform an all-nearest-neighbor search on a variant of the ”MNIST” handwritten digit dataset with 8 million points in 784 dimensions on 16,384 cores of the ”Stampede” system at the Texas Advanced Computing Center, achieving less than one second per PKDT iteration.