Vector search

Level 2 builds every index above from nothing, starting with a picture and eight points on a map.

The everyday picture. You've just moved to a new country and want the house nearest to a landmark. You could walk up to every house in the country and measure the distance. That's always right, but it takes forever. Or you could do what everyone actually does: take the highway to the right region, switch to main roads to the right town, then walk the side streets to the house. You check only a handful of places and still end up at (almost always) the right door.

That is the whole problem of this lesson. A search engine that works on meaning stores every document as a vector (a list of numbers, see primer.ml.embeddings.similarity) and turns the question into a vector too. The best documents are the vectors nearest the question. With ten million documents, "walk to every house" is too slow, so we build a map with highways. Structures like that are called approximate nearest neighbor (ANN) indexes. Approximate because, like the traveller, they very occasionally stop at the second-nearest house instead of the nearest.

Four ideas cover almost every vector database in use:

A few words used throughout, in plain terms:

• Similarity score. How alike two vectors are. Here we use the dot product (multiply the two lists position by position and add up). All vectors are rescaled to length 1 first (normalized), so the dot product equals the cosine similarity, which runs from −1 (opposite) to 1 (identical). Higher means closer.
• Top k. The k best matches, e.g. the top 10.
• Recall@k. Of the true top k (found by exhaustive search), the fraction the index actually returned. Recall@10 = 0.9 means it found 9 of the true 10. It's the standard measure of an ANN index's accuracy.
• Big-O, e.g. O(N). How the work grows with the data. O(N) means doubling the number of vectors N doubles the work; O(log N) means doubling N adds only one more step.

Every index in this lesson is shown first on the same eight points, drawn on a flat map so you can check each step with a pencil. The question (the query) sits at (5, 4). On a flat map "near" is ordinary distance; we keep it squared (dx² + dy²) so every number stays whole. Squaring doesn't change which point is nearest.

Checking all eight rows is flat search. It's exact, and it costs one distance per stored point.

flowchart LR
  Q[Query vector] --> D["Dot product with<br/>EVERY stored vector<br/>(N comparisons)"]
  D --> T[Keep the k highest]
  T --> R[Exact top k]

On the example: with unit vectors q = (0.8, 0.6) and x = (0.6, 0.8): s = 0.8·0.6 + 0.6·0.8 = 0.48 + 0.48 = 0.96, very similar.

In Python:

q = [0.8, 0.6]
x = [0.6, 0.8]
# s(q, x) = Σ q_i x_i
round(sum(q_i * x_i for q_i, x_i in zip(q, x)), 2)  # → 0.96
points = {"A": (0, 0), "B": (2, 1), "C": (4, 0), "D": (1, 3),
          "E": (3, 3), "F": (5, 2), "G": (2, 5), "H": (5, 5)}
query = (5, 4)
def sq_dist(p):
    return (p[0] - query[0]) ** 2 + (p[1] - query[1]) ** 2
# flat search on the map: check all eight
min(points, key=lambda name: sq_dist(points[name]))  # → 'H'

In code: FlatIndex.search scores the query against every stored vector and keeps the best k with top_k; normalize rescales vectors to length 1 first, so the dot product is the cosine.

Why it matters flat search costs N·d multiply-adds per query. At 10 million vectors of 1,536 dimensions that's about 15 billion, far too many for an interactive search, and the raw vectors alone take 10,000,000 × 1,536 × 4 bytes ≈ 61 GB of memory. Below about a million vectors, flat search is often the right answer: simple, exact, and fast enough.

In code: storage_estimate does that memory sum for any n and d.

The everyday picture. A library doesn't search every shelf for a book about sailing. It sorts books into sections ahead of time, and you go to "Sports" and maybe "Travel", the one or two most promising sections, and search only there. A sailing book shelved in the wrong section is missed unless you check that section too.

On the eight points. Split them into two groups (clusters) ahead of time: left = {A, B, C, D} and right = {E, F, G, H}. Each cluster is summarized by its centroid, the average position of its members: left = (1.75, 1), right = (3.75, 3.75). For the query (5, 4):

Scan only E, F, G, H: the nearest is H (distance 1). That's 2 centroid checks + 4 point checks = 6 instead of 8. With a million points in 1,000 clusters the saving is enormous.

flowchart TB
  subgraph BUILD["Ahead of time"]
    V[All vectors] --> KM["k-means: find nlist centroids"]
    KM --> L["Put each vector on the list<br/>of its nearest centroid"]
  end
  subgraph QUERY["At query time"]
    Q[Query] --> C["Compare with the nlist centroids"]
    C --> P["Pick the nprobe nearest"]
    P --> S["Scan only those lists"]
    S --> R[Top k]
  end
  L -.-> S

The set of all points closer to one centroid than to any other is called that centroid's Voronoi cell: the "section" of the library. IVF's blind spot is a true neighbor sitting just across a cell boundary.

On the example: 2 + 8 · 1/2 = 6 comparisons (versus 8). At N = 1,000,000, n_list = 1,000, n_probe = 10: 1,000 + 10,000 = 11,000, about 1% of a flat scan.

In Python:

left = [(0, 0), (2, 1), (4, 0), (1, 3)]
right = [(3, 3), (5, 2), (2, 5), (5, 5)]
def centroid(members):
    return tuple(sum(p[i] for p in members) / len(members) for i in range(2))
centroid(left), centroid(right)  # → ((1.75, 1.0), (3.75, 3.75))
# probe the nearer
[(c[0] - 5) ** 2 + (c[1] - 4) ** 2 for c in (centroid(left), centroid(right))]  # → [19.5625, 1.625]
def comparisons(N, n_list, n_probe):
    # n_list centroids + the vectors in the opened lists
    return n_list + N * n_probe // n_list
comparisons(8, 2, 1)  # → 6
comparisons(1_000_000, 1_000, 10)  # → 11000

In code: IVFIndex.train finds the centroids with kmeans, IVFIndex.add files each vector on its nearest centroid's list, and IVFIndex.search scans only the nprobe nearest lists. tiny_ivf_search replays the eight-point example.

Why it matters IVF is cheap to build and light on memory, and nprobe = nlist gives you back an exact flat search, which is a handy sanity check. Its accuracy depends on how well the clusters fit the data, so retrain the centroids when the data changes a lot.

The everyday picture. A police sketch artist doesn't record every pixel of a face. They pick the closest nose from a catalogue of 256 noses, the closest eyes from 256 eyes, the closest mouth, and so on. The whole face becomes a handful of catalogue numbers: tiny to store, close enough to recognize someone. That's product quantization (PQ). To quantize means to round a value to the nearest entry of a fixed set.

On a 4-number vector. Split x = (0.9, 0.1, −0.2, 0.8) into two halves. Each half is matched against a four-entry catalogue (a codebook), here the four compass directions: 0 = (1, 0), 1 = (0, 1), 2 = (−1, 0), 3 = (0, −1).

The vector is now stored as two codes, [0, 1]. To score the query q = (1, 0, 0, 1) against any stored vector, first build one small table per half: the query's half dotted with each catalogue entry.

The approximate score of codes [0, 1] is T₁[0] + T₂[1] = 1 + 1 = 2. The exact score is 0.9 + 0.8 = 1.7: close, not identical. That's the price of compression.

flowchart LR
  subgraph ENC["Encode (once per vector)"]
    X["x: d numbers"] --> SPLIT["split into m pieces"]
    SPLIT --> NN["each piece → nearest of<br/>256 codebook entries"]
    NN --> CODE["m one-byte codes"]
  end
  subgraph ADC["Score (per query)"]
    Q["query q"] --> TAB["m small tables:<br/>q's piece · every entry"]
    TAB --> SUM["score = sum of m<br/>table lookups"]
  end
  CODE --> SUM

On the example: ŝ = T₁[c₁] + T₂[c₂] = T₁[0] + T₂[1] = 1 + 1 = 2 (exact: 1.7).

In Python:

# the compass codebook, used for both pieces
C = [(1, 0), (0, 1), (-1, 0), (0, -1)]
def dot(a, b):
    return sum(a_i * b_i for a_i, b_i in zip(a, b))
def sq_dist(a, b):
    return sum((a_i - b_i) ** 2 for a_i, b_i in zip(a, b))
x = [0.9, 0.1, -0.2, 0.8]
pieces = [x[0:2], x[2:4]]
# c_j(x)
codes = [min(range(4), key=lambda c: sq_dist(piece, C[c])) for piece in pieces]
codes  # → [0, 1]
q = [1, 0, 0, 1]
# T_j[c] = q^(j) · C_j[c]
T = [[dot(q_j, C[c]) for c in range(4)] for q_j in (q[0:2], q[2:4])]
T  # → [[1, 0, -1, 0], [0, 1, 0, -1]]
# ŝ = Σ_j T_j[c_j(x)]
sum(T[j][codes[j]] for j in range(2))  # → 2
# the exact score, for comparison
round(dot(q, x), 2)  # → 1.7

In code: ProductQuantizer learns the codebooks (ProductQuantizer.train), rounds vectors to codes (ProductQuantizer.encode), builds the Tⱼ tables (ProductQuantizer.lookup_table) and sums the lookups (ProductQuantizer.adc_scores). tiny_pq_example replays the 4-number example by hand-setting the compass codebooks.

Why it matters a 768-dimension float32 vector is 3,072 bytes; with m = 96 it's 96 bytes, 32× smaller, so a billion vectors fit in about 100 GB instead of 3 TB. The lost accuracy is recovered by re-scoring: take the top few hundred by PQ score and recompute their exact scores from the full vectors kept on disk. Real systems combine PQ with IVF (IVF-PQ) and encode each vector's residual (its offset from its cluster centroid), which is smaller and so rounds more precisely.

In code: PQIndex scans every PQ code and can re-score its top candidates with the exact vectors; IVFPQIndex combines IVF lists with PQ-encoded residuals.

The everyday picture. Back to the traveller: highways to get close fast, main roads to get closer, side streets to find the door. HNSW (Hierarchical Navigable Small World) builds exactly that as a graph: a set of points (nodes) joined by links (edges). Every vector is a node on the bottom layer (the side streets). A random few are also placed on layer 1 (main roads), fewer still on layer 2 (highways), and so on.

On the eight points. Layer 1 holds only A, E and G. Layer 0 holds all eight, linked to nearby points (the thin lines in the figure above). A greedy search means "always move to whichever neighbor is closest to the target; stop when none is closer". Start at A:

flowchart TD
  subgraph L2["Top layer: few nodes, long jumps"]
    A2[A] --- D2[D]
  end
  subgraph L1["Middle layer: more nodes"]
    A1[A] --- B1[B] --- D1[D] --- E1[E]
  end
  subgraph L0["Bottom layer: every vector"]
    A0[A] --- B0[B] --- C0[C] --- D0[D] --- E0[E] --- F0[F]
  end
  D2 -.-> D1
  E1 -.-> E0

In code: tiny_hnsw_search replays the greedy walk from the table above on the eight-point map and returns every stop; HNSWIndex is the full index used on real vectors.

The search, step by step

flowchart TD
  S[Start at the entry point<br/>on the top layer] --> G{Is any neighbor<br/>closer to the query?}
  G -->|yes| H[Hop to the closest neighbor] --> G
  G -->|no| B{On the bottom layer?}
  B -->|no| DOWN[Drop one layer,<br/>same node] --> G
  B -->|yes| BEAM["Beam search: keep the best efSearch<br/>candidates, expand the most promising<br/>until nothing better turns up"]
  BEAM --> K[Return the top k]

In code: HNSWIndex.search runs the greedy descent and the bottom-layer beam search; HNSWIndex.search_trace does the same and returns every node it expanded, which is what this figure draws. small_hnsw_map builds the 60-point graph the widget searches, and viz_data hands it to the page.

How the layers are built

flowchart TD
  N[New vector] --> LV["Draw its top layer ℓ at random<br/>(most get 0, a few get more)"]
  LV --> DESC[Greedy descent from the entry point<br/>down to layer ℓ]
  DESC --> FIND["On each layer ≤ ℓ: beam search with<br/>efConstruction to find candidates"]
  FIND --> SEL["Keep up to M diverse neighbors<br/>(the heuristic below)"]
  SEL --> LINK[Link both ways]
  LINK --> PRUNE{Neighbor now has too many links?<br/>more than M, or 2·M on layer 0}
  PRUNE -->|yes| TRIM[Re-select its best-spread links]
  PRUNE -->|no| DONE[Next layer down]
  TRIM --> DONE

The diversity heuristic. When choosing a new node's M neighbors, go through the candidates from nearest to farthest and keep one only if it is closer to the new node than to every neighbor already kept. Otherwise a kept neighbor already "covers" that direction. Without this rule, in clustered data all M links would point into the same clump, and the graph could split into islands a greedy walk can't cross. With it, links fan out in different directions and long "bridges" between clusters survive.

The random layer draw is where the math comes in:

On the example: M = 16, so m_L = 1/ln 16 = 1/2.773 = 0.361. A draw of U = 0.5 gives −ln 0.5 · 0.361 = 0.693 · 0.361 = 0.25 → floor → layer 0. A draw of U = 0.05 gives 2.996 · 0.361 = 1.08 → layer 1. Only 1/16 = 6.25% of nodes reach layer 1, and 1/256 ≈ 0.4% reach layer 2: each layer has about 1/M as many nodes as the one below, which is what makes the upper layers "highways".

In Python:

import math
M = 16
# m_L = 1 / ln M
m_L = 1 / math.log(M)
round(m_L, 3)  # → 0.361
for U in (0.5, 0.05):
    # -ln(U) · m_L
    scaled = -math.log(U) * m_L
    # ... then round down to get ℓ
    print(U, round(scaled, 2), math.floor(scaled))  # → 0.5 0.25 0 0.05 1.08 1
# P(ℓ ≥ l) = M^(-l): 6.25% and about 0.4%
[M ** -l for l in (1, 2)]  # → [0.0625, 0.00390625]

In code: HNSWIndex.add inserts each vector this way: draw its layer, beam-search each layer with efConstruction, keep up to M diverse neighbors and link both ways. HNSWIndex.layer_sizes counts the nodes on each layer, and HNSWIndex.memory_bytes adds up the vectors plus their links.

Why it matters HNSW is the default index in most vector databases because it gives high recall at low latency with no training step. Its costs are memory (every vector plus its links must sit in RAM) and slow builds for very large collections. At billions of vectors, teams switch to IVF-PQ, shard HNSW across machines, or use disk-based graphs (DiskANN).

On the example: if the true top 10 is documents 1 to 10 and the index returns 1 to 9 plus document 42, the overlap is 9, so recall@10 = 9/10 = 0.9.

In Python:

# documents 1 to 10
true_top = set(range(1, 11))
# 1 to 9, plus document 42
returned = set(range(1, 10)) | {42}
k = 10
# |returned ∩ true| / k
len(returned & true_top) / k  # → 0.9

In code: recall_at_k computes this share; ground_truth runs the exact flat search that supplies the true top k, and evaluate reports recall, latency and distance computations per query for any index.

Why it matters recall and latency are traded against each other on every index. The only reliable way to pick a setting is to measure recall@k against a flat index on your own vectors while you turn the dial, then check the 95th-percentile latency.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.