Product Quantization (PQ) Mechanics

Storing and searching through millions or billions of high-dimensional vectors, like those generated by LLM embedding models, presents significant challenges. The sheer memory footprint of these vectors (often float32, requiring $D \times 4$ bytes per vector) can be prohibitive, and calculating exact distances between a query vector and all database vectors becomes computationally expensive. Product Quantization (PQ) is a vector compression technique designed to address both these issues, enabling substantial memory savings and faster approximate distance calculations.

PQ operates on the principle of "divide and conquer." Instead of quantizing the entire high-dimensional vector at once, which often leads to poor resolution, PQ first decomposes the vector into multiple lower-dimensional sub-vectors.

Vector Decomposition

Let's consider an original vector $x$ in a $D$-dimensional space ($x \in \mathbb{R}^D$). PQ splits this vector into $M$ distinct sub-vectors, each of dimension $D/M$. We assume $D$ is divisible by $M$ for simplicity.

$$x = [x_1 | x_2 | \dots | x_M]$$

Here, each $x_m$ ($1 \le m \le M$) is a sub-vector of dimension $D/M$. For example, a 768-dimensional vector could be split into $M=48$ sub-vectors, each of dimension $D/M = 16$.

Subspace Quantization and Codebooks

The core idea is to quantize each of these $M$ subspaces independently. For each subspace $m$, we create a separate codebook $C_m$. This codebook consists of $k^*$ representative vectors, often called centroids, specific to that subspace. These centroids are typically learned using the k-means clustering algorithm on a training set of sub-vectors corresponding to that subspace.

$$C_m = \{c_{m,1}, c_{m,2}, \dots, c_{m,k^*}\} \quad \text{where } c_{m,j} \in \mathbb{R}^{D/M}$$

The number of centroids, $k^*$, is an important parameter. It determines the granularity of the quantization within each subspace. A common choice is $k^*=256$, allowing each centroid index to be represented by a single byte (8 bits).

Once the $M$ codebooks are generated, any original vector $x$ can be encoded. Each sub-vector $x_m$ of $x$ is mapped to the index (ID) of its nearest centroid in the corresponding codebook $C_m$. Let $q_m(x_m)$ be the function that finds the index of the nearest centroid in $C_m$ for sub-vector $x_m$:

$$q_m(x_m) = \underset{j \in \{1, \dots, k^*\}}{\operatorname{argmin}} \| x_m - c_{m,j} \|^2$$

The original vector $x$ is then represented by a sequence of $M$ centroid indices:

$$\text{code}(x) = [q_1(x_1), q_2(x_2), \dots, q_M(x_M)]$$

Each index $q_m(x_m)$ requires $\log_2(k^*)$ bits. If we choose $k^*=256$, each index needs 8 bits (1 byte). Thus, the entire $D$-dimensional vector $x$ is compressed into just $M$ bytes. Compared to storing the original $D$-dimensional float32 vector (requiring $D \times 4$ bytes), this represents a significant compression ratio, typically ranging from 10x to 100x depending on $D$ and $M$. For our 768-dimensional example with $M=48$, the compressed size is 48 bytes, compared to $768 \times 4 = 3072$ bytes originally, a compression factor of 64x.

The process of Product Quantization: A high-dimensional vector is split into sub-vectors. Each subspace undergoes k-means clustering to create a codebook. The original vector is then represented by the sequence of nearest centroid IDs from each subspace.

Approximate Distance Calculation

The primary benefit of PQ, instead of memory compression, is the ability to quickly approximate the distance between a query vector $q$ and a compressed database vector $x$ (represented by its code, $\text{code}(x)$).

Asymmetric Distance Computation (ADC)

The most common method is Asymmetric Distance Computation (ADC). Here, the query vector $q$ remains in its original, uncompressed form. The distance to a database vector $x$ (represented by its PQ code) is approximated by summing the distances between the query's sub-vectors $q_m$ and the centroids corresponding to the codes stored for $x$.

Specifically, let $\text{code}(x) = [id_1, id_2, \dots, id_M]$, where $id_m = q_m(x_m)$. The approximate squared Euclidean distance is calculated as:

$$\hat{d}(q, x)^2 \approx \sum_{m=1}^{M} \| q_m - c_{m, id_m} \|^2$$

The main optimization here is that the distances $\| q_m - c_{m,j} \|^2$ for all $m \in \{1, \dots, M\}$ and $j \in \{1, \dots, k^*\}$ can be precomputed once per query $q$. This creates $M$ lookup tables, each containing $k^*$ precomputed distances. The total distance approximation for a database vector $x$ then reduces to $M$ table lookups and $M-1$ additions. This is significantly faster than calculating the full $D$-dimensional distance.

Symmetric Distance Computation (SDC)

An alternative is Symmetric Distance Computation (SDC), where the query vector $q$ is also quantized using the same codebooks. The distance is then approximated using only the centroids:

$$\tilde{d}(q, x)^2 \approx \sum_{m=1}^{M} \| c_{m, q_m(q_m)} - c_{m, q_m(x_m)} \|^2$$

SDC can be even faster as it operates entirely in the compressed domain, potentially using precomputed distances between all pairs of centroids within each subspace. However, quantizing the query introduces additional approximation error, generally making SDC less accurate than ADC. ADC is usually preferred when accuracy is prioritized.

Trade-offs

Product Quantization offers substantial benefits but comes with inherent trade-offs:

1. Accuracy vs. Compression: PQ is a lossy compression technique. Information is lost when mapping sub-vectors to centroids. This loss directly impacts the accuracy (recall) of the nearest neighbor search. The primary levers controlling this trade-off are:

   • M (Number of Subspaces): Increasing $M$ (reducing subspace dimensionality $D/M$) generally increases accuracy but reduces the compression ratio (since the code length is $M$ bytes, assuming $k^*=256$).
   • k* (Codebook Size): Increasing $k^*$ allows for finer-grained quantization within each subspace, improving accuracy but increasing the memory required for the codebooks and the bit-length needed for each code index. $k^*=256$ is a common practical choice balancing accuracy and byte-alignment.

2. Training Cost: Generating the $M$ codebooks requires running k-means clustering $M$ times on potentially large datasets of sub-vectors derived from a representative training sample. This offline training step can be computationally intensive.

3. Integration with Index Structures: While PQ provides compression and fast approximate distances, it doesn't inherently provide a sub-linear search structure like HNSW or IVF. PQ is most often used in combination with these structures. For instance, in IVFPQ (covered in the next section), PQ is used to compress the vectors stored within the inverted lists of an IVF index, dramatically reducing the index size and speeding up distance calculations during the search refinement stage. Similarly, PQ can compress vectors stored at the nodes of an HNSW graph.

Product Quantization serves as a foundational technique for managing the scale of modern vector datasets. By decomposing vectors and quantizing subspaces independently, it achieves remarkable compression and enables rapid approximate distance calculations, making it an indispensable tool in the construction of large-scale, efficient vector search systems. Understanding its mechanics is essential for optimizing memory usage and query latency in demanding LLM applications.