What it actually does
Once a model has turned a million documents into a million vectors, answering a query means finding the nearest few — and comparing against all million is too slow. So give up exactness: build a structure that looks at a few thousand candidates instead of a million and usually finds the same answer. The word "usually" is the entire subject. An ANN index does not fail by being slow, it fails by silently returning the second-best neighbour, and how often it does so is a dial you set rather than a property you measure once.
- The embedding space is meaningful, so geometric proximity corresponds to whatever relevance means here — the index optimizes proximity and nothing else
- The distance metric matches the one the embedding model was trained under
- Approximate answers are acceptable, and the recall level has been chosen deliberately rather than inherited from a default
- The vector distribution is stable enough that fitted centroids and codebooks remain valid between rebuilds
- Queries at serving time resemble the workload the parameters were tuned on
Loss function
This is a constrained optimization rather than a loss being minimized by gradient descent, and writing it that way makes the trade explicit: maximize recall subject to latency and memory. Every knob in every index moves along that frontier and nothing escapes it. Two consequences follow. Recall is measured against brute force, so you need the exact answer on a sample to know how well the approximation is doing — and teams routinely skip that, which means they have no idea what their retrieval is actually returning. And because the constraint is on expected latency, the tail is not controlled at all: an index tuned to a mean of five milliseconds can have a ninety-ninth percentile many times that.
- the true k nearest neighbours under brute force — the ground truth the index is approximating
- what the index actually returns; recall is the overlap between the two
- the latency budget, which is a hard constraint set by the product rather than by the algorithm
- the memory budget — usually the binding constraint at scale, and what forces quantization
- a query workload representative of production; recall measured on the wrong workload is meaningless
How it is minimized
Graph traversal with a bounded candidate frontier (HNSW), or coarse clustering plus compressed residuals (IVF-PQ)
- the candidate frontier size — the single dial trading recall against latency at query time
- subquantizer count in PQ: the vector is split into M pieces, each coded to one byte
- the k-th centroid of the m-th subspace codebook, fitted by k-means over the training vectors
- which centroid vector x uses in subspace m — one byte, which is the entire compression
- how many coarse clusters IVF actually visits; its equivalent of efSearch
Two families with genuinely different trade profiles, and choosing between them is the decision that matters. HNSW builds a multi-layer proximity graph where upper layers are sparse long-range links and lower ones are dense local links, so a search descends from a coarse hop to a fine walk — greedy traversal with a bounded frontier, giving logarithmic hops and excellent recall, at the cost of holding full vectors plus graph edges in memory. That memory cost is the whole story: it is usually several times the raw vectors, which at a billion vectors is decisive. IVF-PQ attacks memory instead. Coarse k-means partitions the space so a query probes only a few cells, and product quantization splits each vector into subvectors coded to one byte each, compressing a 768-dimensional float vector into perhaps 96 bytes — a 32-fold reduction. Distances are then computed against codebook entries via a precomputed lookup table, so the arithmetic is table lookups rather than floating-point. The price is that the distances are approximate twice over, from the partitioning and from the compression, and recall falls accordingly.
M (HNSW degree)
Edges per node per layer. Drives both memory and recall; the main build-time quality lever (16 to 64)
efConstruction
Frontier size during build. Higher gives a better graph and a slower, one-off build (100 to 500)
efSearch
Frontier size at query time — the dial you actually turn in production, adjustable without a rebuild (50 to 500)
nlist (IVF)
Number of coarse clusters; roughly the square root of the collection size is the usual starting point (4 * sqrt(N) to 16 * sqrt(N))
nprobe
Clusters visited per query. IVF’s recall-latency dial, also adjustable at query time (8 to 128)
PQ subquantizers
Bytes per compressed vector. Must divide the dimension; sets the compression ratio directly (8 to 96)
rerank depth
How many approximate candidates to rescore with exact distances — cheap, and recovers most of the lost recall (2x to 10x k)
Nothing converges, because nothing is trained — which is precisely why the failure modes are unlike the rest of this reference. There is no loss curve and no validation metric that moves; an index either returns good neighbours or quietly does not, and the only way to know is to compare against brute force on a sample. Three characteristic failures. Recall degrades as the collection grows, because parameters tuned at a million vectors are wrong at a hundred million and nothing reports the drift. Deletions in a graph index leave tombstones that fragment connectivity, so a long-lived index with heavy churn slowly loses recall until it is rebuilt. And distribution shift in the embeddings — a new model version, or a new content domain — invalidates IVF centroids and PQ codebooks that were fitted to the old distribution, which shows up as a recall drop with no other symptom anywhere in the system.
HNSW search is O(log N) hops with efSearch distance computations per hop, and memory is the full vectors plus roughly M edges per node — typically several times the raw vector size. IVF-PQ search is O(nprobe · N/nlist) table lookups, with memory of one byte per subquantizer per vector plus the codebooks. Build is the expensive half for both: HNSW is O(N log N) distance computations and is genuinely slow at scale.
How, where, and why it is used
Time-Series Forecasting
Not applicableIt retrieves neighbours in a vector space and has no notion of time, horizon or extrapolation. Nearest-neighbour forecasting exists as a method, but the model there is the k-nearest-neighbours entry and the index would merely be an implementation detail inside it — the index itself forecasts nothing.
Anomaly Detection
AdaptedDistance to the k-th nearest neighbour is a classical anomaly score, and an index makes it affordable on a collection where brute force would not be. Embed everything, index the normal population, and score each new point by how far it sits from its neighbours in that index.
- Near-duplicate and novelty detection over large document or image collections
- Embedding-space outlier detection where the collection is too large for exact k-NN
- Fraud-ring detection on entity embeddings, where isolation in the embedding space is the signal
- Content moderation, where distance from known-good clusters is the first-pass filter
It makes a method that was always sound but computationally infeasible practical at scale — distance-based outlier detection over a hundred million vectors is simply not available without an index. The catch is specific and important: an anomaly is by definition far from everything, which is exactly the regime where approximate search is least reliable, because the graph traversal or cluster probe is guided by proximity that the anomalous point does not have. So the index under-performs precisely on the points that matter, and recall measured on typical queries badly overstates recall on anomalous ones. Measure it separately on known outliers, or the detector is weaker than the benchmark suggests.
- Normalize vectors if using cosine distance, since most indexes assume inner product on unit vectors and silently misrank otherwise
- Measure recall specifically on known anomalous queries, not just on typical ones — the two differ substantially and only one matters here
- Raise efSearch or nprobe for scoring rather than for ordinary retrieval; the extra latency is affordable in a batch detector
- Build the index from a verified-normal population, since an index containing anomalies gives them neighbours and hides them
Precision and recall at a fixed alert budget, and separately the index recall measured on anomalous queries against brute force — that second number is the one that reveals whether the approximation is costing you detections, and it is almost never reported.
- Approximate search failing hardest on exactly the isolated points the detector is looking for
- An index built on data containing anomalies, which gives them neighbours and suppresses their scores
- Embedding drift invalidating the index without any error surfacing
- Quoting recall from typical queries and assuming it holds for outliers, which it does not
Optimization
Primary fitThis entry is an optimization problem wearing a retrieval hat, and it is the cleanest example in the reference of a deliberately chosen accuracy-for-resources trade. Maximize recall subject to latency and memory; every parameter moves along that frontier and none escapes it. The two families are different attacks on the same problem — HNSW spends memory to buy recall, IVF-PQ spends recall to buy memory — and the structures themselves are worth studying independently: a navigable small-world graph is a probabilistic skip list generalized to metric space, and product quantization is vector quantization applied to subspaces so that codebook size stays manageable.
- The recall-latency-memory frontier as an explicit constrained optimization, with efSearch and nprobe as the runtime dials
- Navigable small-world graphs as metric-space skip lists — logarithmic search from a layered structure
- Product quantization as subspace vector quantization, avoiding the exponential codebook a joint quantizer would need
- Two-stage retrieve-then-rerank, the standard pattern for trading a cheap approximate pass against an expensive exact one
Worth studying because the trade is explicit and tunable at runtime rather than baked in at training — you can move along the frontier per query if you want, which almost nothing else in this reference allows. The retrieve-then-rerank pattern generalizes far beyond vector search: use a cheap approximate method to shortlist and an expensive exact one to order, and you recover most of the lost quality for a fraction of the cost. Product quantization is also a genuinely elegant piece of engineering — splitting into subspaces is what keeps the codebook from being exponentially large, and the precomputed distance table turns floating-point arithmetic into table lookups.
- Always measure recall against brute force on a sample; without that number the whole frontier is invisible and you are tuning blind
- Rerank the top candidates with exact distances — it is cheap and recovers most of the recall that quantization costs
- Tune on a query workload that matches production, since recall on random queries differs substantially from recall on real ones
- Track the latency tail rather than the mean; the constraint is expressed as an expectation and the tail is not controlled by it
Plot recall against latency across the parameter sweep and pick a point on the curve deliberately rather than accepting a default. Report the ninety-ninth percentile alongside the mean, and re-measure both after any collection growth, since a point chosen at one scale is not the same point at ten times the size.
- Never measuring recall at all, which is the most common failure and leaves the system silently degraded
- Tuning on random queries rather than production ones, which measures the wrong distribution
- Optimizing the mean latency while the tail is what users actually experience
- Assuming parameters transfer across collection sizes, which they do not
Recommendation & Ranking
Primary fitThe retrieval half of every two-tower system. The item tower embeds the catalogue once into an index; at request time the user tower produces one query vector and the index returns a few hundred candidates, which a heavier cross-feature ranker then orders. Without the index the two-tower architecture has no reason to exist, because scoring the whole catalogue per request is exactly what it was designed to avoid.
- Candidate generation in large-scale recommenders, where the catalogue is millions of items and the budget is milliseconds
- Related-item and similar-content retrieval directly from embedding proximity
- Semantic and hybrid search, combining vector recall with lexical matching
- Retrieval-augmented generation, where the index is the retrieval step and its recall caps the whole system
It is what makes embedding-based retrieval viable at all: a two-tower model exists precisely so that item vectors can be precomputed and searched, and the index is the half that delivers on that promise. The honest framing is that it is infrastructure rather than intelligence — it optimizes geometric proximity and nothing about relevance, so the ranker downstream does the actual relevance work, and a beautiful index over mediocre embeddings retrieves mediocre things quickly. Its recall is also a hard ceiling on everything downstream: a candidate the index misses cannot be ranked, and in RAG that shows up as an answer the model could not have given.
- Normalize vectors and match the index metric to the one the embedding model was trained with — a mismatch silently misranks everything
- Filter before searching where possible; post-filtering a fixed-size result set can return almost nothing after a selective filter
- Rebuild or retrain quantizers whenever the embedding model changes, since codebooks are fitted to the old distribution
- Size the candidate set from measured ranker headroom rather than by convention, since recall at that depth is what caps the system
Recall at the candidate depth against brute force, and separately the end-to-end metric with and without the index — the gap between exact and approximate retrieval is the cost of the approximation, and it is the number that justifies the parameter choice. Measure latency at the ninety-ninth percentile, which is what determines whether the request budget holds.
- A metric mismatch between the embedding model and the index, which misranks everything with no error
- Post-filtering a fixed result set, which after a selective filter can return far fewer items than requested or none
- Stale codebooks after an embedding model upgrade, which degrades recall with nothing logged
- Treating index recall as a solved problem rather than as the ceiling on every downstream metric
Natural Language
Primary fitThe retrieval in retrieval-augmented generation. Documents are chunked, embedded and indexed; a query is embedded and the index returns the chunks that condition the generation. Everything the model says about your data flows through this step, which makes index recall a correctness property of the whole system rather than a performance detail.
- RAG over document collections, where retrieval quality bounds answer quality absolutely
- Semantic search and question answering over enterprise corpora
- Deduplication and clustering of large text collections
- Long-term memory for conversational agents, retrieved by embedding similarity
It is load-bearing in a way that is easy to under-appreciate: in RAG, a chunk the index fails to retrieve is a fact the model cannot use, and the failure looks like the model not knowing rather than like a retrieval bug. That misattribution is the single most common diagnostic error in RAG systems. Hybrid retrieval — combining vector search with lexical matching — consistently outperforms either alone, because embeddings are weak exactly where exact terms matter, such as identifiers, part numbers and names. The other honest limit is that chunking decisions matter more than index parameters in most deployments, and no amount of tuning compensates for chunks that split a fact across a boundary.
- Chunk with overlap and at a size matched to the embedding model’s context, since a fact split across a boundary is unretrievable at any recall
- Combine vector and lexical retrieval, because embeddings underperform badly on exact identifiers and rare terms
- Store metadata alongside vectors and filter before search rather than after
- Re-embed and rebuild whenever the embedding model changes; vectors from two model versions are not comparable at all
Retrieval recall on a labelled query set against brute force, then end-to-end answer quality — measuring only the second makes retrieval failures look like generation failures, which is how RAG systems get debugged in the wrong place for weeks.
- Attributing a retrieval miss to the generator, which is the most common and most expensive misdiagnosis in RAG
- Chunk boundaries splitting facts, which no index parameter can recover
- Vector-only retrieval failing on identifiers and rare terms where lexical matching would have succeeded
- Mixing vectors from two embedding model versions in one index, which is silently meaningless
Computer Vision
ViableVisual search over image embeddings: index the collection once, embed a query image, retrieve the nearest. The same structure and the same trades, with the difference that image collections tend to be larger and the vectors wider, so the memory constraint binds earlier and quantization is usually mandatory rather than optional.
- Reverse image search and visual product discovery over large catalogues
- Near-duplicate detection and copyright matching
- Face and person re-identification against a gallery
- Retrieval-augmented vision, conditioning a model on similar retrieved examples
The scale is what makes it interesting here: billion-image collections with wide embeddings are exactly where IVF-PQ earns its place, because holding full vectors is simply not affordable and a 32-fold compression turns an impossible deployment into a routine one. The caveats are the general ones plus a sharper version of the metric problem — vision embeddings vary widely in whether they are trained for cosine or Euclidean geometry, and a mismatch is both easy to introduce and invisible once introduced.
- Quantize aggressively at billion scale; full vectors are usually not affordable and PQ is the reason the deployment exists
- Rerank the top candidates with exact distances, which recovers most of the recall quantization costs
- Match the metric to how the vision encoder was trained, since the two conventions are both common and the mismatch is silent
- Keep a brute-force sample for continuous recall measurement, since nothing else will tell you the index has drifted
Recall at k against brute force on a held-out sample, plus task metrics for the downstream application. Memory per vector is a first-class number here rather than an implementation detail, because it determines whether the deployment is possible at all.
- A metric mismatch with the vision encoder, which misranks silently
- Over-aggressive quantization collapsing fine-grained distinctions that the task depends on
- Skipping reranking and paying the full quantization recall cost for no reason
What it costs to run
No training in the gradient sense. Building is the cost: HNSW is O(N log N) distance computations and genuinely slow — hours for a hundred million vectors — while IVF-PQ needs a k-means pass and codebook fitting over a sample, which is faster but must be redone whenever the embedding distribution shifts.
Single-digit milliseconds per query for a well-tuned index at millions of vectors, and adjustable at runtime through efSearch or nprobe without rebuilding. Memory is the harder constraint: HNSW typically holds several times the raw vector size, while IVF-PQ can hold a thirty-second of it, and that difference decides which is deployable at scale.
Rebuilds are driven by two independent things: collection churn, since deletions fragment graph connectivity, and embedding model changes, which invalidate every vector and every codebook at once. The second is a full re-embed and rebuild, not an update.
- Measure recall against brute force on a sampled query set continuously — it is the only signal that the index is degrading, and nothing else will surface it
- Track the latency tail rather than the mean, since the constraint is an expectation and the tail is what users experience
- Monitor the deleted-vector fraction on graph indexes; tombstones fragment connectivity and recall falls slowly as they accumulate
- Alert on embedding-distribution shift against the distribution the codebooks were fitted to, which is what invalidates IVF-PQ silently
- The distance metric is a contract between the embedding model and the index. A cosine model on a Euclidean index misranks everything and raises no error anywhere
- Filtered search is the hardest thing these structures do. Post-filtering a fixed-size result set can return almost nothing after a selective filter, and pre-filtering usually needs index support rather than a wrapper
- Deletions are not really deletions. Most graph indexes tombstone, so recall degrades with churn until a rebuild happens
- An embedding model upgrade invalidates the entire index, including codebooks. Mixing vectors from two model versions is silently meaningless rather than merely inaccurate
- Parameters tuned at one collection size do not transfer to ten times the size, and nothing reports that they have stopped being right
- Build memory is often several times steady-state memory, which turns a rebuild into a capacity problem rather than a background job
Pros and cons, in context
Turns an impossible search into a millisecond one
Sub-linear retrieval over hundreds of millions of vectors is what makes every embedding model in this reference usable in production. Without it, two-tower retrieval and RAG have no deployment story at all.
The accuracy-speed trade is explicit and adjustable at runtime
efSearch and nprobe move along the recall-latency frontier per query without a rebuild, which almost nothing else here allows. You can spend more compute on a query you care about.
Quantization makes billion-scale affordable
Product quantization compresses a 768-dimensional float vector into about 96 bytes, which is the difference between a possible and an impossible deployment. It is the reason billion-vector collections are routine rather than exotic.
Nothing is trained, so nothing overfits
No labels, no training run, no generalization gap. The index is deterministic given the vectors and its parameters, which makes it far easier to reason about than the models feeding it.
Fails silently by construction
A degraded index returns plausible but wrong neighbours with no error anywhere. The only defence is continuous recall measurement against brute force, which teams routinely skip and therefore never know what their retrieval is doing.
Filtered search is genuinely hard
Post-filtering a fixed result set can return almost nothing after a selective filter, and pre-filtering needs real index support. This surprises people more than any other property of vector search.
Tightly coupled to the embedding model
A model upgrade invalidates every vector and every codebook simultaneously — a full re-embed and rebuild, not an incremental update. That coupling makes model iteration substantially more expensive than it looks.
Memory is usually the binding constraint
HNSW holds several times the raw vector size, which at a billion vectors decides the architecture for you. Choosing IVF-PQ instead is choosing to pay in recall, and that choice should be made explicitly.
It optimizes proximity, not relevance
A perfect index over mediocre embeddings retrieves mediocre things very quickly. It is infrastructure, and mistaking a retrieval quality problem for an index tuning problem wastes a great deal of time.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""A navigable small-world graph and product quantization, written out.
Two structures, two attacks on the same constrained problem:
maximize recall subject to latency <= tau and memory <= M
HNSW spends MEMORY to buy recall: a layered proximity graph, searched greedily
with a bounded frontier. IVF-PQ spends RECALL to buy memory: coarse clusters
so a query visits few cells, plus per-subspace quantization so a 768-float
vector becomes ~96 bytes.
Brute force is here too, because it is the ground truth every recall number
is measured against - and a system that never measures it has no idea what
its retrieval is returning.
Plain loops, no libraries.
"""
import math
import random
SEED = 3
def squared_distance(a, b):
return sum((x - y) ** 2 for x, y in zip(a, b))
def brute_force(vectors, query, k):
"""The exact answer. Not a fallback - the GROUND TRUTH.
Every recall figure quoted about an index is measured against this on a
sample. Without it the accuracy-speed frontier is invisible and you are
tuning blind.
"""
scored = [(squared_distance(vector, query), index) for index, vector in enumerate(vectors)]
scored.sort()
return [index for _, index in scored[:k]]
def recall_at_k(approximate, exact):
"""The overlap between what the index returned and the truth."""
return len(set(approximate) & set(exact)) / len(exact)
# --------------------------------------------------------------------------
# HNSW: a proximity graph searched greedily with a bounded frontier
# --------------------------------------------------------------------------
def assign_layer(rng, level_multiplier):
"""Geometric layer assignment.
Exactly the skip-list trick, lifted into metric space: most nodes live
only at the bottom, a few reach higher, and the sparse upper layers give
the long-range hops that make the search logarithmic.
"""
return int(-math.log(max(rng.random(), 1e-12)) * level_multiplier)
def search_layer(vectors, graph, query, entry_points, ef, layer):
"""Greedy best-first search with a frontier of at most `ef`.
ef is THE dial. It bounds how many candidates are ever examined, which
bounds latency - and bounds recall at the same time, in the same motion.
Nothing about this structure lets you have both.
"""
visited = set(entry_points)
candidates = [(squared_distance(vectors[p], query), p) for p in entry_points]
candidates.sort()
best = list(candidates)
while candidates:
distance, current = candidates.pop(0)
if best and distance > best[-1][0] and len(best) >= ef:
break # nothing closer can be reached from here
for neighbour in graph[layer].get(current, []):
if neighbour in visited:
continue
visited.add(neighbour)
neighbour_distance = squared_distance(vectors[neighbour], query)
if len(best) < ef or neighbour_distance < best[-1][0]:
candidates.append((neighbour_distance, neighbour))
candidates.sort()
best.append((neighbour_distance, neighbour))
best.sort()
if len(best) > ef:
best.pop()
return best
def build_hnsw(vectors, max_degree=16, ef_construction=100, level_multiplier=1.0):
"""Insert one at a time, descending from the top layer each time."""
rng = random.Random(SEED)
graph = {}
entry_point = None
top_layer = 0
for index, vector in enumerate(vectors):
layer = assign_layer(rng, level_multiplier)
for level in range(layer + 1):
graph.setdefault(level, {}).setdefault(index, [])
if entry_point is None:
entry_point, top_layer = index, layer
continue
current = [entry_point]
# Coarse descent through the sparse upper layers, then a careful
# search at each layer the new node will live on.
for level in range(top_layer, layer, -1):
current = [search_layer(vectors, graph, vector, current, 1, level)[0][1]]
for level in range(min(layer, top_layer), -1, -1):
found = search_layer(vectors, graph, vector, current, ef_construction, level)
current = [node for _, node in found]
for _, neighbour in found[:max_degree]:
graph[level][index].append(neighbour)
graph[level][neighbour].append(index)
# Degree pruning: an unbounded degree destroys both the
# memory budget and the search, so links are capped.
if len(graph[level][neighbour]) > max_degree:
graph[level][neighbour] = sorted(
graph[level][neighbour],
key=lambda n: squared_distance(vectors[n], vectors[neighbour]),
)[:max_degree]
if layer > top_layer:
entry_point, top_layer = index, layer
return graph, entry_point, top_layer
def search_hnsw(vectors, graph, entry_point, top_layer, query, k, ef_search):
"""Descend coarsely, then search the bottom layer with the real frontier."""
current = [entry_point]
for level in range(top_layer, 0, -1):
current = [search_layer(vectors, graph, query, current, 1, level)[0][1]]
found = search_layer(vectors, graph, query, current, ef_search, 0)
return [index for _, index in found[:k]]
# --------------------------------------------------------------------------
# IVF-PQ: coarse clusters plus per-subspace quantization
# --------------------------------------------------------------------------
def kmeans(vectors, num_clusters, iterations=10):
"""Lloyd's algorithm. Used twice: for coarse cells and for codebooks."""
rng = random.Random(SEED)
centroids = [list(vectors[rng.randrange(len(vectors))]) for _ in range(num_clusters)]
for _ in range(iterations):
assignments = [
min(range(num_clusters), key=lambda c: squared_distance(vector, centroids[c]))
for vector in vectors
]
for cluster in range(num_clusters):
members = [vectors[i] for i, a in enumerate(assignments) if a == cluster]
if not members:
continue
centroids[cluster] = [
sum(member[d] for member in members) / len(members)
for d in range(len(members[0]))
]
return centroids
def train_pq(vectors, num_subquantizers, codebook_size=256):
"""Split the vector into M pieces and quantize each SEPARATELY.
The split is the whole trick. A joint quantizer covering the full space
would need an exponentially large codebook; M independent codebooks of
256 entries each represent 256^M distinct vectors using M bytes.
"""
dimension = len(vectors[0])
if dimension % num_subquantizers != 0:
raise ValueError("subquantizer count must divide the dimension")
subspace_width = dimension // num_subquantizers
codebooks = []
for m in range(num_subquantizers):
lo, hi = m * subspace_width, (m + 1) * subspace_width
codebooks.append(kmeans([vector[lo:hi] for vector in vectors], codebook_size))
return codebooks
def encode_pq(vector, codebooks):
"""One byte per subspace. A 768-float vector becomes len(codebooks) bytes."""
subspace_width = len(vector) // len(codebooks)
codes = []
for m, codebook in enumerate(codebooks):
piece = vector[m * subspace_width:(m + 1) * subspace_width]
codes.append(min(range(len(codebook)), key=lambda c: squared_distance(piece, codebook[c])))
return codes
def pq_distance_table(query, codebooks):
"""Precompute query-to-centroid distances ONCE per query.
After this, scoring a compressed vector is M table lookups and M adds -
no floating-point distance computation at all. That substitution is what
makes IVF-PQ fast as well as small.
"""
subspace_width = len(query) // len(codebooks)
return [
[squared_distance(query[m * subspace_width:(m + 1) * subspace_width], centroid)
for centroid in codebook]
for m, codebook in enumerate(codebooks)
]
def search_ivfpq(query, coarse_centroids, inverted_lists, codes, codebooks, k, nprobe):
"""Probe the nearest nprobe cells, score their members by table lookup."""
cell_order = sorted(
range(len(coarse_centroids)),
key=lambda c: squared_distance(query, coarse_centroids[c]),
)
table = pq_distance_table(query, codebooks)
scored = []
for cell in cell_order[:nprobe]:
for index in inverted_lists[cell]:
# M lookups and M adds. No distance is ever computed here.
approximate = sum(table[m][code] for m, code in enumerate(codes[index]))
scored.append((approximate, index))
scored.sort()
return [index for _, index in scored[:k]]
def rerank(vectors, query, candidates, k):
"""Rescore approximate candidates with EXACT distances.
Cheap - it touches only the shortlist - and it recovers most of the
recall that quantization costs. The standard retrieve-then-rerank
pattern, and skipping it pays the full approximation penalty for nothing.
"""
scored = [(squared_distance(vectors[index], query), index) for index in candidates]
scored.sort()
return [index for _, index in scored[:k]]
Brute force is O(N * d) per query; HNSW is O(log N) hops times efSearch distance computations; IVF-PQ is O(nprobe * N/nlist) table lookups. Illustrative, not a measured benchmark: the sort-on-every-insertion in the frontier here makes the literal HNSW search quadratic in efSearch, which a heap fixes in the next stage.
How this language actually wants it written.
"""The same structures, with the frontier as a heap and the metric as a type.
Three things change. The candidate frontier becomes a pair of heaps rather
than a repeatedly sorted list, which is what makes the bounded-frontier
search actually bounded in cost. The distance metric becomes an explicit
contract, because a cosine embedding model against a Euclidean index
misranks everything and raises no error anywhere. And recall measurement
becomes part of the API rather than something a caller is trusted to
remember, since an index that is never measured has no known quality at all.
"""
from __future__ import annotations
import heapq
import math
from dataclasses import dataclass, field
from enum import Enum
from typing import NamedTuple, Sequence
class Metric(Enum):
"""The contract between the embedding model and the index.
A model trained for cosine similarity searched under Euclidean distance
misranks everything, silently and completely. Making this explicit is the
single highest-value piece of typing in the whole structure.
"""
EUCLIDEAN = 'l2'
INNER_PRODUCT = 'ip'
COSINE = 'cosine'
class MetricMismatch(ValueError):
"""Raised when vectors do not satisfy the chosen metric's assumptions."""
class ShapeMismatch(ValueError):
"""Raised on a dimension violation instead of comparing garbage."""
class UnmeasuredIndex(RuntimeError):
"""Raised when an index is queried before its recall has been measured.
Deliberately strict. An ANN index fails by returning plausible wrong
answers, so an unmeasured one has unknown quality - and 'unknown quality'
is the normal state of vector search in production.
"""
@dataclass(frozen=True)
class IndexConfig:
"""Frozen so the metric cannot change under a built index."""
dimension: int
metric: Metric = Metric.COSINE
max_degree: int = 16
ef_construction: int = 100
level_multiplier: float = 1.0
def __post_init__(self) -> None:
if self.dimension < 1:
raise ShapeMismatch('dimension must be positive')
if self.max_degree < 2:
raise ShapeMismatch('a degree below 2 cannot form a navigable graph')
def estimated_bytes_per_vector(self) -> int:
"""Memory is usually the binding constraint, so make it visible.
Full float32 vectors plus roughly max_degree 32-bit edges per node,
doubled because links are bidirectional. This number is what decides
whether a graph index is deployable at all at a given scale.
"""
return 4 * self.dimension + 2 * 4 * self.max_degree
class Neighbour(NamedTuple):
"""Distance first, so heap ordering is the natural one."""
distance: float
index: int
def normalize(vector: Sequence[float]) -> list[float]:
norm = math.sqrt(math.fsum(value * value for value in vector))
if norm == 0.0:
raise MetricMismatch('a zero vector has no direction; cosine is undefined')
return [value / norm for value in vector]
def distance(a: Sequence[float], b: Sequence[float], metric: Metric) -> float:
"""Lower is nearer, for every metric — so one comparison works throughout.
Inner product is negated for exactly that reason: mixing a
higher-is-better metric into a lower-is-better search is a classic and
completely silent ranking inversion.
"""
if metric is Metric.EUCLIDEAN:
return math.fsum((x - y) ** 2 for x, y in zip(a, b))
dot = math.fsum(x * y for x, y in zip(a, b))
return -dot
class Frontier:
"""The bounded candidate set, as two heaps.
A min-heap of candidates still to expand and a max-heap of the best found
so far, capped at ef. The literal version re-sorted a list on every
insertion, which made a bounded frontier quadratic in its own bound; this
makes each insertion logarithmic, which is what 'bounded' should mean.
"""
__slots__ = ('_ef', '_candidates', '_best')
def __init__(self, ef: int) -> None:
if ef < 1:
raise ShapeMismatch('the frontier must admit at least one candidate')
self._ef = ef
self._candidates: list[Neighbour] = []
self._best: list[tuple[float, int]] = [] # negated distance, for a max-heap
def push(self, neighbour: Neighbour) -> None:
heapq.heappush(self._candidates, neighbour)
heapq.heappush(self._best, (-neighbour.distance, neighbour.index))
if len(self._best) > self._ef:
heapq.heappop(self._best)
def pop_nearest(self) -> Neighbour | None:
return heapq.heappop(self._candidates) if self._candidates else None
@property
def worst_kept(self) -> float:
return -self._best[0][0] if self._best else math.inf
@property
def full(self) -> bool:
return len(self._best) >= self._ef
def results(self) -> list[Neighbour]:
return sorted(Neighbour(-d, i) for d, i in self._best)
@dataclass
class RecallReport:
"""What the index actually returns, measured against brute force.
A dataclass rather than a float because the tail matters: mean recall
hides that a subset of queries is served badly, and in retrieval that
subset is often the interesting one.
"""
mean_recall: float
worst_recall: float
queries_measured: int
ef_search: int
def acceptable(self, floor: float) -> bool:
return self.mean_recall >= floor
class HnswIndex:
"""A layered proximity graph with an explicit metric and measured recall."""
def __init__(self, config: IndexConfig) -> None:
self._config = config
self._vectors: list[list[float]] = []
self._graph: dict[int, dict[int, list[int]]] = {}
self._entry_point: int | None = None
self._top_layer = 0
self._recall: RecallReport | None = None
def add(self, vector: Sequence[float], layer: int) -> int:
"""Layer is supplied rather than drawn, so builds are reproducible."""
if len(vector) != self._config.dimension:
raise ShapeMismatch(
f'expected {self._config.dimension} dimensions, got {len(vector)}'
)
stored = (
normalize(vector)
if self._config.metric is Metric.COSINE
else list(vector)
)
index = len(self._vectors)
self._vectors.append(stored)
for level in range(layer + 1):
self._graph.setdefault(level, {})[index] = []
if self._entry_point is None:
self._entry_point, self._top_layer = index, layer
return index
self._link(index, stored, layer)
if layer > self._top_layer:
self._entry_point, self._top_layer = index, layer
# Any structural change invalidates the measured recall.
self._recall = None
return index
def _link(self, index: int, vector: Sequence[float], layer: int) -> None:
entry = self._entry_point
assert entry is not None
current = [entry]
for level in range(self._top_layer, layer, -1):
current = [n.index for n in self._search_layer(vector, current, 1, level)]
for level in range(min(layer, self._top_layer), -1, -1):
found = self._search_layer(vector, current, self._config.ef_construction, level)
current = [n.index for n in found]
for neighbour in found[:self._config.max_degree]:
self._graph[level][index].append(neighbour.index)
self._graph[level][neighbour.index].append(index)
self._prune(neighbour.index, level)
def _prune(self, node: int, level: int) -> None:
"""Cap the degree. Unbounded degree destroys the memory budget AND
the search, because a hub node makes every traversal expensive."""
links = self._graph[level][node]
if len(links) <= self._config.max_degree:
return
anchor = self._vectors[node]
links.sort(key=lambda n: distance(self._vectors[n], anchor, self._config.metric))
self._graph[level][node] = links[:self._config.max_degree]
def _search_layer(
self, query: Sequence[float], entry_points: Sequence[int], ef: int, level: int
) -> list[Neighbour]:
frontier = Frontier(ef)
visited = set(entry_points)
for point in entry_points:
frontier.push(Neighbour(distance(self._vectors[point], query, self._config.metric), point))
while True:
current = frontier.pop_nearest()
if current is None:
break
# Guard clause: once the nearest unexpanded candidate is further
# than the worst kept result, nothing better is reachable.
if frontier.full and current.distance > frontier.worst_kept:
break
for neighbour in self._graph[level].get(current.index, ()):
if neighbour in visited:
continue
visited.add(neighbour)
frontier.push(
Neighbour(
distance(self._vectors[neighbour], query, self._config.metric),
neighbour,
)
)
return frontier.results()
def search(self, query: Sequence[float], k: int, ef_search: int) -> list[int]:
if self._recall is None:
raise UnmeasuredIndex(
'call measure_recall before querying; an unmeasured index has '
'unknown quality and fails by returning plausible wrong answers'
)
return self._raw_search(query, k, ef_search)
def _raw_search(self, query: Sequence[float], k: int, ef_search: int) -> list[int]:
if self._entry_point is None:
return []
prepared = (
normalize(query) if self._config.metric is Metric.COSINE else list(query)
)
current = [self._entry_point]
for level in range(self._top_layer, 0, -1):
current = [n.index for n in self._search_layer(prepared, current, 1, level)]
return [n.index for n in self._search_layer(prepared, current, ef_search, 0)[:k]]
def brute_force(self, query: Sequence[float], k: int) -> list[int]:
"""The ground truth. Not a fallback — what recall is measured against."""
prepared = (
normalize(query) if self._config.metric is Metric.COSINE else list(query)
)
scored = [
Neighbour(distance(vector, prepared, self._config.metric), index)
for index, vector in enumerate(self._vectors)
]
# nsmallest, not a sort: k is tiny relative to the collection.
return [n.index for n in heapq.nsmallest(k, scored)]
def measure_recall(
self, queries: Sequence[Sequence[float]], k: int, ef_search: int
) -> RecallReport:
"""Mandatory before querying, and worth re-running as the index grows.
Parameters tuned at a million vectors are wrong at a hundred million,
and nothing else in the system will report that they have drifted.
"""
if not queries:
raise ShapeMismatch('recall cannot be measured with no queries')
recalls: list[float] = []
for query in queries:
exact = set(self.brute_force(query, k))
approximate = set(self._raw_search(query, k, ef_search))
recalls.append(len(exact & approximate) / k)
self._recall = RecallReport(
mean_recall=math.fsum(recalls) / len(recalls),
worst_recall=min(recalls),
queries_measured=len(recalls),
ef_search=ef_search,
)
return self._recall
What changed: Three changes. The candidate frontier becomes a pair of heaps instead of a repeatedly sorted list — the literal version re-sorted on every insertion, which made a bounded frontier quadratic in its own bound, so "bounded" meant nothing operationally. The metric becomes an explicit enum, because a model trained for cosine similarity searched under Euclidean distance misranks everything with no error anywhere, and that is the highest-value piece of typing in the whole structure; inner product is negated so that lower is nearer under every metric, since mixing a higher-is-better metric into a lower-is-better search is a classic silent ranking inversion. The third is that recall measurement becomes part of the API rather than a caller’s good intention: querying an unmeasured index raises, which is deliberately strict because an ANN index fails by returning plausible wrong answers and unknown quality is the normal state of vector search in production. Alongside those, the recall report keeps the worst case as well as the mean, memory per vector is exposed as a method since it is usually the binding constraint, the layer is passed in rather than drawn so builds are reproducible, and any structural change invalidates the measurement.
The frontier drops from O(ef^2) per layer search to O(ef log ef), and brute force uses a bounded heap rather than a full sort. Illustrative, not a measured benchmark: the heap change is what makes a large efSearch usable, which is exactly the regime where recall is worth buying.
Faster, and the justification for each change.
"""Batched queries and table-lookup scoring. The distance loop disappears.
The structural observation: a distance computation is a dot product, and a
batch of queries against a batch of candidates is a MATRIX PRODUCT. So the
inner loop that dominates every ANN structure becomes BLAS - and separately,
product quantization removes the floating-point distance entirely by
precomputing a lookup table per query.
Three changes:
1. Brute force becomes one GEMM. It is the ground truth every recall number
is measured against, so making it fast is what makes measurement
affordable enough to actually do.
2. PQ scoring becomes a table gather and a sum along the subquantizer axis:
no distance arithmetic at all, just M lookups per candidate.
3. Codes are stored subquantizer-major, so scanning one subquantizer across
every candidate is a contiguous read - which is what the gather needs.
"""
from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
FLOAT = np.float32
CODE = np.uint8
def brute_force_batch(
vectors: NDArray[np.float32],
queries: NDArray[np.float32],
k: int,
) -> NDArray[np.int32]:
"""Exact k-NN for a batch of queries, as one matrix product.
||q - x||^2 = ||q||^2 - 2 q.x + ||x||^2, and only the middle term depends
on both - so the whole distance matrix is one GEMM plus two broadcasts.
The query norm is constant per row and cannot change the ranking, so it
is dropped entirely.
"""
vector_norms = np.einsum('ij,ij->i', vectors, vectors, optimize=True)
cross = queries @ vectors.T
cross *= -2.0
cross += vector_norms # broadcasts over the query axis
# argpartition, not argsort: k is tiny relative to N, so a selection is
# O(N) where ordering the rest is O(N log N) of discarded work.
candidates = np.argpartition(cross, k, axis=1)[:, :k]
ordered = np.take_along_axis(cross, candidates, axis=1).argsort(axis=1)
return np.take_along_axis(candidates, ordered, axis=1).astype(np.int32)
class ProductQuantizer:
"""Per-subspace quantization with table-lookup scoring.
The split is the whole trick: a joint quantizer over the full space would
need an exponentially large codebook, while M independent codebooks of
256 entries represent 256^M distinct vectors in M bytes.
"""
def __init__(self, dimension: int, num_subquantizers: int, codebook_size: int = 256) -> None:
if dimension % num_subquantizers:
raise ValueError('subquantizer count must divide the dimension')
if codebook_size > 256:
raise ValueError('a codebook above 256 entries no longer fits in a byte')
self._m = num_subquantizers
self._width = dimension // num_subquantizers
# (M, codebook_size, subspace_width), contiguous so the per-subspace
# distance computation below is a dense product.
self._codebooks = np.zeros((num_subquantizers, codebook_size, self._width), dtype=FLOAT)
@property
def bytes_per_vector(self) -> int:
"""The number that decides whether the deployment is possible.
768 float32 dimensions is 3072 bytes; 96 subquantizers is 96. That
32-fold reduction is the entire reason IVF-PQ exists.
"""
return self._m
def encode(self, vectors: NDArray[np.float32]) -> NDArray[np.uint8]:
"""(n, d) -> (M, n) uint8 codes.
SUBQUANTIZER-MAJOR on purpose: scoring scans one subquantizer across
every candidate, and that is a contiguous read only in this layout.
The natural (n, M) layout makes the hot loop strided.
"""
reshaped = vectors.reshape(len(vectors), self._m, self._width)
codes = np.empty((self._m, len(vectors)), dtype=CODE)
for m in range(self._m):
# One GEMM per subspace against its codebook, same expansion as
# brute force: the centroid norm plus the cross term.
centroids = self._codebooks[m]
cross = reshaped[:, m, :] @ centroids.T
cross *= -2.0
cross += np.einsum('ij,ij->i', centroids, centroids, optimize=True)
codes[m] = cross.argmin(axis=1).astype(CODE)
return codes
def distance_table(self, query: NDArray[np.float32]) -> NDArray[np.float32]:
"""(M, codebook_size) of query-to-centroid distances, once per query.
After this, scoring a compressed vector is M lookups and M adds — no
floating-point distance computation at all. That substitution is what
makes IVF-PQ fast as well as small.
"""
pieces = query.reshape(self._m, 1, self._width)
difference = self._codebooks - pieces
return np.einsum('mkd,mkd->mk', difference, difference, optimize=True)
def score(
self, table: NDArray[np.float32], codes: NDArray[np.uint8]
) -> NDArray[np.float32]:
"""codes is (M, n). Returns (n,) approximate distances.
A fancy-index gather along the subquantizer axis followed by a sum:
the whole scoring pass is one gather and one reduction, with no
arithmetic on the vectors themselves anywhere.
"""
rows = np.arange(self._m)[:, None]
return table[rows, codes].sum(axis=0, dtype=FLOAT)
class IvfPqIndex:
"""Coarse cells plus compressed residuals."""
def __init__(self, centroids: NDArray[np.float32], quantizer: ProductQuantizer) -> None:
self._centroids = np.ascontiguousarray(centroids, dtype=FLOAT)
self._quantizer = quantizer
self._lists: list[NDArray[np.int32]] = []
self._codes: list[NDArray[np.uint8]] = []
def search(
self,
query: NDArray[np.float32],
k: int,
nprobe: int,
rerank_depth: int = 0,
vectors: NDArray[np.float32] | None = None,
) -> NDArray[np.int32]:
"""Probe nprobe cells, score by table lookup, optionally rerank."""
cell_distances = self._centroids @ query
cell_distances *= -2.0
cell_distances += np.einsum('ij,ij->i', self._centroids, self._centroids, optimize=True)
cells = np.argpartition(cell_distances, nprobe)[:nprobe]
table = self._quantizer.distance_table(query)
# One scoring pass per probed cell, each a gather over a contiguous
# block of subquantizer-major codes.
scores: list[NDArray[np.float32]] = []
ids: list[NDArray[np.int32]] = []
for cell in cells:
scores.append(self._quantizer.score(table, self._codes[cell]))
ids.append(self._lists[cell])
all_scores = np.concatenate(scores)
all_ids = np.concatenate(ids)
depth = max(k, rerank_depth)
if depth >= len(all_scores):
shortlist = all_ids
else:
shortlist = all_ids[np.argpartition(all_scores, depth)[:depth]]
if rerank_depth and vectors is not None:
# Rescore the shortlist with EXACT distances. Cheap, since it
# touches only the shortlist, and it recovers most of the recall
# that quantization costs.
exact = vectors[shortlist] - query
exact = np.einsum('ij,ij->i', exact, exact, optimize=True)
return shortlist[np.argsort(exact)[:k]].astype(np.int32)
return shortlist[:k].astype(np.int32)
def recall_at_k(
approximate: NDArray[np.int32], exact: NDArray[np.int32]
) -> tuple[float, float]:
"""Mean AND worst recall over a batch of queries.
The worst case matters: mean recall hides that a subset of queries is
served badly, and in retrieval that subset is frequently the interesting
one. Returning only the mean is how a degraded index passes review.
"""
k = exact.shape[1]
# Broadcast comparison rather than a Python set per row: (n, k, k) of
# matches, reduced along both candidate axes in one pass.
matches = (approximate[:, :, None] == exact[:, None, :]).any(axis=2).sum(axis=1)
per_query = matches / k
return float(per_query.mean()), float(per_query.min())
What changed: The structural observation is that a distance computation is a dot product, so a batch of queries against a collection is a matrix product — and the expansion of squared Euclidean distance means only the cross term depends on both operands, with the query norm constant per row and therefore droppable since it cannot change a ranking. That turns brute force into one GEMM, which matters more than it first appears: brute force is the ground truth every recall number is measured against, so making it cheap is what makes measurement affordable enough that teams actually do it. Product quantization then removes floating-point distance entirely: one table per query, after which scoring a compressed vector is a gather and a sum along the subquantizer axis. The codes are stored subquantizer-major rather than vector-major precisely because scoring scans one subquantizer across every candidate, which is contiguous only in that layout. Selections use argpartition rather than argsort throughout, since k is tiny relative to the collection and ordering the rest is discarded work, and recall returns the worst case alongside the mean because a mean hides the subset of queries being served badly.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Exact distances for a whole query batch become one GEMM via the norm expansion, and the per-subspace encoding is one GEMM per subquantizer against its codebook.
Trade-off: The distance matrix is queries times collection size in float32, so a large batch against a large collection cannot be materialized at all and has to be chunked — which is exactly the case where you most wanted the batch.
Ensure contiguous memory layout and a single dtype
Subquantizer-major uint8 codes make the scoring gather a contiguous read per subquantizer, and float32 codebooks keep the encoding products on the BLAS fast path.
Trade-off: Subquantizer-major is the wrong layout for adding or removing a single vector, since one vector’s code is strided across M arrays — so inserts become substantially more expensive than queries.
Fuse adjacent operations so intermediates are never materialized
The norm expansion scales and adds in place over the cross-product matrix, and the squared norms are contracted with einsum rather than squaring into a temporary and summing it.
Trade-off: The cross-product buffer is destroyed as it becomes the distance matrix, so nothing can inspect the raw similarities afterwards — which is what you want when debugging a metric mismatch.
Batch work to amortize interpreter overhead
Queries, candidates and subquantizers all become array axes, so measuring recall over a thousand queries costs a constant number of interpreter operations rather than a thousand searches.
Trade-off: Batching only helps offline evaluation and bulk encoding; production serves one query at a time, so the batched path and the served path diverge and the measured latency is not the served latency.
Brute force becomes one GEMM per batch instead of N * d scalar operations per query; PQ scoring becomes M lookups per candidate with no distance arithmetic. Illustrative, not a measured benchmark: a fast brute force is what makes continuous recall measurement affordable, which is the practice that separates a tuned index from an unmeasured one.