Residual Vector Quantization — Visual Experiments¶

Goal: See how stacking small codebooks on the residual gives exponential resolution from linear parameters — the core trick of neural audio codecs.

What we'll explore¶

#	Experiment	Question answered
1	Coarse-to-fine	Does the residual really shrink stage by stage, with $\hat z$ closing in on $z$?
2	Exponential expressivity	At an equal code budget, does RVQ beat a single big codebook?
3	Bitrate scaling	Can decoding with fewer stages give a usable lower-rate reconstruction?
4	Why later stages are harder	How does the residual distribution change across stages?

---

Linked theory: research/05-residual-vq.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import warnings
warnings.filterwarnings('ignore')

np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})


def quantize(z, codebook):
    d2 = ((z[:, None, :] - codebook[None, :, :]) ** 2).sum(-1)
    idx = d2.argmin(1)
    return idx, codebook[idx]


def fit_rvq(data, n_stages, K, seed=0):
    """Fit N_q codebooks; each stage k-means-fits the residual of the previous."""
    residual = data.copy()
    codebooks = []
    for i in range(n_stages):
        km = KMeans(n_clusters=K, n_init=3, random_state=seed + i).fit(residual)
        cb = km.cluster_centers_
        codebooks.append(cb)
        _, q = quantize(residual, cb)
        residual = residual - q
    return codebooks


def encode_rvq(data, codebooks):
    """Return cumulative reconstruction after each stage and the residual norms."""
    residual = data.copy()
    z_hat = np.zeros_like(data)
    cum_recon, res_norms, all_residuals = [], [], []
    for cb in codebooks:
        _, q = quantize(residual, cb)
        residual = residual - q
        z_hat = z_hat + q
        cum_recon.append(z_hat.copy())
        res_norms.append(np.linalg.norm(residual, axis=1).mean())
        all_residuals.append(residual.copy())
    return cum_recon, np.array(res_norms), all_residuals


print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt from sklearn.cluster import KMeans import warnings warnings.filterwarnings('ignore') np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) def quantize(z, codebook): d2 = ((z[:, None, :] - codebook[None, :, :]) ** 2).sum(-1) idx = d2.argmin(1) return idx, codebook[idx] def fit_rvq(data, n_stages, K, seed=0): """Fit N_q codebooks; each stage k-means-fits the residual of the previous.""" residual = data.copy() codebooks = [] for i in range(n_stages): km = KMeans(n_clusters=K, n_init=3, random_state=seed + i).fit(residual) cb = km.cluster_centers_ codebooks.append(cb) _, q = quantize(residual, cb) residual = residual - q return codebooks def encode_rvq(data, codebooks): """Return cumulative reconstruction after each stage and the residual norms.""" residual = data.copy() z_hat = np.zeros_like(data) cum_recon, res_norms, all_residuals = [], [], [] for cb in codebooks: _, q = quantize(residual, cb) residual = residual - q z_hat = z_hat + q cum_recon.append(z_hat.copy()) res_norms.append(np.linalg.norm(residual, axis=1).mean()) all_residuals.append(residual.copy()) return cum_recon, np.array(res_norms), all_residuals print('Setup complete.')

Setup complete.

---

Experiment 1: Coarse-to-fine¶

Question: RVQ claims each stage quantizes the leftover residual of the last. Does the residual actually shrink stage by stage, and does the cumulative sum $\hat z = \sum_i e^{(i)}$ converge to $z$?

We fit an 8-stage RVQ on a 2D dataset and watch both the mean residual norm and one example point being refined.

In [2]:

np.random.seed(1)

# data: a filled annulus -> needs fine resolution to cover
ang = np.random.uniform(0, 2 * np.pi, 3000)
rad = np.sqrt(np.random.uniform(0.5, 4, 3000))
data = np.stack([rad * np.cos(ang), rad * np.sin(ang)], axis=1)

n_stages, K = 8, 4
codebooks = fit_rvq(data, n_stages, K, seed=1)
cum_recon, res_norms, _ = encode_rvq(data, codebooks)
print('mean residual norm after each stage:')
print(np.round(res_norms, 4))

np.random.seed(1) # data: a filled annulus -> needs fine resolution to cover ang = np.random.uniform(0, 2 * np.pi, 3000) rad = np.sqrt(np.random.uniform(0.5, 4, 3000)) data = np.stack([rad * np.cos(ang), rad * np.sin(ang)], axis=1) n_stages, K = 8, 4 codebooks = fit_rvq(data, n_stages, K, seed=1) cum_recon, res_norms, _ = encode_rvq(data, codebooks) print('mean residual norm after each stage:') print(np.round(res_norms, 4))

mean residual norm after each stage:
[0.6732 0.36   0.1815 0.0957 0.0508 0.0277 0.018  0.0118]

In [3]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5.3))
fig.suptitle('Each stage shrinks the residual; $\\hat z$ closes in on $z$', fontsize=13, fontweight='bold')

ax = axes[0]
ax.plot(range(1, n_stages + 1), res_norms, '-o', color='darkorange', lw=2)
ax.set_yscale('log')
ax.set_xlabel('stage'); ax.set_ylabel('mean residual norm $\\|r_i\\|$ (log)')
ax.set_title(f'Residual decay ({n_stages} stages, K={K} each)')

ax2 = axes[1]
j = 0   # one example point
ax2.scatter(data[:, 0], data[:, 1], s=4, color='lightgray', alpha=0.4)
ax2.scatter(*data[j], c='crimson', s=120, marker='X', zorder=6, label='target $z$')
approach = np.array([r[j] for r in cum_recon])
ax2.plot(approach[:, 0], approach[:, 1], '-o', ms=5, color='seagreen', zorder=5,
         label='cumulative $\\hat z$ per stage')
ax2.scatter(0, 0, c='black', s=40, zorder=4, label='start (0)')
ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1')
ax2.set_title('One point refined stage by stage')
ax2.legend(loc='upper left', fontsize=9)
ax2.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5.3)) fig.suptitle('Each stage shrinks the residual; $\\hat z$ closes in on $z$', fontsize=13, fontweight='bold') ax = axes[0] ax.plot(range(1, n_stages + 1), res_norms, '-o', color='darkorange', lw=2) ax.set_yscale('log') ax.set_xlabel('stage'); ax.set_ylabel('mean residual norm $\\|r_i\\|$ (log)') ax.set_title(f'Residual decay ({n_stages} stages, K={K} each)') ax2 = axes[1] j = 0 # one example point ax2.scatter(data[:, 0], data[:, 1], s=4, color='lightgray', alpha=0.4) ax2.scatter(*data[j], c='crimson', s=120, marker='X', zorder=6, label='target $z$') approach = np.array([r[j] for r in cum_recon]) ax2.plot(approach[:, 0], approach[:, 1], '-o', ms=5, color='seagreen', zorder=5, label='cumulative $\\hat z$ per stage') ax2.scatter(0, 0, c='black', s=40, zorder=4, label='start (0)') ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1') ax2.set_title('One point refined stage by stage') ax2.legend(loc='upper left', fontsize=9) ax2.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The left curve falls steadily (log scale) — the average residual norm drops with every added stage. The right panel traces one point: the green cumulative estimate starts at the origin and zig-zags ever closer to the red target, each stage a smaller correction.
2. Why it looks this way: Stage $i$ quantizes $r_{i-1}$, so it can only add a vector roughly the size of the remaining error. As error shrinks, each new code is a finer adjustment — exactly the coarse-to-fine refinement RVQ is built for.
3. Key takeaway: RVQ turns one hard quantization into a sequence of easy corrections; the sum of small codes reconstructs $z$ far better than any single small codebook could.

---

Experiment 2: Exponential expressivity¶

Question: RVQ with $N_q$ codebooks of size $K$ reaches $K^{N_q}$ combinations using only $N_q \cdot K$ codes. At an equal code budget, does this beat a single codebook with all $N_q \cdot K$ codes?

We compare distortion for a single VQ with $M$ codes against an RVQ whose stages also total $M$ codes, across budgets.

In [4]:

np.random.seed(2)

# uniform square: resolution-hungry, so combinatorial reach pays off
sq = np.random.uniform(-1, 1, size=(4000, 2))

stage_sizes = [2, 3, 4, 5, 6]   # number of RVQ stages
K = 4
single_d, rvq_d, budgets = [], [], []
for nq in stage_sizes:
    M = nq * K                               # equal code budget
    budgets.append(M)
    # single big codebook
    km = KMeans(n_clusters=M, n_init=3, random_state=0).fit(sq)
    _, q = quantize(sq, km.cluster_centers_)
    single_d.append(((sq - q) ** 2).sum(1).mean())
    # RVQ
    cbs = fit_rvq(sq, nq, K, seed=2)
    cum, _, _ = encode_rvq(sq, cbs)
    rvq_d.append(((sq - cum[-1]) ** 2).sum(1).mean())
print('budget M :', budgets)
print('single   :', np.round(single_d, 4))
print('rvq      :', np.round(rvq_d, 4))

np.random.seed(2) # uniform square: resolution-hungry, so combinatorial reach pays off sq = np.random.uniform(-1, 1, size=(4000, 2)) stage_sizes = [2, 3, 4, 5, 6] # number of RVQ stages K = 4 single_d, rvq_d, budgets = [], [], [] for nq in stage_sizes: M = nq * K # equal code budget budgets.append(M) # single big codebook km = KMeans(n_clusters=M, n_init=3, random_state=0).fit(sq) _, q = quantize(sq, km.cluster_centers_) single_d.append(((sq - q) ** 2).sum(1).mean()) # RVQ cbs = fit_rvq(sq, nq, K, seed=2) cum, _, _ = encode_rvq(sq, cbs) rvq_d.append(((sq - cum[-1]) ** 2).sum(1).mean()) print('budget M :', budgets) print('single :', np.round(single_d, 4)) print('rvq :', np.round(rvq_d, 4))

budget M : [8, 12, 16, 20, 24]
single   : [0.0845 0.0566 0.0405 0.0329 0.0267]
rvq      : [0.0417 0.0105 0.0027 0.0007 0.0002]

In [5]:

fig, ax = plt.subplots(figsize=(8.5, 5.3))
fig.suptitle('Equal code budget: RVQ reaches far more effective cells', fontsize=13, fontweight='bold')

ax.plot(budgets, single_d, '-o', color='royalblue', lw=2, label='single codebook (M codes)')
ax.plot(budgets, rvq_d, '-s', color='seagreen', lw=2, label=f'RVQ (M/{K} stages × {K} codes)')
ax.set_xlabel('total code budget M = $N_q \\cdot K$')
ax.set_ylabel('distortion (mean squared error)')
ax.set_yscale('log')
ax.set_title('Lower is better')
ax.legend()
for nq, b in zip(stage_sizes, budgets):
    ax.annotate(f'$K^{{{nq}}}$={K**nq}', (b, rvq_d[stage_sizes.index(nq)]),
                textcoords='offset points', xytext=(4, -14), fontsize=8, color='seagreen')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8.5, 5.3)) fig.suptitle('Equal code budget: RVQ reaches far more effective cells', fontsize=13, fontweight='bold') ax.plot(budgets, single_d, '-o', color='royalblue', lw=2, label='single codebook (M codes)') ax.plot(budgets, rvq_d, '-s', color='seagreen', lw=2, label=f'RVQ (M/{K} stages × {K} codes)') ax.set_xlabel('total code budget M = $N_q \\cdot K$') ax.set_ylabel('distortion (mean squared error)') ax.set_yscale('log') ax.set_title('Lower is better') ax.legend() for nq, b in zip(stage_sizes, budgets): ax.annotate(f'$K^{{{nq}}}$={K**nq}', (b, rvq_d[stage_sizes.index(nq)]), textcoords='offset points', xytext=(4, -14), fontsize=8, color='seagreen') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Both distortions fall with budget, but the green RVQ curve drops far faster than the blue single-codebook curve. The annotations show RVQ's effective number of cells $K^{N_q}$ exploding (256, 1024, ...) while spending the same number of codes.
2. Why it looks this way: A single codebook with $M$ codes can only carve $M$ Voronoi cells. RVQ combines $N_q$ independent stage-choices, so $M$ codes address $K^{N_q}$ summed locations — exponential reach from linear parameters. On resolution-hungry data this combinatorial coverage wins decisively.
3. Key takeaway: This is why neural audio codecs use RVQ (e.g. 8 codebooks of 1024) instead of one impossibly large codebook of $1024^8$ entries.

---

Experiment 3: Bitrate scaling¶

Question: Because stages go coarse→fine, decoding with only the first few stages should give a lower-bitrate reconstruction. Does truncating stages produce a graceful quality ladder?

We RVQ-quantize the RGB colours of a synthetic image and reconstruct using 1, 2, 4, and all 8 stages — the basis of quantizer dropout for variable-bitrate codecs.

In [6]:

np.random.seed(3)

H = W = 64
yy, xx = np.mgrid[0:H, 0:W] / H
img = np.stack([
    0.5 + 0.5 * np.sin(5 * xx),
    0.5 + 0.5 * yy,
    0.5 + 0.4 * np.cos(4 * (xx - yy)),
], axis=-1)
img = np.clip(img, 0, 1)
pixels = img.reshape(-1, 3)

n_stages, K = 8, 4
cbs = fit_rvq(pixels, n_stages, K, seed=3)
cum, _, _ = encode_rvq(pixels, cbs)
shown = [1, 2, 4, 8]
recons = {s: np.clip(cum[s - 1], 0, 1).reshape(H, W, 3) for s in shown}
psnr = {s: ((pixels - cum[s - 1]) ** 2).mean() for s in shown}
print('MSE by #stages:', {s: round(psnr[s], 4) for s in shown})

np.random.seed(3) H = W = 64 yy, xx = np.mgrid[0:H, 0:W] / H img = np.stack([ 0.5 + 0.5 * np.sin(5 * xx), 0.5 + 0.5 * yy, 0.5 + 0.4 * np.cos(4 * (xx - yy)), ], axis=-1) img = np.clip(img, 0, 1) pixels = img.reshape(-1, 3) n_stages, K = 8, 4 cbs = fit_rvq(pixels, n_stages, K, seed=3) cum, _, _ = encode_rvq(pixels, cbs) shown = [1, 2, 4, 8] recons = {s: np.clip(cum[s - 1], 0, 1).reshape(H, W, 3) for s in shown} psnr = {s: ((pixels - cum[s - 1]) ** 2).mean() for s in shown} print('MSE by #stages:', {s: round(psnr[s], 4) for s in shown})

MSE by #stages: {1: np.float64(0.0167), 2: np.float64(0.0069), 4: np.float64(0.0016), 8: np.float64(0.0001)}

In [7]:

fig, axes = plt.subplots(1, 5, figsize=(15, 3.4))
fig.suptitle('More stages = higher bitrate = finer reconstruction', fontsize=13, fontweight='bold')

axes[0].imshow(img)
axes[0].set_title('original')
axes[0].set_xticks([]); axes[0].set_yticks([])
for ax, s in zip(axes[1:], shown):
    ax.imshow(recons[s])
    ax.set_title(f'{s} stage(s)\n{s * K} bits/px-ish, MSE={psnr[s]:.3f}', fontsize=9)
    ax.set_xticks([]); ax.set_yticks([])

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 5, figsize=(15, 3.4)) fig.suptitle('More stages = higher bitrate = finer reconstruction', fontsize=13, fontweight='bold') axes[0].imshow(img) axes[0].set_title('original') axes[0].set_xticks([]); axes[0].set_yticks([]) for ax, s in zip(axes[1:], shown): ax.imshow(recons[s]) ax.set_title(f'{s} stage(s)\n{s * K} bits/px-ish, MSE={psnr[s]:.3f}', fontsize=9) ax.set_xticks([]); ax.set_yticks([]) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: With 1 stage the image is crudely posterised; each additional stage visibly sharpens the colours, and by 8 stages it closely matches the original, with MSE dropping monotonically.
2. Why it looks this way: The first stage captures the dominant colour structure; later stages add finer colour corrections (the residual). Truncating after $s$ stages keeps only the coarse part — a lower-bitrate, lower-fidelity version of the same encoding.
3. Key takeaway: One RVQ encoding supports a whole ladder of bitrates just by choosing how many stages to decode — exactly the scalability SoundStream exploits with quantizer dropout.

---

Experiment 4: Why later stages are harder¶

Question: The theory warns that later stages are more collapse-prone. What happens to the residual distribution that each stage must quantize?

We compare the distribution of residuals entering an early stage versus a late stage of the annulus RVQ from Experiment 1.

In [8]:

np.random.seed(4)

codebooks = fit_rvq(data, 8, 4, seed=1)
_, _, all_res = encode_rvq(data, codebooks)
# residual ENTERING a stage = residual leaving the previous one
entering = [data] + all_res[:-1]
norm_stage1 = np.linalg.norm(entering[0], axis=1)
norm_stage6 = np.linalg.norm(entering[5], axis=1)
print(f'spread entering stage 1: std={norm_stage1.std():.3f}, mean={norm_stage1.mean():.3f}')
print(f'spread entering stage 6: std={norm_stage6.std():.3f}, mean={norm_stage6.mean():.3f}')

np.random.seed(4) codebooks = fit_rvq(data, 8, 4, seed=1) _, _, all_res = encode_rvq(data, codebooks) # residual ENTERING a stage = residual leaving the previous one entering = [data] + all_res[:-1] norm_stage1 = np.linalg.norm(entering[0], axis=1) norm_stage6 = np.linalg.norm(entering[5], axis=1) print(f'spread entering stage 1: std={norm_stage1.std():.3f}, mean={norm_stage1.mean():.3f}') print(f'spread entering stage 6: std={norm_stage6.std():.3f}, mean={norm_stage6.mean():.3f}')

spread entering stage 1: std=0.357, mean=1.462
spread entering stage 6: std=0.028, mean=0.051

In [9]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Late-stage residuals collapse toward the origin', fontsize=13, fontweight='bold')

ax = axes[0]
r1, r6 = entering[0], entering[5]
ax.scatter(r1[:, 0], r1[:, 1], s=5, alpha=0.3, color='royalblue', label='entering stage 1')
ax.scatter(r6[:, 0], r6[:, 1], s=5, alpha=0.5, color='crimson', label='entering stage 6')
ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1')
ax.set_title('Residual cloud (stage 1 vs stage 6)')
ax.legend(fontsize=9); ax.set_aspect('equal')

ax2 = axes[1]
ax2.hist(norm_stage1, bins=40, alpha=0.6, color='royalblue', label='stage 1', density=True)
ax2.hist(norm_stage6, bins=40, alpha=0.6, color='crimson', label='stage 6', density=True)
ax2.set_xlabel('residual norm'); ax2.set_ylabel('density')
ax2.set_title('Residual norm distribution')
ax2.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('Late-stage residuals collapse toward the origin', fontsize=13, fontweight='bold') ax = axes[0] r1, r6 = entering[0], entering[5] ax.scatter(r1[:, 0], r1[:, 1], s=5, alpha=0.3, color='royalblue', label='entering stage 1') ax.scatter(r6[:, 0], r6[:, 1], s=5, alpha=0.5, color='crimson', label='entering stage 6') ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1') ax.set_title('Residual cloud (stage 1 vs stage 6)') ax.legend(fontsize=9); ax.set_aspect('equal') ax2 = axes[1] ax2.hist(norm_stage1, bins=40, alpha=0.6, color='royalblue', label='stage 1', density=True) ax2.hist(norm_stage6, bins=40, alpha=0.6, color='crimson', label='stage 6', density=True) ax2.set_xlabel('residual norm'); ax2.set_ylabel('density') ax2.set_title('Residual norm distribution') ax2.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The blue stage-1 residual cloud is large and spread out (it is essentially the data); the red stage-6 cloud is a tight blob near the origin. The histogram confirms it: stage-6 residual norms are far smaller and more concentrated.
2. Why it looks this way: Every stage removes most of the remaining magnitude, so the input to a late stage is a small, dense cluster around zero. A codebook trying to tile that tight blob has little spread to work with, so many codes sit unused.
3. Key takeaway: Concentrated late-stage residuals are exactly the condition that breeds codebook collapse, which is why each RVQ stage needs its own EMA and restart bookkeeping.

---

Summary¶

Experiment	Key insight
1. Coarse-to-fine	Each stage quantizes the shrinking residual; the cumulative sum converges to $z$.
2. Exponential expressivity	At equal code budget, RVQ's $K^{N_q}$ combinations crush a single $M$-code codebook.
3. Bitrate scaling	Truncating stages yields a graceful quality ladder from one encoding (quantizer dropout).
4. Why later stages are harder	Late-stage residuals concentrate near zero, breeding codebook collapse.

Connection to the broader theory¶

RVQ scales the precision of vector quantization by stacking codebooks, at the cost of $N_q$ tokens per position and $N_q$ chances of codebook collapse. It is the additive, more-codebooks answer to the resolution problem. The next topic, FSQ, takes the opposite route — it removes the learned codebook entirely and quantizes each dimension on a fixed grid, sidestepping collapse by construction.