Distributed Approximate Nearest Neighbors

"They asked me for the closest vector among a billion. I found the eighth-closest in a millisecond and called it close enough. Nobody complained."

An Index Shard With a Relaxed Definition of Truth

Every earlier section of this chapter shared one shape: split the data across workers, compute a partial summary on each, and combine the summaries into the exact global answer. Linear models combined gradient sums, decision trees combined histogram counts, k-means combined cluster sums. Nearest-neighbor search breaks that shape in an instructive way. The answer we want, the $k$ stored vectors closest to a query, has no small sufficient statistic; in the worst case you must look at every vector to be certain you have the closest. So the distributed story changes from "combine partial summaries into the exact answer" to "search partitions of the data and merge their best candidates into a very good answer." That shift, from exactness to controlled approximation, is what makes search at scale tractable, and it is the natural close to a chapter about doing classical machine learning on more data than one machine can hold.

1. Why Exact Nearest Neighbor Does Not Scale Beginner

Given a query vector $q$ and a database of $N$ vectors $\{x_1, \dots, x_N\}$ in $\mathbb{R}^d$, the nearest-neighbor problem asks for the $k$ database vectors with the smallest distance to $q$, usually Euclidean distance or its cosine cousin. The obvious algorithm computes all $N$ distances and keeps the $k$ smallest. Its cost is

which is fine when $N$ is a few thousand and ruinous when $N$ is a billion. A single query over a billion 768-dimensional float vectors touches roughly $10^{9} \times 768$ multiply-add operations and reads several terabytes from memory; multiply by thousands of queries per second and no single machine, however large, can keep up. The vectors themselves, at four bytes per dimension, occupy about three terabytes, which already exceeds the memory of one server. Both ceilings from Chapter 1, the memory ceiling and the throughput ceiling, bind at once.

The classical escape, tree structures such as k-d trees that prune away most of the database, works beautifully in low dimensions and collapses in high ones. As $d$ grows past a few dozen, distances concentrate: every pair of points becomes nearly equidistant, the triangle-inequality pruning that trees rely on stops eliminating branches, and an exact tree search degrades to scanning almost everything. This is the curse of dimensionality, and embeddings live squarely inside it at $d$ in the hundreds or thousands. The response is not a better exact structure but a deliberate retreat from exactness.

2. The Three Index Families Intermediate

Production ANN systems are built from three ideas, often combined. The first is IVF, the inverted file. You cluster the database into $\texttt{nlist}$ cells with k-means (the very algorithm of the previous section, reused as a coarse quantizer), and store each vector in the inverted list of its nearest centroid. At query time you find the $\texttt{nprobe}$ centroids closest to $q$ and scan only the vectors in those few cells, ignoring the rest of the database. If a query touches $\texttt{nprobe}$ of $\texttt{nlist}$ cells holding $N/\texttt{nlist}$ vectors each, the work falls from $N$ to roughly $\texttt{nprobe} \cdot N / \texttt{nlist}$, a large constant-factor win whose recall depends on how often the true neighbors live in the probed cells.

The second idea is graph-based HNSW, hierarchical navigable small-world graphs. Each vector becomes a node linked to a handful of nearby vectors, with a few long-range links layered on top, exactly the small-world structure that makes social networks navigable. A search starts at an entry node and greedily walks toward the query, hopping closer at each step, reaching a high-recall neighborhood after a logarithmic number of hops rather than a linear scan. HNSW gives the best recall-versus-latency trade-off in memory and dominates in-RAM benchmarks, at the cost of a graph that can be larger than the vectors themselves and is awkward to update incrementally.

The third idea attacks memory rather than time: product quantization (PQ). Split each $d$-dimensional vector into $m$ sub-vectors, run k-means on each subspace to learn a small codebook of, say, 256 entries, and store each sub-vector as a single byte naming its nearest codebook entry. A 768-dimensional float vector (3072 bytes) compresses to $m$ bytes, often 96 or fewer, a 30-fold or larger reduction, and distances are approximated by table lookups over the codebooks. PQ rarely stands alone; the workhorse production index is IVF-PQ, coarse IVF cells to narrow the search plus PQ codes to make each cell cheap to scan and small enough to keep in memory. Table 12.7.1 summarizes how the three families place themselves on the recall, latency, and memory axes.

3. Building an IVF Index From Scratch Intermediate

The cleanest way to feel the recall-versus-speed trade is to build the simplest of the three, an IVF index, in a few lines of NumPy and measure both numbers on the same data. The code below clusters a synthetic database into cells, stores an inverted list per cell, then answers each query by probing only its nearest few cells. It also runs an exhaustive search to get ground truth, so it can report true recall@10 alongside the speedup over scanning everything.

import numpy as np, time

rng = np.random.default_rng(7)
N, d, Q = 200_000, 64, 500       # database vectors, dimension, queries
nlist, nprobe, topk = 256, 8, 10 # cells, cells probed per query, neighbors

# A clustered database so neighborhoods are real.
centers = rng.standard_normal((40, d)) * 6.0
base    = centers[rng.integers(0, 40, N)] + rng.standard_normal((N, d))
queries = centers[rng.integers(0, 40, Q)] + rng.standard_normal((Q, d))

# Train the coarse quantizer: k-means lite -> nlist cell centroids.
cent = base[rng.choice(N, nlist, replace=False)].copy()
for _ in range(8):
    cell_of = np.empty(N, np.int64)
    for s in range(0, N, 20000):                 # chunked nearest-centroid assign
        e = min(s + 20000, N)
        cell_of[s:e] = ((base[s:e, None] - cent[None]) ** 2).sum(-1).argmin(1)
    for c in range(nlist):                        # recompute centroids (cluster means)
        m = cell_of == c
        if m.any(): cent[c] = base[m].mean(0)
inv = [np.where(cell_of == c)[0] for c in range(nlist)]   # inverted lists

# Exact search = ground truth and the cost we want to beat.
t0 = time.perf_counter()
exact = np.array([np.argpartition(((base - q) ** 2).sum(1), topk)[:topk] for q in queries])
t_exact = time.perf_counter() - t0

# IVF search: probe only the nprobe nearest cells.
t0 = time.perf_counter(); approx = []; scanned = 0
for q in queries:
    probe = np.argpartition(((cent - q) ** 2).sum(1), nprobe)[:nprobe]
    cand  = np.concatenate([inv[c] for c in probe]); scanned += cand.size
    dd    = ((base[cand] - q) ** 2).sum(1)
    approx.append(cand[np.argpartition(dd, topk)[:topk]])
t_ivf = time.perf_counter() - t0
approx = np.array(approx)

recall = np.mean([len(set(exact[i]) & set(approx[i])) / topk for i in range(Q)])
print(f"avg vectors scanned  : {scanned/Q:.0f} of {N}  ({100*scanned/Q/N:.2f}%)")
print(f"recall@{topk}           : {recall:.3f}")
print(f"speedup over exact   : {t_exact/t_ivf:.1f}x")

avg vectors scanned  : 6660 of 200000  (3.33%)
recall@10           : 0.987
speedup over exact   : 17.0x

The numbers tell the whole story of ANN in three lines: scanning one part in thirty of the data still found almost every true neighbor, because clustering puts a query's real neighbors overwhelmingly in the cells nearest its own centroid. The 1.3% of neighbors missed are the price, and nprobe is the knob that sets that price. Everything more sophisticated, HNSW graphs and PQ compression, is a way to push the same curve further: more recall per unit of latency, or the same recall in far less memory.

4. Distributing the Index: Shard, Scatter, Gather Advanced

One machine holds the whole IVF index in Code 12.7.1, which works at two hundred thousand vectors and fails at two hundred billion. The distribution pattern is the one this book has used since Section 4.7: partition the data across machines, send the query to all of them, and merge their partial answers. Each machine, a shard, holds a disjoint slice of the vectors and its own complete index over that slice. A query is broadcast to every shard (scatter), each shard returns its local top-$k$ (a partial answer), and a coordinator merges the $S$ shard results into the global top-$k$ (gather). Because Euclidean and cosine distance are computed independently per vector, the union of every shard's local top-$k$ is guaranteed to contain the global top-$k$: merging cannot lose a true neighbor that some shard found. Figure 12.7.1 shows the full path.

Sharding answers the memory ceiling: each machine indexes only $N/S$ vectors. Throughput is the other ceiling, and it is answered by the orthogonal move of replication. Each shard is copied onto several machines, and incoming queries are load-balanced across the replicas, so a system serving ten thousand queries per second simply adds replicas until the rate is met. Sharding and replication compose into a grid: $S$ shards wide for capacity, $R$ replicas deep for throughput, $S \times R$ machines in total. The query latency is set by the slowest shard to respond, which makes ANN serving vulnerable to the stragglers of Section 4.7: one slow replica delays every query routed to it, so the same tail-latency tactics, hedged requests and replica selection, return here.

import faiss, numpy as np
# base, queries: float32 arrays of shape (N, d) and (Q, d)
index = faiss.index_factory(d, "IVF256,Flat")   # IVF with 256 cells; "IVF256,PQ16" to compress
index.train(base)                               # learns the coarse centroids
index.add(base)                                 # builds the inverted lists
index.nprobe = 8                                # the recall-vs-latency knob
D, I = index.search(queries, 10)                # top-10 neighbor ids per query

5. Where ANN Connects, Backward and Forward Intermediate

Approximate nearest neighbor did not appear from nowhere; it is the high-dimensional, learned-embedding descendant of ideas already in this book. Section 6.7 introduced locality-sensitive hashing (LSH), the original sublinear trick for similarity at scale: hash vectors so that close vectors collide and far ones do not, then search only colliding buckets. IVF is LSH with a learned, data-adaptive partition (k-means cells) in place of random hyperplanes, which is why it achieves higher recall per probe on real embeddings. The vectors being searched are exactly the embeddings whose distributed storage Section 11.6 built: sharded embedding tables produce the vectors, sharded ANN indexes retrieve over them, and the two sit on opposite ends of the same pipeline.

Forward, ANN is the engine inside the vector databases and retrieval-augmented generation systems of Part V. Chapter 25 takes the sharded scatter-gather index of Figure 12.7.1 and wraps it in the durability, filtering, incremental-update, and consistency machinery that turns an index into a database, then puts it on the critical path of a language model that retrieves context before it generates. The 1.3% of neighbors an IVF index misses becomes, in that setting, a passage of context the model never sees, which is why recall tuning stops being academic and starts shaping answer quality.

6. Chapter Summary: One Pattern, Five Algorithms Beginner

This section closes Chapter 12, and the chapter has a single spine worth stating plainly. Classical machine learning distributes by the same move every time: partition the training data across workers, compute a small sufficient statistic on each shard, and combine those statistics with an all-reduce or a histogram merge into a result that is exact, or in the case of search, very nearly so. Linear and logistic regression combined gradient sums. Decision trees and gradient boosting combined per-bin histogram counts to find the best split. K-means combined per-cluster point sums and counts. ANN, the one method whose answer has no compact sufficient statistic, combined per-shard candidate lists by a top-$k$ merge. The collective from Chapter 4 is the common engine beneath all five, which is the deeper reason this chapter sits where it does: classical machine learning, distributed, is mostly the art of finding the right small summary and reducing it across the cluster.