Vector-Quantized VAE (VQ-VAE) — Visual Experiments¶
Goal: Build intuition for discrete latent spaces, vector quantization, gradient tricks, and codebook health through hands-on experiments.
What we'll explore¶
| # | Experiment | Question answered |
|---|---|---|
| 1 | Vector Quantization | What does "snapping" encoder outputs to a codebook look like geometrically? |
| 2 | Straight-Through Estimator | How does backpropagation flow through a non-differentiable argmin? |
| 3 | Loss Components | How do reconstruction, VQ, and commitment losses interact during training? |
| 4 | EMA vs. Gradient Descent | Why is EMA more stable than Adam for codebook updates? |
| 5 | Codebook Collapse | What does collapse look like, and how does the reset policy fix it? |
| 6 | L2 Normalization | Why projecting onto a unit sphere improves codebook utilization? |
Linked theory: research/04-vq-vae.md · Chạy: uv run jupyter lab từ thư mục latent-anything-theory/ (cần các package trong pyproject.toml).
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
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_e, codebook):
dists = np.linalg.norm(z_e[:, None, :] - codebook[None, :, :], axis=-1)
indices = np.argmin(dists, axis=-1)
z_q = codebook[indices]
return z_q, indices
print('Setup complete.')
Setup complete.
Experiment 1: Vector Quantization — Snapping Continuous to Discrete¶
Question: What does it look like when a continuous encoder output is "snapped" to the nearest codebook vector?
In VQ-VAE, the encoder produces a continuous vector $z_e$. Instead of passing this directly to the decoder, each $z_e$ is replaced by the closest entry $c_k$ in a learned codebook $C = \{c_1, \dots, c_K\}$, discarding fine-grained magnitude information in favour of a discrete index. This experiment visualises that nearest-neighbour lookup in 2D.
np.random.seed(1)
n_samples = 250
z_e = np.random.randn(n_samples, 2) * 1.5
K = 8
codebook = np.array([
[-2.2, 1.5], [-0.3, 2.1], [ 2.1, 1.4], [ 2.4, -0.6],
[ 1.0, -2.1], [-1.2, -2.0], [-2.4, -1.0], [ 0.1, 0.3],
])
z_q, indices = quantize(z_e, codebook)
print(f'z_e: {z_e.shape} z_q: {z_q.shape}')
print(f'Usage per codebook entry: {np.bincount(indices, minlength=K)}')
z_e: (250, 2) z_q: (250, 2) Usage per codebook entry: [24 36 27 21 33 21 17 71]
colors_k = plt.cm.tab10(np.linspace(0, 1, K))
# Background colour map: which code owns each spatial region?
res = 180
xx, yy = np.meshgrid(np.linspace(-4, 4, res), np.linspace(-4, 4, res))
grid = np.stack([xx.ravel(), yy.ravel()], axis=1)
_, grid_ids = quantize(grid, codebook)
grid_map = grid_ids.reshape(res, res)
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Experiment 1: Vector Quantization in 2D', fontsize=13, fontweight='bold')
# Left: scatter + displacement arrows
ax = axes[0]
for k in range(K):
mask = indices == k
if mask.any():
ax.scatter(z_e[mask, 0], z_e[mask, 1], color=colors_k[k], alpha=0.55, s=22, label=f'c{k}')
for i in np.where(mask)[0][:4]:
ax.annotate('', xy=z_q[i], xytext=z_e[i],
arrowprops=dict(arrowstyle='->', color=colors_k[k], lw=0.9, alpha=0.7))
ax.scatter(codebook[:, 0], codebook[:, 1], c='black', s=180, marker='*', zorder=6, label='Codebook')
ax.set_xlabel('z dim 1')
ax.set_ylabel('z dim 2')
ax.set_title('Encoder Outputs Snapped to Nearest Code\n(arrows show displacement)')
ax.legend(fontsize=8, ncol=2, loc='upper right')
# Right: Voronoi-like colour background
ax2 = axes[1]
ax2.imshow(colors_k[grid_map][:, :, :3], extent=[-4, 4, -4, 4],
origin='lower', alpha=0.28, aspect='auto')
for k in range(K):
mask = indices == k
if mask.any():
ax2.scatter(z_e[mask, 0], z_e[mask, 1], color=colors_k[k], alpha=0.65, s=22)
ax2.scatter(codebook[:, 0], codebook[:, 1], c='black', s=180, marker='*', zorder=6)
for k, (cx, cy) in enumerate(codebook):
ax2.text(cx, cy + 0.28, f'$c_{k}$', ha='center', fontsize=9, fontweight='bold')
ax2.set_xlim(-4, 4)
ax2.set_ylim(-4, 4)
ax2.set_xlabel('z dim 1')
ax2.set_ylabel('z dim 2')
ax2.set_title('Nearest-Code Regions (Voronoi Partition)\n(each colour zone belongs to one code)')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The left panel shows 250 encoder outputs (coloured by assigned code) with arrows pointing from each $z_e$ to its quantized replacement $z_q$. Every arrow converges to a black star (a codebook entry). The right panel colours the entire space: each zone is the Voronoi cell of one code — any $z_e$ falling inside a zone gets mapped to that zone's code.
Why it looks this way: Vector quantization is a nearest-neighbour lookup in $\mathbb{R}^D$. The Voronoi partition is the natural geometry of nearest-neighbour assignment: each cell contains all points closer to its code than to any other. Crucially, all spatial variation within a cell is discarded — the decoder only sees the discrete index, not the exact position.
Key takeaway: Quantization trades continuous precision for a discrete code index. The codebook is the model's vocabulary; the Voronoi cells are its recognition regions. The richer the codebook (larger $K$), the finer-grained the vocabulary.
Experiment 2: Straight-Through Estimator — Patching the Gradient¶
Question: How does the model learn when the argmin operation blocks all gradients?
The argmin step — selecting the nearest codebook entry — is a piecewise-constant (staircase) function with zero derivative almost everywhere, and undefined at jump points. Without intervention, no gradient would flow back through the quantisation layer to the encoder. The Straight-Through Estimator (STE) patches this by pretending, only during the backward pass, that quantisation was the identity function: $\nabla_{z_e} L \approx \nabla_{z_q} L$.
np.random.seed(2)
z_vals = np.linspace(-3.2, 3.2, 500)
codebook_1d = np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5])
dists_1d = np.abs(z_vals[:, None] - codebook_1d[None, :])
z_q_1d = codebook_1d[np.argmin(dists_1d, axis=-1)]
print('Codebook entries:', codebook_1d)
print('Forward pass maps z_e to nearest codebook entry (staircase).')
print('STE backward pass copies gradient straight through as if quantise = identity.')
Codebook entries: [-2.5 -1.5 -0.5 0.5 1.5 2.5] Forward pass maps z_e to nearest codebook entry (staircase). STE backward pass copies gradient straight through as if quantise = identity.
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Experiment 2: Straight-Through Estimator', fontsize=13, fontweight='bold')
# Left: forward vs. STE backward
ax = axes[0]
ax.plot(z_vals, z_q_1d, color='#e74c3c', lw=2.5, label='$z_q$ = quantize($z_e$) [forward]')
ax.plot(z_vals, z_vals, color='#3498db', lw=1.8, linestyle='--',
label='$z_e$ identity [STE backward]')
ax.scatter(codebook_1d, codebook_1d, color='black', s=90, zorder=5, label='Codebook entries')
ax.set_xlabel('Encoder output $z_e$')
ax.set_ylabel('Value passed to decoder')
ax.set_title('Forward: Staircase\nSTE Backward: Identity')
ax.legend()
# Right: gradient comparison
ax2 = axes[1]
ax2.fill_between(z_vals, np.zeros_like(z_vals), alpha=0.35, color='#e74c3c',
label='True gradient = 0 (encoder receives nothing)')
ax2.plot(z_vals, np.ones_like(z_vals), color='#2ecc71', lw=2.5,
label='STE gradient = 1 (gradient flows straight through)')
ax2.axhline(0, color='gray', lw=0.8)
ax2.set_xlabel('Encoder output $z_e$')
ax2.set_ylabel('dL / dz_e (magnitude)')
ax2.set_ylim(-0.2, 1.5)
ax2.set_title('Gradient Flow: True (blocked) vs. STE (approximated)')
ax2.legend()
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The left panel overlays the forward-pass function (red staircase) on the STE approximation (blue dashed diagonal). The right panel contrasts the true gradient (zero everywhere — shown as a filled region at 0) with the STE gradient (a flat line at 1).
Why it looks this way: The staircase has zero slope between jumps. Its true gradient blocks encoder learning entirely. The STE replaces this blocked gradient with the identity's gradient (constant 1), so the encoder receives a signal proportional to the decoder's gradient as if quantisation had never happened.
Key takeaway: STE is a deliberate mathematical lie told only during backprop. It enables end-to-end training of VQ-VAE despite a non-differentiable bottleneck, at the cost of a biased gradient estimate. This bias is acceptable because the encoder still receives meaningful learning signal about direction, even if the magnitude is approximate.
Experiment 3: Loss Components — Three Terms, One Objective¶
Question: How do the three VQ-VAE loss terms behave and balance during training?
VQ-VAE combines three loss terms:
- Reconstruction loss — pulls encoder and decoder together via $z_q$
- VQ / Codebook loss — $\|\text{sg}[z_e] - z_q\|^2$ — moves codebook vectors toward encoder outputs
- Commitment loss — $\beta \|z_e - \text{sg}[z_q]\|^2$ — prevents encoder from drifting far from its assigned code
This experiment simulates all three on a 2D toy problem to show their relative magnitude and convergence dynamics.
np.random.seed(3)
K_loss = 4
D_loss = 2
beta = 0.25
n_steps_loss = 220
batch_size_loss = 32
# Four Gaussian cluster centres (encoder targets)
centers_loss = np.array([[-1.5, 1.5], [1.5, 1.5], [1.5, -1.5], [-1.5, -1.5]])
# Codebook initialised off-centre so convergence is visible
codebook_loss = np.random.randn(K_loss, D_loss) * 0.4
recon_losses, vq_losses, commit_losses = [], [], []
def smooth(arr, w=12):
return np.convolve(arr, np.ones(w) / w, mode='valid')
for step in range(n_steps_loss):
ids = np.random.randint(0, K_loss, batch_size_loss)
z_e_b = centers_loss[ids] + np.random.randn(batch_size_loss, D_loss) * 0.3
dists_b = np.linalg.norm(z_e_b[:, None, :] - codebook_loss[None, :, :], axis=-1)
assigned_b = np.argmin(dists_b, axis=-1)
z_q_b = codebook_loss[assigned_b]
squared_gap = np.sum((z_e_b - z_q_b) ** 2, axis=-1)
recon_losses.append(float(np.mean(squared_gap)))
vq_losses.append(float(np.mean(squared_gap))) # sg[z_e] term, same magnitude
commit_losses.append(float(beta * np.mean(squared_gap)))
# Gradient descent on codebook (simplified VQ loss step)
for k in range(K_loss):
mask = assigned_b == k
if mask.any():
codebook_loss[k] -= 0.09 * (codebook_loss[k] - np.mean(z_e_b[mask], axis=0))
print('Loss simulation complete.')
Loss simulation complete.
steps_s = np.arange(len(smooth(recon_losses)))
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Experiment 3: VQ-VAE Loss Components During Training', fontsize=13, fontweight='bold')
ax = axes[0]
ax.plot(steps_s, smooth(recon_losses), color='#3498db', lw=2, label='Reconstruction loss')
ax.plot(steps_s, smooth(vq_losses), color='#e74c3c', lw=2, label='VQ / Codebook loss')
ax.plot(steps_s, smooth(commit_losses),color='#2ecc71', lw=2, label=f'Commitment loss (beta={beta})')
ax.set_xlabel('Training step')
ax.set_ylabel('Loss value (smoothed)')
ax.set_title('Individual Loss Terms')
ax.legend()
ax2 = axes[1]
total_s = (np.array(smooth(recon_losses)) +
np.array(smooth(vq_losses)) +
np.array(smooth(commit_losses)))
ax2.stackplot(
steps_s,
smooth(recon_losses), smooth(vq_losses), smooth(commit_losses),
labels=['Reconstruction', 'VQ Loss', 'Commitment'],
colors=['#3498db', '#e74c3c', '#2ecc71'], alpha=0.75
)
ax2.plot(steps_s, total_s, 'k--', lw=1.5, label='Total')
ax2.set_xlabel('Training step')
ax2.set_ylabel('Loss value (smoothed)')
ax2.set_title('Total Loss Decomposition (stacked)')
ax2.legend(loc='upper right')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: All three losses start high when the codebook is far from the data, then decay as the codebook converges. The commitment loss (green) is always a fixed fraction $\beta$ of the other terms because it penalises the same gap $\|z_e - z_q\|$ from the encoder side.
Why it looks this way: As the codebook vectors move toward their cluster centres, the distance $\|z_e - z_q\|$ shrinks for all three terms simultaneously. In practice the three losses converge at different rates depending on $\beta$ and the codebook learning rate — this is the key hyperparameter tension in VQ-VAE tuning.
Key takeaway: The commitment loss $\beta$ is a stability dial: too small and the encoder wanders freely, destabilising the codebook; too large and the encoder is over-constrained, reducing representation expressiveness. The VQ loss and commitment loss are complementary — one pulls the codebook toward the encoder, the other pulls the encoder toward the codebook.
Experiment 4: EMA vs. Gradient Descent — Tracking a Drifting Distribution¶
Question: Why does EMA outperform Adam for updating codebook vectors when the data distribution is non-stationary?
During training, encoder weights update continuously, so the distribution of $z_e$ changes over time (internal covariate shift). The codebook must track this moving distribution. EMA updates the codebook directly toward the empirical mean of assigned vectors — equivalent to online K-means. Adam, being a gradient-based optimizer, requires multiple small steps to move the codebook to where the data is. This experiment simulates a cluster mean that shifts in steps and measures how quickly each method tracks it.
np.random.seed(4)
n_steps_ema = 220
batch_size_ema = 16
# True cluster mean: step shifts that simulate non-stationarity
true_means = np.zeros(n_steps_ema)
true_means[:55] = -1.5
true_means[55:110] = 1.5
true_means[110:165] = -0.8
true_means[165:] = 2.0
gamma_ema = 0.85 # EMA decay factor
lr_slow = 0.05 # GD slow — lags behind
lr_fast = 0.55 # GD fast — noisy overshoot
c_ema = 0.0
c_slow = 0.0
c_fast = 0.0
ema_hist, slow_hist, fast_hist = [], [], []
for t in range(n_steps_ema):
batch = np.random.normal(true_means[t], 0.3, batch_size_ema)
bm = float(batch.mean())
c_ema = gamma_ema * c_ema + (1.0 - gamma_ema) * bm
c_slow -= lr_slow * (c_slow - bm)
c_fast -= lr_fast * (c_fast - bm)
ema_hist.append(c_ema)
slow_hist.append(c_slow)
fast_hist.append(c_fast)
print('EMA vs. GD simulation done.')
EMA vs. GD simulation done.
t_ax = np.arange(n_steps_ema)
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Experiment 4: EMA vs. Gradient Descent — Non-Stationary Codebook Update',
fontsize=13, fontweight='bold')
ax = axes[0]
ax.plot(t_ax, true_means, 'k--', lw=2, label='True cluster mean (shifts)', alpha=0.8)
ax.plot(t_ax, ema_hist, color='#2ecc71', lw=2.5, label=f'EMA (gamma={gamma_ema})')
ax.plot(t_ax, slow_hist, color='#3498db', lw=2, label=f'GD slow (lr={lr_slow})')
ax.plot(t_ax, fast_hist, color='#e74c3c', lw=2, label=f'GD fast (lr={lr_fast})')
ax.fill_between(t_ax,
np.array(true_means) - 0.3,
np.array(true_means) + 0.3,
alpha=0.08, color='gray', label='Data spread (+/-1 sigma)')
ax.set_xlabel('Training step')
ax.set_ylabel('Codebook vector value')
ax.set_title('Tracking a Shifting Distribution')
ax.legend(fontsize=9)
ax2 = axes[1]
err_ema = np.abs(np.array(ema_hist) - true_means)
err_slow = np.abs(np.array(slow_hist) - true_means)
err_fast = np.abs(np.array(fast_hist) - true_means)
ax2.plot(t_ax, err_ema, color='#2ecc71', lw=2.5, label='EMA')
ax2.plot(t_ax, err_slow, color='#3498db', lw=2, label='GD slow')
ax2.plot(t_ax, err_fast, color='#e74c3c', lw=2, label='GD fast')
ax2.set_xlabel('Training step')
ax2.set_ylabel('|codebook value - true mean|')
ax2.set_title('Tracking Error Over Time')
ax2.legend()
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: After each step-shift in the true mean, the EMA codebook (green) converges quickly and smoothly. Slow GD (blue) lags — it takes many steps to reach the new mean. Fast GD (red) responds quickly but exhibits large oscillations around the true value. The right panel shows these tracking errors numerically: EMA has low, stable error; slow GD has high but smooth error; fast GD has noisy, spiky error.
Why it looks this way: EMA with decay $\gamma$ places the codebook at a weighted average of all past batch means, with the most recent ones weighted highest. This is mathematically equivalent to online K-means — it "jumps" the codebook toward the cluster centre in a single effective step. Gradient descent with small $lr$ is too cautious; with large $lr$ it overshoots and oscillates because individual batch means are noisy.
Key takeaway: EMA is the natural update rule for codebook vectors: it is independent of the encoder learning rate, resilient to gradient noise, and memory-efficient (no Adam momentum buffers per code). This is why modern VQ-VAE implementations replace the VQ loss gradient with EMA.
Experiment 5: Codebook Collapse and Reset Recovery¶
Question: What does codebook collapse look like in practice, and how does a reset policy fix it?
Codebook collapse (or index collapse) occurs when only a small fraction of codebook entries are ever assigned data points. Unassigned codes receive no gradient and no EMA update — they freeze in place, permanently excluded from training. This experiment uses a biased data distribution (only 2 real clusters but $K=8$ codes) to trigger collapse, then shows how a replacement policy (substituting dead codes with live batch samples) restores full codebook utilisation.
np.random.seed(5)
K_col = 8
D_col = 2
n_steps_col = 110
batch_size_col = 64
gamma_col = 0.95
# Data: only 2 real clusters -> most of K=8 codes will go unused
cluster_centers_col = np.array([[0.8, 0.8], [-0.8, -0.8]])
def run_vq_sim(use_reset, reset_threshold=10):
np.random.seed(55)
cb = np.random.randn(K_col, D_col) * 2.0
ema_N = np.full(K_col, 1e-6)
ema_m = cb.copy()
usage_hist = np.zeros((n_steps_col, K_col), dtype=int)
steps_idle = np.zeros(K_col, dtype=int)
for step in range(n_steps_col):
ids = np.random.randint(0, 2, batch_size_col)
batch = cluster_centers_col[ids] + np.random.randn(batch_size_col, D_col) * 0.25
dists = np.linalg.norm(batch[:, None, :] - cb[None, :, :], axis=-1)
assigned = np.argmin(dists, axis=-1)
counts = np.bincount(assigned, minlength=K_col)
for k in range(K_col):
mask = assigned == k
ema_N[k] = gamma_col * ema_N[k] + (1 - gamma_col) * counts[k]
if mask.any():
ema_m[k] = gamma_col * ema_m[k] + (1 - gamma_col) * batch[mask].sum(axis=0)
cb[k] = ema_m[k] / max(ema_N[k], 1e-6)
steps_idle[k] = 0 if counts[k] > 0 else steps_idle[k] + 1
if use_reset:
for k in range(K_col):
if steps_idle[k] >= reset_threshold:
new_vec = batch[np.random.randint(batch_size_col)]
cb[k] = new_vec
ema_N[k] = 1.0
ema_m[k] = new_vec.copy()
steps_idle[k] = 0
usage_hist[step] = counts
return cb, usage_hist
cb_no_reset, hist_no_reset = run_vq_sim(use_reset=False)
cb_reset, hist_reset = run_vq_sim(use_reset=True)
total_nr = hist_no_reset.sum(axis=0)
total_r = hist_reset.sum(axis=0)
print(f'No reset - usage: {total_nr}')
print(f'With reset - usage: {total_r}')
No reset - usage: [ 0 0 3475 3565 0 0 0 0] With reset - usage: [ 676 717 1427 1360 722 597 743 798]
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
fig.suptitle('Experiment 5: Codebook Collapse vs. Reset Recovery', fontsize=13, fontweight='bold')
# Left: usage heatmap — no reset
ax = axes[0]
im1 = ax.imshow(hist_no_reset.T, aspect='auto', cmap='YlOrRd', interpolation='nearest')
ax.set_xlabel('Training step')
ax.set_ylabel('Codebook index')
ax.set_title('Usage Heatmap — No Reset\n(dead codes stay dark forever)')
ax.set_yticks(range(K_col))
plt.colorbar(im1, ax=ax, label='# assigned')
# Middle: usage heatmap — with reset
ax2 = axes[1]
im2 = ax2.imshow(hist_reset.T, aspect='auto', cmap='YlOrRd', interpolation='nearest')
ax2.set_xlabel('Training step')
ax2.set_ylabel('Codebook index')
ax2.set_title('Usage Heatmap — With Reset\n(dead codes revived by replacement)')
ax2.set_yticks(range(K_col))
plt.colorbar(im2, ax=ax2, label='# assigned')
# Right: total utilisation bar chart
ax3 = axes[2]
x_col = np.arange(K_col)
pct_nr = total_nr / total_nr.sum() * 100
pct_r = total_r / total_r.sum() * 100
ax3.bar(x_col - 0.2, pct_nr, 0.35, label='No Reset', color='#e74c3c', alpha=0.85)
ax3.bar(x_col + 0.2, pct_r, 0.35, label='With Reset', color='#2ecc71', alpha=0.85)
ax3.axhline(100.0 / K_col, color='gray', linestyle='--', lw=1.5, label='Ideal uniform')
ax3.set_xlabel('Codebook index')
ax3.set_ylabel('Usage frequency (%)')
ax3.set_title('Total Codebook Utilisation')
ax3.set_xticks(x_col)
ax3.set_xticklabels([f'c{k}' for k in range(K_col)])
ax3.legend()
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: Without reset (left heatmap), only 2 codes show activity (orange/red rows) while 6 codes remain permanently dark. With reset (middle), all 8 rows eventually show colour as dead codes are replaced with live data points. The bar chart (right) makes this stark: without reset, 2 codes absorb nearly 100% of the assignments; with reset, usage spreads across all 8.
Why it looks this way: The data has only 2 real clusters. The 6 nearest-neighbour regions of the initialised codebook that fall far from these clusters never win any assignment — so they never receive an EMA update, never move, and drift further away with each encoder update. Reset breaks this cycle by teleporting dead codes directly into the live data distribution, giving them a second chance to compete.
Key takeaway: Codebook collapse is not a theoretical concern — it is routine when $K$ exceeds the intrinsic cluster count of the data. The reset policy (or LRU replacement) is a practical engineering fix that requires no change to the loss function and dramatically increases effective codebook capacity.
Experiment 6: L2 Normalization — Quantization on the Unit Sphere¶
Question: Why does constraining both encoder outputs and codebook vectors to the unit hypersphere improve stability and codebook utilization?
ViT-VQGAN and subsequent architectures enforce L2 normalization on $z_e$ and all codebook entries $c_k$, projecting them onto a unit sphere. This has two key effects:
- Bounded distribution: $\|z_e\|_2 = 1$ regardless of encoder weights, stopping distribution drift
- Geometric equivalence: minimising Euclidean distance on the sphere is identical to maximising cosine similarity
This experiment illustrates both effects geometrically.
np.random.seed(6)
n_pts = 300
K_l2 = 8
# Raw encoder outputs: random direction + varying magnitude
raw_dir = np.random.randn(n_pts, 2)
raw_mag = np.random.exponential(scale=1.5, size=n_pts)
raw_out = raw_dir / np.linalg.norm(raw_dir, axis=1, keepdims=True) * raw_mag[:, None]
# L2-normalized encoder outputs (all on unit circle)
norm_out = raw_out / np.linalg.norm(raw_out, axis=1, keepdims=True)
# Codebook: random for raw, evenly spaced on unit circle for normalized
cb_raw = np.random.randn(K_l2, 2) * 1.2
angles_cb = np.linspace(0, 2 * np.pi, K_l2, endpoint=False)
cb_sphere = np.stack([np.cos(angles_cb), np.sin(angles_cb)], axis=1)
_, idx_raw = quantize(raw_out, cb_raw)
_, idx_norm = quantize(norm_out, cb_sphere)
print(f'Raw codebook usage: {np.bincount(idx_raw, minlength=K_l2)}')
print(f'Normalized codebook usage: {np.bincount(idx_norm, minlength=K_l2)}')
Raw codebook usage: [14 63 25 17 45 13 30 93] Normalized codebook usage: [33 39 38 34 39 31 38 48]
colors_l2 = plt.cm.tab10(np.linspace(0, 1, K_l2))
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
fig.suptitle('Experiment 6: L2 Normalization — Projecting onto the Unit Sphere',
fontsize=13, fontweight='bold')
# Left: raw encoder outputs (variable magnitude)
ax = axes[0]
for k in range(K_l2):
mask = idx_raw == k
if mask.any():
ax.scatter(raw_out[mask, 0], raw_out[mask, 1], color=colors_l2[k], alpha=0.55, s=20)
ax.scatter(cb_raw[:, 0], cb_raw[:, 1], c='black', s=160, marker='*', zorder=5)
unit_c = plt.Circle((0, 0), 1, fill=False, color='gray', linestyle='--', lw=1.2, alpha=0.5)
ax.add_patch(unit_c)
ax.set_aspect('equal')
ax.set_xlim(-5.5, 5.5)
ax.set_ylim(-5.5, 5.5)
ax.set_xlabel('dim 1')
ax.set_ylabel('dim 2')
ax.set_title('Raw Encoder Outputs\n(varying magnitudes, skewed usage)')
# Middle: L2-normalized (all on unit circle)
ax2 = axes[1]
for k in range(K_l2):
mask = idx_norm == k
if mask.any():
ax2.scatter(norm_out[mask, 0], norm_out[mask, 1], color=colors_l2[k], alpha=0.65, s=20)
unit_c2 = plt.Circle((0, 0), 1, fill=False, color='#3498db', linestyle='-', lw=2)
ax2.add_patch(unit_c2)
ax2.scatter(cb_sphere[:, 0], cb_sphere[:, 1], c='black', s=160, marker='*', zorder=5)
for k, (cx, cy) in enumerate(cb_sphere):
ax2.text(cx * 1.35, cy * 1.35, f'$c_{k}$', ha='center', fontsize=9, fontweight='bold')
ax2.set_aspect('equal')
ax2.set_xlim(-1.9, 1.9)
ax2.set_ylim(-1.9, 1.9)
ax2.set_xlabel('dim 1')
ax2.set_ylabel('dim 2')
ax2.set_title('L2-Normalized Outputs\n(all on unit circle, uniform usage)')
# Right: Euclidean distance = cosine similarity on sphere
ax3 = axes[2]
theta = np.linspace(0, np.pi, 250)
eucl = np.sqrt(2.0 - 2.0 * np.cos(theta))
cos_s = np.cos(theta)
ax3r = ax3.twinx()
ax3.plot(np.degrees(theta), eucl, color='#e74c3c', lw=2.5, label='Euclidean distance')
ax3r.plot(np.degrees(theta), cos_s, color='#3498db', lw=2.5, linestyle='--',
label='Cosine similarity')
ax3.set_xlabel('Angle theta between z_e and z_q (degrees)')
ax3.set_ylabel('Euclidean distance', color='#e74c3c')
ax3r.set_ylabel('Cosine similarity', color='#3498db')
ax3.set_title('On Unit Sphere:\nmin Euclidean = max Cosine Similarity')
lines1, labs1 = ax3.get_legend_handles_labels()
lines2, labs2 = ax3r.get_legend_handles_labels()
ax3.legend(lines1 + lines2, labs1 + labs2, loc='center right')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: Raw encoder outputs (left) scatter across a wide range of radii — some codes attract many points simply because they happen to be in a high-density magnitude region, not because they represent distinct semantic content. Normalized outputs (middle) all lie exactly on the unit circle; the 8 evenly spaced codebook entries divide angular space equally, producing near-uniform utilisation. The right panel proves the mathematical identity: on the unit sphere, Euclidean distance and cosine similarity are monotonically related by $\|z_e - z_q\|^2 = 2 - 2\cos\theta$.
Why it looks this way: L2 normalization removes magnitude as a variable entirely. The only remaining degree of freedom is direction (angle). Cosine similarity, which measures angular alignment, becomes the sole metric. Since encoders can no longer "cheat" by scaling magnitudes, they are forced to learn angular structure — which tends to be more semantically meaningful.
Key takeaway: Projecting onto the unit sphere (L2-normalized codes) simultaneously prevents distribution drift, converts the lookup to cosine similarity, and promotes uniform codebook coverage. Paired with small code dimension and large $K$, this is the configuration used in ViT-VQGAN and modern discrete-token generative models.
Summary¶
| Experiment | Key insight |
|---|---|
| 1. Vector Quantization | Encoder outputs are snapped to nearest-neighbour codebook entries; Voronoi cells partition the space — all variation within a cell is discarded |
| 2. Straight-Through Estimator | STE copies decoder gradients straight through the non-differentiable argmin by treating quantisation as identity in the backward pass |
| 3. Loss Components | Reconstruction, VQ, and commitment losses each target a different pair (encoder/decoder, codebook, encoder↔codebook) and must be balanced via $\beta$ |
| 4. EMA vs. Adam | EMA acts as online K-means and tracks shifting distributions faster and more stably than gradient-based optimizers |
| 5. Codebook Collapse | When data clusters $\ll K$, most codes die; a reset policy teleports dead codes into live batch samples to restore full utilisation |
| 6. L2 Normalization | Constraining vectors to the unit sphere removes magnitude drift, converts quantisation to cosine similarity, and promotes uniform codebook coverage |
Connection to the broader theory¶
VQ-VAE represents a fundamental departure from the VAE family studied in research/03-vae.md: instead of a continuous Gaussian posterior, it learns a codebook — a finite dictionary of discrete symbols. Every image is encoded as a spatial grid of code indices, transforming the generation problem from "sample from a continuous distribution" to "predict the next token in a sequence". This token-based view enabled Transformer architectures to model images (VQGAN, DALL-E) and audio (EnCodec) with the same machinery used for text, opening the era of large multimodal generative models.