Codebooks that add up · Close Enough

1The sum trick

\kvec{x} \;\approx\; \kdec{\hat x} \;=\; \kbook{c^{1}}_{\kcode{i_1}} + \kbook{c^{2}}_{\kcode{i_2}} + \cdots + \kbook{c^{M}}_{\kcode{i_M}}, \qquad \kcode{i_m} \in \{0, \dots, K-1\}

The code for $x$ is the tuple $(\kcode{i_1}, \dots, \kcode{i_M})$ : $M \log_2 K$ bits. With $M = 8$ and $K = 256$ that is 8 bytes.

2,048 vectors stored ( $M \cdot K$ )
2⁶⁴ points reachable ( $K^M = 256^8$ )
8 B per code

Fig 4.2.1Three codebooks that add up

Drag the yellow point Three codebooks of four arrows each: big, medium, small. Pick one arrow from each and chain them. Only 12 arrows are stored, yet 4 × 4 × 4 = 64 grey points are reachable.

Now the smallest interesting case with data: two codebooks, A and B, four arrows each.

Fig 4.2.2Sums of codewords

Drag the arrow tips Each mint dot is Aᵢ + Bⱼ, so moving one arrow moves four dots at once. Then switch to the product constraint and press Fit: the sums are forced onto an axis-aligned grid, which wastes points on a tilted cloud.

2Two things to take from playing with it

Decoding is trivial. Given the code, $\hat x$ is $M$ table lookups and $M-1$ vector additions. The decoder never sees the giant implicit codebook.

Encoding got hard. To find the best code we would have to consider all $K^M$ combinations: the same $2^{64}$ we called impossible. The big search has not gone away. It has moved from the decoder to the encoder.

3Restrict: product quantization

The Product button is product quantization (PQ; Jégou, Douze and Schmid, 2011). Split the $d$ dimensions into $M$ blocks. Codebook $m$ only has non-zeros in block $m$ . Codewords from different codebooks are then perpendicular, and the error separates:

\Big\lVert \kvec{x} - \sum_m \kbook{c^m}_{i_m} \Big\rVert^2 \;=\; \sum_{m=1}^{M} \big\lVert \kvec{x^{(m)}} - \kbook{c^{m}}_{i_m} \big\rVert^2

Each term depends on one index only, so each is minimised separately: $M$ independent nearest-neighbour searches over $K$ entries. That is the whole encoder. Training is k-means in each block.

Product quantizationnumpy

def pq_encode(X, codebooks):        # M arrays of shape (K, d/M)
    blocks = np.split(X, len(codebooks), axis=1)
    return np.stack([nearest(B, C) for B, C in zip(blocks, codebooks)], axis=1)

def pq_decode(codes, codebooks):
    return np.concatenate([C[codes[:, m]] for m, C in enumerate(codebooks)], axis=1)

The price is what you saw. The implicit codebook is a grid aligned with the blocks. When dimensions are correlated across blocks, and in embeddings they always are, much of the grid sits where no data lives.

OPQ (Ge et al., 2013) learns a rotation first, tilting the grid toward the data. It helps a great deal; the grid is still a grid.

PQ / OPQdistortion mediumencode K·d, trivialsearch M lookupstraining k-means per block