Optimized Product Quantization and Score-Aware Quantization

Overview & motivation

The previous topic left a clean cliffhanger. Product quantization splits a vector into $m$ disjoint subvectors and quantizes each subspace independently, and it works because squared Euclidean distance separates over disjoint coordinate blocks. But that separation is also its weakness: the quantizer treats the subspaces as if they were independent and gives each the same number of bits. On a deliberately variance-imbalanced cloud — coordinates whose variance decays geometrically — raw product quantization paid a reconstruction distortion of about $85.6$ , because the equal bit budget over-quantizes the low-variance tail and starves the high-variance head. A rotation that balanced the variance across subspaces cut that distortion to about $15.4$ . We showed the direction of the effect and named the fix; this topic is the fix.

There are two independent ways to make a product quantizer better, and they are orthogonal in the literal sense. The first changes the axes: rotate the space before the split so each subspace carries a balanced share of the data’s information. This is optimized product quantization (OPQ), and it is the first half of the topic. The second changes the loss: when the downstream task is maximum-inner-product search rather than reconstruction, the error that matters is not how far the reconstruction sits from the vector, but how much it disturbs the inner products that high-scoring queries see. This is score-aware, or anisotropic, quantization — the idea behind ScaNN — and it is the second half. The laboratory below lets you switch between the four rotations, watch the alternating optimization converge, and raise the anisotropy to see a quantizer trade reconstruction for inner-product fidelity.

PQ distortion ∑‖x−Q(x)‖²

85.58

vs raw

0% lower

per-subspace variance (m = 4) — OPQ balances the product

subspace 11.35

subspace 20.17

subspace 30.02

subspace 40.00

distortion across rotations

raw (native axes)85.6

PCA-only (naive)83.3

balanced heuristic15.4

learned OPQ14.4

A naive PCA-alignment barely helps (85.6 → 83.3): it concentrates variance into subspace 1 and starves the rest. Balancing variance across subspaces collapses the distortion (→ 15.4), and the learned OPQ rotation refines it further (→ 14.4) via the alternating optimization in the next panel. The balanced and OPQ variance bars look alike because OPQ's extra gain is a within-block rotation, not a different split.

The first panel is the on-ramp made precise: switch among the raw axes, a naive PCA-alignment, the previous topic’s balancing heuristic, and the learned OPQ rotation, and read the distortion drop and the per-subspace variance bars. The second panel runs the alternating optimization and shows its distortion descending monotonically below the heuristic. The third panel is the score-aware story: a datapoint, its candidate codewords, and the residual split into the part that moves the inner product and the part that does not.

The rotation is a free isometry

Throughout, $x \in \mathbb{R}^d$ is a database vector, $R \in O(d)$ is an orthogonal matrix ( $R^{\mathsf T} R = I$ ), and the data splits into $m$ contiguous subspaces of width $d/m$ . A rotated product quantizer first maps $x \mapsto Rx$ , then product-quantizes the rotated vector.

Definition 1 (Rotated product quantizer).

Fix an orthogonal $R \in O(d)$ and a product codebook $\mathcal{C} = \mathcal{C}^1 \times \cdots \times \mathcal{C}^m$ with $k^\*$ centroids per subspace. The rotated product quantizer encodes $x$ by the tuple of sub-centroid indices of $Rx$ and reconstructs

Q_R(x) \;=\; R^{\mathsf T}\, c(Rx), \qquad c(Rx) \in \mathcal{C},

where $c(\cdot)$ is the ordinary product-quantization codeword. Optimized product quantization minimizes the reconstruction distortion jointly over the rotation and the codebooks:

\min_{R^{\mathsf T} R = I,\;\mathcal{C}^1,\dots,\mathcal{C}^m}\ \sum_{i} \big\lVert R x_i - c(R x_i) \big\rVert^2 .

Why is it legitimate to rotate the data before quantizing — does it not change the retrieval problem? It does not, and this is the whole license for the method.

Theorem 1 (Rotation is a distance- and inner-product-preserving isometry).

For any orthogonal $R$ and any $x, y \in \mathbb{R}^d$ ,

\lVert Rx - Ry \rVert = \lVert x - y \rVert, \qquad \langle Rx, Ry \rangle = \langle x, y \rangle .

Consequently, rotating every database vector and every query by the same $R$ leaves all pairwise distances and all inner products unchanged, so the true nearest neighbors and the true maximum-inner-product results are identical. Since ordinary product quantization is the special case $R = I$ , and the feasible set of optimized product quantization strictly contains it, the optimized distortion never exceeds the unoptimized one:

E_{\mathrm{OPQ}} \;\le\; E_{\mathrm{PQ}} .

Proof.

Orthogonality gives $\lVert Rx \rVert^2 = (Rx)^{\mathsf T}(Rx) = x^{\mathsf T} R^{\mathsf T} R x = x^{\mathsf T} x = \lVert x \rVert^2$ , and the same computation applied to $x - y$ and to the polarization identity $\langle x, y \rangle = \tfrac12(\lVert x \rVert^2 + \lVert y \rVert^2 - \lVert x - y \rVert^2)$ gives both equalities. Encoding $x$ as the nearest codeword of $Rx$ and decoding through $R^{\mathsf T}$ is therefore a quantizer on the original space whose distortion equals that of ordinary product quantization applied to the rotated data. Optimizing over $R$ includes $R = I$ , so the minimum can only decrease. $\blacksquare$

∎

The rotation is, in this exact sense, free: it cannot move a single true neighbor, but it adds $d(d-1)/2$ degrees of freedom the quantizer can spend to reduce distortion. The companion code verifies the isometry directly and confirms that rotating both the database and the queries leaves every true nearest neighbor in place. The only question left is which $R$ to choose.

Parametric OPQ: balance the product, not the sum

When the data is well modeled by a Gaussian, the optimal rotation has a closed form. Write $\Sigma$ for the covariance of $x$ , and after rotation let $\hat\Sigma = R \Sigma R^{\mathsf T}$ , partitioned into $m$ diagonal blocks $\hat\Sigma_{jj}$ of size $(d/m) \times (d/m)$ — one per subspace. High-rate quantization theory bounds the distortion of a $k$ -centroid quantizer of a Gaussian by a constant times $|\Sigma|^{1/d}$ , the determinant to the $1/d$ power. Applying that bound per subspace and summing gives the quantity OPQ must minimize.

Theorem 2 (Parametric optimum: decorrelate and balance the determinant).

Under $x \sim \mathcal{N}(0, \Sigma)$ and the high-rate bound, the product-quantization distortion after rotation satisfies

E(R) \;\ge\; (k^\*)^{-2m/d}\,\frac{d}{m}\,\sum_{j=1}^{m} \big| \hat\Sigma_{jj} \big|^{m/d},

and minimizing the right-hand side over $R$ is equivalent to minimizing $\sum_{j=1}^{m} |\hat\Sigma_{jj}|^{m/d}$ . By the arithmetic–geometric-mean inequality and Fischer’s inequality,

\sum_{j=1}^{m} \big| \hat\Sigma_{jj} \big|^{m/d} \;\ge\; m \Big(\prod_{j=1}^{m} \big|\hat\Sigma_{jj}\big|\Big)^{1/d} \;\ge\; m\,|\Sigma|^{1/d},

a bound that is constant in $R$ . It is attained if and only if (i) the off-diagonal blocks of $\hat\Sigma$ vanish — the subspaces are decorrelated, achieved by aligning to the principal axes — and (ii) the block determinants $|\hat\Sigma_{jj}|$ are equal across subspaces, i.e. the product of the eigenvalues assigned to each subspace is balanced.

The two conditions are distinct, and both are needed. Diagonalizing $\Sigma$ by PCA makes Fischer’s inequality tight — it removes cross-subspace correlation — but it does not say which eigen-axis belongs to which subspace. The arithmetic–geometric-mean step is tight only when the products of eigenvalues match across subspaces, and that is a question of allocation: distribute the $d$ eigenvalues into $m$ groups so the products are as equal as possible. The greedy eigenvalue allocation does exactly this — it repeatedly places the largest unassigned eigen-axis into the subspace with the smallest current product (equivalently, the smallest sum of log-eigenvalues).

This is where the previous topic’s heuristic and the true optimum part ways, and the difference is provable rather than cosmetic. The balancing heuristic equalized the sum of variances across subspaces — the trace of each block. The parametric optimum equalizes the product — the determinant. Under the high-rate bound the distortion is governed by the geometric mean of the variances, not the arithmetic mean, so balancing the sum is only an approximation while balancing the product is exact. In the companion code, the spread of per-subspace log-variances on the imbalanced cloud falls from about $12.2$ under the raw contiguous blocks to about $0.3$ under the eigenvalue allocation: the determinant is balanced almost perfectly, which the sum-balancing heuristic does not achieve.

The Gaussian assumption is doing real work here, and we flag it: on data that is not Gaussian, the eigenvalue allocation is a principled construction and an excellent starting point, but not a certified optimum. For that we need an algorithm that makes no distributional assumption at all.

Non-parametric OPQ: alternating optimization

Drop the Gaussian model and return to the raw objective of Definition 1. It couples the rotation and the codebooks, and it is not jointly convex — but it splits into two subproblems, each of which we can solve to its global optimum, and alternating between them is a descent.

Theorem 3 (The rotation step is an Orthogonal Procrustes problem).

Hold the codebooks fixed, and let $Y$ collect the current reconstructions $c(Rx_i)$ — one per training vector — with the data $X$ and reconstructions $Y$ stacked as columns. The rotation that minimizes the reconstruction error,

\min_{R^{\mathsf T} R = I}\ \big\lVert R X - Y \big\rVert_F^2,

has the closed form

X Y^{\mathsf T} = U \Sigma V^{\mathsf T} \quad\Longrightarrow\quad R = V U^{\mathsf T},

read directly off the singular value decomposition of the cross-covariance $X Y^{\mathsf T}$ . Holding the rotation fixed instead, the codebook subproblem is ordinary product quantization on the rotated data $RX$ — $m$ independent Lloyd problems. Alternating the two steps produces a sequence of distortions that is monotone non-increasing and converges.

Proof.

Expand the objective: $\lVert RX - Y \rVert_F^2 = \lVert X \rVert_F^2 + \lVert Y \rVert_F^2 - 2\,\mathrm{tr}(R X Y^{\mathsf T})$ , so minimizing it is maximizing $\mathrm{tr}(R\, X Y^{\mathsf T})$ . With $X Y^{\mathsf T} = U \Sigma V^{\mathsf T}$ and $Z = V^{\mathsf T} R U$ orthogonal, $\mathrm{tr}(R X Y^{\mathsf T}) = \mathrm{tr}(R U \Sigma V^{\mathsf T}) = \mathrm{tr}(Z^{\mathsf T} \Sigma) = \sum_i \Sigma_{ii} Z_{ii} \le \sum_i \Sigma_{ii}$ , since the singular values are non-negative and the diagonal of an orthogonal matrix is bounded by one in magnitude. Equality holds when $Z = I$ , i.e. $R = V U^{\mathsf T}$ . Each step is the global optimum of its subproblem — the codebook step by Lloyd’s monotonicity, the rotation step by this closed form — so the joint distortion cannot increase. Bounded below by zero, the sequence converges. $\blacksquare$

∎

The guarantee is monotone descent, not global optimality: alternating optimization over a non-convex coupled objective reaches a local optimum that depends on the initialization, and the standard remedy is to warm-start from the parametric eigenvalue-allocation rotation, which is what the companion code does. On the imbalanced cloud the trajectory starts at the heuristic’s distortion of about $15.4$ and descends monotonically to about $14.4$ — the second panel of the laboratory traces it. The extra gain over the balancing heuristic comes from the within-block rotation the Procrustes step finds; the per-subspace variance sums barely move, which is why the variance bars for the heuristic and for OPQ look almost identical even though the distortion differs. We cross-check the rotation step against scipy.linalg.orthogonal_procrustes rather than trust an opaque library, and assert the monotone descent as a test.

Score-aware quantization: protect the inner product

Everything so far minimized reconstruction error, $\lVert x - Q(x) \rVert^2$ . For maximum-inner-product search that is the wrong objective. What a search actually cares about is the inner-product error — how much quantization perturbs $\langle q, x \rangle$ — and, more sharply, it cares about that error only for the queries $q$ that score $x$ highly enough to retrieve it. A reconstruction that is excellent on average but distorts the high-scoring inner products is worse than one that is mediocre on average but preserves them.

Make the residual $r = x - \tilde x$ explicit and split it relative to the datapoint’s own direction.

Definition 2 (Parallel and orthogonal residual).

For a datapoint $x \ne 0$ and a reconstruction $\tilde x$ , decompose the residual $r = x - \tilde x$ into its component parallel to $x$ and the orthogonal remainder:

r_\parallel \;=\; \frac{\langle r, x \rangle}{\lVert x \rVert^2}\, x, \qquad r_\perp \;=\; r - r_\parallel, \qquad \lVert r \rVert^2 = \lVert r_\parallel \rVert^2 + \lVert r_\perp \rVert^2 .

The parallel component is the one that moves the inner product. For a query $q$ roughly aligned with $x$ — which is what a high-scoring query is — the inner-product error $\langle q, x - \tilde x \rangle$ is dominated by the part of the residual along $x$ , while the orthogonal part is largely invisible to it. Averaging over queries drawn uniformly from the unit sphere makes this exact.

Theorem 4 (Anisotropic decomposition of the score-aware loss).

Let $q$ be uniform on the unit sphere in $\mathbb{R}^d$ , and weight each query by a score-dependent weight $w(\langle q, x \rangle)$ . The expected weighted inner-product error decomposes as an anisotropic sum of the two residual magnitudes,

\mathbb{E}_q\big[\, w(\langle q, x \rangle)\,\langle q,\, x - \tilde x \rangle^2 \,\big] \;=\; h_\parallel\,\lVert r_\parallel \rVert^2 \;+\; h_\perp\,\lVert r_\perp \rVert^2,

and for any weight $w$ that is zero on negative scores and non-decreasing on positive ones, $h_\parallel \ge h_\perp$ . For the threshold weight $w(t) = \mathbf{1}[t \ge T]$ , the ratio $\eta = h_\parallel / h_\perp$ grows with dimension and, for $T$ bounded away from $0$ , has the large- $d$ form

\eta \;\approx\; (d-1)\,\frac{(T/\lVert x \rVert)^2}{1 - (T/\lVert x \rVert)^2} .

The ratio $\eta \ge 1$ is the entire content of the result: the parallel error is penalized at least as heavily as the orthogonal error, and strictly more as the score threshold $T$ rises. At $T = 0$ every query counts equally and the exact ratio is $\eta = 1$ — the isotropic endpoint that the displayed large- $d$ form, which tracks only how $\eta$ grows as $T$ rises, does not itself capture — and the loss collapses back to ordinary reconstruction. As $T \to \lVert x \rVert$ only near-perfectly-aligned queries count and $\eta \to \infty$ , penalizing only the parallel error. So the score-aware quantizer minimizes

\ell(x, \tilde x, \eta) \;=\; \eta\,\lVert r_\parallel \rVert^2 \;+\; \lVert r_\perp \rVert^2, \qquad \eta \ge 1,

and chooses, for each datapoint, the codeword that best preserves its magnitude along its own direction even at the cost of larger orthogonal error. The third panel of the laboratory shows the swap: at $\eta = 1$ the quantizer picks the Euclidean-nearest codeword; raise $\eta$ and it switches to a codeword lying exactly on the datapoint’s direction, with zero parallel residual.

The closed-form weight is a large-dimension limit, the threshold $T$ (equivalently $\eta$ ) is a tuned knob, and the additive objective above is cleanest when the database has constant norm — all of which we flag. What survives those caveats is the mechanism, and it is robust: protecting the parallel component protects the inner products a query sees.

§ Finance case study text

On the synthetic finance cloud, we train two codebooks of the same size — an isotropic $k$ -means codebook and a score-aware one that minimizes the anisotropic loss at $\eta = 8$ — and compare them on the inner-product error of each query’s true top-ten pairs, the high-scoring pairs a maximum-inner-product search would surface. The score-aware codebook lowers that error by about a third (from $0.00439$ to $0.00284$ in the companion code) and raises the MIPS recall@10 from about $12.0\%$ to $13.1\%$ , at exactly the same code length. The two improvements are the same phenomenon: a codebook that spends its representational budget on the magnitude-along-the-vector preserves the inner products that ranking depends on. For a production multimodal retrieval index whose embeddings are compared by inner product, this is the difference between a quantizer tuned for a reconstruction objective nobody queries and one tuned for the search that actually runs. The gain is a direction, not a guarantee — its size depends on the cloud, the code size, and the anisotropy — but its sign is the point.

Honest accounting

⚠ Rigor flag

Both optimality claims in this topic are conditional, and the conditions are load-bearing. Parametric OPQ — decorrelate and balance the product of per-subspace variances — is optimal only under a high-rate rate-distortion bound for a zero-mean Gaussian, and only when the dimension divides evenly into subspaces; off that model it is a principled initializer, not a certified optimum. Non-parametric OPQ is block coordinate descent: the rotation step (Orthogonal Procrustes) and the codebook step (Lloyd per subspace) are each globally optimal, so the distortion is monotone non-increasing and converges, but only to an initialization-dependent local optimum — never a certified global one, which is why the parametric rotation is used to warm-start it. The rotation is a strict isometry, so it changes no true neighbor; its only effect is lower quantization distortion at a fixed rate, measured here on the same synthetic clouds as the previous topics. Score-aware quantization’s optimality holds for a uniform-on-the-sphere query distribution; the convenient weight ratio is a large-dimension limit of a ratio of integrals; the score threshold is a tuned hyperparameter; and the clean additive training objective assumes a constant-norm database. Every numeric result — the rotation distortions, the monotone trajectory, the score-aware error reduction — is a property of these clouds and these hyperparameters, asserted as a direction in the companion code and never promised as a universal margin.

The rotated, score-aware codebook is the atom the production indexes assemble. An inverted-file index quantizes the residual after a coarse partition; a multi-vector index quantizes each token embedding; and both inherit, unchanged, the rotation that balances the determinant and the loss that protects the inner product. Those are the topics that follow.