Nearest neighbors at scale.

Vector search is how a computer finds, among millions of stored passages, the few whose meaning sits closest to your question, when each passage has been boiled down to a long list of numbers.

It’s the step that runs before your cross-encoder ever sees anything, and it quietly sets a ceiling on how good your ranker can look. When a search result is bad, the fault is often here, in a stage that’s easy to treat as plumbing. It’s also a favourite interview question, because the honest answer is a page of arithmetic about bytes and memory speed.

It leans on two things you already have: the napkin math from How does it scale?, and the idea from Why ~17 tokens per second? that a chip is often stuck waiting for numbers to arrive from memory rather than working on them. After it, Ranking metrics and eval sets asks how sure you can be that one search setup really beats another, and the “match by meaning” cache from Caching and Redis turns out to be this same machinery.

Today

• Picture each passage’s list of numbers as a point in space, and say when the three common ways of measuring “how close” two points are all agree on the order.
• Work out on a napkin what it costs to compare one question against 100 million passages, in bytes, arithmetic and milliseconds, and why the answer is set by how fast memory can be read.
• Explain, in plain words, the three classic shortcuts: sort passages into rooms and only search a few rooms, shrink each passage’s numbers to a few bytes, and walk a road map of links between neighbours. Name the dial on each that trades accuracy for speed or memory.
• Build all three on real search data on the Spark, check how often each finds the same top 10 as the slow exact search, and plot that against speed and memory.
• Put a shortcut search in front of your cross-encoder, the two-step fetch-then-reorder setup you work on, and measure what each step costs and how much it improves the final ranking.

You’ll want the napkin math from How does it scale? and the “waiting on memory versus working” idea from Why ~17 tokens per second?; Why is NumPy 100x faster? helps too. Plan on one evening to read and one long session on the Spark.

The Problem

Your cross-encoder is good. Give it a query and a passage and it reads them together, every word of one attending to every word of the other, and outputs one relevance score. It is also slow on purpose: one forward pass of a small transformer per pair. Call it a millisecond per pair on a GPU when batched, which is generous.

Now the index has 100 million passages. Scoring all of them for one query is 100 million forward passes. At a millisecond each that is 28 hours. Per query.

So nobody does that. Search systems run in two stages. First, a cheap stage pulls a thousand or so plausible candidates out of the hundred million. Then the expensive model, your model, reorders only those. The ranker can only reorder what it is handed: if the right passage isn’t in the thousand, no amount of training on your side brings it back.

The cheap stage, for a growing share of systems, works like this. A small bi-encoderbi-encoderA model that turns a query and a document into vectors separately, so documents can be turned into vectors once, ahead of time. Relevance is then a simple vector comparison. Contrast with a cross-encoder, which must see both texts together. turns every passage into a list of numbers ahead of time, say 768 of them. At query time it turns the query into 768 numbers too, and the job becomes: which of the 100 million stored lists is closest to this one?

That sounds easy. It’s one dot product per passage. Let’s see what “one dot product per passage” costs when there are 100 million passages, and then spend the rest of the lesson cheating our way out of it.

The Idea

An embedding is a point

An embeddingembeddingA fixed-length list of numbers that a model outputs for a piece of text (or an image, or a user). The model is trained so that things with similar meaning get similar lists. is a list of numbers, like 768 floats. Read the list as coordinates, and it is a point in a 768-dimensional space. Two numbers would be a point on a map. Three would be a point in a room. 768 is a point in a space you can’t picture, but every rule you know from the map still holds: points have distances, and some points are closer than others.

The embedding modelembedding modelThe neural network that produces embeddings, for example a small sentence-transformer. Trained so that a query lands near the passages that answer it. was trained to arrange that space for you. A query like “how do I reset my router” lands near passages about routers and modem lights and far from passages about pasta. Search becomes geometry: find the stored points nearest the query point. That task is called k-nearest neighbourskNNk-nearest neighbours: given a query point, find the k stored points closest to it. “Nearest neighbour search” usually means this..

Three ways to say “near”

There are three common measures, and you will see all three in library flags.

• dot productdot productMultiply the two vectors position by position and add up the results. a·b = a₁b₁ + a₂b₂ + … + a₇₆₈b₇₆₈. Bigger means more similar.: 768 multiplies and 768 adds. Big when the vectors point the same way and are long.
• cosine similaritycosine similarityThe dot product divided by both vectors’ lengths. It is the cosine of the angle between them: 1 for the same direction, 0 for unrelated (at right angles), −1 for opposite.: the dot product with length divided out, so it only measures the angle.
• Euclidean distanceL2 distanceThe straight-line distance between two points: subtract position by position, square, add, take the square root. Smaller means more similar. Called L2 after the mathematical name of this norm.: the straight-line gap between the points. Small means similar.

Here is the identity that ties them together. Expand the squared L2 distance:

|a − b|²  =  |a|²  +  |b|²  −  2 (a · b)

Now suppose every vector has length 1 (you normalizenormalizeDivide a vector by its own length so it has length exactly 1. It keeps the direction and throws away the size.d them). Then |a|² = |b|² = 1, the dot product is the cosine, and

|a − b|²  =  2 − 2 cos θ

The squared distance is a fixed, decreasing function of the cosine. So sorting by smallest L2, by biggest dot product, and by biggest cosine gives the exact same order. That’s why a lot of systems normalize everything and then use whichever operation their hardware or library does fastest, which is usually a plain dot product.

When the vectors aren’t normalized, they disagree. Dot product rewards length: a long vector that points roughly the right way can beat a short one that points exactly the right way. Some retrieval models are trained to use that (length can encode “this passage is generally useful”), and for those you must search by dot product, not cosine. The rule: use the similarity the model was trained with. The model card tells you. The small sentence-transformer you’ll use in Build It was trained with cosine and outputs normalized vectors, so all three agree.

Brute force, and what it costs

The honest way to find the nearest neighbours: compute the similarity between the query and every stored vector, keep the top k. This is brute forcebrute-force kNNExact nearest neighbour search by comparing the query to every stored vector. Always correct. Cost grows linearly with the number of vectors. Also called flat or exhaustive search. (FAISS calls it a “flat” index). It’s exact. Let’s price it, back-of-envelope style.

For N vectors of d dimensions, one query costs N × d multiply-adds. And to do them, the machine has to read every stored number once.

For 100 million passages at 768 dimensions in 32-bit floats:

• Multiply-adds per query: 10⁸ × 768 ≈ 7.7 × 10¹⁰. About 77 billion.
• Bytes: 10⁸ × 768 × 4 bytes = 307 GB. That’s the size of the index, and every query reads all of it.
• Does it fit on the Spark? The spec sheet says 128 GB. No. Even at 16 bits per number (154 GB) it doesn’t.
• If it did fit, how fast? The Spark’s spec-sheet memory bandwidth is about 273 GB/s. Reading 307 GB takes 307 / 273 ≈ 1.1 seconds. Per query.

Play with the knobs. Note how the read time scales with every factor, and which ones you can actually change.

Two ways out appear.

1. Read fewer bytes per vector. Store them smaller. fp16 halves it, int8 quarters it, and product quantization (below) shrinks each vector to tens of bytes.
2. Read fewer vectors. Don’t look at most of the index. Somehow jump to the right neighbourhood and only check points there.

Every serious vector index does one or both. And every one of them gives up the guarantee of the exact answer to do it.

Approximate, and the triangle

Approximate nearest neighbourANNApproximate nearest neighbour search: return points that are probably the nearest, or very close to the nearest, in exchange for being much faster than brute force. search means: return the true nearest neighbours most of the time, much faster. The way you score “most of the time” is recall@krecall@kFor ANN: of the true k nearest neighbours (found by brute force), what fraction did the index return in its top k? 0.95 means it found 9.5 of the true 10, on average over queries..

Be careful, because you already know a different recall from ranking. In ranking metrics, recall asks “did we find the relevant documents?”, judged by human labels. ANN recall asks “did we find the same top 10 that brute force would?”, judged by the exact search over the same embeddings. The first measures the embedding model. The second measures only the index. If the embedding model ranks the right passage 400th, a perfect index still returns it 400th. Keep the two apart: ANN recall@10 = 0.98 can sit next to a retrieval recall@10 of 0.6.

Every index lives somewhere in a triangle with three corners:

• Recall. How often you get the true neighbours.
• Latency (or its flip side, queries per second). How much work per query.
• Memory. How many bytes per vector, plus the index’s own structures.

You can buy any one with the other two. Look at more of the index: recall up, speed down. Compress the vectors: memory down, recall down. Add a graph of links between vectors: speed and recall up, memory up. Index build time is a fourth cost that hides off to the side: some indexes take hours to build and can’t easily absorb new vectors.

The rest of the Idea is three tricks. Each one is a different bet about which corner to give up.

IVF: sort the library into rooms

Picture a library with 100 million books and no catalogue. Brute force reads every spine. A smarter librarian first sorts the books into a thousand rooms by topic. When you ask a question, she figures out which few rooms your question belongs to, and only searches those.

That’s an inverted file indexIVFInverted file index. Split the vectors into clusters ahead of time; store, for each cluster, the list of vectors in it. At query time, find the few clusters nearest the query and only search their lists.. Mechanically:

1. Train. Run k-meansk-meansA clustering algorithm: pick k centre points, assign every vector to its nearest centre, move each centre to the average of its vectors, repeat. The result splits the space into k regions. on (a sample of) your vectors to get nlist cluster centres, called centroidcentroidThe centre point of a cluster: the average of the vectors assigned to it.s. Each centroid owns the region of space closer to it than to any other centroid. Those regions are called Voronoi cellsVoronoi cellsThe regions you get by giving every point in space to its nearest centre. Draw the borders and the space looks like a cracked tile floor..
2. Add. Put every vector in the list of its nearest centroid.
3. Search. Compare the query to all nlist centroids (cheap: a thousand, not a hundred million). Pick the nprobe nearest. Brute-force only the vectors in those cells.

The cost per query is roughly nlist + nprobe × (N / nlist) distance computations. For 100 million vectors, nlist = 65,536 (2¹⁶) and nprobe = 64: that’s 65,536 + 64 × 1,526 ≈ 160,000 distances instead of 100 million. About 600× less work.

Where’s the catch? The border. The query’s true nearest neighbour can sit just across the line in a cell you didn’t open. In 768 dimensions, a cell has a huge number of neighbouring cells, and the query often lands near several borders at once. nprobe is the dial: probe 1 cell and you miss a lot; probe more and recall climbs toward 1 while speed falls linearly. The FAISS docs suggest nlist somewhere between about 4√N and 16√N; then you sweep nprobe.

Product quantization: describe each vector in a few bytes

IVF reads fewer vectors. It doesn’t make them smaller, so 100 million × 768 floats still needs 307 GB. product quantizationPQCompress a vector by cutting it into m short pieces and replacing each piece with the ID (one byte) of the closest entry in a small learned codebook for that piece. A 768-float vector becomes m bytes. fixes the memory corner.

The analogy: describing a face to a sketch artist. You don’t send a photo. You say “eyes: type 17, nose: type 203, mouth: type 4”, picking each from a book of 256 standard shapes. The artist rebuilds an approximate face from a handful of numbers.

PQ does exactly that to a vector:

1. Cut the 768 numbers into m = 96 consecutive pieces of 8 numbers each.
2. For piece 1, run k-means with 256 centres on piece 1 of all your vectors. That’s piece 1’s codebookcodebookFor one piece of the vector, a list of 256 representative short vectors learned by k-means. A piece is stored as the index (0-255) of the closest one, which fits in one byte.. Same for each of the 96 pieces.
3. Store each vector as 96 bytes: for each piece, the ID of its nearest codebook entry.

The memory math: 768 × 4 = 3,072 bytes becomes 96 bytes. 32× smaller. 100 million vectors go from 307 GB to 9.6 GB, plus the IDs that IVF stores alongside (8 bytes each in FAISS, another 0.8 GB). That fits on the Spark with room to spare. Use m = 48 and it’s 64× smaller, at a further cost in accuracy.

The clever part is the bottom half of the figure. To score a compressed vector you don’t rebuild it. You split the query into the same 96 pieces and compute, once, the similarity of each query piece to each of the 256 codebook entries for that piece. That’s a 96 × 256 table, about 25,000 numbers, small enough to sit in cache. Now any stored vector’s (approximate) dot product with the query is the sum of 96 table lookups, one per code. This is called asymmetric distance computationADCScoring compressed database vectors against an uncompressed query, using a precomputed table. “Asymmetric” because only one side is compressed, which keeps more accuracy than compressing both.. Fewer bytes read, cheap arithmetic, and in a bandwidth-bound problem fewer bytes is exactly what you need.

The price is accuracy. Each piece is snapped to the nearest of 256 shapes, so every score is a little wrong, and the ranking near the top gets shuffled. The standard fix is cheap: ask the PQ index for the top 100, then re-score those 100 with the full-precision vectors (kept on disk or in slower memory) and keep the best 10. This “compressed search, exact re-rankre-rankTake a short list of candidates from a cheap scorer and re-order it with a more accurate (and more expensive) scorer.” pattern is the whole lesson in miniature, and you’ll see it again at the end.

In FAISS the combination of the two ideas is IVF65536,PQ96: room-sorting to read fewer vectors, PQ to make each one small.

HNSW: highways, then local roads

The third trick throws out the rooms entirely. Instead, build a graph: connect each vector to a handful of its near neighbours. To search, start somewhere, look at your neighbours, step to whichever is closest to the query, and repeat until no neighbour is closer. That’s greedy searchgreedy searchAt each step, move to whichever neighbour is closest to the target, and stop when none is closer. Fast, but can get stuck in a spot that is locally best and not globally best. on a graph, and it’s how you’d find a friend’s house in a town with no map: keep taking whichever street heads most directly toward it.

On its own this is slow for a big graph. If every link is short, crossing the whole space takes thousands of tiny steps. The fix is the one road networks use. You don’t drive across the country on residential streets. You take local roads to a highway, the highway across most of the distance, and local roads again at the end.

A HNSWHNSWHierarchical Navigable Small World graph. Vectors are linked into several layers of graphs: the top layers hold few vectors with long links, the bottom layer holds every vector with short links. Search walks down the layers. index builds exactly that, in layers:

• Layer 0 contains every vector, each linked to its nearest neighbours. Local roads.
• Layer 1 contains a random subset, about one in M of them, linked to their nearest neighbours within that subset. Because the subset is sparse, those links are longer. A regional highway.
• Layer 2 keeps one in M of those, with even longer links. And so on up, until the top layer has a handful of vectors.

Each vector’s top layer is chosen at random when it’s inserted, with exponentially falling odds, like flipping coins. If that sounds like a skip listskip listA sorted linked list with extra “express” lists on top that skip over many elements, built by coin flips. Search runs along the express lane, then drops down. Lookups take about log N steps., it is: HNSW is the skip list idea with graphs instead of sorted lists.

Three parameters control it, and they are the ones you’ll tune:

• M: how many links each node keeps per layer (the bottom layer usually gets 2M). More links means more ways to route around a dead end, so better recall, more memory (each link is a 4-byte ID: at M = 32 that’s about 64 × 4 = 256 extra bytes per vector on layer 0), and slower builds. Typical: 16 to 64.
• efConstruction: when inserting a vector, how wide a search to run to find its neighbours. Bigger means a better-wired graph and a slower build. Typical: 100 to 500. Set once; you can’t change it after building.
• efSearch (often just ef): the width of the search on the bottom layer at query time. Instead of keeping only the single best node, the search keeps a running list of the best ef nodes it has seen, and keeps expanding the closest unexplored one until nothing in the queue can beat the worst of its list. Bigger ef explores more, so recall up, speed down. This is the dial you turn per query, live, with no rebuild.

Two more details from the paper that matter. The search on the upper layers is plain greedy (ef = 1): those layers only need to get you to the right neighbourhood. And when choosing a node’s links, HNSW doesn’t just take the M nearest. It skips a candidate if it’s closer to an already-chosen neighbour than to the node itself. That keeps a few links pointing away, toward other clusters, instead of M links all pointing into the same tight huddle. Without it, clustered data breaks the graph into islands that greedy search can’t leave.

Now try both indexes on a small map. Click to place a query anywhere. In HNSW mode, step through and watch the big dark dots (the upper-layer nodes) carry the search across the map in a few hops, then the bottom layer fill in locally. In IVF mode, watch it measure the cell centres first and then scan whole cells.

Things to try:

1. HNSW, query in a far corner. The first hops are huge. Count the hops per layer: roughly constant, a handful each, which is why the cost grows roughly with log N instead of N.
2. HNSW, slide ef from 5 to 40. The distance count rises and the found set stays perfect. This 2D map is too easy to show misses. On real 768-dimension data, low ef misses a lot (you’ll see recall around 0.8-0.9 at small ef).
3. IVF, nprobe = 1, put the query right on a dotted border. Some true neighbours sit in the cell next door, and IVF misses them. Bump nprobe to 2.
4. Compare the counters. On this map HNSW checks about 40-70 points, IVF with nprobe = 2 about 120. Brute force always 600. At 100 million points the gap is enormous.

Why not always HNSW?

HNSW is the default in most vector databases because its recall-vs-speed curve is usually the best at a given memory budget, and it handles inserts without retraining. But look at the memory: it stores full vectors plus the links. 100 million × (3,072 + ~300) bytes is about 340 GB. It doesn’t compress, and its access pattern hops randomly through memory, so it really wants the whole thing in RAM.

A rough guide:

At the 100 million scale on one box, IVF-PQ is often the only thing that fits. At a few million, HNSW in RAM is hard to beat. At billions, people combine the ideas (graph indexes over compressed vectors, with full vectors on SSD, as in Microsoft’s DiskANN).

Filters make everything harder

Real queries come with conditions: “nearest passages, but only English, only from the last month, only this customer’s documents”. This sounds trivial and is one of the hardest problems in the field.

• Post-filter: search the index for the top 1,000, then throw away the ones that fail the filter. If 1% of the corpus passes the filter, about 10 of your 1,000 survive, and you wanted 100. You can raise the 1,000, but for a very selective filter you’d need to search nearly everything.
• Pre-filter: first find every vector that passes, then brute force over just those. Great when the filter keeps few vectors (brute force over 10,000 is fast). Terrible when it keeps 40% of 100 million.
• Filter during the search: skip failing vectors as you go. For IVF that’s fine: you just skip entries while scanning a cell. For HNSW it’s subtle: if you refuse to step through failing nodes, the graph falls apart into disconnected islands and greedy search gets stuck. If you step through them but don’t return them, you waste work on a selective filter. Libraries let you pass a filter function (hnswlib has one; FAISS has an IDSelector), and newer research such as ACORN (2024) builds denser graphs so filtered searches stay connected.

Production systems usually decide per query, using an estimate of how selective the filter is: pre-filter plus brute force when few vectors pass, filtered graph search otherwise. Keep this in mind when someone says “just add a WHERE clause to the vector search”.

Hybrid search: words and vectors

Embeddings are good at meaning (“cheap flights” finds “low-cost airfare”) and bad at exact strings: product codes, error messages, rare names, a specific version number. The decades-old keyword method is the opposite. BM25BM25A classic keyword ranking formula (Robertson and others, 1990s). Scores a document by how often it contains the query’s words, weighting rare words more and dampening repeated ones and long documents. Runs on an inverted index of words. scores documents by the query words they contain, with rare words counting more. It finds CUDA error 700 instantly, and has no idea that “car” and “automobile” mean the same thing.

So many systems run both and merge the lists: hybrid searchhybrid searchRunning a keyword search (like BM25) and a vector search for the same query and merging their result lists, to get the strengths of both.. The scores aren’t comparable (BM25 is unbounded, cosine sits between −1 and 1), so the usual merge ignores scores and uses ranks: reciprocal rank fusionRRFMerge several ranked lists by giving each document the sum over lists of 1 / (k + its rank in that list), with k usually 60. Documents near the top of several lists win.. Each document gets 1/(60 + rank) from each list it appears in, summed. Simple, robust, and hard to beat without training a model to do the merge.

The two-stage pipeline you work on

Put it all together and you get the shape of most modern search stacks:

The division of labour is the important idea. The first stage’s job is recall: get the right passage somewhere into the candidate list. It can be sloppy about order. The second stage’s job is precision: put the right passage at the top. It can be slow per item, because there are only a hundred items.

That gives you a budget to reason about. Say the whole request must finish in 100 ms. If the reranker scores 100 pairs in 40 ms, the retrieval stage gets what’s left. If you rerank 300 pairs instead of 100, you might catch a few more good passages that retrieval ranked low, but you pay 3× on the reranker. And an ANN index at recall@10 = 0.95 instead of 0.99 might cost you nothing at all end to end, because the passages it misses are the near-duplicates at ranks 8-10 that the reranker would have reordered anyway. Or it might cost you a lot. You can’t reason your way to the answer. You measure it, end to end, in NDCGNDCGNormalized discounted cumulative gain. Each relevant result earns credit that shrinks with its position (1/log₂(rank + 1)), summed and divided by the best possible score, so 1.0 means a perfect ordering. NDCG@10 looks only at the top 10., which is what the Build is for.

Build It

Everything runs on the Spark. Make a folder, say ~/workbench/vecsearch/, and a virtual environment. The libraries: faiss-cpu (Meta’s vector search library; it has ARM64 Linux wheels on PyPI, but check, and the GPU build is usually installed through conda, so check its install docs for ARM64), hnswlib (the reference HNSW implementation from the paper’s authors), sentence-transformers for the models, and datasets for MS MARCO.

# make_slice.py
import random, json, pyarrow.parquet as pq
from huggingface_hub import hf_hub_download
from datasets import load_dataset

corpus_path = hf_hub_download("BeIR/msmarco", "corpus/corpus-00000-of-00001.parquet", repo_type="dataset")
query_path = hf_hub_download("BeIR/msmarco", "queries/queries-00000-of-00001.parquet", repo_type="dataset")
qrels = load_dataset("BeIR/msmarco-qrels", split="validation")  # the dev small labels

random.seed(0)
rel: dict[str, set[str]] = {}
for r in qrels:
    rel.setdefault(str(r["query-id"]), set()).add(str(r["corpus-id"]))
qids = random.sample(sorted(rel), 1000)
need = set().union(*(rel[q] for q in qids))

corpus = pq.read_table(corpus_path, columns=["_id", "text"]).to_pandas()
is_rel = corpus["_id"].isin(need)
others = corpus[~is_rel].sample(500_000 - int(is_rel.sum()), random_state=0)
sl = corpus[is_rel]._append(others).sample(frac=1, random_state=0).reset_index(drop=True)
sl.to_parquet("passages.parquet")

queries = pq.read_table(query_path).to_pandas().set_index("_id")["text"]
json.dump({q: {"text": queries[q], "rel": sorted(rel[q])} for q in qids}, open("queries.json", "w"))
print(len(sl), "passages,", len(qids), "queries")

# embed.py
import json, time, numpy as np, pandas as pd
from sentence_transformers import SentenceTransformer

m = SentenceTransformer("sentence-transformers/msmarco-MiniLM-L6-cos-v5", device="cuda")
p = pd.read_parquet("passages.parquet")
qs = json.load(open("queries.json"))

t = time.perf_counter()
xb = m.encode(p["text"].tolist(), batch_size=512, normalize_embeddings=True,
              convert_to_numpy=True, show_progress_bar=True).astype("float32")
print(f"{len(xb)} passages in {time.perf_counter() - t:.0f}s")
xq = m.encode([v["text"] for v in qs.values()], normalize_embeddings=True).astype("float32")
np.save("xb.npy", xb); np.save("xq.npy", xq)
print(xb.shape, xq.shape, "norms:", np.linalg.norm(xb[:3], axis=1))

# brute.py
import time, numpy as np, torch

xb = np.load("xb.npy"); xq = np.load("xq.npy"); k = 10

def topk_numpy_one(q: np.ndarray, k: int) -> np.ndarray:
    """One query (shape [d]) against all of xb. Return indices of the k highest dot products, best first.
    Hint: xb @ q, then np.argpartition for the top k, then sort just those k."""
    ...

def topk_numpy_batch(Q: np.ndarray, k: int) -> np.ndarray:
    """All queries at once (shape [nq, d]) -> [nq, k]. One matrix multiply."""
    ...

def topk_torch_gpu(Q: torch.Tensor, B: torch.Tensor, k: int) -> torch.Tensor:
    """Same, on the GPU with torch.topk. Time it with torch.cuda.synchronize() before stopping the clock."""
    ...

if __name__ == "__main__":
    # 1. time topk_numpy_one over 100 queries -> ms per query
    # 2. time topk_numpy_batch over all 1000 -> ms per query
    # 3. time topk_torch_gpu, batch of 1 and batch of 1000 (warm up once first)
    # 4. save the exact top-100 for every query as gt.npy (you'll need 100, not just 10, later)
    ...

import faiss, hnswlib

d = xb.shape[1]
flat = faiss.IndexFlatIP(d); flat.add(xb)                               # brute force, for sanity

quant = faiss.IndexFlatIP(d)                                            # IVF: k-means cells
ivf = faiss.IndexIVFFlat(quant, d, nlist, faiss.METRIC_INNER_PRODUCT)
ivf.train(xb); ivf.add(xb); ivf.nprobe = 16
D, I = ivf.search(xq, 10)

quant2 = faiss.IndexFlatIP(d)                                           # IVF + PQ: m bytes per vector
ivfpq = faiss.IndexIVFPQ(quant2, d, nlist, m, 8, faiss.METRIC_INNER_PRODUCT)

hnsw = faiss.IndexHNSWFlat(d, 32, faiss.METRIC_INNER_PRODUCT)          # HNSW, M = 32
hnsw.hnsw.efConstruction = 200; hnsw.add(xb); hnsw.hnsw.efSearch = 64

h = hnswlib.Index(space="ip", dim=d)                                    # or hnswlib directly
h.init_index(max_elements=len(xb), M=32, ef_construction=200)
h.add_items(xb); h.set_ef(64); labels, dists = h.knn_query(xq, k=10)

# two_stage.py
from sentence_transformers import CrossEncoder
ce = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device="cuda")  # or your model

def retrieve(qvec, index, depth: int) -> list[int]:
    """Top `depth` passage row ids from the ANN index."""
    ...

def rerank(qtext: str, ids: list[int], top: int) -> list[int]:
    """Score (qtext, passage) for every id with ce.predict(pairs, batch_size=...), return the best `top` ids."""
    ...

def ndcg_at_10(ranked_ids: list[int], relevant: set[int]) -> float:
    """Binary relevance: DCG = sum over hits of 1/log2(rank + 1), divided by the best possible DCG."""
    ...

Measure

Fill these in on the Spark. Predict each blank cell before you run it.

Brute force (500k × 384, fp32)

Effective GB/s for the one-at-a-time rows is index bytes / time per query. Compare it to the 273 GB/s spec. You won’t reach it from one CPU thread; the GPU path should get much closer.

Indexes (1 thread, one query at a time)

Two stages (1,000 dev queries)

For reference, I ran the same pipeline on an M4 Pro MacBook (spec-sheet bandwidth 273 GB/s, same as the Spark), embedding on the Mac’s GPU, FAISS 1.15 on the CPU, one thread for the index rows. The embedding took 699 s for 500,000 passages, about 715 per second. Yours will differ; the shapes should not.

Two-stage, first 300 of the 1,000 queries, public MiniLM cross-encoder on the Mac GPU, HNSW32:

What the numbers should tell you:

Brute force is bandwidth-bound one query at a time and compute-bound in a batch. One query at a time, the best path read 768 MB in 3.8 ms, about 200 GB/s, close to the spec. A batch of 1,000 on one CPU thread was 39× faster per query than one at a time on that same thread, because each byte read got reused 1,000 times. That’s the Predict answer, measured. (If your numpy batch row looks slow, check you used argpartition, not a full sort of 500,000 scores per query: my first try did, and the sort cost more than the matrix multiply.)

Real embeddings are hard for IVF. With 4,096 cells of about 120 vectors each, nprobe = 8 found only 68% of the true top 10. The neighbours of a 384-dimension query are scattered across many cells. HNSW at ef = 64 got 0.976 at 0.38 ms, about 28× faster than single-thread brute force. That’s why HNSW is the default when it fits in RAM.

PQ trades recall for 27× less memory, and exact re-ranking buys most of it back. 48 bytes per vector (plus the 8-byte ID) dropped recall from 0.82 to 0.62 at nprobe = 32. Re-scoring the top 100 with the full vectors brought it to 0.82: exactly IVF-Flat’s recall at the same nprobe. The re-rank fixes PQ’s shuffling of the order; it can’t fix what IVF never looked at.

HNSW builds slowly. 71 s for half a million vectors on 10 threads, against 16 s for IVF. At 100 million, multiply by more than 200 (insertion cost grows faster than linearly). Plan for hours.

End to end, the approximate index cost nothing. After reranking the top 100, HNSW and brute force both scored NDCG@10 0.841. The few true neighbours HNSW misses rarely matter once a cross-encoder reorders the list. Reranking added 0.07 NDCG over retrieval alone, which is the reason your job exists.

Rerank depth past 50 bought nothing here and cost linearly. 50, 100 and 200 all scored 0.841, while rerank time went 61 → 125 → 233 ms on the Mac’s GPU. On the Spark your cross-encoder will be faster, but the shape holds: find the knee of NDCG against depth and stop there. One caveat: these scores are high (full MS MARCO MRR@10 for small bi-encoders is closer to 0.3) because a 500k slice has 17× fewer distractors than the real 8.8M corpus. Compare configurations against each other, not against leaderboards.

Break It

1. Forget to normalize. Re-embed the slice without normalize_embeddings=True (or multiply a random 10% of xb rows by 3) and search with IndexFlatIP. Recall against the cosine ground truth drops, and the top results fill up with long vectors. Then normalize and search with IndexFlatL2: rankings match IndexFlatIP exactly, which is the identity from the Idea, checked. Lesson: the metric is part of the model’s contract.

2. Make IVF’s training sample lie. Train the IVF quantizer on only the first 10,000 vectors after sorting the passages by length (or by any field correlated with topic), then add all 500k. Many cells end up huge and many nearly empty, and at the same nprobe your recall and speed both get worse, because a few giant cells are scanned every time. Print the sizes of the 20 biggest inverted lists (ivf.invlists.list_size(i)). This is what happens in production when the corpus drifts after the index was trained, and why IVF indexes get periodically retrained.

3. A selective filter. Tag each passage with a random “language” so that 1% are “de”. Ask for the top 10 German passages three ways: post-filter HNSW top 100 (count how often you get fewer than 10), post-filter top 5,000, and pre-filter then brute force over the ~5,000 matching vectors. The pre-filtered brute force is both exact and fast here. Now make the filter keep 50% and redo it. The winner flips.

4. Starve the reranker. Set HNSW efSearch = 10 and retrieve only 10 candidates, then rerank them. The cross-encoder can’t find a passage it was never shown: NDCG@10 is capped by retrieval’s recall@10. Push the depth to 100 and watch NDCG climb even though the ANN index itself didn’t change. Your ranker’s ceiling is set one stage earlier.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Yu. A. Malkov and D. A. Yashunin, “Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs” (arXiv 2016, IEEE TPAMI 2018). The HNSW paper. Short, readable, and section 4’s neighbour-selection heuristic is the detail most explanations skip.
• Hervé Jégou, Matthijs Douze, Cordelia Schmid, “Product Quantization for Nearest Neighbor Search” (IEEE TPAMI 2011). Introduces PQ, asymmetric distance computation and the IVF + PQ combination. Figures worth staring at.
• Jeff Johnson, Matthijs Douze, Hervé Jégou, “Billion-scale similarity search with GPUs” (2017), and Matthijs Douze et al., “The Faiss library” (2024). How FAISS is built, and a survey of the whole index zoo from the people who wrote it.
• The FAISS wiki on GitHub, especially “Guidelines to choose an index” and “The index factory”. The practical rules of thumb for nlist, m and when to use what.
• Martin Aumüller, Erik Bernhardsson, Alexander Faithfull, “ANN-Benchmarks” (ann-benchmarks.com). The standard recall-vs-QPS plots for every library on common datasets. Compare your plot’s shape to theirs.
• Suhas Jayaram Subramanya et al., “DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node” (NeurIPS 2019). What to do when the vectors don’t fit in RAM: a graph built for SSDs, with compressed vectors in memory.
• Rodrigo Nogueira and Kyunghyun Cho, “Passage Re-ranking with BERT” (2019), and Vladimir Karpukhin et al., “Dense Passage Retrieval for Open-Domain Question Answering” (EMNLP 2020). The two papers behind the two stages you work on: a cross-encoder reranker on MS MARCO, and a bi-encoder retriever that beat BM25.
• Gordon Cormack, Charles Clarke, Stefan Büttcher, “Reciprocal Rank Fusion outperforms Condorcet and individual Rank Learning Methods” (SIGIR 2009). Two pages, and the reason hybrid search is a one-liner.

Connections

This lesson is How does it scale? and Why ~17 tokens per second? applied to search: brute force is one division (bytes over bandwidth), and the fix is the same as in decoding, read fewer bytes or reuse them across a batch. Product quantization is a cousin of Quantization: fewer bits per number, accuracy traded for bandwidth. IVF’s inverted lists and BM25’s word lists are both indexes in the sense of Indexes and B-trees: precomputed structure that lets a read skip most of the data, paid for at write time.

Next, Ranking metrics and eval sets goes deep on the NDCG and MRR you measured in the two-stage table, and on the harder question of whether 1,000 queries are enough to trust a difference of 0.01. And p99 and load testing explains why the retrieval stage’s per-query latency, not its batch throughput, is the number your search service lives or dies by.