VQGAN — Visual Experiments¶
Goal: Build intuition for why VQGAN produces sharper images than VQ-VAE, how it turns images into discrete token sequences, and why those tokens unlock Transformer-based generation.
What we'll explore¶
| # | Experiment | Question answered |
|---|---|---|
| 1 | The MSE Blur Problem | Why does pixel-level reconstruction loss produce blurry outputs? |
| 2 | Perceptual vs. Pixel Distance | How does feature-space distance differ from pixel-space distance? |
| 3 | PatchGAN: Local Authenticity | How does a patch discriminator force sharp local textures? |
| 4 | Adaptive Loss Weight λ | How does the gradient-based λ balance reconstruction and adversarial losses? |
| 5 | Image Tokenization | How does VQGAN compress an image into a short sequence of discrete tokens? |
| 6 | Autoregressive Token Generation | How does treating image tokens like text tokens enable Transformer-based generation? |
Linked theory: research/05-vqgan.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,
})
print('Setup complete.')
Setup complete.
Experiment 1: The MSE Blur Problem¶
Question: Why does minimising pixel-level MSE loss cause the decoder to produce blurry outputs?
VQ-VAE uses an L2/MSE reconstruction loss which compares each output pixel directly to the ground truth. When the data distribution is multimodal — for example, a patch of texture that could plausibly be sharp either way — the MSE-optimal decoder outputs the mean of all valid completions. If those completions are structurally opposite (e.g., bright stripes vs. dark stripes), their mean is a flat, gray, low-contrast image. VQGAN replaces this pixel comparison with a perceptual loss and an adversarial loss that force the decoder to commit to one sharp mode rather than averaging over all of them.
np.random.seed(1)
size = 48
x = np.linspace(0, 4 * np.pi, size)
# Two valid sharp textures (opposite phases)
stripe_A = (np.sin(3 * x) > 0).astype(float)
stripe_B = (np.sin(3 * x) < 0).astype(float)
# 2D images (tile along rows)
img_A = np.tile(stripe_A, (size, 1))
img_B = np.tile(stripe_B, (size, 1))
# MSE-optimal output: average of both modes -> flat gray
img_mse = 0.5 * (img_A + img_B)
# GAN output: commits to one sharp mode (picks A)
img_gan = img_A.copy()
# 1D version for frequency analysis
signal_A = np.sin(3 * x)
signal_B = np.sin(3 * x + np.pi)
signal_mse = 0.5 * (signal_A + signal_B)
print(f'Image shape: {img_A.shape}')
print(f'Pixel variance -- A: {img_A.var():.3f}, MSE: {img_mse.var():.3f}, GAN: {img_gan.var():.3f}')
print(f'Mean pixel value -- A: {img_A.mean():.3f}, MSE: {img_mse.mean():.3f}, GAN: {img_gan.mean():.3f}')
Image shape: (48, 48) Pixel variance -- A: 0.250, MSE: 0.005, GAN: 0.250 Mean pixel value -- A: 0.479, MSE: 0.490, GAN: 0.479
fig = plt.figure(figsize=(16, 9))
fig.suptitle('Experiment 1: MSE Averages Over Modes → Blurry Output', fontsize=13, fontweight='bold')
# Top row: 2D texture images
ax1 = fig.add_subplot(2, 4, 1)
ax1.imshow(img_A, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax1.set_title('Mode A\n(sharp stripes, phase 0)')
ax1.set_xlabel('pixel x')
ax1.set_ylabel('pixel y')
ax1.grid(False)
ax2 = fig.add_subplot(2, 4, 2)
ax2.imshow(img_B, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax2.set_title('Mode B\n(sharp stripes, phase π)')
ax2.set_xlabel('pixel x')
ax2.set_ylabel('pixel y')
ax2.grid(False)
ax3 = fig.add_subplot(2, 4, 3)
ax3.imshow(img_mse, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax3.set_title(f'MSE output (mean of A+B)\nvar={img_mse.var():.3f} ← blurry!')
ax3.set_xlabel('pixel x')
ax3.set_ylabel('pixel y')
ax3.grid(False)
ax4 = fig.add_subplot(2, 4, 4)
ax4.imshow(img_gan, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax4.set_title(f'GAN output (commits to A)\nvar={img_gan.var():.3f} ← sharp!')
ax4.set_xlabel('pixel x')
ax4.set_ylabel('pixel y')
ax4.grid(False)
# Bottom row: 1D signals
ax5 = fig.add_subplot(2, 1, 2)
ax5.plot(x, signal_A, color='#3498db', lw=2, label='Mode A (sin 3x)', alpha=0.9)
ax5.plot(x, signal_B, color='#e74c3c', lw=2, label='Mode B (sin 3x + π)', alpha=0.9)
ax5.plot(x, signal_mse, color='#2ecc71', lw=3, label='MSE mean = 0 (complete cancellation)', linestyle='--')
ax5.axhline(0, color='gray', lw=0.8, alpha=0.5)
ax5.set_xlabel('x')
ax5.set_ylabel('signal value')
ax5.set_title('1D View: MSE Minimisation Cancels High-Frequency Detail')
ax5.legend(loc='upper right')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The top row has four images: two valid sharp stripe patterns (A and B), the MSE-optimal output (uniform gray, variance ≈ 0), and the GAN output (identical to A, full-contrast sharp). The bottom 1D plot makes the cancellation explicit — modes A and B are opposite-phase sinusoids whose mean is exactly zero.
Why it looks this way: MSE loss penalises the expected squared error. When two outcomes are equally likely, the minimiser of $\mathbb{E}[(\hat{x} - x)^2]$ is $\hat{x} = \mathbb{E}[x]$. If the distribution of real textures is bimodal and symmetric, this expectation is a flat constant — the decoder has no incentive to produce sharp edges it might get wrong.
Key takeaway: MSE is a mode-averaging loss: it learns the mean, not a sample. The adversarial loss in VQGAN breaks this symmetry — the discriminator penalises any output that looks like an average, forcing the decoder to commit to a single plausible sharp reconstruction.
Experiment 2: Perceptual vs. Pixel Distance¶
Question: How does measuring distance in feature space (as VQGAN's perceptual loss does) differ from measuring it at the pixel level?
Perceptual loss (LPIPS) compares images via the activations of a pretrained network (e.g. VGG-16), not their raw pixels. The intuition is that a network trained on natural images builds internal representations that capture structure — edges, textures, shapes — rather than absolute pixel values. Two images that look structurally similar should have nearby feature maps even if their pixel values differ (e.g., a slight spatial shift). Conversely, a blurry version of an image can have low pixel distance from the sharp original yet be far away in feature space because it lacks the high-frequency edge responses.
np.random.seed(2)
img_size = 40
# Base image: horizontal stripes
base = np.zeros((img_size, img_size))
base[::4, :] = 1.0
base[1::4, :] = 0.8
# Image B: shifted by 1 pixel (perceptually similar, high pixel distance)
shifted = np.roll(base, 2, axis=0)
# Image C: blurred (perceptually very different — lacks edges, but close in pixel value)
def box_blur(img, k=3):
out = img.copy()
for _ in range(k):
out = 0.25 * (out + np.roll(out, 1, 0) + np.roll(out, -1, 0) + np.roll(out, 1, 1))
return out
blurred = box_blur(base, k=6)
# Simple proxy for "feature map": vertical gradient magnitude (edge strength)
def edge_feature(img):
grad = np.abs(np.diff(img, axis=0, prepend=img[:1, :]))
return grad
feat_base = edge_feature(base)
feat_shifted = edge_feature(shifted)
feat_blurred = edge_feature(blurred)
# Distances
pixel_dist_bs = np.mean((base - shifted) ** 2)
pixel_dist_bb = np.mean((base - blurred) ** 2)
feat_dist_bs = np.mean((feat_base - feat_shifted) ** 2)
feat_dist_bb = np.mean((feat_base - feat_blurred) ** 2)
print(' Pixel MSE Feature MSE')
print(f'Base vs. Shifted: {pixel_dist_bs:.4f} {feat_dist_bs:.4f}')
print(f'Base vs. Blurred: {pixel_dist_bb:.4f} {feat_dist_bb:.4f}')
Pixel MSE Feature MSE Base vs. Shifted: 0.8200 0.0390 Base vs. Blurred: 0.2011 0.3830
fig, axes = plt.subplots(2, 3, figsize=(14, 8))
fig.suptitle('Experiment 2: Perceptual (Feature) Distance vs. Pixel Distance', fontsize=13, fontweight='bold')
images = [base, shifted, blurred]
feats = [feat_base, feat_shifted, feat_blurred]
titles = ['Base image', 'Shifted (+2 px)\nhigh pixel dist, low feature dist', 'Blurred\nlow pixel dist, high feature dist']
for col, (img, feat, title) in enumerate(zip(images, feats, titles)):
ax_img = axes[0, col]
ax_img.imshow(img, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax_img.set_title(title)
ax_img.set_xlabel('x')
ax_img.set_ylabel('y')
ax_img.grid(False)
ax_feat = axes[1, col]
ax_feat.imshow(feat, cmap='plasma', aspect='auto')
ax_feat.set_title(f'Edge feature map')
ax_feat.set_xlabel('x')
ax_feat.set_ylabel('y')
ax_feat.grid(False)
# Add distance annotations as text below
axes[1, 1].set_xlabel(
f'Pixel MSE vs base: {pixel_dist_bs:.4f} | Feature MSE: {feat_dist_bs:.4f}',
fontsize=9, color='#2c3e50'
)
axes[1, 2].set_xlabel(
f'Pixel MSE vs base: {pixel_dist_bb:.4f} | Feature MSE: {feat_dist_bb:.4f}',
fontsize=9, color='#2c3e50'
)
plt.tight_layout()
plt.show()
# Summary bar chart
fig2, axes2 = plt.subplots(1, 2, figsize=(11, 4))
fig2.suptitle('Distance Summary: Pixel vs. Feature Space', fontsize=12, fontweight='bold')
labels = ['Base vs.\nShifted', 'Base vs.\nBlurred']
pixel_vals = [pixel_dist_bs, pixel_dist_bb]
feat_vals = [feat_dist_bs, feat_dist_bb]
x_pos = np.arange(2)
axes2[0].bar(x_pos, pixel_vals, color=['#3498db', '#e74c3c'], alpha=0.85)
axes2[0].set_xticks(x_pos)
axes2[0].set_xticklabels(labels)
axes2[0].set_ylabel('MSE distance')
axes2[0].set_title('Pixel-Space Distance')
for i, v in enumerate(pixel_vals):
axes2[0].text(i, v + 0.001, f'{v:.4f}', ha='center', fontweight='bold')
axes2[1].bar(x_pos, feat_vals, color=['#3498db', '#e74c3c'], alpha=0.85)
axes2[1].set_xticks(x_pos)
axes2[1].set_xticklabels(labels)
axes2[1].set_ylabel('MSE distance')
axes2[1].set_title('Feature-Space Distance (Edge Maps)')
for i, v in enumerate(feat_vals):
axes2[1].text(i, v + 0.0005, f'{v:.4f}', ha='center', fontweight='bold')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: Three images (base stripes, 2-pixel shift, blurred) and their edge feature maps. The bar charts quantify the paradox: the shifted image has a higher pixel MSE than the blurred image (its stripes are in different rows), yet its feature-map MSE is lower (both edge maps show the same periodic edge structure, just offset).
Why it looks this way: A 2-pixel shift moves individual pixel values dramatically, but the structural content — where the edges are, how strong they are — is almost unchanged. Blurring destroys that edge information entirely, collapsing the feature maps toward zero, yet many pixels end up at intermediate values close to the original. Pixel MSE is sensitive to spatial alignment; feature MSE is sensitive to structural content.
Key takeaway: Perceptual loss measures what an image contains, not where each pixel fell. A decoder trained with perceptual loss is rewarded for reconstructing the right edges and textures, not for memorising exact pixel positions — which is why VQGAN reconstructions look sharp and natural even at high compression.
Experiment 3: PatchGAN — Local Authenticity Scoring¶
Question: How does a PatchGAN discriminator detect local texture flaws that a global discriminator would miss?
A standard discriminator reduces an entire image to a single real/fake score. VQGAN uses a PatchGAN discriminator that outputs a grid of scores — one per non-overlapping patch of the image. Each patch is evaluated independently: is this local region realistic? Forcing the decoder to fool every patch prevents it from producing a globally plausible image that hides local artifacts behind a good overall composition. It is much harder to cheat a bank of local critics than one global judge.
np.random.seed(3)
img_h, img_w = 64, 64
patch_size = 8
n_patches = img_h // patch_size # 8x8 = 64 patches
# --- Real image: regular horizontal stripes (authentic structured texture) ---
real_img = np.zeros((img_h, img_w))
for i in range(0, img_h, 4):
real_img[i, :] = 1.0
real_img[i+1, :] = 0.7
# --- Fake image: correct structure on left half, random noise on right half ---
# Globally it looks "mostly" structured, but local patches on the right are wrong
fake_img = real_img.copy()
np.random.seed(33)
fake_img[:, img_w//2:] = np.random.rand(img_h, img_w//2) * 0.4 + 0.3 # noisy region
# --- Authenticity scoring function (proxy for a trained discriminator) ---
def patch_authenticity(patch):
"""Score 0..1 measuring how closely a patch matches the expected stripe texture."""
expected = np.zeros_like(patch)
for i in range(0, patch.shape[0], 4):
if i < patch.shape[0]: expected[i, :] = 1.0
if i+1 < patch.shape[0]: expected[i+1, :] = 0.7
return float(1.0 - np.mean((patch - expected) ** 2) / 0.25)
# Compute per-patch scores
real_scores = np.zeros((n_patches, n_patches))
fake_scores = np.zeros((n_patches, n_patches))
for i in range(n_patches):
for j in range(n_patches):
r_patch = real_img[i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size]
f_patch = fake_img[i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size]
real_scores[i, j] = patch_authenticity(r_patch)
fake_scores[i, j] = patch_authenticity(f_patch)
global_real = real_img.mean()
global_fake = fake_img.mean()
print(f'Global mean pixel -- real: {global_real:.3f}, fake: {global_fake:.3f}')
print(f'PatchGAN avg score -- real: {real_scores.mean():.3f}, fake: {fake_scores.mean():.3f}')
print(f'Patches with score < 0.4 -- real: {(real_scores < 0.4).sum()}, fake: {(fake_scores < 0.4).sum()}')
Global mean pixel -- real: 0.425, fake: 0.463 PatchGAN avg score -- real: 1.000, fake: 0.576 Patches with score < 0.4 -- real: 0, fake: 32
fig, axes = plt.subplots(2, 2, figsize=(13, 11))
fig.suptitle('Experiment 3: PatchGAN — Local Authenticity Scoring', fontsize=13, fontweight='bold')
# Image panels
ax1 = axes[0, 0]
ax1.imshow(real_img, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax1.set_title(f'Real image (mean={global_real:.3f})')
ax1.set_xlabel('pixel x')
ax1.set_ylabel('pixel y')
ax1.grid(False)
ax2 = axes[0, 1]
ax2.imshow(fake_img, cmap='gray', vmin=0, vmax=1, aspect='auto')
ax2.set_title(f'Fake image (mean={global_fake:.3f})\nRight half: random noise')
ax2.set_xlabel('pixel x')
ax2.set_ylabel('pixel y')
ax2.axvline(img_w // 2, color='red', lw=2, label='Noise boundary')
ax2.legend()
ax2.grid(False)
# PatchGAN score heatmaps
ax3 = axes[1, 0]
im3 = ax3.imshow(real_scores, cmap='RdBu_r', vmin=0, vmax=1, aspect='auto')
ax3.set_title(f'PatchGAN scores — Real\n(mean={real_scores.mean():.3f}, all patches high)')
ax3.set_xlabel('patch column')
ax3.set_ylabel('patch row')
ax3.grid(False)
plt.colorbar(im3, ax=ax3, label='Authenticity (0=fake, 1=real)')
ax4 = axes[1, 1]
im4 = ax4.imshow(fake_scores, cmap='RdBu_r', vmin=0, vmax=1, aspect='auto')
ax4.set_title(f'PatchGAN scores — Fake\n(mean={fake_scores.mean():.3f}, right half exposed!)')
ax4.set_xlabel('patch column')
ax4.set_ylabel('patch row')
ax4.grid(False)
plt.colorbar(im4, ax=ax4, label='Authenticity (0=fake, 1=real)')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The real and fake images look superficially similar in global statistics (both have roughly the same mean pixel value). But the PatchGAN score heatmaps are starkly different: real patches are uniformly high (blue = authentic), while the fake image's right half patches score near zero (red = not authentic).
Why it looks this way: The global mean cannot distinguish a structured texture from random noise if both have the same average intensity. PatchGAN divides the image into independent sub-problems — each patch must independently convince a local critic. The noisy right half of the fake image fails every one of those local tests, showing up as a red column in the score heatmap. A global discriminator would see an acceptable overall mean and possibly classify the fake as real.
Key takeaway: PatchGAN provides dense gradient feedback: every poorly-rendered patch contributes a separate loss signal. This forces the decoder to produce convincing textures everywhere in the image, not just in aggregate — which is why VQGAN decoders generate sharp, detailed reconstructions across the full spatial extent of an image.
Experiment 4: Adaptive Loss Weight λ — Gradient-Based Balancing¶
Question: How does the adaptive λ automatically balance the reconstruction/perceptual loss against the adversarial loss throughout training?
VQGAN computes λ at each step as the ratio of gradient norms at the final decoder layer: $$\lambda = \frac{\|\nabla_{G_L}[\mathcal{L}_{\text{rec}}]\|}{\|\nabla_{G_L}[\mathcal{L}_{\text{GAN}}]\| + \delta}$$ Early in training, the reconstruction loss has large gradients (the decoder is still learning coarse structure) and the GAN loss contributes small gradients (the discriminator has not yet learned to distinguish real from fake). The ratio λ is large, downweighting the GAN. Once the coarse structure is established, the reconstruction gradient shrinks and the GAN gradient grows — λ falls, allowing the adversarial signal to sharpen fine details.
np.random.seed(4)
n_steps = 400
t = np.arange(n_steps)
# Reconstruction gradient norm: large at start, decays as coarse structure is learned
grad_rec = 8.0 * np.exp(-t / 100) + 0.3
grad_rec += 0.4 * np.abs(np.random.randn(n_steps))
# GAN gradient norm: tiny at start (discriminator not yet trained), grows mid-training
# Then stabilises once adversarial training converges
grad_gan = 0.05 + 3.5 * (1 - np.exp(-np.maximum(0, t - 60) / 80))
grad_gan += 0.25 * np.abs(np.random.randn(n_steps))
# Adaptive lambda
delta = 1e-4
lambda_adaptive = np.clip(grad_rec / (grad_gan + delta), 0, 20)
# Smoothing helper
def smooth(arr, w=15):
return np.convolve(arr, np.ones(w) / w, mode='same')
# Three training phases for annotation
phase_boundaries = [0, 80, 200, n_steps]
phase_labels = [
'Phase 1: Coarse structure learning\n(large λ, GAN suppressed)',
'Phase 2: Transition\n(λ decreasing)',
'Phase 3: Detail sharpening\n(small λ, GAN active)',
]
phase_colors = ['#3498db', '#f39c12', '#e74c3c']
print(f'Lambda range: {lambda_adaptive.min():.2f} – {lambda_adaptive.max():.2f}')
print(f'Lambda at step 0: {lambda_adaptive[0]:.2f} (GAN suppressed)')
print(f'Lambda at step 200: {lambda_adaptive[200]:.2f} (balanced)')
print(f'Lambda at step 350: {lambda_adaptive[350]:.2f} (GAN active)')
Lambda range: 0.13 – 20.00 Lambda at step 0: 20.00 (GAN suppressed) Lambda at step 200: 0.48 (balanced) Lambda at step 350: 0.16 (GAN active)
fig, axes = plt.subplots(2, 1, figsize=(14, 9))
fig.suptitle('Experiment 4: Adaptive Loss Weight λ — Self-Regulating Gradient Balance',
fontsize=13, fontweight='bold')
# Top: gradient norms
ax1 = axes[0]
ax1.plot(t, smooth(grad_rec), color='#3498db', lw=2.5, label='||∇ reconstruction loss||')
ax1.plot(t, smooth(grad_gan), color='#e74c3c', lw=2.5, label='||∇ GAN loss||')
for i, (start, end) in enumerate(zip(phase_boundaries[:-1], phase_boundaries[1:])):
ax1.axvspan(start, end, alpha=0.07, color=phase_colors[i])
ax1.set_xlabel('Training step')
ax1.set_ylabel('Gradient norm (smoothed)')
ax1.set_title('Gradient Norms at the Final Decoder Layer')
ax1.legend()
# Bottom: adaptive lambda
ax2 = axes[1]
ax2.plot(t, smooth(lambda_adaptive), color='#8e44ad', lw=3, label='λ (adaptive weight)')
ax2.axhline(1.0, color='gray', lw=1.2, linestyle='--', alpha=0.6, label='λ = 1 (equal contribution)')
phase_x = [(0 + 80) / 2, (80 + 200) / 2, (200 + n_steps) / 2]
for i, (px, label, col) in enumerate(zip(phase_x, phase_labels, phase_colors)):
ax2.axvspan(phase_boundaries[i], phase_boundaries[i+1], alpha=0.07, color=col)
ax2.text(px, 15, label, ha='center', va='top', fontsize=8.5,
color=col, fontweight='bold', wrap=True)
ax2.set_xlabel('Training step')
ax2.set_ylabel('λ value')
ax2.set_title('Adaptive λ: Ratio of Gradient Norms')
ax2.set_ylim(-0.5, 20)
ax2.legend()
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The top panel shows two gradient norm curves over training: reconstruction gradients (blue) start large and decay; GAN gradients (red) start near zero and grow. The bottom panel shows λ — which starts very high (≫ 1), falls through a transition phase, and settles near a small value where GAN gradients dominate.
Why it looks this way: λ =
grad_rec / grad_gan. Early in training, the decoder is still learning to reconstruct coarse shapes — the reconstruction gradient is huge, making λ large and effectively turning off the GAN term. Once the discriminator is trained and starts distinguishing real from fake,grad_gangrows, pulling λ down. This is an automatic curriculum: the model learns structure first, then uses the GAN to sharpen it.Key takeaway: The adaptive λ is an elegant self-pacing mechanism that requires no hand-tuning. It prevents the GAN from destabilising early training (when the decoder needs uncontested gradient flow to learn basic reconstruction) while ensuring it becomes active once the decoder is ready to focus on perceptual quality.
Experiment 5: Image Tokenization — VQGAN as Visual Vocabulary¶
Question: How does VQGAN compress a high-resolution image into a compact sequence of discrete integer tokens?
VQGAN's encoder maps an image to a spatial feature grid (e.g., a 16×16 grid from a 256×256 image — a 16× spatial downscale). Each spatial position holds a D-dimensional feature vector. Vector quantization then replaces each vector with the index of its nearest codebook entry. The result is a 2D grid of integers — the image's token map. Flattened to 1D, this is a sequence of discrete indices that can be fed directly to a Transformer. The critical property is the dramatic compression: a 256×256 image with 3 channels (196 608 floats) becomes 256 integers.
np.random.seed(5)
# Simulate encoder output: spatial grid of D-dim feature vectors
H_feat, W_feat, D_feat = 8, 8, 16 # 8x8 grid, each position is 16-dim
K_vocab = 32 # codebook (vocabulary) size
# Structured encoder output: 4 natural clusters (simulating semantic regions)
cluster_centres = np.array([
[2.0, -1.5, 0.5, -2.0] + [0.0] * (D_feat - 4), # cluster 0
[-2.0, 1.5, -0.5, 2.0] + [0.0] * (D_feat - 4), # cluster 1
[1.5, 1.5, -1.5, -1.5] + [0.0] * (D_feat - 4), # cluster 2
[-1.5, -1.5, 1.5, 1.5] + [0.0] * (D_feat - 4), # cluster 3
])
# Assign spatial regions to clusters (4 quadrants)
region_map = np.zeros((H_feat, W_feat), dtype=int)
region_map[:H_feat//2, :W_feat//2] = 0
region_map[:H_feat//2, W_feat//2:] = 1
region_map[H_feat//2:, :W_feat//2] = 2
region_map[H_feat//2:, W_feat//2:] = 3
z_e_grid = cluster_centres[region_map] + 0.4 * np.random.randn(H_feat, W_feat, D_feat)
# Codebook: random entries, but 4 of them are close to the cluster centres
codebook = np.random.randn(K_vocab, D_feat) * 1.5
codebook[:4] = cluster_centres + 0.1 * np.random.randn(4, D_feat)
# Quantize: find nearest codebook entry for each spatial position
z_flat = z_e_grid.reshape(-1, D_feat) # (H*W, D)
dists = np.linalg.norm(z_flat[:, None, :] - codebook[None, :, :], axis=-1) # (H*W, K)
token_ids = np.argmin(dists, axis=-1).reshape(H_feat, W_feat)
token_seq = token_ids.flatten()
n_unique = len(np.unique(token_seq))
orig_floats = H_feat * W_feat * D_feat # before quantisation
n_tokens = H_feat * W_feat
print(f'Encoder output: {H_feat}×{W_feat}×{D_feat} = {orig_floats} float32 values')
print(f'Token sequence: {n_tokens} integers, vocabulary size K={K_vocab}')
print(f'Float compression: {orig_floats / n_tokens:.0f}× reduction')
print(f'Unique tokens used: {n_unique} / {K_vocab}')
print(f'Token grid:\n{token_ids}')
Encoder output: 8×8×16 = 1024 float32 values Token sequence: 64 integers, vocabulary size K=32 Float compression: 16× reduction Unique tokens used: 4 / 32 Token grid: [[0 0 0 0 1 1 1 1] [0 0 0 0 1 1 1 1] [0 0 0 0 1 1 1 1] [0 0 0 0 1 1 1 1] [2 2 2 2 3 3 3 3] [2 2 2 2 3 3 3 3] [2 2 2 2 3 3 3 3] [2 2 2 2 3 3 3 3]]
# PCA on D_feat -> 2D for visualisation
def pca_2d(X):
X_c = X - X.mean(axis=0)
cov = X_c.T @ X_c / len(X)
vals, vecs = np.linalg.eigh(cov)
idx = np.argsort(vals)[::-1]
return X_c @ vecs[:, idx[:2]]
z_2d = pca_2d(z_flat)
cb_2d = pca_2d(codebook)
colors_tok = plt.cm.tab20(np.linspace(0, 1, K_vocab))
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
fig.suptitle('Experiment 5: VQGAN Image Tokenization Pipeline', fontsize=13, fontweight='bold')
# Panel 1: Token ID grid (spatial)
ax1 = axes[0]
im1 = ax1.imshow(token_ids, cmap='tab20', vmin=0, vmax=K_vocab - 1, aspect='auto')
for i in range(H_feat):
for j in range(W_feat):
ax1.text(j, i, str(token_ids[i, j]), ha='center', va='center', fontsize=8, color='white',
fontweight='bold')
ax1.set_title(f'Token ID Map ({H_feat}×{W_feat} grid)\nEach cell = one codebook index')
ax1.set_xlabel('spatial column')
ax1.set_ylabel('spatial row')
ax1.grid(False)
plt.colorbar(im1, ax=ax1, label='Token ID')
# Panel 2: Feature vectors in PCA space, coloured by assigned token
ax2 = axes[1]
for k in range(K_vocab):
mask = token_ids.flatten() == k
if mask.any():
ax2.scatter(z_2d[mask, 0], z_2d[mask, 1], color=colors_tok[k], s=60, alpha=0.8, label=f'token {k}')
ax2.scatter(cb_2d[:, 0], cb_2d[:, 1], color='black', s=120, marker='*',
zorder=5, label='Codebook entries')
for k in range(K_vocab):
ax2.scatter(cb_2d[k, 0], cb_2d[k, 1], color=colors_tok[k], s=80, marker='D',
edgecolors='black', lw=0.8, zorder=6)
ax2.set_xlabel('PCA dimension 1')
ax2.set_ylabel('PCA dimension 2')
ax2.set_title(f'Encoder Outputs in PCA Space\n(coloured by assigned token, stars=codebook)')
# Panel 3: Flat token sequence visualisation
ax3 = axes[2]
seq_img = token_seq.reshape(1, -1)
ax3.imshow(seq_img, cmap='tab20', vmin=0, vmax=K_vocab - 1, aspect='auto')
for j, tok in enumerate(token_seq):
ax3.text(j, 0, str(tok), ha='center', va='center', fontsize=7.5, color='white', fontweight='bold')
ax3.set_title(f'Flattened Token Sequence\n({n_tokens} tokens, K={K_vocab} vocabulary)')
ax3.set_xlabel('position in sequence')
ax3.set_yticks([])
ax3.grid(False)
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: Three panels: a spatial grid of token IDs (each colour = one codebook entry), the encoder feature vectors projected to 2D via PCA (coloured by their assigned token, black stars = codebook positions), and the flattened 1D token sequence. The spatial grid clearly shows four coherent quadrant regions, each dominated by a cluster of nearby codebook entries.
Why it looks this way: The encoder learned to map structurally similar image regions to nearby points in feature space. Vector quantization then assigns the same token to all encoder outputs that fall in the same codebook Voronoi cell. The result is a spatially coherent token map — adjacent patches that depict similar content share the same token ID. The flat sequence is just this 2D map read in row-major order.
Key takeaway: VQGAN turns an image into a compact integer sequence by: (1) encoding it to a much smaller spatial grid, (2) looking up the nearest codebook entry at each position. The final sequence of 64 integers (for an 8×8 grid) fully characterises the visual content in a form that a Transformer can process — this is VQGAN's most important contribution to generative modelling.
Experiment 6: Autoregressive Token Generation — Images as Language¶
Question: How can a Transformer model image generation as next-token prediction, exactly like language modelling?
Once VQGAN tokenizes an image into a 1D integer sequence $[s_1, s_2, \ldots, s_T]$, generation becomes trivially identical to language modelling. A Transformer is trained to predict $p(s_i \mid s_1, \ldots, s_{i-1})$ — the probability distribution over the next token given all previous ones. To generate a new image, the model samples sequentially: it picks $s_1$ from $p(s_1)$, then $s_2$ from $p(s_2 \mid s_1)$, and so on. After all $T$ tokens are generated, VQGAN's decoder reconstructs the image from those indices. Conditioning on a text prompt is straightforward: prepend the text tokens to the image token sequence and let the Transformer attend to them.
np.random.seed(6)
H_ar, W_ar = 6, 6
T = H_ar * W_ar # sequence length
K_ar = 24 # vocabulary size
# Simulate a structured image token sequence
# Rows 0-2: tokens cluster around IDs 2,3,4 (a semantic region)
# Rows 3-5: tokens cluster around IDs 10,11,12 (another region)
true_seq = np.array([
2, 3, 4, 3, 2, 4, # row 0
3, 2, 4, 2, 3, 4, # row 1
4, 3, 2, 4, 2, 3, # row 2
10, 11, 12, 11, 10, 12, # row 3
11, 10, 12, 10, 11, 12, # row 4
12, 11, 10, 12, 10, 11, # row 5
])
# Simulate conditional probability distributions at 6 positions in the sequence
positions = [0, 5, 10, 18, 23, 30]
cond_probs = []
for pos in positions:
probs = np.random.dirichlet(np.ones(K_ar) * 0.15)
# Concentrate around the true token and its neighbours
for offset in [-1, 0, 1]:
tok = true_seq[pos] + offset
if 0 <= tok < K_ar:
probs[tok] += 2.5 if offset == 0 else 1.0
probs /= probs.sum()
cond_probs.append(probs)
# Simulate sampling (greedy = argmax)
sampled_seq = true_seq.copy()
sampled_seq[-5:] = [np.argmax(c) for c in cond_probs[-5:]] # last 5 tokens are predicted
print(f'Image: {H_ar}×{W_ar} = {T} tokens, vocabulary K={K_ar}')
print(f'True sequence: {true_seq}')
print(f'Sampled sequence: {sampled_seq}')
Image: 6×6 = 36 tokens, vocabulary K=24 True sequence: [ 2 3 4 3 2 4 3 2 4 2 3 4 4 3 2 4 2 3 10 11 12 11 10 12 11 10 12 10 11 12 12 11 10 12 10 11] Sampled sequence: [ 2 3 4 3 2 4 3 2 4 2 3 4 4 3 2 4 2 3 10 11 12 11 10 12 11 10 12 10 11 12 12 4 3 10 12 12]
fig = plt.figure(figsize=(17, 10))
fig.suptitle('Experiment 6: Autoregressive Image Generation — Images as Language',
fontsize=13, fontweight='bold')
# Top-left: token ID grid
ax1 = fig.add_subplot(2, 3, 1)
ax1.imshow(true_seq.reshape(H_ar, W_ar), cmap='viridis', vmin=0, vmax=K_ar-1, aspect='auto')
for i in range(H_ar):
for j in range(W_ar):
ax1.text(j, i, str(true_seq[i*W_ar + j]), ha='center', va='center',
fontsize=9, color='white', fontweight='bold')
ax1.set_title(f'Token Grid ({H_ar}×{W_ar})\n2D image → integer map')
ax1.set_xlabel('column')
ax1.set_ylabel('row')
ax1.grid(False)
# Top-centre: flat sequence with generation arrow
ax2 = fig.add_subplot(2, 3, 2)
seq_display = true_seq.reshape(1, -1)
ax2.imshow(seq_display, cmap='viridis', vmin=0, vmax=K_ar-1, aspect='auto')
for j, tok in enumerate(true_seq):
color = 'yellow' if j >= T - 5 else 'white'
ax2.text(j, 0, str(tok), ha='center', va='center', fontsize=7.5,
color=color, fontweight='bold')
ax2.axvline(T - 5.5, color='yellow', lw=2, linestyle='--')
ax2.set_title('Flattened Sequence (row-major)\nyellow = tokens being predicted')
ax2.set_xlabel('position in sequence')
ax2.set_yticks([])
ax2.grid(False)
# Top-right: token frequency distribution
ax3 = fig.add_subplot(2, 3, 3)
counts = np.bincount(true_seq, minlength=K_ar)
active_ids = np.where(counts > 0)[0]
ax3.bar(active_ids, counts[active_ids], color=plt.cm.viridis(active_ids / K_ar), alpha=0.85)
ax3.set_xlabel('Token ID')
ax3.set_ylabel('Frequency')
ax3.set_title('Token Frequency Distribution\n(vocabulary usage)')
ax3.set_xlim(-1, K_ar)
# Bottom row: conditional probability distributions at 3 selected positions
bottom_positions = [0, 17, 30]
bottom_probs = [cond_probs[0], cond_probs[3], cond_probs[5]]
bottom_titles = [
'p(s₁) — first token\n(unconditional)',
f'p(s₁₈ | s₁..s₁₇)\n(start of 2nd region)',
f'p(s₃₁ | s₁..s₃₀)\n(near end of sequence)',
]
for k, (pos, probs, title) in enumerate(zip(bottom_positions, bottom_probs, bottom_titles)):
ax = fig.add_subplot(2, 3, 4 + k)
bar_colors = [plt.cm.viridis(i / K_ar) for i in range(K_ar)]
ax.bar(range(K_ar), probs, color=bar_colors, alpha=0.85)
ax.axvline(true_seq[pos], color='red', lw=2.5, linestyle='--',
label=f'True token = {true_seq[pos]}')
ax.set_xlabel('Token ID')
ax.set_ylabel('Probability')
ax.set_title(title)
ax.set_xlim(-1, K_ar)
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The top row shows an image as: a 2D token grid (colour-coded by token ID), the same sequence flattened to 1D (yellow = tokens still to be predicted), and the token frequency histogram (the visual vocabulary in use). The bottom row shows conditional probability distributions $p(s_i \mid s_{<i})$ at three positions — earlier positions are less peaked (more uncertainty), later positions become sharply concentrated around the correct token (the context has established which region of the image we're in).
Why it looks this way: The two semantic regions (tokens 2–4 in the top half, tokens 10–12 in the bottom half) create a clear structure in the conditional distributions. After seeing 17 tokens all from the {2,3,4} cluster, the model should be highly confident the 18th token is also from {2,3,4}. When the sequence transitions to row 3 (bottom region), the distribution shifts to peak at {10,11,12}. This context-dependent sharpening is exactly what a Transformer learns via attention.
Key takeaway: An image token sequence is structurally identical to a text token sequence — both are integer arrays with long-range dependencies. Any Transformer architecture (GPT, LLaMA, etc.) can model images without modification by simply treating image token IDs as "words". This is the conceptual bridge that enabled models like DALL-E and Parti to generate images with language models, and ultimately led to unified multimodal architectures.
Summary¶
| Experiment | Key insight |
|---|---|
| 1. MSE Blur Problem | MSE minimisation learns the mean of a multimodal distribution — which, for opposing textures, is a flat gray output with no high-frequency detail |
| 2. Perceptual Distance | Feature-space distance (proxy for LPIPS) captures structural similarity that pixel MSE misses: a shift has high pixel error but low feature error; blur has low pixel error but high feature error |
| 3. PatchGAN Scoring | A bank of local patch critics detects spatially localised artifacts that a global discriminator would miss, forcing the decoder to produce convincing textures everywhere |
| 4. Adaptive λ | The gradient-ratio λ automatically schedules a curriculum: suppress GAN during coarse learning, then let it sharpen details once reconstruction gradients have decayed |
| 5. Image Tokenization | VQGAN encodes a spatial feature grid into a compact sequence of integer codebook indices — orders-of-magnitude compression of the raw float representation |
| 6. Autoregressive Generation | Image tokens are structurally identical to text tokens; next-token prediction on the flattened sequence is all a Transformer needs to generate images |
Connection to the broader theory¶
VQGAN is the culmination of the representation-learning thread that began with autoencoders (research/02-autoencoder.md) and progressed through VAE (research/03-vae.md) and VQ-VAE (research/04-vq-vae.md). Each step added one capability: VAE added probabilistic latents; VQ-VAE made latents discrete; VQGAN made discrete latents perceptually sharp and transformer-compatible. The integer token sequences produced by VQGAN are the bridge between the continuous-latent world studied in this section and the large-scale autoregressive generation models that define the current frontier of AI — models that treat pixels, audio, and text as interchangeable streams of discrete tokens.