Fast k Nearest Neighbor Search using GPU

In plain words

Finding the nearest stored items to a query is a common step in vision and statistics: a query point is compared with every reference point and the closest few are kept. Done naively, the work grows with the number of points on both sides, so large sets get slow.

The paper's idea is to do this exhaustive comparison on a graphics card. Every query-to-reference distance is independent of the others, so a GPU can work on thousands of them at once. The authors find that this makes the simple, exhaustive method faster than CPU versions of the same method and faster than a tree-based shortcut in their tests.

The same search also feeds an entropy estimator, which needs only the distance to the k-th neighbour of each point. The paper's experiments measure timing only; they do not test accuracy against a tree-based index.

TL;DR

• Exact brute-force k-NN search, with the distance and sort steps run as CUDA kernels §2.1.
• Table 1 times reference and query sets of 1200 to 38400 points in 8 to 96 dimensions, against Matlab, C, a kd-tree (ANN) and the GPU version tab. 1.
• The headline "up to 120x" is author-reported and is measured against Matlab brute force; against the C baseline the table gives about 103x at most tab. 1 (derived by us).
• Timings only: the paper reports no accuracy or recall numbers for the GPU method.
• Relevant if you need exact k-NN on moderate sets, or the k-th-neighbour distance for entropy estimates.

Problem

Given m reference points and n query points in d dimensions, find the k nearest references for each query under a chosen distance §2.1. The brute-force method computes all n·m distances, sorts them, and selects the k smallest. Its cost is about 2·n·m·d additions and n·m·d multiplications for the distances, plus about n·m·log m comparisons for the sorts §2.1. The cost therefore grows with both set sizes. The paper says these searches are often the bottleneck in the applications it lists §2.2.

Assumptions: Euclidean or Manhattan distance are the usual choices; other distances are allowed §2.1. The paper's experiments use k = 20 §3.

Method

Brute force on the GPU. The paper's method is brute force with every step parallel: all distances are computed at once, each query's column is sorted independently, and the k smallest are kept §2.1. The paper does not describe the kernels; the design below is taken from the official code, which is the 2008 version kept in legacy/. Code citations in the rest of this section use the file legacy/KNN.CUDA.without.indexes/knn_cuda_without_indexes.cu, shortened to its basename inside the pseudocode.

• Distances. A tiled kernel loads 16x16 blocks of reference and query coordinates into shared memory and accumulates squared differences over the dimensions knn_cuda_without_indexes.cu. A variant reads the reference points from texture memory and is chosen when they fit knn_cuda_without_indexes.cu.
• Selection. Each query column is handled by one thread. An insertion routine keeps only the k smallest values at the top of the column, not a full sort knn_cuda_without_indexes.cu.
• Root. Square roots are taken only for the k kept values, so the output is Euclidean distances knn_cuda_without_indexes.cu.
• Batching. Queries are split into chunks sized by free GPU memory, with a 0.9 fraction of free memory used knn_cuda_without_indexes.cu.

# Brute-force k-NN, all queries in parallel  [§2.1]
# Input:  ref (m, d) reference points; query (n, d) query points; k
# Output: dist (k, n) Euclidean distances, ascending per query

def knn_brute_force(ref, query, k):
    # Squared distance for every reference-query pair  [§2.1]
    # Sums squared differences per dimension  [knn_cuda_without_indexes.cu#L153-L158]
    sq_dist = squared_distances(ref, query)         # (m, n)

    # Keep the k smallest values per query column  [§2.1 step 2]
    # Code uses insertion, not a full sort  [knn_cuda_without_indexes.cu#L180-L218]
    for each query column j in parallel:
        sq_dist[:, j] = keep_k_smallest_first(sq_dist[:, j], k)

    # Square roots of the k kept values only  [knn_cuda_without_indexes.cu#L229-L233]
    return sqrt(sq_dist[:k, :])                     # (k, n)

Application 1: entropy. The Kozachenko–Leonenko style estimator uses the distance from each point to its k-th nearest neighbour §2.2. The paper gives the Shannon entropy definition eq. 1 eq. 2, the estimator eq. 3, the digamma function eq. 4 and the unit-ball volume eq. 5. Note that the paper does not state a log base.

\widehat{H}_{n,k}(Y) = \frac{1}{n} \sum_{i=1}^{n} \big[ \log((n-1)\,\rho_k(y_i)) + \log(c_1(d)) - \Psi(k) \big] eq. 3

\Psi(k) = \frac{\Gamma'(k)}{\Gamma(k)} = \int_0^\infty \left[ \frac{e^{-t}}{t} - \frac{e^{-kt}}{1-e^{-t}} \right] dt eq. 4

c_1(d) = \frac{2\,\pi^{d/2}}{d\,\Gamma(d/2)} eq. 5

Here ρ_k(y_i) is the distance from y_i to its k-th nearest neighbour, and c_1(d) is the volume of the unit ball in R^d §2.2.

# k-NN entropy estimate  [§2.2] [eq. 3]
# Input:  points (n, d) sample set Y; k
# Output: scalar entropy estimate, log base not stated

def knn_entropy(points, k):
    # Distance rho_k to the k-th nearest neighbour of each point  [§2.2]
    rho = kth_neighbor_distance(points, k)      # (n,)
    n, d = points.shape

    # Constant term from the unit-ball volume  [eq. 5]
    log_c1 = log(2 * pi ** (d / 2) / (d * gamma(d / 2)))

    # Digamma term  [eq. 4]
    psi_k = digamma(k)

    # Average the per-point terms  [eq. 3]
    terms = log((n - 1) * rho) + log_c1 - psi_k  # (n,)
    return mean(terms)

Application 2: classification and CBIR. The paper also names KNN classification §2.2 and KNN-based content-based image retrieval §2.2, where descriptors are matched by nearest-descriptor search and votes are counted per image. It gives no equations or parameters for these uses.

Hyperparameters

Results

Author-reported. Setup: brute-force CUDA against Matlab brute force, C brute force and the ANN kd-tree (KDT-C), on the hardware above §3. Data source, data type, timing scope and repetitions are not stated §3. Times are in seconds, and bold marks the fastest entry per cell, as in the paper tab. 1.

Source: tab. 1. The caption gives the BF-CUDA factors as roughly 120 over BF-Matlab, 100 over BF-C and 40 over KDT-C tab. 1.

Ratios we computed from the table (not stated in the paper): the largest BF-Matlab / BF-CUDA ratio is about 118 (D=96, N=4800); the largest BF-C / BF-CUDA ratio is about 103 (D=96, N=9600); the largest KDT-C / BF-CUDA ratio is about 41 (D=96, N=4800). KDT-C is faster than BF-CUDA at D=8 with N=19200 and 38400 tab. 1 §3.

Other reported results:

• The worked example: 38400 reference and query points in 96 dimensions take about one hour for BF-Matlab and BF-C, about 20 minutes for KDT-C, and 43 seconds for BF-CUDA §3.
• Time grows about linearly with dimension for every method, and is nearly flat for BF-CUDA §3 fig. 2.
• At D=8, most of the BF-CUDA time is host-to-device copies, according to the CUDA profiler §3.

Paper vs code

The official repository is vincentfpgarcia/kNN-CUDA at commit d135b79391f5ffbfd45ef1c53d2217d9fa53c299. Its root code/ holds a 2018 rewrite with a different API, and legacy/ holds the 2008 code. The paper's experiments match the legacy code. The paper does not say which version produced Table 1, and the legacy file header is dated 03/07/2009 knn_cuda_without_indexes.cu.

Differences:

• Output. Section 2.1 defines k-NN as finding the k nearest neighbours of each query §2.1. The legacy code returns only the distance to the k-th neighbour, not the neighbours or their indices knn_cuda_without_indexes.cu.
• Selection. Section 2.1 describes a full sort, with about n·m·log m comparisons §2.1. The legacy code keeps the k smallest with insertion knn_cuda_without_indexes.cu.
• Distance. The paper allows other distances §2.1. The code computes squared Euclidean distance and takes the root only at the end knn_cuda_without_indexes.cu.
• Speed-up baseline. Section 2.1 gives the factors as 120 over a C version and 40 over the kd-tree §2.1. Section 3 and the Table 1 caption give 120 over Matlab and 100 over C §3 tab. 1. The abstract gives a factor of 120 without naming the baseline.
• Slopes at N=38400. Section 3 gives slopes of 34 for BF-C and 31 for BF-Matlab §3. A straight-line slope between D=8 and D=96 in Table 1 gives about 29.7 for BF-C and about 33.8 for BF-Matlab (derived by us). The paper's two values look swapped.
• Texture path. The code uses texture memory only when the reference set is at most 16384 points and the dimension is at most 8192 knn_cuda_without_indexes.cu. Table 1 at N=19200 and 38400 would therefore use the global-memory path, if this code produced it. The paper does not mention the switch.
• Batching. The code splits queries into batches sized by free GPU memory knn_cuda_without_indexes.cu. The paper does not mention batching.
• 2018 rewrite. The root code/knncuda.cu computes distances with the expansion |q|² + |r|² − 2·q·r through cuBLAS, using a cublasSgemm call knncuda.cu knncuda.cu. It also adds a texture variant knncuda.cu. It does not match the 2008 experiments.
• Setup per call. Every knn call initialises CUDA and creates and detaches a context knn_cuda_without_indexes.cu. The timing loop calls knn 100 times knn_cuda_without_indexes.cu. The paper does not say whether Table 1 includes this per-call overhead.

Pitfalls

• Output is a single distance. The legacy code returns the k-th neighbour's distance per query, not the neighbour list. Adapt it before using it for classification or retrieval knn_cuda_without_indexes.cu.
• Headline speed-up. "Up to 120" is measured against Matlab brute force. Against C it is about 100 in the table, and about 41 against the kd-tree at best tab. 1 (derived by us). Do not quote one figure without its baseline.
• Timing scope. Table 1 is computation time in seconds tab. 1. The paper does not say whether host-to-device copies, context creation or allocation are included. The code's per-call setup suggests they are (see Paper vs code).
• Accuracy is not measured. The paper reports only times. The GPU method is exact brute force, so its results should match the C baseline, but the paper does not check this §2.1 §3.
• Small-D regime. At D=8 the kd-tree is faster than BF-CUDA for N=19200 and 38400 tab. 1 §3. The paper's explanation is host-to-device copies §3.
• Entropy. The log base is not stated eq. 3. The paper does not say whether the query point is excluded from its own neighbour list. The (n − 1) term suggests it is excluded, but this is our reading eq. 3.
• Numerical range. The code computes c_1(d) and digamma from the formulas eq. 4 eq. 5 with no numerical details. For large d, compute the log of c_1(d) with log-gamma to avoid overflow. This is our suggestion, not from the paper.
• Memory limits. The code allocates from free GPU memory and reserves 90% of it knn_cuda_without_indexes.cu. If free memory is too small, the computed batch size can drop to 0. The batch loop then steps by 0 and never ends knn_cuda_without_indexes.cu. This is our reading of the code; it is not tested here.
• Sizes. The reference set must fit the texture limits to use the texture path knn_cuda_without_indexes.cu. The global-memory path has its own pitch limit, checked only in the Matlab build knn_cuda_without_indexes.cu.
• Section and equation numbers. These are counted from the LaTeX source. The book class has no chapters here, so the compiled numbers may differ. The compiled PDF was not checked.
• Licence of the code. The repository is CC BY-NC-SA 3.0, which is non-commercial LICENSE.md. This card cites it by path and line, does not quote it, and is not a licence to reuse the code in commercial products.

Lineage

• Builds on: the exhaustive brute-force k-NN search §2.1, with the sort replaced by a partial insertion selection in the code knn_cuda_without_indexes.cu.
• Compared against: Matlab brute force, C brute force, and the ANN kd-tree library §3 tab. 1.
• Related later work: the code's README lists a 2010 ICIP paper by the same authors on GPU k-NN README.md. That paper is not on this wiki.