Product Quantization for Fast AI Vector Search

Product Quantization (PQ) is a powerful vector compression technique used in artificial intelligence databases to dramatically speed up Approximate Nearest Neighbor (ANN) searches while drastically reducing memory consumption. This article explains how PQ works by splitting high-dimensional embedding vectors into smaller sub-vectors, clustering those segments, and replacing heavy floating-point data with lightweight centroid IDs. Furthermore, it covers how AI databases use these compressed codes alongside precomputed distance tables to execute rapid similarity searches across billions of vectors.

Splitting High-Dimensional Vectors into Sub-Vectors

Modern AI models generate embeddings that often contain hundreds or thousands of dimensions. Storing and comparing these raw vectors requires massive amounts of RAM and compute power. Product Quantization solves this by dividing a single high-dimensional vector into several smaller, equal-sized chunks called sub-vectors.

If you have a vector with \(D\) dimensions, PQ splits it into \(M\) distinct sub-vectors. Each sub-vector therefore has a dimension of \(D^* = D / M\).

For example, take a 128-dimensional vector containing 32-bit floating-point numbers. If we set \(M = 8\), the vector is partitioned into 8 separate sub-vectors. Each sub-vector now contains just 16 dimensions. What was once a single complex space is now broken down into 8 independent, lower-dimensional subspaces.

Independent Clustering and Codebook Generation

Once the database's training vectors are split into \(M\) subspaces, PQ applies the standard \(k\)-means clustering algorithm to each subspace independently.

If we choose a cluster count of \(K = 256\) for each subspace, the \(k\)-means algorithm will find 256 centroids (representative average points) for the first subspace, 256 centroids for the second subspace, and so on, up to the \(M\)-th subspace.

The collection of centroids for a single subspace is called a codebook. Because there are \(M\) subspaces, the system generates \(M\) distinct codebooks. Even though each codebook only contains 256 centroids, the Cartesian product of these subspaces allows the system to represent \(K^M\) (or \(256^8\), which is \(1.84 \times 10^{19}\)) possible vector approximations. This is where "Product" Quantization gets its name.

Quantization: Encoding Vectors into Short Codes

After the codebooks are trained, the database compresses the actual high-dimensional vectors. Each sub-vector of an original embedding is compared against the centroids in its corresponding subspace codebook. The sub-vector is then replaced by the numerical ID of the single closest centroid.

Because we used \(K = 256\) centroids per subspace, every centroid ID can be represented as an integer between 0 and 255. In binary, a number from 0 to 255 requires exactly 8 bits (1 byte) of storage.

Looking back at our 128-dimensional vector example:

  • Original Size: 128 dimensions \(\times\) 4 bytes (32-bit float) = 512 bytes.
  • PQ Compressed Size: 8 sub-vectors \(\times\) 1 byte (centroid ID) = 8 bytes.

The original high-dimensional vector is now stored as a compact 8-byte array of integers (the PQ code). This achieves a 64x reduction in memory, allowing AI databases to keep massive vector indexes entirely in fast system RAM rather than on slow disk storage.

Fast Approximate Nearest Neighbor Search via Lookup Tables

When a user queries the AI database, the system must find the stored vectors that are closest to the new query vector. Instead of decompressing the vectors, PQ uses a technique called Asymmetric Distance Computation (ADC) to calculate distances directly on the compressed codes.

  1. Split the Query: The incoming high-dimensional query vector is split into the same \(M\) sub-vectors using the exact same partitioning rules as the database.
  2. Build Distance Tables: For each of the \(M\) query sub-vectors, the system calculates the exact squared Euclidean distance to all 256 centroids in the corresponding codebook. This creates \(M\) lookup tables, each containing 256 precomputed distances.
  3. Table Lookups and Addition: To find the approximate distance between the query and a compressed database vector, the system reads the 1-byte centroid IDs of the database vector. It uses these IDs as indexes to fetch the precomputed distances from the lookup tables and adds them together.

Instead of performing 128 complex floating-point subtractions and multiplications per vector, the CPU simply performs 8 fast memory lookups and 8 simple additions. Because the vectors are tiny and the math is reduced to basic table lookups, the AI database can scan millions of vectors in milliseconds, returning the highest-scoring approximate nearest neighbors with minimal loss in accuracy.