Optimized Product Quantization

PQ's blind spot: the fixed split

Recall base PQ. A vector x ∈ R^D is split into M contiguous subvectors of dimension D/M; each subvector is quantized against its own k-entry codebook (typically k = 256, so 8 bits per subquantizer), and the code is the concatenation of the M nearest-codeword indices. The reconstruction is q(x) = [c^1, c^2, …, c^M], and training minimizes total quantization error Σ_i ||x_i − q(x_i)||^2 by running k-means independently in each subspace.

The weak point is the word contiguous. PQ takes dimensions 1…D/M as subspace 1, the next block as subspace 2, and so on — an arbitrary, data-blind partition. If the embedding puts 90% of its energy in a handful of directions (as PCA spectra of real features routinely show), those high-variance dimensions crowd into a few subspaces while others get near-constant slices. A codebook spent on a low-variance slice is wasted; a high-variance slice is under-resourced. Correlation across subspace boundaries compounds the loss — PQ’s error is hostage to the coordinate system it was handed.

The key insight: a rotation is free

OPQ’s move rests on one fact: an orthogonal matrix R (with R^T R = I) preserves Euclidean distance. For any two vectors, ||Rx − Ry|| = ||x − y||, since (x−y)^T R^T R (x−y) = (x−y)^T (x−y). Rotating the entire dataset therefore leaves the nearest-neighbor structure — the thing we search — completely intact.

But while distances are invariant, how those distances distribute across the coordinate axes is not. A rotation can take an ellipsoidal cloud whose long axis lies at 45° to the axes and re-orient it so the variance lines up with the coordinates we are about to slice along. So OPQ has a free parameter to exploit: pick the R that makes the post-rotation coordinate split as quantization-friendly as possible. We quantize Rx, store the same compact codes, and rotate the query the same way at search time. Retrieval semantics are unchanged; only the reconstruction error shrinks. That is the whole idea — the rest is how to find a good R.

The OPQ objective

Where base PQ optimizes codebooks alone, OPQ optimizes the rotation and the codebooks jointly. Writing q(·) for the PQ quantizer (concatenate the nearest codeword from each of the M subspaces), the objective is:

minimize    Σ_i  || R x_i  −  q(R x_i) ||^2
  R, C

subject to  R^T R = I         (R is D x D, orthogonal)
            q(y) = [ c^1(y_1), ..., c^M(y_M) ]   (PQ on the rotated y)

Two knobs, one loss. The codebooks C = {c^m} are the usual PQ centroids; the new variable is the orthogonal R. Setting R = I recovers ordinary PQ, so OPQ can only do at least as well — the optimum over a larger feasible set cannot be worse. The constraint R^T R = I keeps the transform distance-preserving; drop it and you could collapse dimensions to cheat the loss while destroying the geometry. The difficulty is that R and C are coupled: the best split depends on the codebooks, and the best codebooks depend on the split. OPQ breaks that coupling two ways — a parametric closed form and a non-parametric alternating loop.

Why rotation helps: the distortion math

To see what a good R should achieve, model the rotated data as zero-mean Gaussian with covariance Σ. High-rate quantization theory says the minimum achievable distortion of a d-dimensional subspace, at a fixed bit budget, is proportional to the geometric mean of its variances — equivalently to (∏ λ_j)^(1/d), the d-th root of the product of the eigenvalues assigned to that subspace. Total distortion is the sum over the M subspaces:

D_total  ∝  Σ_m  ( ∏_{j ∈ subspace m} λ_j )^(1 / (D/M))

Minimizing this reveals two design rules. First, decorrelate: R should diagonalize Σ — align the axes with the principal directions so no correlation leaks across slice boundaries. Second, balance the variance: for a fixed total, a sum of geometric means is smallest when the per-subspace products ∏ λ_j are all equal (a consequence of the AM–GM inequality). So the ideal R both removes correlation and spreads variance evenly across the M groups — exactly the two things the contiguous split fails to do.

The parametric (Gaussian) solution

Those two rules give a direct recipe when the data is roughly Gaussian. Step one, decorrelate: take the eigendecomposition of the covariance, Σ = U Λ U^T, and use the eigenvectors as a PCA rotation — this diagonalizes the data. Step two, balance: assign the D eigenvalue-directions to the M subspaces (each of size D/M) so the product of eigenvalues in every subspace is as equal as possible. A greedy pass — sort eigenvalues descending, drop each into the subspace with the currently smallest product — does the job. The final R is the PCA rotation composed with the permutation that realizes this assignment.

Worked example. Say D = 4, M = 2 (two 2-D subspaces), with eigenvalues λ = [100, 50, 2, 1]. The naive contiguous split after PCA pairs the two largest and the two smallest: products 100×50 = 5000 and 2×1 = 2 — wildly imbalanced. The balanced assignment pairs largest-with-smallest: {100, 1} and {50, 2}, giving products 100 and 100 — perfectly equal. Same eigenvalues, same PCA, but the balanced grouping slashes the summed geometric-mean distortion.

The non-parametric solution: alternating optimization

Real embeddings are not exactly Gaussian, so OPQ also offers a data-driven solver that optimizes the actual objective directly by alternating between the two coupled variables, holding one fixed while solving the other. Initialize R from the parametric solution above (a strong warm start), then iterate until the loss stops dropping:

repeat:
  (1) fix R  ->  train codebooks C on the rotated data { R x_i }
               = standard PQ: run k-means in each of the M subspaces

  (2) fix C  ->  solve for the best orthogonal R given the current
               reconstruction targets y_i = q(R x_i)   [Procrustes]

Step (1) is just plain PQ applied to Rx: quantize the rotated vectors and update the per-subspace centroids. Step (2) is the new piece. With the codebooks and each point’s codeword assignment frozen, every point has a fixed target y_i = q(R x_i), and we ask which orthogonal rotation best maps the raw x_i onto those targets. Each half-step can only lower (or hold) the shared loss, so the alternation converges monotonically to a local optimum — the same guarantee that underwrites Lloyd’s k-means.

Solving for R: the Orthogonal Procrustes problem

Step (2) has a clean closed form. With targets fixed, we minimize

minimize  Σ_i || R x_i − y_i ||^2   subject to  R^T R = I

Expand the squared norm. Because R is orthogonal, ||R x_i||^2 = ||x_i||^2 is constant, and ||y_i||^2 does not involve R, so only the cross term survives. Minimizing the loss is equivalent to maximizing Σ_i y_i^T R x_i = tr(R · Σ_i x_i y_i^T) = tr(R B), where B = X Y^T = Σ_i x_i y_i^T is the D × D correlation between raw data and targets. This is the classic Orthogonal Procrustes problem, solved by one SVD:

B = X Y^T          (D x D)
SVD:  B = U Σ V^T
R* = V U^T         (orthogonal; maximizes tr(R B))

Intuitively, B’s singular vectors are the correlated directions between inputs and quantized targets, and R* = V U^T rotates the inputs to line up with them. One SVD of a D × D matrix per outer iteration — cheap, exact, and why OPQ needs no gradient descent over the orthogonal manifold.

Shapes, complexity, and cost

OPQ adds exactly one object to a PQ index: the rotation matrix R, shape [D, D]. At build and query time it costs one extra matrix–vector product Rx, which is O(D^2) per vector — negligible next to PQ’s distance-table lookups, and done once per query, not per database vector. The compressed codes are identical in size to PQ’s: M bytes per vector at k = 256. Memory overhead is the single D×D matrix shared across the index — a rounding error against millions of codes.

Training is where the extra work lives: each outer iteration runs a full round of per-subspace k-means (step 1) plus one D×D SVD (step 2), converging in a couple of dozen iterations. For the common case where D is a few hundred to a thousand, the SVD is milliseconds and k-means dominates — so OPQ training is a small constant factor over PQ, paid once, offline. At serving time the marginal cost is essentially zero: rotate the query, then run the exact same asymmetric-distance lookups PQ uses.

Recall gains, practice, and pitfalls

The payoff is higher recall at the same bit budget. On real feature sets — SIFT/GIST descriptors and deep embeddings alike — OPQ measurably reduces quantization error and lifts recall@k over plain PQ, with the gap widest where PQ is weakest: anisotropic, correlated data whose energy is concentrated in a few directions. The dividend is largest under aggressive compression (small M), where every wasted codebook hurts. OPQ ships as a first-class option in FAISS (the OPQ pre-transform, e.g. OPQ64,PQ64) and in vector databases built on it, precisely because it is a near-free upgrade — same codes, same query path, one added matmul.

Pitfalls to respect. The rotation is part of the model: you must persist R and apply it to every query, or the codes decode into nonsense. Warm-start the alternation from the parametric PCA-plus-balance solution; a random R converges slowly to a worse local optimum. Gains shrink when the data is already near-isotropic (little for a rotation to fix) and, in an IVF pipeline, partly overlap with structure the coarse quantizer already removes — so measure recall on your own data rather than assuming a fixed uplift.