PQ Math — Product Quantization

The problem: exact vectors don't fit

Approximate nearest-neighbour (ANN) search over embeddings has one stubborn enemy: memory. A dataset of N vectors in D dimensions, stored as 32-bit floats, occupies N · D · 4 bytes. For N = 10^9 and D = 128 that is roughly 512 GB — before any index structure. A brute-force scan is also O(N · D) per query, since every candidate needs a full D-dimensional distance.

Quantization attacks both costs at once. If we can represent each vector with a few bytes instead of 4D, the dataset shrinks by an order of magnitude or two and fits in RAM, where random access is fast. And if distances can be computed on the compact codes rather than the raw floats, the per-candidate cost collapses too. The catch is fidelity: a scalar quantizer that rounds each dimension to a few bits destroys the geometry similarity search depends on. PQ’s contribution is a quantizer with an enormous effective codebook that still stores and searches cheaply — the ‘product’ in the name is how it manufactures that huge codebook from tiny parts.

The core idea: split into m sub-vectors

Take a vector x ∈ R^D and cut it into m contiguous sub-vectors of equal length D* = D/m (PQ assumes m divides D):

x = [ x^1 | x^2 | ... | x^m ],   each x^j ∈ R^(D/m)

Each slice x^j lives in its own low-dimensional subspace. The key assumption is that these subspaces are roughly independent, so we can quantize each one separately and recombine. Quantizing a full D-dim vector well would need a codebook with astronomically many centroids; quantizing a D/m-dim slice needs only a small one. Because the slices combine as a Cartesian product, m codebooks of size k yield an effective codebook of k^m distinct reconstructions — with m = 8 and k = 256 that is 256^8 ≈ 1.8 × 10^19 possible codes, all expressible in 8 bytes: a titanic codebook you never have to store explicitly.

Codebooks: k-means per subspace

For each subspace j we learn a codebook C^j = {c^j_1, ..., c^j_k} of k centroids by running k-means on the j-th slices of a training set. This happens once, offline. The quantizer for subspace j maps a slice to its nearest centroid:

q^j(x^j) = argmin_{i ∈ 1..k}  || x^j − c^j_i ||^2

The centroids are the reusable ‘prototype’ slices; every database vector’s j-th piece gets snapped to one of them. Choosing k = 256 is near-universal, and not by accident: an index in 0..255 fits in exactly one byte, so log2(k) = 8 bits per sub-quantizer. Larger k gives finer resolution per slice but costs more bits and a bigger, slower-to-train codebook; the codebooks themselves occupy m · k · (D/m) = k · D floats total, tiny next to the dataset. The reconstruction of x is just the concatenation of the chosen centroids, and the residual it leaves behind is the quantization error PQ trades accuracy for.

Encoding: m indices, m·log2(k) bits

Encoding a vector is now trivial: run each slice through its sub-quantizer and keep the winning index, not the centroid. The code for x is the tuple

code(x) = ( i_1, i_2, ..., i_m ),   where i_j = q^j(x^j) ∈ {1..k}

Each index needs log2(k) bits, so the whole code is m · log2(k) bits. That single expression is the compression knob. With the canonical k = 256, an index is 1 byte and a PQ code is exactly m bytes — independent of D. A vector is stored as, say, [37, 209, 4, 128, 91, 17, 250, 63]: eight bytes that point into eight codebooks. Nothing about the original floats is kept. Decoding, if ever needed, looks up c^j_{i_j} in each subspace and concatenates — an approximation of x, never the original — but in practice you rarely decode, because distances are computed straight from the codes.

A worked compression example

Make it concrete with D = 128, m = 8, k = 256.

original : 128 dims × 4 bytes (float32) = 512 bytes
sub-vector length D* = 128 / 8            = 16 dims each
bits per index = log2(256)               = 8 bits = 1 byte
PQ code size   = m × 1 byte             = 8 bytes
compression    = 512 / 8                  = 64×

A billion such vectors drop from ~512 GB to ~8 GB — from impossible-on-one-box to comfortably in RAM. For higher fidelity, m = 16 gives 16-byte codes (32× compression, 8-dim slices) and m = 32 gives 64-byte codes (8×) with 4-dim slices. The codebooks add a fixed k · D = 256 × 128 = 32,768 floats (128 KB), negligible at scale. The dial is stark: every doubling of m halves the compression ratio while shrinking each slice, so the centroids track the data more tightly — more bytes, less error.

Searching the codes: ADC and distance tables

The elegant part is querying. Given an uncompressed query y, we want ||y − x||^2 for every compressed database vector x. Because the squared Euclidean distance decomposes over the independent subspaces, and x’s j-th slice is approximated by c^j_{i_j}:

||y − x||^2  ≈  Σ_{j=1..m} || y^j − c^j_{i_j} ||^2

This is Asymmetric Distance Computation (ADC): only the database side is quantized, the query stays exact. The trick is to precompute, once per query, a lookup table for each subspace — m tables of k entries:

T^j[i] = || y^j − c^j_i ||^2    for j = 1..m,  i = 1..k
d(y, x)^2 ≈ T^1[i_1] + T^2[i_2] + ... + T^m[i_m]

Now each database distance is just m table reads and m−1 adds — no multiplies, no D-dim work. Building the tables costs m · k sub-distances in D/m dims, i.e. O(k · D); scanning is O(N · m). Total O(kD + Nm) beats brute-force O(ND) handily because m « D.

ADC vs SDC, and why ADC is more accurate

There are two ways to compute PQ distances. Symmetric (SDC) quantizes the query too, then reads distances from a single precomputed k × k table of centroid-to-centroid distances per subspace — so both query and database contribute quantization error. Asymmetric (ADC), above, keeps the query exact and builds m query-specific tables of size k. Do not confuse the two shapes: SDC’s table is k × k and query-independent; ADC’s tables are m × k and rebuilt per query.

ADC is the default because it is strictly more accurate for the same code size. The distance estimate has only one source of error — the database vector’s quantization — instead of two, so its error is bounded by the database distortion alone rather than that plus the query’s. SDC’s only edge is that its table is computed once ever, not per query, which matters when queries are themselves stored codes; for live search with a fresh query vector, the O(kD) table build is cheap and ADC wins.

The memory / accuracy trade

PQ exposes two dials, m and k, and they set a clean frontier. Code size is m · log2(k) bits; accuracy rises with both. Increasing m shrinks each slice, so its k centroids blanket a lower-dimensional space more densely — less distortion, but proportionally more bytes. Increasing k adds resolution within each slice; beyond k = 256 you lose the one-byte alignment and gain little.

The error that matters is the quantization distortion Σ_j ||x^j − c^j_{i_j}||^2, the squared gap between a vector and its reconstruction, averaged over the dataset. Every recalled distance inherits this error, which is why high compression eventually degrades recall: two genuinely close vectors can look far apart once both are coarsely snapped. The practical regime — m from 8 to 32, k = 256 — keeps codes at 8–32 bytes and recall high enough that a cheap exact re-ranking of the top candidates recovers most of the lost precision.

IVFPQ: a coarse quantizer plus PQ residuals

PQ compresses vectors but, alone, still scans all N codes per query. IVFPQ adds a coarse pruning stage. First a coarse quantizer — k-means with k_c centroids (an inverted file, IVF) — partitions the space into Voronoi cells and assigns each vector to its nearest coarse centroid c(x). At query time you probe only the nprobe cells nearest the query, scanning a small fraction of the database.

The sharper idea is what PQ encodes. Instead of the raw vector, IVFPQ quantizes the residual r = x − c(x) — the vector relative to its cell centre:

x  →  cell id c(x)  +  PQ_code( x − c(x) )

Residuals are smaller and lower-variance than the originals, so the same m, k budget quantizes them with markedly less distortion — better accuracy at identical code size. ADC then runs on the residual: the query’s distance tables are built against the probed cell’s centroid. IVFPQ — coarse filtering plus residual PQ — is the workhorse behind billion-scale indexes.

How PQ differs from OPQ

PQ’s independence assumption is its weak point: it splits the vector along fixed, arbitrary coordinate blocks and hopes the subspaces are balanced and uncorrelated. Real embeddings rarely oblige — variance clusters in some dimensions, and correlations straddle the block boundaries — so a naive split wastes codebook capacity on some slices while starving others.

Optimized Product Quantization (OPQ), the sibling covered separately, fixes exactly this. It learns an orthogonal rotation matrix R and applies it before splitting: PQ(R · x). The rotation decorrelates dimensions and balances variance evenly across the m subspaces, and it is trained jointly with the codebooks by alternating between updating R and re-running k-means — minimizing the same quantization distortion. The boundary is clean: PQ = fixed dimension split; OPQ = a learned orthogonal rotation before the same split, optimized together with the codebooks. OPQ buys lower distortion at equal code size for one extra matrix-multiply per encode and query — a near-free upgrade that leaves PQ’s ADC machinery, IVF pruning, and byte layout intact.