I wonder if there is any study that compares the performance of kd-tree vs brute-force nearest neighbor search on GPU. Post #4 on this page suggests that kd-tree may not be the optimal algorithm for GPU but I wonder if there is any data supporting this claim?
3 Answers
Ultimately, naive brute-force KNN is an $O(n^2)$ algorithm, while kd-tree is $O(n \log n)$, so at least in theory, kd-tree will eventually win out for a large enough $n$. In practice, the leading constants for a GPU implementation may be vastly different --- we may be comparing $0.0001n^2$ vs $1000n\log n$ --- so it may indeed be the case that the former wins out for practical problem sizes.
Now, the reason for the difference in leading constants come down to parallelizability and memory access patterns. Naive KNN is embarrassingly parallelizable, and can be implemented entirely using vectorized sequential memory access. By contrast, kd-tree is naturally serial, and requires extensive conditional statements and random memory access. Algorithms like the former enjoy a massive speed-up on a GPU, while those like the latter often run slower on a GPU than a CPU.
Additionally to what @richard-zang said, instead of a "naive brute-force" search, you can often use some refinement, e.g. a locality-based hashing or if you have fixed neighbor distance radius, a common approach is to pre-grid/sort the search space to limit the lookup to neighboring cells (commonly used in molecular simulation and referred to as linked cell list.
- $k$: number of dimensions.
- $n$: number of points.
The naive KNN with brute-force linear scan is $O(kn)$ for both worst-case and average case. When ran on GPU, the speedup can be (roughly) one thousand ($10^3$).
kd-tree's worst case is $O(n)$ for a completely unbalanced tree and $O(n^{1−1/k})$ for the average case. It can be barely parallelized and does not support online/incremental updates well (although amortized time is good enough).
Let's take the average case of kd-tree with a single thread. When $n=2^{20}$ (about 1 million) and $k=100$, (this is a good average case for text-based classification problem) then kd-tree needs about 1 million operations in the average case, where linear scan needs 100 million operations. Considering the GPU speedup, then linear is 10 times faster than kd-tree.
We can see that only when $k$ is very small that kd-tree can beat linear scan, but that is pretty rare in real-world problems. This problem is called the dimension curse which led to the research of Approximate Nearest Neighbor (ANN) search algorithms for practical implementations: e.g. https://arxiv.org/pdf/1702.05911.pdf.