Product Quantization and the Asymmetric Distance Computation

An inverted file reduces how many vectors a query compares, as covered in Coarse Quantization with Inverted Files and the nprobe Setting. Product quantization reduces what each comparison costs. The two compose, and together they are the basis of most systems that serve hundreds of millions of vectors on commodity hardware. The idea is a factorization: instead of encoding a 768-dimensional vector as one of $2^{768}$ possibilities, split it into contiguous subvectors and encode each subvector separately against a codebook of a few hundred entries. The resulting code is a handful of bytes, and the distance between a query and a code becomes a sum of table lookups rather than a sequence of multiply-accumulates.

The codebook construction

Choose a number of subspaces $m$ that divides the dimension $d$. With $d = 768$ and $m = 96$, each subvector has 8 components. Choose a codebook size $k^\ast$, conventionally 256 so that each sub-code fits in one byte. For each of the $m$ subspaces independently, run k-means over the corresponding 8-component slices of a training sample, producing 256 centroids per subspace.

The index stores $m$ codebooks of $256 \times (d/m)$ floats each. At $d = 768$, $m = 96$, that is 96 × 256 × 8 × 4 bytes = 786,432 bytes, under a megabyte. Each record is then encoded by assigning each of its $m$ subvectors to the nearest centroid in that subspace’s codebook, yielding an $m$-byte code. A 3072-byte float32 vector becomes a 96-byte code, a factor of 32.

The compression ratio follows from the byte counts: bytes per record go from $4d$ to $m$, so the ratio is $4d/m$. With $d = 768$ and $m = 96$ the ratio is 32. The codebook memory is $m \cdot k^\ast \cdot (d/m) \cdot 4 = 4 d k^\ast$ bytes, independent of $m$: 786 kilobytes regardless of how the dimension is partitioned. Only the code length and the granularity change with $m$.

That independence is a useful design fact. Splitting into more subspaces does not cost more codebook memory; it costs shorter subvectors, which are harder to quantize accurately because there are fewer dimensions in which to distinguish 256 centroids.

Why the split is what makes it work

Quantizing a 768-dimensional vector directly into 256 centroids would need 256 centroids to cover a space of dimension 768. The number of centroids required to maintain a given distortion grows exponentially with dimension, so 256 centroids in 768 dimensions are spread so thinly that the average distance from a vector to its assigned centroid is comparable to the spread of the data. The quantization would carry almost no information.

Splitting into 96 independent 8-dimensional problems changes the exponent. In 8 dimensions, 256 centroids are dense: $256^{1/8} = 2$ centroids per axis if laid on a grid, which is enough to resolve each coordinate into a few levels. The distortion per subspace is small, and the total distortion is the sum over subspaces. This is the mechanism: product quantization converts one hopeless high-dimensional quantization into many easy low-dimensional ones, and the total code length grows linearly with the number of subspaces while the distortion falls with the per-subspace dimension.

The cost of the independence assumption is that correlations between subspaces are lost. If components 1 and 9 are strongly correlated, and they fall in different subvectors, the code cannot represent the constraint that one implies the other. The reconstruction error then has a component along the correlation direction that no amount of codebook training removes. The remedy is a rotation applied before splitting: a fixed orthogonal transform, such as a random rotation or one derived from the data’s covariance, spreads correlated components across subspaces so that each subspace carries a mixture. Optimized product quantization learns the rotation jointly with the codebooks, and the gain over a random rotation is measurable but modest.

Asymmetric distance computation

The query is not quantized. For a query $q$, precompute for each subspace $j$ and each code value $c$ the squared distance between the query’s subvector $q_j$ and centroid $c$ of codebook $j$:

$$T[j][c] = \lVert q_j - C_j[c] \rVert^2$$

This table has $m \cdot k^\ast$ entries: 96 × 256 = 24,576 floats, about 98 kilobytes. It is computed once per query with $m \cdot k^\ast \cdot (d/m) = d \cdot k^\ast$ multiply-accumulates, which at $d = 768$ and $k^\ast = 256$ is about 196,608 operations, roughly the cost of 256 exact distance computations, amortized over every record scanned.

Then, for a record with code $(c_1, \ldots, c_m)$, the approximate squared distance to the query is

$$\hat{d}(q, x)^2 = \sum_{j=1}^{m} T[j][c_j]$$

The sum is $m$ additions and $m$ table lookups: 96 of each, against 768 multiply-accumulates for the exact computation, and the bytes read are 96 rather than 3072. The speedup is roughly a factor of 32 in bytes and a factor of several in operations, with the exact ratio depending on how well the table lookups overlap.

The asymmetry is the point. Quantizing the query as well would make both sides codes, and the distance between two codes cannot be decomposed into a per-subspace sum of table lookups unless a second table indexed by both codes is built, with $k^{\ast 2}$ entries per subspace, or 65,536 per subspace, which at 96 subspaces is far too large. Keeping the query exact lets the table be one-dimensional in the code, which is the whole trick. Symmetric computation exists and is used when the queries are also stored codes, at a modest loss of accuracy.

The distance being computed is not the distance between the query and the original vector. It is the distance between the query and the vector reconstructed from the code, which is the concatenation of the assigned centroids. The error is bounded by the quantization error of the record itself, and that error is not random with respect to the ranking: records in sparse regions of a subspace quantize badly, so their distances are systematically less accurate, and the ordering among near-tied candidates is where the loss appears.

import numpy as np

def train_pq(vectors: np.ndarray, m: int, k: int = 256, iters: int = 20, seed: int = 0):
    """Independent k-means per contiguous subspace. Returns codebooks and codes."""
    rng = np.random.default_rng(seed)
    n, d = vectors.shape
    sub = d // m
    books = np.empty((m, k, sub), dtype=np.float32)
    for j in range(m):
        block = vectors[:, j * sub:(j + 1) * sub]
        cent = block[rng.choice(n, size=k, replace=False)].copy()
        for _ in range(iters):
            assign = np.argmin(((block[:, None, :] - cent[None]) ** 2).sum(-1), axis=1)
            for c in range(k):
                members = block[assign == c]
                if len(members):
                    cent[c] = members.mean(axis=0)
        books[j] = cent
    return books

def encode(vectors: np.ndarray, books: np.ndarray) -> np.ndarray:
    m, k, sub = books.shape
    codes = np.empty((len(vectors), m), dtype=np.uint8)
    for j in range(m):
        block = vectors[:, j * sub:(j + 1) * sub]
        codes[:, j] = np.argmin(((block[:, None, :] - books[j][None]) ** 2).sum(-1), axis=1)
    return codes

def adc_table(query: np.ndarray, books: np.ndarray) -> np.ndarray:
    m, k, sub = books.shape
    table = np.empty((m, k), dtype=np.float32)
    for j in range(m):
        diff = books[j] - query[j * sub:(j + 1) * sub]
        table[j] = (diff ** 2).sum(1)
    return table

def adc_distances(table: np.ndarray, codes: np.ndarray) -> np.ndarray:
    m = codes.shape[1]
    return table[np.arange(m)[None, :], codes].sum(1)

The last function is the hot loop. It gathers $m$ values per record and sums them; in a tight implementation the gather is vectorized with a per-subspace loop over a block of records so that the table row stays in registers.

The parameters and what each one moves

Three numbers determine the behavior: $m$, $k^\ast$, and whether a rotation is applied.

Increasing $m$ lengthens the code, since each subspace contributes one byte, while shortening each subvector. More subspaces means less compression and lower distortion, because each subspace’s quantization error is smaller, but it also means a longer table and a longer sum. At the extreme $m = d$, each subspace is one scalar and the scheme degenerates to scalar quantization: one byte per component, with per-component centroids, which is a valid and sometimes preferable configuration.

Decreasing $m$ compresses harder and loses more. With $m = 8$ and $d = 768$, each subvector is 96 dimensions and 256 centroids are sparse in that space, so distortion rises sharply. Codes become 8 bytes per record, an impressive ratio, and recall at fixed candidate count falls.

$k^\ast$ is almost always 256 because it fits a byte. Using $k^\ast = 16$ allows two codes per byte and halves the code length at a cost in accuracy; using $k^\ast = 65{,}536$ requires two bytes per code and doubles the length for a small accuracy gain, since 256 centroids already cover an 8-dimensional subspace densely. The byte boundary dominates the choice.

The rule of thumb from the original paper is that $m$ between 8 and 32 with 8 to 16 components per subvector gives the best accuracy per byte for typical embedding widths. For $d = 768$ that suggests $m$ in the range of 48 to 96.

Where the error becomes a ranking error

Quantization error is not uniform across the collection. Records whose subvectors lie near a centroid reconstruct almost exactly; records between centroids reconstruct poorly. The distortion is therefore correlated with local density: dense clusters, which are exactly where many near-duplicates and near-ties live, quantize well while sparse outliers quantize badly.

The practical consequence is that PQ is a good fit for candidate generation and a poor fit for final ranking. It is used to scan a large collection cheaply and produce a shortlist, which is then rescored exactly against the stored vectors. That two-stage structure is the subject of Over-Fetch and Exact Reranking in Two-Stage Vector Retrieval, and the width of the over-fetch is set by how much the code’s ranking diverges from the exact one.

Two failure modes deserve explicit handling. The first is codes trained on a sample that does not represent the collection: a codebook trained on one language or one document type reconstructs vectors from another poorly, and the degradation is invisible in aggregate metrics computed on the training distribution. Retraining after a corpus shift is required, and the retraining invalidates every stored code, which makes it a full rebuild.

The second is empty or near-empty subspaces. If a model emits a component that is constant across the collection, the corresponding subspace has zero variance, k-means produces 256 identical centroids, and every record gets code zero for that subspace. The table row is constant, contributing the same value to every distance, which is harmless. But if the constant is large, the constant offset can dominate the comparison in floating-point terms and mask the real differences. Centering the data before training removes the issue.

Within an inverted file, PQ codes replace the stored vectors in the posting lists, and the scan cost per candidate drops from $4d$ bytes to $m$ bytes. The combination is the standard billion-scale configuration, and its remaining weakness is the coarse assignment: the code decides which cell a record belongs to as well as how it is encoded, so an error in either role compounds. On storage that cannot hold the full vectors, the same codes serve a different purpose as a routing layer, which is the arrangement examined in Disk-Resident Graph Indexes and the Memory Bandwidth Wall. Methods that keep a small exact residual alongside the code, or that use a finer second-level quantizer, exist precisely to bound that compounding. What does not change is the underlying arithmetic: the distance is a sum of independent per-subspace terms, and the quality of the result is the quality of the subspaces.