Nearest neighbour search: brute force, heaps and clustering

Brute force with a heap

Compute the similarity between the query and every vector, then keep the top k. Sorting all n scores costs O(n log n), and a heap of size k costs O(n log k), which matters when n is large and k is 5.

import heapq

import numpy as np

rng = np.random.default_rng(1)
n, dim, k = 20000, 64, 5
vectors = rng.normal(size=(n, dim)).astype(np.float32)
vectors /= np.linalg.norm(vectors, axis=1, keepdims=True)
query = vectors[42] + 0.05 * rng.normal(size=dim).astype(np.float32)
query /= np.linalg.norm(query)

scores = vectors @ query
exact = [i for _, i in heapq.nlargest(k, ((s, i) for i, s in enumerate(scores)))]
print("exact top-k:", exact)

n_clusters = 50
centroids = vectors[rng.choice(n, n_clusters, replace=False)]
for _ in range(5):
    assign = np.argmax(vectors @ centroids.T, axis=1)
    for c in range(n_clusters):
        members = vectors[assign == c]
        if len(members):
            centroids[c] = members.mean(axis=0) / np.linalg.norm(members.mean(axis=0))

for n_probe in (1, 3, 10):
    probe = np.argsort(-(centroids @ query))[:n_probe]
    candidates = np.where(np.isin(assign, probe))[0]
    approx = candidates[np.argsort(-(vectors[candidates] @ query))[:k]]
    recall = len(set(approx) & set(exact)) / k
    print(f"probe {n_probe} clusters: searched {len(candidates)} vectors, recall@{k} {recall:.1f}")

With normalised vectors, the dot product is cosine similarity, so one matrix multiplication scores everything.

Approximate search: clustering (IVF)

Split the vectors into clusters with k-means, and at query time search only the closest few clusters. More clusters probed means higher recall and slower queries. The example above shows the trade-off: probing 1 cluster searches a small fraction of the vectors, and probing more brings recall back up. The vectors in the example are random, so they have no natural clusters and recall is lower than you'd see with real embeddings, which cluster by topic.

Approximate search: graphs (HNSW)

HNSW builds a graph where each vector links to its close neighbours, with sparse upper layers for long jumps and dense lower layers for local search. A query starts at the top, greedily moves toward closer nodes, and drops a layer each time it can't improve. It's the default index in most vector databases, and it uses more memory than IVF. See vector index internals.

Measuring recall

Recall@k is the fraction of the true top k (from brute force) that the approximate index returns. Measure it on a sample of real queries before you tune anything else, because a fast index that misses the right chunk makes your RAG system wrong in a way the LLM can't fix.

When brute force is enough

Vectors	Usually fine
Up to about 100,000	Brute force, or pgvector without an index
100,000 to tens of millions	HNSW or IVF in a vector database
More, or tight memory	IVF with compressed vectors (product quantization)

These ranges are rough and depend on dimensions, latency targets and hardware, so time it on your data.

Check your understanding

Change n_clusters to 200 and explain what happens to the number of vectors searched and to recall.

From the floor

In the AI data structures lecture, I placed similarity search right after embeddings in the sequence of stages, because once meaning is a vector, finding related items becomes a question of distance, and I tied it back to search: a query like "can you get medicine for someone" needs the system to find pages about buying medicine on someone's behalf, not pages that contain the same words.

Related

• Vector index internals
• Vector databases
• Embedding tables
• Hybrid search

Sources

• Efficient and robust approximate nearest neighbor search using HNSW graphs, Malkov and Yashunin (2016)
• Faiss documentation, Meta
• ANN Benchmarks