Why product quantization wants a rotation first

Product quantization cuts a vector into contiguous slices and replaces each slice with a codebook index. The compression is dramatic and the mechanism is simple, which conceals an assumption sitting inside the very first step: that cutting the vector at positions 0–7, 8–15, 16–23 and so on is a reasonable way to divide it. That assumption is usually false, the cost is unevenly distributed error, and the fix is a rotation applied before encoding — a change that costs nothing at query time and can matter a great deal.

The naive approach

Split the d-dimensional vector into m contiguous slices of d/m components each. Train a codebook per slice position. Encode each slice to its nearest codebook entry.

Each slice gets the same number of codebook entries — usually 256, one byte — and therefore the same share of the bit budget. That is the assumption: every slice is allocated equal representational capacity, so every slice had better carry a comparable amount of information.

How it fails: embedding coordinates are neither equally variable nor independent. Some coordinates have far larger variance than others. Groups of coordinates are correlated with each other, and those groups have no reason to be contiguous — the ordering of components in an embedding is an artifact of the model’s arbitrary basis and carries no meaning at all, as the space can be rotated freely without changing any distance.

So the contiguous split distributes the variance arbitrarily across slices. One slice may contain several high-variance coordinates and another almost none. And the total quantization error is the sum of the per-slice errors, which produces a specific, avoidable waste.

Why unequal variance across slices is expensive

Take a concrete illustrative case with numbers chosen to make the arithmetic obvious. Suppose two slices, each getting 256 codebook entries, and suppose slice A’s data has ten times the spread of slice B’s.

K-means places 256 representatives within whatever region its data occupies. In slice A, 256 points must cover a large region, so the typical distance from a point to its nearest representative is large. In slice B, 256 points cover a small region and fit it snugly, so the residual error is small.

Total error is the sum. Slice A dominates it — and slice B’s bits are largely wasted, buying precision finer than the data’s own variation. We spent equal capacity on unequal problems, and the error is set by the worst slice.

Two independent repairs exist for this. Allocate unequal bits per slice — more codebook entries for high-variance slices — which works, and complicates both the storage layout and the distance table. Or make the slices equally hard, which is what the rotation does, and it leaves the format untouched.

There is a second cost, subtler and just as real. PQ quantizes each slice independently, so it cannot represent any correlation that spans a slice boundary. If coordinates 7 and 8 are strongly correlated — they always move together — then a joint codebook over both would need very few entries to capture their behaviour, because they only ever occupy a thin diagonal band of their two-dimensional space. Split between slices, each is quantized as if it were free to take any value independently, and the codebooks spend entries on combinations that never occur. Correlation across a cut is capacity thrown away.

The idea: rotate before you cut

The vector can be rotated without changing any distance between vectors — that is the defining property of a rotation. So we are free to apply any rotation R we like to every vector before encoding, and any distance computed among rotated vectors equals the distance among the originals.

That gives us a free variable, and we can choose R to make the subsequent slicing better. Two goals, both achieved by the same move:

Balance the variance across slices. Choose R so that after rotation, each group of d/m coordinates carries a comparable share of the total variance. Then every codebook faces a problem of comparable difficulty, no slice dominates the error, and no slice’s bits are wasted.

Decorrelate within slices, or at least across the cuts. Choose R to remove correlation that spans slice boundaries, so that the independence assumption PQ makes becomes closer to true. Rotating to a basis where the coordinates are uncorrelated — the principal directions of the data — does exactly this globally, and then the remaining question is just how to allocate those directions to slices.

Encoding becomes:

train(sample, m, bits):
    R ← learn_rotation(sample, m)
    rotated ← { R·v : v in sample }
    for each slice position p in 1..m:
        codebook[p] ← k-means(rotated slice p, k = 2^bits)

encode(v):
    r ← R·v
    return [ nearest entry in codebook[p] to slice p of r  for p in 1..m ]

query_tables(q):
    r ← R·q                                  # the query is rotated too, and stays exact
    for p, for each entry e in codebook[p]:
        table[p][e] ← d(slice p of r, e)²

Note where the rotation lands at query time. It is one matrix multiply against the query, once per query, and then the search proceeds exactly as before: build the distance tables, sum table lookups per candidate. The per-candidate cost — the part multiplied by millions — is completely unchanged. This is why the technique is close to free: the rotation is applied to the one vector in flight, not to any of the stored ones, which were rotated at build time.

Choosing the rotation

Three approaches in increasing order of ambition, all of them producing an orthogonal matrix.

A random rotation. Draw a random orthogonal matrix. This sounds crude and it captures a surprising share of the benefit, because a random rotation mixes coordinates: each output coordinate becomes a combination of many input ones, so variance and correlation are spread roughly evenly across all of them by averaging. It removes the pathological cases — a slice of near-constant coordinates, a correlated pair straddling a cut — without knowing anything about the data. It requires no training and it is a reasonable default.

Principal-direction rotation with balanced allocation. Rotate to the data’s principal axes, which makes the coordinates uncorrelated and orders them by variance. Then the allocation problem is explicit: assign the principal directions to slices so each slice’s total variance is comparable — a greedy assignment of directions to the currently-lightest slice works. This is deterministic, cheap, and it attacks both goals directly.

A jointly learned rotation. Treat R and the codebooks as parameters of one optimisation: minimise total quantization error over both. Alternate — fix R, retrain codebooks; fix codebooks, solve for the best R given them. This is the “optimised product quantization” family, and it produces the lowest error of the three because it optimises the actual objective rather than a proxy for it. The cost is training time, and the extra storage is one d × d matrix, which is negligible next to the collection.

The trade-off among them is entirely a build-time one. All three have identical query cost and identical storage per vector.

What it gives up

Nothing at query time, beyond a single matrix multiply on the query, which is one dense operation against a cost dominated by millions of table lookups.

A dependence on the training distribution. A learned rotation is fitted to a sample, and if the collection’s distribution shifts, the rotation is balancing variance according to a picture that is no longer accurate. The failure is graceful — a mismatched rotation is not worse than no rotation, it is merely less helpful — but it is one more thing that is quietly stale after a large distribution change, alongside the codebooks and the centroids.

Interpretability of the stored codes. After rotation the codes correspond to nothing in the original coordinate space. This costs nothing real, since embedding coordinates were never interpretable anyway, but it does mean you cannot inspect a code and relate it to a slice of the original vector.

It doesn’t fix the fundamental loss. The vector is still being replaced by a concatenation of codebook entries, and the position within each cell is still discarded. Rotation improves how efficiently the bit budget is spent; it does not change the fact that a budget exists. The reranking stage is still what recovers exact ordering.

The general principle

The reason to know this goes beyond PQ. Any scheme that divides a vector into parts and allocates capacity per part is making an implicit claim about how information is distributed across those parts — and for a learned embedding, that claim is about an arbitrary basis, so it is almost certainly wrong. The same reasoning applies to any per-dimension scheme: scalar quantization computes a range per dimension and is therefore sensitive to which dimensions have outliers, and binary quantization takes the sign along each coordinate axis, which is a strictly worse choice of projection directions than random ones would be. Whenever you see a compression scheme treating coordinates as meaningful units, a rotation is a candidate improvement, and it is usually a cheap one.