Chapter 6: Scaling Retrieval - Velmurugan Narayanasamy

Every search in this series so far, Chapter 2’s three documents, Chapter 3’s six sections, Chapter 4 and 5’s same handbook, used the same strategy: loop through every chunk, compute cosine similarity, keep the best one. That’s a linear scan, and it has an honest limitation that’s been sitting in the background since Chapter 4: it compares the question against everything, every single time.

This chapter is about what changes when “everything” stops being six chunks and starts being millions.

Chapter 6 of 10 · Better retrieval

Where we left off: Every search so far, three documents, six sections, one handbook, used a linear scan: compare the question against every single chunk, every time.

What we fix now: Cluster the chunks ahead of time so a query only has to check a small, promising subset instead of everything, and measure the speed/recall trade-off that comes with going approximate.

By the end, you can:

explain why linear scan stops being viable as a chunk store grows into the millions
build an inverted-file index from scratch using k-means clustering
measure the real trade-off between nprobe, query speed, and recall

Making the problem concrete first¶

This chapter is about the mechanics of search at scale, not retrieval quality on real text, so instead of the employee handbook, we’ll use synthetic embeddings standing in for a large chunk store: thousands to millions of vectors, generated with some real cluster structure baked in (documents about the same topic tend to land near each other), the same property Chapter 2’s real embeddings had.

import numpy as np
import time

rng = np.random.default_rng(42)

def make_dataset(n, d=25, n_true_clusters=40, noise=0.15):
    # n_true_clusters simulates "topics": vectors are scattered
    # around a handful of centers, like real chunk embeddings tend to be.
    centers = rng.normal(size=(n_true_clusters, d))
    centers /= np.linalg.norm(centers, axis=1, keepdims=True)
    assign = rng.integers(0, n_true_clusters, size=n)
    return centers[assign] + rng.normal(scale=noise, size=(n, d))

def brute_force_search(query, vectors):
    q = query / np.linalg.norm(query)
    v = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
    sims = v @ q
    best = np.argmax(sims)
    return best, sims[best]

print("Ready.")

# Time a single query against a growing number of chunks.
for n in [1_000, 10_000, 100_000, 1_000_000]:
    vectors = make_dataset(n)
    query = vectors[0] + rng.normal(scale=0.05, size=vectors.shape[1])

    start = time.perf_counter()
    reps = 5
    for _ in range(reps):
        brute_force_search(query, vectors)
    elapsed_ms = (time.perf_counter() - start) / reps * 1000

    print(f"N={n:>9,}  {elapsed_ms:8.3f} ms per query")

Roughly 10 times more chunks costs roughly 10 times longer, every single query, forever. That’s what “linear scan” means in practice: there’s no cleverness to fall back on, the only way to be sure you found the best match is to actually check every candidate. Fine at six. Not fine at six million, especially if many people are asking questions at once.

The fix: don’t compare against everything¶

The core idea behind real vector databases: organize the vectors ahead of time so that most of them can be skipped entirely at query time. One real, widely used technique, simple enough to build from scratch: cluster the vectors into groups first (an inverted file index, if you want the name real vector databases use for this). At query time, figure out which few groups are actually relevant, and only search inside those.

The clustering itself uses k-means, an algorithm simple enough to write in a few lines: pick some starting points, assign every vector to its nearest one, move each point to the average of what got assigned to it, repeat.

def kmeans(vectors, k, n_iters=8, seed=0):
    rng2 = np.random.default_rng(seed)
    n = vectors.shape[0]
    centroid_idx = rng2.choice(n, size=k, replace=False)
    centroids = vectors[centroid_idx].copy()

    v_sq = (vectors ** 2).sum(axis=1, keepdims=True)
    for _ in range(n_iters):
        # squared distance from every vector to every centroid,
        # computed as one matrix multiply instead of N*K individual
        # comparisons, this is what keeps k-means fast at scale.
        c_sq = (centroids ** 2).sum(axis=1)
        dists = v_sq - 2 * vectors @ centroids.T + c_sq[None, :]
        assignments = np.argmin(dists, axis=1)

        for c in range(k):
            members = vectors[assignments == c]
            if len(members) > 0:
                centroids[c] = members.mean(axis=0)

    return centroids, assignments

N = 20_000
K = 50
vectors = make_dataset(N)

start = time.perf_counter()
centroids, assignments = kmeans(vectors, K)
elapsed = time.perf_counter() - start

cluster_sizes = np.bincount(assignments, minlength=K)
print(f"Clustered {N:,} vectors into {K} groups in {elapsed:.2f}s")
print(f"Cluster sizes: min={cluster_sizes.min()}, "
      f"mean={cluster_sizes.mean():.0f}, max={cluster_sizes.max()}")

Searching the index instead of everything¶

At query time: first compare the question against the (small number of) cluster centroids, not the whole dataset, to find which clusters are actually promising. Then only run the expensive comparison inside those clusters. nprobe controls how many clusters to check, more means slower but more thorough.

def indexed_search(query, vectors, centroids, assignments, nprobe):
    # Step 1: cheap check against K centroids, not N vectors.
    q = query / np.linalg.norm(query)
    c = centroids / np.linalg.norm(centroids, axis=1, keepdims=True)
    nearest_clusters = np.argsort(c @ q)[::-1][:nprobe]

    # Step 2: only look inside the chosen clusters.
    candidate_idx = np.nonzero(np.isin(assignments, nearest_clusters))[0]
    candidates = vectors[candidate_idx]

    local_best, score = brute_force_search(query, candidates)
    return candidate_idx[local_best], score, len(candidate_idx)

query = vectors[123] + rng.normal(scale=0.05, size=vectors.shape[1])

start = time.perf_counter()
for _ in range(30):
    brute_force_search(query, vectors)
bf_ms = (time.perf_counter() - start) / 30 * 1000
print(f"brute-force:       {bf_ms:.3f} ms/query  (checks all {N:,})")

for nprobe in [1, 3, 5, 10]:
    start = time.perf_counter()
    for _ in range(30):
        result = indexed_search(query, vectors, centroids, assignments, nprobe)
    idx_ms = (time.perf_counter() - start) / 30 * 1000
    print(f"indexed nprobe={nprobe:<3} {idx_ms:.3f} ms/query  (checks ~{result[2]:,})")

A real speedup, several times faster, by skipping most of the dataset outright. But speed alone isn’t the whole story: an approximate search that’s fast and wrong isn’t actually useful. The real question is whether indexed_search still finds the same answer brute_force_search would.

# Recall: across many test queries, how often does the indexed
# search find the exact same answer brute force would?
n_test = 150
test_queries = [
    vectors[i] + rng.normal(scale=0.05, size=vectors.shape[1])
    for i in rng.integers(0, N, n_test)
]
true_best = [brute_force_search(q, vectors)[0] for q in test_queries]

for nprobe in [1, 2, 3, 5, 10]:
    matches = sum(
        1 for q, tb in zip(test_queries, true_best)
        if indexed_search(q, vectors, centroids, assignments, nprobe)[0] == tb
    )
    print(f"nprobe={nprobe:<3} recall={matches / n_test:.3f}")

This is the real trade-off every vector database asks you to make, laid bare: nprobe=1 is the fastest and already right most of the time, but not always, the true best match occasionally lives in a cluster that didn’t get checked. Raising nprobe searches more clusters, costs more time, and recovers the missed cases. There’s no setting that’s simply “correct,” only a speed/accuracy trade-off you choose deliberately, based on how much a wrong answer costs you versus how much a slow answer costs you.

MY_K = 50
MY_NPROBE = 2

my_centroids, my_assignments = kmeans(vectors, MY_K)

matches = sum(
    1 for q, tb in zip(test_queries, true_best)
    if indexed_search(q, vectors, my_centroids, my_assignments, MY_NPROBE)[0] == tb
)
print(f"K={MY_K}, nprobe={MY_NPROBE}: recall={matches / n_test:.3f}")

Quick check¶

Before moving on, answer for yourself (no peeking ahead):

In your own words, why does comparing the question against K centroids first make the rest of the search so much cheaper?

In the recall experiment, nprobe=1 already got most queries right. Why does it ever get one wrong, what has to happen for the true best match to end up outside the single cluster that got checked?

If you were building a real system and a wrong answer was costly (say, a medical or legal question) but a slow answer was merely annoying, how would that change the nprobe you’d choose, compared to a system where speed mattered more than occasional misses?

This chapter’s index is built once, from a fixed dataset. What do you think has to happen differently if new chunks need to be added to a live system without rebuilding the whole index from scratch?

Chapter 5 built citations so a wrong retrieval becomes visible to a reader. Now that retrieval is approximate (nprobe less than the total clusters), does the citation mechanism from Chapter 5 still work the same way when the “wrong” chunk is a genuine approximation miss rather than an embedding mistake?

You can now: cluster millions of chunks so search only checks a small, promising subset instead of everything, and you’ve measured the real speed/recall trade-off that comes with going approximate.

But one problem remains: clustering by meaning solves the scale problem, but it still relies entirely on embeddings, and there are questions, like one that’s really just a code or a number, where meaning-based search loses to plain keyword matching outright.

Next → Chapter 7: Hybrid Search