Product Quantization Tessellation

• Product Quantization Tessellation is a geometric partitioning strategy for high-dimensional spaces that enables efficient vector compression and rapid similarity search.
• It divides the space into structured cells using methods like axis-aligned, projective, and Voronoi tessellations, leveraging k-means and directional clustering.
• The approach underpins applications in nearest neighbor search and large-model inference, offering GPU-optimized solutions and end-to-end differentiable quantization.

Product Quantization Tessellation refers to a class of geometric partitioning strategies for high-dimensional vector spaces that underlie modern quantization schemes for vector compression, fast similarity search, and quantized neural network operations. These methods tessellate $\mathbb{R}^d$ into cells—either axis-aligned Cartesian blocks, cones defined by directional clustering, or Voronoi polytopes—each associated with quantization centers or codes. Product quantization tessellation thus encapsulates several closely related frameworks, including classical product quantization (PQ), its advanced projective clustering variants, and differentiable tessellation approaches for end-to-end learning. The factorizations and cell constructions in these schemes determine quantization error, encoding and search complexity, and eventual system accuracy in tasks such as nearest neighbor search, max inner product search (MIPS), and large-model inference.

1. Classical Product Quantization Tessellation

Classical product quantization (PQ) decomposes $\mathbb{R}^d$ into a Cartesian product of $M$ subspaces. Each input vector $x \in \mathbb{R}^d$ is divided into $M$ disjoint, equal-sized blocks $(x^{(1)},\ldots,x^{(M)})$, with each subvector $x^{(i)} \in \mathbb{R}^{d/M}$. For each subspace $i$, a $k$-means codebook $\mathcal{C}_i = \{c_{i,1}, \ldots, c_{i,K}\}$ is learned, typically via Lloyd’s algorithm:

$$\min_{C_i} \sum_{n=1}^N \min_{j=1, \ldots, K} \| x_n^{(i)} - c_{i,j} \|_2^2,$$

where $K=2^b$ quantization centers per subspace are specified by $b$ bits.

The tessellation consists of axis-aligned (Cartesian) cells, each defined as the cross-product of a subspace cell, and each point is encoded by the index tuple of its closest centroid in every subspace. Quantization then amounts to table lookups for encoding and decoding; reconstruction error is additive over subspaces:

$$\|x - \hat x\|_2^2 = \sum_{i=1}^M \|x^{(i)} - c_{i,k_i}\|_2^2.$$

Because each subspace is quantized independently, the product quantization tessellation partitions $\mathbb{R}^d$ into $K^M$ hyperrectangular cells (Wang et al., 12 Mar 2025).

2. Projective Clustering Product Quantization (PCPQ) and Anisotropic Tessellation

Projective Clustering Product Quantization (PCPQ) generalizes classical PQ by using directional (projective) clustering in each block of coordinates, forming high-resolution, anisotropic partitions of $\mathbb{R}^d$ adapted to the data and query distribution. For a data vector $x\in\mathbb{R}^d$, divided into $m$ contiguous blocks $x = [x^{(1)},...,x^{(m)}]$, PCPQ fits $k$ unit direction vectors $c_1, ..., c_k$ per block. Each $x^{(j)}$ is assigned to a direction $c_{j,\phi_j}$ and scaling coefficient $\alpha$ minimizing:

$$L_\text{proj}(C;X) = \sum_{i=1}^n \min_{j \in [k]} \| x_i - ( \langle x_i, c_j \rangle / \|c_j\|_2^2 ) c_j \|_2^2.$$

This is equivalent to projecting $x_i$ onto the span of $c_j$; the section induces a cone-like tessellation, where each cell contains vectors closest in projection to a given direction.

Anisotropic PCPQ (APCQ) adapts resolution by using error weights $(h_\parallel, h_\perp)$ and redefining loss as a weighted sum of parallel and orthogonal errors, focused on maximizing inner product accuracy (MIPS):

$$L_\text{aniso}(C,\alpha;X) = \sum_{i=1}^n \min_{j \in [k]} \left[ h_\parallel(\|x_i\|)\|r_\parallel\|_2^2 + h_\perp(\|x_i\|)\|r_\perp\|_2^2 \right],$$

where $r_\parallel$ and $r_\perp$ are component errors along and orthogonal to $c_j$ (Krishnan et al., 2021).

3. Quantized Variants and Encoding Schemes

To reduce storage beyond full-precision scaling coefficients $\alpha_i$, quantized variants such as Q-PCPQ and Q-APCPQ constrain all scale values in each section to a set of $s$ quantization points $\{\lambda_1, ..., \lambda_s\}$. The minimization solves a 1D $k$-means on the projection coefficients, alternately optimizing cluster assignments and quantized values:

$$\min_{\bar u_j: |\text{Im}(\bar u_j)| \leq s} \sum_{j=1}^k \| X_j - \bar u_j v_j^\top \|_F^2.$$

After index training, each data point is encoded per block by

• $\phi_j(i) \in [k]$ (center/direction index),
• $\gamma_j(i) \in [s]$ (quantized scalar index).

Reconstruction in $\mathbb{R}^d$ is:

$$\tilde x_i = [ \lambda_{\gamma_1(i)} c_{\phi_1(i)}^{(1)}, ..., \lambda_{\gamma_m(i)} c_{\phi_m(i)}^{(m)} ].$$

Dot product approximation for query $q = [q_1,...,q_m]$ utilizes precomputed inner products in each block, summed over the appropriate indices (Krishnan et al., 2021).

4. GPU-Optimized PQ Tessellation for Large-Scale Inference

Applications in LLMs require efficient partitioning for compressed key-value (KV) cache storage. The MILLION framework employs PQ tessellation to split each KV vector $x \in \mathbb{R}^{d_k}$ into $M$ subspaces, learning $2^b$-centroid codebooks per subspace via $k$-means. The quantized codes $(k_1,...,k_M)$ encode $x$; dequantization reconstructs via table lookup.

MILLION’s PQ is non-uniform: $k$-means clustering densifies centroids in subspaces/channels with high variance, "immunizing" quantization against amplitude and standard deviation outliers common in KV caches. Explicit sparse outlier handling yields negligible accuracy gains (\textless1% reduction), but incurs additional memory and indexing overhead, so PQ's flexible centroid allocation alone is preferred (Wang et al., 12 Mar 2025).

On supported GPUs, fused CUDA kernels integrate table lookup and attention kernel computation, eliminating explicit global dequantization. Asynchronous streams quantize new keys/values in the background, overlapping compute and quantization for end-to-end speedup. Empirically, 4-bit PQ delivers over 2x inference speedup with less than 1 point of accuracy or perplexity degradation in 32K+ context LLMs (Wang et al., 12 Mar 2025).

5. Differentiable Tessellation with Voronoi Quantization

Recent approaches generalize PQ via differentiable tessellation using learnable anchor points, as in the differentiable Voronoi tessellation framework (Chen et al., 2022). Here, the codebook consists of $K$ learnable anchor vectors $\{x_k\}$ in $\mathbb R^D$, and each Voronoi cell $V_k$ is defined by:

$$V_k = \{ z \in \mathbb R^D : \|z - x_k\|^2 < \|z - x_i\|^2 \ \forall i\neq k,\ c_\ell < z < c_r \}.$$

A bijective, analytic mapping $f_k : \mathbb R^D \rightarrow V_k$ with tractable Jacobian ensures full differentiability and enables end-to-end optimization of quantization boundaries. This semi-discrete approach can learn non-axis-aligned, flexible polytopes, supporting structured quantization and facilitating "Voronoi dequantization" for discrete-to-continuous latent mappings in flows.

Encoding is by nearest-anchor search, and unlike PQ, the number of codebook entries $K$ and dimension $D$ can be decoupled. The per-point cost is dominated by nearest-anchor lookup plus a constraint solve, remaining practical for GPU acceleration (Chen et al., 2022).

6. Theoretical Guarantees, Complexity, and Empirical Results

Product quantization tessellation schemes exhibit several theoretical properties:

• For (Q-)PCPQ, quantization loss is bounded by the optimal projective clustering loss plus the 1D $k$-means error on the $\alpha$'s; dot-product approximation error is bounded pointwise by $\|q\|_2\|x_i-\tilde{x}_i\|_2$ (Krishnan et al., 2021).
• Alternating minimization in clustering converges after a few iterations, each costing $O(nk\bar{d})$ per section ($\bar{d} = d/m$).

For large-scale PQ:

• Query-time complexity is $O(kd + skm + nm)$ for $k$ clusters, $s$ scalars, $m$ blocks, $n$ database points.
• Storage per vector is $m(\log_2 k+\log_2 s)$ bits (Krishnan et al., 2021).
• MILLION achieves $2\times$ to $3\times$ inference speedup over half-precision baselines without significant accuracy loss at 4 bits (Wang et al., 12 Mar 2025).

Differentiable tessellation in normalizing flow applications realizes test set log-likelihood and bits-per-character gains across diverse structured datasets, strictly improving upon previous dequantization methods (Chen et al., 2022).

7. Comparison and Applications Across Domains

Tessellation Type	Cell Geometry	Learning Paradigm
Classical PQ	Axis-aligned boxes	Non-differentiable k-means
Projective Clustering PQ	Cones, lines	Alternating minimization
Voronoi Tessellation	General polytopes	End-to-end backprop

Classical and projective PQ tessellations dominate in fast similarity search, inner-product estimation, and LLM inference due to their tractable encoding/decoding and efficient implementation. Differentiable tessellation methods enable flexible, data-adaptive quantization boundaries critical for semi-discrete normalizing flows and structured generative models. Product quantization tessellation thus forms the foundation for high-performance, scalable vector quantization in both classical and modern neural data processing pipelines (Krishnan et al., 2021, Wang et al., 12 Mar 2025, Chen et al., 2022).