Product Quantization: Vector Compression for ANN Search › Codebook Learning

The quality of PQ depends entirely on how well the codebooks capture the structure of your data. A good codebook should have its centroids positioned where data points cluster in each subspace. Poorly positioned centroids (for instance, those uniformly distributed in a region where real data is sparse) lead to high quantization error and degraded recall. This is why PQ always requires a training phase: the codebooks must be fitted to the distribution of your embeddings before you can compress and search them effectively.

The training process is conceptually simple: for each subspace, collect all the corresponding subvectors from your training data, and run k-means clustering to find $K$ representative centroids. The tricky part is scale. You need enough training vectors to populate all $K$ centroids with reasonable coverage, which in practice means at least several hundred times $K$ training examples per subspace. For $K = 256$, this means at least 50,000 to 100,000 training vectors, a number easily obtained from a typical corpus but worth planning for in advance.

Think of codebook training as photographing a neighborhood to build a visual catalog. If you photograph only one street, your catalog will misrepresent the full variety of buildings. If you photograph a representative sample of every street and neighborhood, your catalog will cover the full range of architectural styles that a visitor might encounter. PQ codebooks work the same way: a training set drawn from the same distribution as your final index will produce centroids that are well-positioned for the actual data, while a mismatched training set will leave gaps in the coverage and inflate quantization error.

An important practical consideration is that the training set should match the distribution of vectors you will index. If your index will contain embeddings from a news corpus, train the codebooks on a sample from that corpus, not from a generic Wikipedia dump. The closer the match between training and indexing distributions, the lower the average quantization error will be, and the better your recall.

For each subspace $m \in \{1, \ldots, M\}$, you extract the $m$-th segment from every training vector to form a dataset of $N$ subvectors in $\mathbb{R}^{d^*}$. The goal is to find a compact set of centroids that together describe the typical shapes this subvector takes across your data. You then run k-means clustering on this dataset, seeking $K$ centroids $\mathcal{C}_m = \{c_{m,1}, c_{m,2}, \ldots, c_{m,K}\}$ that minimize the total squared distance between each subvector and its nearest centroid, called the quantization error:

$$\mathcal{L}_m = \sum_{i=1}^{N} \min_{k \in \{1,\ldots,K\}} \left\| x_i^{(m)} - c_{m,k} \right\|^2$$

where:

• $\mathcal{L}_m$: the quantization loss (total reconstruction error) for subspace $m$
• $N$: the number of training vectors
• $x_i^{(m)}$: the $m$-th subvector of training vector $i$, a point in $\mathbb{R}^{d^*}$
• $c_{m,k}$: the $k$-th centroid in the codebook for subspace $m$
• $\min_{k \in \{1,\ldots,K\}}$: the assignment step, selecting the nearest centroid for each subvector

This is standard k-means, as discussed in the context of IVF coarse quantization in the previous chapter. The difference here is that you run it $M$ separate times, each time on a lower-dimensional subspace, which makes the clustering much faster and more reliable than trying to cluster in the full 768-dimensional space. In high dimensions, k-means is notoriously difficult: the curse of dimensionality makes all points appear roughly equidistant from one another, and finding meaningful cluster boundaries becomes unreliable. By restricting each k-means run to $d^*$ dimensions, PQ turns a hard high-dimensional clustering problem into $M$ tractable low-dimensional ones. This is one of the most practically important aspects of the design: the decomposition is a compression trick and a strategy for making the learning problem statistically well-posed.

The k-means algorithm iterates between two steps: (1) assign each training subvector to the nearest centroid, and (2) update each centroid to the mean of all subvectors assigned to it. These two steps alternate until the centroid positions converge. The loss $\mathcal{L}_m$ is guaranteed to decrease or stay the same on every iteration, so the algorithm always terminates. However, k-means can converge to local optima, so it is common to run it multiple times with different random initializations and keep the result with the lowest total loss. In practice, one good initialization strategy is k-means++, which places the initial centroids far from one another to improve convergence speed and quality.

After training, the codebooks $\{\mathcal{C}_1, \mathcal{C}_2, \ldots, \mathcal{C}_M\}$ are fixed. They are the only data structure you need to store alongside the compressed codes. At search time, you do not need to keep the original training vectors in memory at all. The codebooks, which total only $M \times K \times d^*$ floating-point values, are loaded once and held in cache throughout the search session. For typical parameters, this is on the order of a few megabytes, an entirely negligible overhead compared to the index itself.

Once the codebooks are trained, encoding a database vector $x$ is straightforward. The process visits each subspace in order, extracts the corresponding slice of the vector, and finds the nearest centroid in that subspace's codebook. Split $x$ into $M$ subvectors, then for each subspace find the nearest centroid:

$$k_m = \underset{k}{\operatorname{argmin}} \left\| x^{(m)} - c_{m,k} \right\|^2$$

where:

• $k_m$: the centroid index assigned to the $m$-th subvector, i.e., the value stored in the compressed code
• $x^{(m)}$: the $m$-th subvector of database vector $x$
• $c_{m,k}$: the $k$-th centroid in the codebook for subspace $m$
• $\underset{k}{\operatorname{argmin}}$: the operator that returns the index $k$ minimizing the squared distance, selecting the nearest centroid to subvector $x^{(m)}$

The encoded vector is the tuple $(k_1, k_2, \ldots, k_M)$, storing $M$ integers each in the range $[0, K-1]$. With $K = 256$, each fits in one byte (uint8). The intuition is that instead of storing a full floating-point subvector, you store only the address of the nearest representative in a shared codebook. This is a drastic reduction from $d^* \times 4$ bytes down to 1 byte per subspace.

This encoding step only happens once at index build time. After encoding, the original floating-point vectors can be discarded entirely if memory is the primary concern, and only the compact integer codes and the codebook centroids need to be retained. At query time, every distance computation will refer exclusively to these codes and the precomputed lookup tables derived from the codebooks. The encoding process scales linearly with the number of database vectors: encoding 100 million vectors requires 100 million independent nearest-centroid lookups across all subspaces, and these can be parallelized trivially since each vector is encoded independently of all others.

The time complexity of encoding a single vector is $O(M \cdot K \cdot d^*)$: for each of the $M$ subspaces, you compute the distance from the subvector to each of the $K$ centroids and select the minimum. With $d^* = d/M$, this simplifies to $O(K \cdot d)$, which is the same order as a single full-precision distance computation multiplied by $K$. In practice, the encoding is fast enough that it is never the bottleneck in an indexing pipeline; the bottleneck is almost always generating the embeddings themselves.