Learning the codebooks

1The decoder update is least squares

Fix the codes for $N$ training vectors. Build a sparse $N \times MK$ matrix $B$ with one 1 in each of the $M$ blocks of every row, marking the entry that row uses. Stack all codebooks into one $MK \times d$ matrix $C$.

The reconstruction of the whole training set is then $\kdec{\hat X} = B\,\kbook{C}$, and the best codebooks for those codes solve a plain least-squares problem:

def update_codebooks(X, codes, K):
    N, M = codes.shape
    B = np.zeros((N, M * K))                  # sparse in practice
    B[np.arange(N)[:, None], np.arange(M) * K + codes] = 1
    C = np.linalg.lstsq(B, X, rcond=None)[0]  # (M*K, d)
    return np.split(C, M)

codebooks = rq_train(X)                       # the starting point matters a lot
for epoch in range(10):
    codes = encode(X, codebooks)              # beam / ICM: expensive
    codebooks = update_codebooks(X, codes, K) # least squares: cheap

This is k-means with a different assignment step.

Training inherits the encoder’s cost.

Every epoch re-encodes the whole training set with the expensive encoder. That is why LSQ-style methods spent much of their engineering on fast GPU encoding.

2The overfitting trap

Three defences, in order of preference:

1. More data than parameters. Marginalia has 300 million vectors to train on; the experiment here had 63,120, about one sample per ten parameters.
2. Regularise. A ridge term in the least-squares solve keeps rarely used codewords from chasing noise.
3. Always score on held-out vectors. It is the only way to notice.