Nearest-neighbor search

Overview

Given reference points and query points, find for each query the k reference points nearest to it under a chosen distance 0804.1448 §2.1. The exhaustive method computes every query-to-reference distance, orders them, and keeps the k smallest. Its cost grows with both set sizes, and the paper says these searches are often the bottleneck in the applications it lists 0804.1448 §2.2.

Nine cards cover this concept. They differ in three ways: whether the search is exact or approximate, whether the full table of distances is built in memory, and whether the index is compressed. In one card, the neighbour sets are an input to a visualisation method rather than the answer a caller wants 2012.04456 §7. In another, distances to the nearest archive members decide what enters an evolutionary collection, so the neighbours are not returned to a caller either 1708.09251 §3.1.2. Three are exact over all reference points 0804.1448 §2.1 2004.11127 §2 2112.01511 §4.8. Faiss has an exact flat index as well as compressed approximate ones 2401.08281 §3. Two approximate methods return neighbours at a recall the authors report only in part: about 90% of the true top-k for Memorizing Transformers 2203.08913 §3.3, and an unstated level for RETRO's SCaNN search 2112.04426 §2.3. Faiss is the only card that measures accuracy against exact search 2401.08281 §3.4 2401.08281 §5.5. The other exactness claims are stated by the methods and are not checked against a reference result in the papers.

A 2024 survey of retrieval-augmented generation lists IVFPQ, HNSW and tree-based indexes such as Annoy as the ANN structures used in these systems 2407.13193 §3.1.3. It gives no parameters or measurements for them.

Variants

Exact brute-force search on a GPU

Every query-reference distance is independent, so all of them are computed in parallel on a graphics card, and each query's k smallest values are kept 0804.1448 §2.1. The paper uses Euclidean distance in its experiments and fixes k = 20 0804.1448 §3. Its reported speed-up against Matlab brute force is up to 120x, which the authors state as their headline result 0804.1448 tab. 1. Card: Fast k Nearest Neighbor Search using GPU.

Exact GPU reduction without a stored distance table

KeOps writes each pairwise value as a symbolic formula and reduces it over one point set for every point of the other. Its reductions include argKmin(K), which returns the indices of the K smallest entries in each row 2004.11127 §2. Memory stays linear in the number of points, and gradients come from automatic differentiation 2004.11127 §2. The paper times a Gaussian matrix-vector product, not a k-nearest-neighbour query. At D = 3 on one RTX 2080 Ti it takes 1.39 s for 1M points, while PyTorch runs out of memory from 100k points 2004.11127 §3. The k-NN reduction itself is not timed. Card: KeOps.

Linear lookup in a visual imitation policy

VINN embeds each demonstration frame once. At test time it finds the k nearest demonstration embeddings under Euclidean distance and returns a softmin-weighted average of their actions 2112.01511 §3.2. The paper uses k around 10 and reports no gain beyond k = 20 2112.01511 §4.7. The search is brute force, so its cost is linear in the number of demonstration frames 2112.01511 §4.8. The paper reports about 0.038 s per query, with the hardware not stated 2112.01511 §4.8. It names FAISS as a faster option and does not test it 2112.01511 §4.8. Card: The Surprising Effectiveness of Representation Learning for Visual Imitation.

Approximate top-k as external memory in a transformer

Memorizing Transformers adds a store of past (key, value) pairs to one attention layer. Each query retrieves its top 32 stored pairs by approximate kNN, which it reaches about 90% recall on TPUs, and blends the result with local attention through a learned per-head gate 2203.08913 §3.1 2203.08913 §3.3 2203.08913 §4.2. The stored entries are not trained by gradients. Author-reported: at context 512 with an XL cache, test perplexity on arXiv falls from 2.67 to 2.31 when the memory grows to 65K tokens 2203.08913 §4.3. The paper reports that the approximation did not noticeably hurt quality, but it gives no sweep over recall 2203.08913 §4.2 2203.08913 §3.3. Card: Memorizing Transformers.

Compressed inverted-file and graph indexes in a library

Faiss builds an index by chaining a vector codec with a step that skips most of the database, and optionally re-ranks the shortlist exactly [arxiv:2401.08281 §3.5] 2401.08281 §5. The codecs include product quantizers, which split each vector into M sub-vectors with a codebook of K entries each, and scalar quantizers down to 4 bits per component 2401.08281 §4.1. Search by default compares the raw query with the decoded codes, which loses no query accuracy 2401.08281 §4.

The skipping step is either an inverted file or a graph. An inverted file has nlist centroids, and a query scans only the nprobe nearest lists 2401.08281 §5.1. A graph index (HNSW or NSG) is searched by a traversal whose length is the knob; HNSW accepts new vectors later, while NSG is frozen after its first batch 2401.08281 §5.2 2401.08281 §6.1. The paper advises graphs when memory is not the limit and the set is below about 1M vectors, and compressed inverted files above that 2401.08281 §5.3.

The headline result is author-reported: a 1.5-trillion-vector index with 72-dimensional PCA and 6-bit scalar codes, 54 bytes per vector, at about 1 s per query. That setup used 20 intermediate servers, and the paper says network bandwidth was the limit 2401.08281 §7.1. On the Deep10M benchmark, the paper reports it faster than SCANN and about on par with Vamana, depending on the operating point 2401.08281 §5.5. Card: The Faiss library.

Approximate kNN over precomputed keys for retrieval in a language model

RETRO splits each sequence into 64-token chunks and retrieves neighbours for each chunk from a database of up to about 2 trillion tokens. Training uses two neighbours per chunk 2112.04426 §4.1. Each key is the time-averaged embedding of a frozen BERT encoder, and the search uses squared L2 distance with SCaNN as the approximate index 2112.04426 §2.3. Neighbours from the same document as the training sequence are removed 2112.04426 §2.3.

Training neighbours are computed ahead of time, so training does not search on the fly. Sampling retrieves again after each generated chunk, so the cost grows with the number of chunks 2112.04426 §2.4. Storage is large: about 215 GB for a Wikipedia index and about 93 TB for MassiveText 2112.04426 §4.1. The paper does not state the SCaNN settings or the recall they reach 2112.04426 §2.3. Card: Improving language models by retrieving from trillions of tokens.

Neighbour graph inside a dimension-reduction method

PaCMAP needs the nearest neighbours of each point before its optimisation starts. It takes the closest Euclidean candidates, a pool of min(n_NB + 50, N), re-ranks them by a distance scaled with each point's local spread, and keeps n_NB = 10 2012.04456 §7 2012.04456 eq. 2 2012.04456 alg. 1. The neighbour pairs are built once, before the loop 2012.04456 alg. 1. Its mid-near pairs are random draws, not a search 2012.04456 §7.

The paper's timing setup names Annoy as the neighbour backend 2012.04456 §8.1. The released repository defaults to an approximate FAISS HNSW index instead, so the neighbour sets, and the scale values taken from them, can differ from exact search. The paper does not discuss this (checked in the PaCMAP card). The reported local and global scores are author-reported on 12 datasets 2012.04456 §8.1 2012.04456 §8.2. Card: PaCMAP.

Nearest-member distance as a density test in a quality-diversity archive

The archive in the quality-diversity framework has no fixed grid. A candidate is judged by its Euclidean distance to the existing members, and the layout emerges from the solutions found 1708.09251 §3.1.2. The insertion rule looks at the two nearest members. If the nearest is farther than the density threshold l, the candidate is added. If the second-nearest is within l, it is rejected. Otherwise it may replace the nearest member, provided it exclusively ε-dominates that member on novelty and quality 1708.09251 §3.1.2 fig. 2. Novelty is the mean distance to the k nearest members, with k = 15 in all three experiments 1708.09251 §3.1.2 1708.09251 tab. 2.

The threshold l is in normalised descriptor units. The paper uses 0.01 for its two 2-D tasks and 0.25 for the 6-D gait task, with ε = 0.1 throughout, so l has to be set for the normalisation in use 1708.09251 tab. 2. After each batch, the algorithm recomputes novelty and local quality for every member 1708.09251 alg. 1. The arm collection reached about 6,000 solutions, so this step grows with the archive 1708.09251 §4.2.2. The paper's results are author-reported and come from simulation only. Its arm archive collections look similar to its grid collections 1708.09251 §4.2.2 fig. 3. Card: Quality and Diversity Optimization: A Unifying Modular Framework.

When to use which

Timing trends: for every method, time grows roughly linearly with dimension, while the GPU time stays nearly flat 0804.1448 §3 fig. 2. Its largest settings take about an hour with Matlab or C brute force, about 20 minutes with the kd-tree, and 43 seconds on the GPU at N = 38400 and D = 96 0804.1448 §3 tab. 1. These figures are on a 2008 Pentium 4 host and a GeForce 8800 GTX 0804.1448 §3, so they show the trend, not what a current GPU would do.

The cards use different hardware and settings: the 2008 GPU above, an RTX 2080 Ti for KeOps 2004.11127 §3, TPU hardware for the approximate search 2203.08913 §3.3, an unstated machine for VINN 2112.01511 §4.8, and no stated hardware for the quality-diversity archive 1708.09251 §4. Faiss reports its trillion-scale timing on a sharded cluster 2401.08281 §7.1. Their timings show trends within each paper and are not a ranking across them.

Caution for the transformer variant: the authors' repository also contains a dense memory path that does not do kNN. Check which path you run before you compare numbers 2203.08913 §3.1.

Implementation sketch

Language-agnostic pseudocode of the exact GPU brute-force core, in our own words, following the steps in the paper's method section. The paper does not say how the selection is implemented; the sort here is one option. For the compressed inverted-file search, see the Faiss card's pipeline 2401.08281 §5.1 2401.08281 §4.

# Exact brute-force k-NN, all pairs in parallel  [arxiv:0804.1448 §2.1]
# Input:  ref (m, d) reference points; query (n, d) query points; k
# Output: (k, n) distances to the k nearest references, per query

def knn_brute_force(ref, query, k):
    # Squared Euclidean distance for every query-reference pair  [§2.1]
    # Pairs are independent, so the GPU computes them in parallel
    sq = pairwise_sq_distances(ref, query)        # (m, n)

    # Keep the k smallest values in each query's column  [§2.1]
    # Columns are independent, so each one can run on its own thread
    for each column j in parallel:
        sq[:, j] = smallest_k_sorted(sq[:, j], k)

    # The square root is monotone, so it can be taken last (our note)
    return sqrt(sq[:k, :])                        # (k, n)

Contradictions

Across papers, the descriptions of tree-based search differ. The survey lists tree-based indexes such as Annoy among the ANN structures, with no dimension limit given 2407.13193 §3.1.3. Faiss says LSH and tree-based pruning do not scale well above about 10 dimensions 2401.08281 §5. The survey gives no measurement, so the two statements cannot be reconciled from these cards. The kd-tree results in 0804.1448 are consistent with a low-dimension limit: at D = 8 it beats the GPU method on large sets 0804.1448 tab. 1.

Within the papers, the reports disagree in these places:

• Speed-up baselines (0804.1448). Section 2.1 gives the factors as about 120 over a C version and about 40 over the kd-tree 0804.1448 §2.1. Section 3 and the Table 1 caption give about 120 over Matlab and about 100 over C 0804.1448 §3 0804.1448 tab. 1. The ratios from Table 1 match the second set (largest BF-Matlab to BF-CUDA ratio about 118; largest BF-C to BF-CUDA about 103; largest KDT-C to BF-CUDA about 41), as computed in the source card (derived by us).
• Slopes at N = 38400 (0804.1448). Section 3 gives slopes of 34 for BF-C and 31 for BF-Matlab 0804.1448 §3. A straight line through the Table 1 values at D = 8 and D = 96 gives about 29.7 for BF-C and about 33.8 for BF-Matlab (derived by us) 0804.1448 tab. 1.
• Neighbour backend (2012.04456). The timing setup names Annoy 2012.04456 §8.1, while the repository's default is an approximate FAISS HNSW index. The paper does not say which backend its results used. This is paper versus code, as checked in the PaCMAP card.
• Recall versus quality (2203.08913). The paper says the approximation did not noticeably hurt model quality 2203.08913 §4.2, but it gives no sweep over recall, so the size of any loss is not measured 2203.08913 §3.3.
• Density test, paper versus code (1708.09251). The paper rejects a candidate when its distance to the second-nearest member is within l 1708.09251 §3.1.2. The released code uses (1 − ε)·l as that bound, so a candidate between 0.9l and l is accepted by the code and rejected by the paper. This is checked in the quality-diversity card, not a disagreement between two papers.

Papers

• Fast k Nearest Neighbor Search using GPU: exact brute-force k-NN with the distance and selection steps on CUDA, timed against Matlab, C and a kd-tree on sets of up to 38400 points in up to 96 dimensions.
• KeOps: symbolic kernel and distance reductions on GPU, including k-smallest selection, without storing the full matrix, with autodiff.
• The Surprising Effectiveness of Representation Learning for Visual Imitation: VINN, a brute-force k-NN lookup over self-supervised embeddings to predict robot actions.
• Memorizing Transformers: approximate top-k retrieval from an external memory of past keys and values inside one attention layer.
• The Faiss library: a C++ and Python toolkit for compressed inverted-file and graph indexes, with exact re-ranking and a parameter sweep that keeps Pareto-optimal settings.
• Improving language models by retrieving from trillions of tokens: RETRO, approximate SCaNN retrieval of neighbours for every chunk from a trillion-token database, with keys computed ahead of time.
• Retrieval-Augmented Generation for Natural Language Processing: A Survey: a taxonomy of retrieval components, including the ANN index types above. It runs no experiments of its own.
• PaCMAP: a 2-D visualisation method that builds scaled-distance nearest-neighbour pairs, plus mid-near and random pairs, and optimises them with Adam.
• Quality and Diversity Optimization: A Unifying Modular Framework: one quality-diversity loop with a swappable archive, whose insertion rule uses distances to the nearest members, plus a curiosity-based parent selection.