Approximate Nearest Neighbor Search in Production: From Brute-Force to IVF+PQ

Exact similarity search over a billion items is correct, thorough, and hopelessly slow. Here’s how giving up 3% accuracy bought a 12x speedup.

This one landed on my plate as a “more like this” feature, the kind of thing that sounds like a one-sprint job right up until you read the fine print. For any product in our catalog, show the handful of others most similar to it. I’d built that before on small datasets, where it genuinely is an afternoon’s work. The fine print this time was the catalog. We had about a billion products.

The trick to comparing two products is to turn each one into a list of 128 numbers, called an embedding. A good embedding has one magic property that does all the heavy lifting: similar products end up with similar numbers. So “find me similar products” quietly becomes “find the lists of numbers closest to this one,” and “closest” is just a score you compute between two lists. Higher score, more alike. The model that produces those numbers was someone else’s problem for the length of this project. I just had to search them.

What I actually had to produce was the 20 closest matches for 450,000 products, each compared against all billion in the catalog. Do that the obvious way, every one of the 450,000 against every one of the billion, and you are looking at about 450 trillion comparisons per run. Our cluster could do it. It could also, I suspect, heat the building. What it couldn’t do was finish quickly or cheaply, and “it’ll be ready tomorrow” stops being charming the moment someone asks to run it every day.

The fingerprints came from a language model we’d fine-tuned (XLM-Roberta, if you like the name on the box). The same model runs on both sides of every comparison, which is what keeps the scores honest. That part I wasn’t allowed to touch. The search was where I had room to move, and room to make mistakes.

What follows is the whole trip: from the slow, embarrassingly simple version that actually worked, to the fast one I could run on a schedule, with every wall I walked into in between.

The short version, before the long version

• Start with the slow, exact version. It finds the true best matches every time, which gives you something honest to grade the fast versions against.
• Step one (IVF) buys roughly 8 times the speed for almost no accuracy. You stop searching the whole catalog and look only at the promising parts.
• Step two (PQ) shrinks each fingerprint by about 16 times (20.5 GB down to 1.3 GB per machine) and trims a little more time, for about 3 points of accuracy.
• Right-size the machines. Fewer, beefier machines that each use all their cores beat a swarm of tiny ones. We landed on 25 machines with 16 cores and 64 GB each.
• The default Spark timeout will quietly kill your job. A single piece of work runs for hours, and Spark gives up after two minutes. Raise the limit.
• Cheap search plus a second-opinion step beats expensive search alone. If something downstream double-checks the results, you can afford to be a little less precise up front.

The cluster, and why it looks like this

Before any search, a word on the hardware, because the shape of the cluster decides every number after this. It’s also where I made my most expensive mistake, so we may as well start with the confession.

The full catalog, stored as raw fingerprints, is a billion times 128 numbers times 4 bytes, about 512 GB. No single machine holds that, so it has to be split. The only real questions are how many machines, and how big each slice should be.

My first instinct, which I do not recommend, was lots of small machines. It felt appropriately “big data.” In practice it was a great way to watch a hundred boxes fall over one at a time. So here is the version that actually works.

We settled on 25 machines, each with 16 CPU cores and 64 GB of memory. The reasoning:

• Split the catalog into 25 slices, 40 million items each. One slice as raw fingerprints is 40M times 512 bytes, about 20.5 GB, which sits comfortably in 64 GB with room for everything else a machine is juggling.
• Why 64 GB and not less? The index itself is 20.5 GB, but add Spark’s own appetite, the working space the search needs, and a little headroom, and peak usage lands near 35 GB. A 32 GB box would spend its life one bad garbage collection away from the out-of-memory killer. 64 GB buys peace of mind for not much money.
• Why 16 cores? This is the one that actually mattered. FAISS will happily spread a single search across every core you give it, so the more cores a machine has, the faster it tears through its slice. A swarm of single-core machines leaves all that speed on the floor: every core is starved, and the fixed setup costs get paid over and over. Fewer, beefier machines is both faster and cheaper, which is the rare win that comes with no asterisk. I learned it the expensive way.

So: 25 machines, 400 cores total, each owning a 40-million-item slice of the catalog. Hold onto those numbers; everything below is built on them.

Phase 1: brute force with IndexFlatIP

My first working version wasn’t clever, and I didn’t try to make it clever. Each machine built a plain index from its 40-million-item slice and compared every product against all of it. It’s the algorithmic equivalent of finding a word by reading the dictionary from “aardvark.” Slow, but gloriously, unarguably correct. I had good reasons to start here even knowing it wouldn’t last, and I’ll get to them before you accuse me of wasting a cluster.

import faiss
import numpy as np

# Build index: add this machine's 40M catalog vectors
index = faiss.IndexFlatIP(128)            # 128-dimensional inner product
index = faiss.IndexIDMap(index)           # Wrap to support custom IDs
index.add_with_ids(catalog_embeddings, catalog_ids)

# Search: find top-20 for each source-product vector
distances, indices = index.search(source_embeddings, k=20)

# distances: (N_source, 20), the inner product scores
# indices:   (N_source, 20), the catalog IDs of the nearest neighbors

IndexFlatIP does the simplest possible thing: compare your products against every catalog vector, keep the best scores. Exact, deterministic, nothing to tune, nothing to argue about in code review. Just slow once the catalog grows past “demo.”

The performance profile

A quick word on how to read that time. Spark, the tool we use to spread work across many machines, hands each machine its 40-million-item slice. We mail a copy of all 450,000 source products to every machine, and each one searches its slice for all 450,000 at the same time. Since the 25 machines run in parallel, the whole run takes about as long as one machine’s share, which lands in the four-to-five-hour range.

The runtime threw me at first, because the napkin math promises better. The whole job is about 115 quadrillion small arithmetic operations, and 400 cores doing this kind of math well should get through it in under three hours. We saw four to five. The missing hour or two goes to the things napkin math always forgets: picking the top 20 out of 40 million scores for every single product, dragging each slice in from storage, and merging everyone’s answers at the end. The arithmetic is never the whole bill. It’s just the part that’s easy to count.

The bigger lesson is the one I confessed earlier, and it’s worth the repeat because it cost real money. The first version ran on hundreds of small, single-core machines. It was both slower and more expensive: each core was starved, because FAISS wants to spread one search across many cores, and the tiny slices meant the fixed costs got paid hundreds of times over. Moving to 25 machines that each use all 16 cores was a large and nearly free win, the kind you wish you’d found in week one instead of week six.

Getting FAISS to run on Spark

FAISS is a C++ library with a Python wrapper and absolutely no idea Spark exists. The pattern that worked is to mail a copy of the source fingerprints to every machine, then have each one build a FAISS index from its own slice of the catalog and search it, using all its cores.

# Broadcast all 450K source vectors once (450K × 128D = 230 MB)
bc_source_ids = spark.sparkContext.broadcast(source_ids)
bc_source_embs = spark.sparkContext.broadcast(source_embeddings)

# Each machine builds an index from its own catalog slice, multi-threaded
def search_shard(catalog_rows):
    faiss.omp_set_num_threads(16)               # use all of the machine's cores
    index = faiss.IndexFlatIP(128)
    index = faiss.IndexIDMap(index)
    ids, embs = prepare_catalog(catalog_rows)   # this machine's ~40M items
    index.add_with_ids(embs, ids)
    # Search all 450K source products in a single pass
    D, I = index.search(bc_source_embs.value, k=20)
    yield from format_results(D, I, bc_source_ids.value)

# Run across the 25 slices, then merge to a global top-20
results = catalog_rdd.mapPartitions(search_shard)

# ... global top-20 deduplication via Window function

Notice the source products get mailed out once, not in pieces. The whole source set is only 230 MB (450,000 fingerprints at 512 bytes each), which is a rounding error on a 64 GB machine. So each machine builds its index a single time and searches all 450,000 in one go. No looping, no rebuilding.

The heartbeat trap. While a machine is heads-down searching, it can’t pause to check in. FAISS keeps each one busy for hours, but Spark expects a hello every two minutes, and hearing nothing, decides the machine has died and shoots it. The machine was working harder than ever. It just hadn’t answered its messages. The fix is gloriously boring: raise spark.executor.heartbeatInterval and spark.network.timeout well past the longest a single piece of work could take. I lost a couple of early runs to these phantom deaths before the penny dropped. Boring fixes are the best fixes.

Why bother starting with the slow version

I was itching to skip straight to the clever version. I’m glad I didn’t, and not only because a wiser engineer than me once warned that premature optimization is the root of all evil. Three concrete reasons. The first is that the exact run gives you the true answer to grade against. You can’t claim a faster method finds “99% of the right matches” without first knowing what 100% looks like. The second is that it shakes out the whole pipeline: mailing data to every machine, splitting the catalog, merging results at the end, all of it has to work before you start fiddling with search settings. The third, and the one I appreciated most later, is debugging. When the fast version eventually returned something weird, I could re-run that one product the slow way and tell instantly whether the search was broken or the embeddings were.

That third reason is what bites people who skip this step. Without an exact baseline to check against, “the new index is broken” and “the embeddings are bad” look identical from the outside, and you can lose a week chasing the wrong one. I have lost that week, in a previous life. Not again.

Phase 2: IVF, searching only the nearby clusters

So exact search hit the wall I always knew it would. Every product gets compared against every item in its slice, and at 40 million items per slice, that’s a lot of comparing. The first real speedup I reached for was to stop looking everywhere. Group the catalog into clusters, and for each product, only search the clusters that look promising.

That’s the whole idea behind IVF (it stands for Inverted File index, but the name has never helped anyone understand it). Think of a library. You don’t read every book on every shelf to find one title. You walk to the few shelves most likely to have it, and look only there.

There’s one trap worth flagging, because I nearly fell in it: the grouping step. Do not redo it on every machine. You group the catalog into clusters once, training on a random sample (FAISS quietly subsamples it anyway, so this takes seconds to a couple of minutes), then ship that finished grouping to every machine. Each machine just drops its own items into the right buckets. At search time, a product checks which cluster centers are nearest, picks a handful (a setting called nprobe), and searches only inside those. Redoing the grouping on all 25 machines would be 25 copies of the same work for the privilege of waiting longer.

Here’s the payoff in numbers. We make 16,384 clusters and open the nearest 256 of them per product. That means each product gets compared against 40M × (256/16384) = 625,000 items, about 1.6% of its slice. On paper that’s about 64 times less work. In reality the run is about 8 times faster, not 64, and the gap is worth understanding. Two costs don’t shrink. First, to pick which clusters to open, every product still has to be checked against all the cluster centers. Second, the fixed costs from before, reading the slice from storage and merging results at the end, don’t care how clever the search got. So roughly 8 times is the ceiling for IVF on its own. That ceiling is also the setup for the next trick: compressing the vectors cuts both the comparison cost and the memory shuffling, which is what lets the total drop further.

Picking nlist and nprobe

IVF on its own takes the run from 4–5 hours down to roughly 30–40 minutes. About 8 times faster, which felt like magic right up until I looked at the memory and realized each machine was still hauling around the full, uncompressed fingerprints. Its 20.5 GB hadn’t budged. That’s the next thing to fix.

Phase 3: product quantization, shrinking the vectors

IVF cut down how many comparisons we make. PQ goes after the other lever: how expensive each comparison is, and how much memory the fingerprints eat. It’s lossy compression for vectors, a JPEG for your fingerprints. You trade a little fidelity for a lot of size. The idea is to replace every 128-number fingerprint with a tiny 32-byte stand-in, small enough to be cheap, without throwing away so much detail that the search falls apart.

How PQ works

You chop each 128-number vector into 32 little pieces of 4 numbers each. For each piece, you’ve pre-learned a “menu” of 256 common patterns, and you replace the actual piece with the menu number it most resembles (a single byte, since 256 options fit in one byte). So a fingerprint that used to be 128 numbers becomes 32 menu picks.

Comparing two products then becomes looking up pre-computed scores in a small table instead of doing the full math. Where the original comparison needed 128 multiplications, the compressed one needs 32 quick lookups and 32 additions. Faster, and a lot lighter on memory.

What it saves

That’s about 19 GB saved on every machine, or roughly 480 GB across all 25. Put another way: once the catalog is compressed, the whole billion items take only about 32 GB. Which raises a question I didn’t think to ask at first. If the data suddenly fits almost anywhere, do we still need 25 machines?

How small a cluster can you get away with

The only reason we ever needed 25 machines was memory. 512 GB of float32 vectors has to live somewhere, and no single box holds that. Once PQ shrinks the catalog to about 32 GB of codes (call it 40 GB with the bookkeeping), the entire compressed index fits on one or two machines. The fleet size stops being dictated by the data.

That frees you to pick machines by how fast you want the run to finish rather than by how much data you’re holding. And since IVF+PQ also does far less work per query, a small cluster goes a surprisingly long way. The trade is the usual one: fewer cores, longer run.

All three hold the compressed catalog in memory without breaking a sweat. The only thing that changes as you go down the rows is how much patience you need.

One catch decides how far you can really push this. Training the PQ menus and encoding the catalog reads all 512 GB of float32 vectors, so if you rebuild the index from scratch on every run, a tiny fleet just spends its time doing that and the savings cap out. The fix is to build the compressed index once and store the 32 GB of codes. After that, every run loads the small thing instead of the big one, the search itself is cheap, and a handful of machines genuinely does what 25 did. (If you put much bigger slices on each machine, nudge nlist up to match the larger slice. The recall doesn’t change.)

Putting it together: IVF+PQ

With both tricks stacked, the pipeline lands somewhere completely different.

The catch is that roughly 3-point drop in accuracy. IVF+PQ misses a few of the true top-20 matches for some products, and in some projects that would be a deal-breaker. In mine it was fine, and it’s worth being honest about why. That ~97% isn’t a number you get for free, either: compressing real product embeddings, which tend to be more tangled than the tidy test data in benchmarks, can cost more than this, and the usual way to firm it up is a smarter rotation before compressing plus a quick exact re-check of the shortlist (FAISS calls that IndexRefineFlat). I didn’t even need those, for three reasons. A later step in our system uses a language model to re-rank the candidates and catch borderline cases the embeddings miss. We return the top 20, not the single best, so missing match number 18 rarely changes anything. And I kept watch: anything that drops below 95% overlap with the exact answer gets flagged before it ships.

That last bit earned its own section.

Checking I didn’t break recall

Switching from exact to approximate is a one-way door, so before I flipped the switch I measured the gap carefully on a real sample. Measure twice, delete the flat index once.

# Sample 5,000 source products across volume tiers
sample_source = stratified_sample(all_source, n=5000)

# Run exact search (ground truth)
exact_results = faiss_flat_search(sample_source, catalog)  # top-20 per source product

# Run approximate search
approx_results = faiss_ivfpq_search(sample_source, catalog)  # top-20 per source product

# Compute recall@20: how many of the true top-20 appear in approximate top-20?
for p in sample_source:
    exact_set = set(exact_results[p][:20])
    approx_set = set(approx_results[p][:20])
    recall = len(exact_set & approx_set) / 20
    recalls.append(recall)

mean_recall = np.mean(recalls)  # Target: ≥ 0.95

The shape of that curve is the whole point. Opening a few more clusters (nprobe) buys you a lot of accuracy at first, then less and less, while the cost keeps climbing at a steady rate. So there’s a sweet spot, a knee in the curve, where you’ve bought most of the accuracy you’re going to get without paying through the nose for the last sliver. For us that knee sits around nprobe=256. Chasing the final fraction of a percent past it is a great way to spend an afternoon and a cloud budget on nothing.

Build the index once, search in one pass

One choice quietly saves a lot of time, and the obvious instinct gets it wrong, so it’s worth being explicit. The instinct is to feed the source products through in batches, rebuilding the search each round. Don’t. Rebuilding the index every batch is a time-honored way to turn a fast job slow. The whole source set is 230 MB, so you mail it out once. And building the index, grouping into clusters and then dropping in 40 million items, is the expensive setup step, so you do it once per machine and reuse it for the whole run.

# Train the cluster grouping ONCE on a sample, then ship the empty
# trained index to every machine.
quantizer = train_ivfpq(sample, nlist=16384, m=32, nbits=8)  # seconds to minutes
bc_index = spark.sparkContext.broadcast(serialize(quantizer))

# Send all 450K source products out once (230 MB)
bc_src = spark.sparkContext.broadcast(source_embeddings)
def search_shard(catalog_rows):
    faiss.omp_set_num_threads(16)
    index = deserialize(bc_index.value)     # empty, already grouped
    ids, embs = prepare_catalog(catalog_rows)
    index.add_with_ids(embs, ids)           # built once, no retraining
    index.nprobe = 256
    D, I = index.search(bc_src.value, k=20) # all 450K in a single pass
    yield from format_results(D, I)

results = catalog_rdd.mapPartitions(search_shard)   # one pass, then merge

The batch-and-rebuild version multiplies the most expensive setup step by the number of batches for exactly zero benefit. The merge at the end is cheap, too: each of the 450,000 products collects 20 candidates from each of the 25 machines, which is about 225 million rows to sift into a final top-20. Nothing here needs breaking up.

Where the embeddings come from

Stepping back to the start of the pipeline for a moment. The fingerprints come from a single fine-tuned XLM-Roberta model. Both products in any comparison run through the same model, which reads the text out to 768 numbers, squeezes that down to 128, and normalizes the result so all the fingerprints sit on the same scale. Same model on both sides, same number space, so a plain score between any two fingerprints actually means something.

For production, we export the model to ONNX, a portable format you can run without dragging the whole training setup along:

# Export product encoder to ONNX (dynamo exporter, dynamic batch)
torch.onnx.export(
    product_encoder,
    dummy_input,
    "product_encoder.onnx",
    opset_version=18,                # opset 18+ recommended with the dynamo exporter
    dynamo=True,                     # default in PyTorch 2.9+; explicit for clarity
    dynamic_axes={"input_ids": {0: "batch"}, "attention_mask": {0: "batch"}},
)

# Inference: ~50 sequences/second per core on CPU
session = onnxruntime.InferenceSession("product_encoder.onnx")
embeddings = session.run(None, {"input_ids": ids, "attention_mask": mask})[0]
# → (batch_size, 128) float32 array, L2-normalized

Going through ONNX instead of calling the model directly was deliberate. There’s no GPU needed, so the embeddings run on the same ordinary machines as the rest of the job, and we skip standing up a separate, pricey GPU cluster just for this. A 16-core machine handles several hundred products a second, and we only re-embed the products that actually changed rather than the whole billion every run, so keeping up with daily edits is easy. ONNX also keeps the peace between the model’s libraries and Spark’s Python environment, two things that otherwise get along like cats in a sack. (That fight cost me an afternoon twice before I gave up and switched.)

Lessons, mostly learned the hard way

If you scrolled straight here, which is exactly what I’d have done, this is the whole post squeezed into the things I’d actually say to the next person handed this ticket.

Start with the embarrassingly slow version. Exact brute force is the one nobody brags about, and it’s the only thing that tells you the truth: what 100% looks like, whether the pipeline holds together end to end, and later, whether a weird result is the search’s fault or the embeddings’. Skip it and you debug blind. I’ve debugged blind. Zero stars, would not recommend.

Buy cores, not boxes. Fewer, beefier machines that each use all their cores beat a swarm of starved single-core ones every time. Letting FAISS stretch across all 16 cores was the biggest speedup I found and, naturally, the last one I found. The 32 GB boxes I started with were neither cheap enough nor useful enough. They were just wrong.

Spark has trust issues. A machine elbow-deep in a multi-hour search can’t stop to check in, and Spark, hearing nothing for two minutes, declares it dead and shoots it. Raise the heartbeat and network timeouts above your longest piece of work and the phantom funerals stop. It’s a how-long-the-work-takes problem, not a threading one, since FAISS lets go of Python’s lock while it searches.

nprobe is a dial, not a light switch. Turning it up buys accuracy fast, then slowly, then barely at all, while the bill rises the whole way. Find the knee, around 256 for us, and walk away. The last half a percent always costs more than it’s worth.

Do the expensive thing exactly once. The source set is tiny, so broadcast it whole. Building the index is the pricey part, so build it once per machine and reuse it. Rebuilding it inside a batch loop is a time-honored way to take a fast job and lovingly hand-craft it into a slow one.

Let something downstream cover for you. Our results feed a language model that re-ranks them, so 97% recall in the search is plenty and the next stage mops up the misses. Without a safety net like that, you’d be chasing a much higher number and paying for every point of it.

Memory is money. Squeezing each machine’s index from 20.5 GB down to about 1.3 GB didn’t just speed up the search, it freed roughly 480 GB across the cluster and shrank the whole catalog small enough to serve from a box or two. Cheaper hardware is a feature, not a footnote.

That’s the arc. Exact search over a billion items was correct and unbearably slow. IVF made it roughly 8 times faster by politely ignoring the 98% of the catalog that was never going to match. PQ pushed it to about 12 times faster overall and 16 times smaller in memory, all for giving up around 3 points of accuracy I genuinely never missed. Once you’ve proven that 97% lands the same business result as 100%, you stop walking back through that door. The measuring was a one-time tax. The savings show up on every single run after, which is the closest thing to a free lunch this job has ever served me.