Be greedy

1Quantize what is left

Now let every codebook use all $d$ dimensions. The simplest encoder:

Train the same way: k-means on $x$, then k-means on the residuals, and so on. This is residual quantization (RQ), an old idea from signal coding (Juang and Gray, 1982) that came back for vector search and for neural audio codecs.

def rq_encode(X, codebooks):        # M arrays of shape (K, d)
    R, codes = X.copy(), []
    for C in codebooks:
        i = nearest(R, C)
        codes.append(i)
        R -= C[i]                   # the next stage sees only what's left
    return np.stack(codes, axis=1)

def additive_decode(codes, codebooks):
    return sum(C[codes[:, m]] for m, C in enumerate(codebooks))

Greedy is cheap: $O(MKd)$, the same order as PQ. But it is not optimal.

2The shape of the failure

Stage $m$ chooses without knowing what stages $m+1, \dots, M$ can do. The readout under the figure says how often that matters in the toy; press New codebooks a few times to see it vary.

3A side effect worth remembering

This is used in three places later:

• One model, many bitrates. Neural audio codecs randomly drop later stages during training so that one model serves every bitrate (4.9 Bells and whistles).
• Depth chosen afterwards. For lossless coding you can train 24 stages and decide how many to use later, even per vector (6.7 Where should the bits go?).
• Refinement on the wire. Send the first stages now and the rest if the client asks (6.9 Two codecs).

4Measured