Vector Quantization — Visual Experiments¶

Goal: Build intuition for how a continuous vector becomes a discrete token — the operation that lets a latent space be modelled like language.

What we'll explore¶

#	Experiment	Question answered
1	Voronoi lookup	How does argmin over a codebook partition the space and snap vectors to codes?
2	Straight-through estimator	Why does training stall through a raw argmin, and how does STE fix it?
3	Rate–distortion	How does quantization error fall as the codebook grows from K codes?
4	Image → token grid	What does it mean to turn a continuous image into a grid of integer tokens?

---

Linked theory: research/01-vector-quantization.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from scipy.spatial import Voronoi, voronoi_plot_2d
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_e, codebook):
    """Hard nearest-neighbour quantization.

    z_e:      (N, D) encoder outputs
    codebook: (K, D) code vectors
    returns:  (indices (N,), z_q (N, D)) -- token id and the chosen code.
    """
    # squared Euclidean distance via the |a-b|^2 = |a|^2 - 2 a.b + |b|^2 expansion
    d2 = (
        (z_e ** 2).sum(1, keepdims=True)
        - 2 * z_e @ codebook.T
        + (codebook ** 2).sum(1)[None, :]
    )
    idx = np.argmin(d2, axis=1)
    return idx, codebook[idx]


print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as mpatches from scipy.spatial import Voronoi, voronoi_plot_2d 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_e, codebook): """Hard nearest-neighbour quantization. z_e: (N, D) encoder outputs codebook: (K, D) code vectors returns: (indices (N,), z_q (N, D)) -- token id and the chosen code. """ # squared Euclidean distance via the |a-b|^2 = |a|^2 - 2 a.b + |b|^2 expansion d2 = ( (z_e ** 2).sum(1, keepdims=True) - 2 * z_e @ codebook.T + (codebook ** 2).sum(1)[None, :] ) idx = np.argmin(d2, axis=1) return idx, codebook[idx] print('Setup complete.')

Setup complete.

---

Experiment 1: Voronoi lookup¶

Question: How does the argmin nearest-neighbour rule carve up the latent space, and what does "snapping to the nearest code" look like geometrically?

Vector quantization replaces each encoder output $z_e$ by the closest codebook vector $e_{k^\star}$, where $k^\star = \arg\min_j \lVert z_e - e_j \rVert_2$. Every code therefore owns a Voronoi cell — the region of space closer to it than to any other code — and quantization is just "which cell am I in?".

In [2]:

np.random.seed(1)

K = 12          # codebook size
n = 600         # number of encoder outputs

codebook = np.random.uniform(-3, 3, size=(K, 2))
# encoder outputs: a few blobs scattered around the plane
centers = np.random.uniform(-2.5, 2.5, size=(5, 2))
z_e = np.repeat(centers, n // 5, axis=0) + 0.7 * np.random.randn(n, 2)

idx, z_q = quantize(z_e, codebook)
usage = np.bincount(idx, minlength=K)
print(f'{K} codes, {n} points. Codes actually used: {(usage > 0).sum()}/{K}')

np.random.seed(1) K = 12 # codebook size n = 600 # number of encoder outputs codebook = np.random.uniform(-3, 3, size=(K, 2)) # encoder outputs: a few blobs scattered around the plane centers = np.random.uniform(-2.5, 2.5, size=(5, 2)) z_e = np.repeat(centers, n // 5, axis=0) + 0.7 * np.random.randn(n, 2) idx, z_q = quantize(z_e, codebook) usage = np.bincount(idx, minlength=K) print(f'{K} codes, {n} points. Codes actually used: {(usage > 0).sum()}/{K}')

12 codes, 600 points. Codes actually used: 12/12

In [3]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))
fig.suptitle('Nearest-neighbour quantization partitions the plane', fontsize=13, fontweight='bold')

vor = Voronoi(codebook)
ax = axes[0]
voronoi_plot_2d(vor, ax=ax, show_vertices=False, show_points=False,
                line_colors='gray', line_width=1.0, line_alpha=0.7)
ax.scatter(z_e[:, 0], z_e[:, 1], c=idx, cmap='tab20', s=14, alpha=0.7)
ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=240,
           edgecolor='white', linewidth=1.2, label='codebook $e_k$', zorder=5)
ax.set_xlim(-4, 4); ax.set_ylim(-4, 4)
ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1')
ax.set_title('Encoder outputs coloured by assigned code\n(Voronoi cells = code ownership)')
ax.legend(loc='upper right')

ax2 = axes[1]
# draw the snap: arrow from z_e to its code
for i in range(0, n, 6):
    ax2.plot([z_e[i, 0], z_q[i, 0]], [z_e[i, 1], z_q[i, 1]],
             color='steelblue', alpha=0.25, linewidth=0.6)
ax2.scatter(z_e[:, 0], z_e[:, 1], c='lightgray', s=10, alpha=0.6, label='$z_e$ (continuous)')
ax2.scatter(codebook[:, 0], codebook[:, 1], c='crimson', marker='*', s=240,
            edgecolor='white', linewidth=1.2, label='$z_q$ (snapped)', zorder=5)
ax2.set_xlim(-4, 4); ax2.set_ylim(-4, 4)
ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1')
ax2.set_title('Quantization collapses every point onto its code')
ax2.legend(loc='upper right')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5.5)) fig.suptitle('Nearest-neighbour quantization partitions the plane', fontsize=13, fontweight='bold') vor = Voronoi(codebook) ax = axes[0] voronoi_plot_2d(vor, ax=ax, show_vertices=False, show_points=False, line_colors='gray', line_width=1.0, line_alpha=0.7) ax.scatter(z_e[:, 0], z_e[:, 1], c=idx, cmap='tab20', s=14, alpha=0.7) ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=240, edgecolor='white', linewidth=1.2, label='codebook $e_k$', zorder=5) ax.set_xlim(-4, 4); ax.set_ylim(-4, 4) ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1') ax.set_title('Encoder outputs coloured by assigned code\n(Voronoi cells = code ownership)') ax.legend(loc='upper right') ax2 = axes[1] # draw the snap: arrow from z_e to its code for i in range(0, n, 6): ax2.plot([z_e[i, 0], z_q[i, 0]], [z_e[i, 1], z_q[i, 1]], color='steelblue', alpha=0.25, linewidth=0.6) ax2.scatter(z_e[:, 0], z_e[:, 1], c='lightgray', s=10, alpha=0.6, label='$z_e$ (continuous)') ax2.scatter(codebook[:, 0], codebook[:, 1], c='crimson', marker='*', s=240, edgecolor='white', linewidth=1.2, label='$z_q$ (snapped)', zorder=5) ax2.set_xlim(-4, 4); ax2.set_ylim(-4, 4) ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1') ax2.set_title('Quantization collapses every point onto its code') ax2.legend(loc='upper right') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The left panel draws the Voronoi diagram of the 12 codes (gray cell boundaries) and colours each encoder point by the code it lands on — every point inside a cell takes that cell's colour. The right panel shows the result of quantization: the thin blue lines are each $z_e$ being pulled onto its code, and all the continuous gray cloud collapses onto just 12 red stars.
2. Why it looks this way: argmin over Euclidean distance is exactly the rule that defines a Voronoi tessellation — the cell boundaries are the perpendicular bisectors between neighbouring codes. After the lookup, the only information that survives is which cell, i.e. the integer index.
3. Key takeaway: Quantization is a lossy, piecewise-constant map: an entire region of continuous space is replaced by a single discrete symbol. The error you pay is the length of those blue snap-lines.

---

Experiment 2: Straight-through estimator¶

Question: The forward pass uses argmin, which has zero gradient almost everywhere. So how does the encoder ever learn — and what does the straight-through estimator actually do?

We optimise an encoder output $z_e$ so that its quantized version $z_q$ lands near a target. With the true gradient through argmin (which is 0), nothing moves. With the STE trick $z_q = z_e + \mathrm{sg}[e_{k^\star} - z_e]$, the gradient at $z_q$ is copied straight back to $z_e$, so the encoder gets a usable signal.

In [4]:

np.random.seed(2)

# fixed codebook on a line so we can watch z_e hop between cells
codebook2 = np.array([[-2.0, 0.0], [-0.7, 0.0], [0.7, 0.0], [2.0, 0.0]])
target = np.array([1.9, 0.0])   # we want the quantized output to reach the rightmost code

lr = 0.15
steps = 60


def run(use_ste):
    z = np.array([-1.8, 1.2])      # start near the leftmost code
    traj, loss_hist = [z.copy()], []
    for _ in range(steps):
        _, zq = quantize(z[None, :], codebook2)
        zq = zq[0]
        loss = ((zq - target) ** 2).sum()
        loss_hist.append(loss)
        if use_ste:
            grad = 2 * (zq - target)   # STE: dL/dz_e := dL/dz_q
        else:
            grad = np.zeros(2)         # true gradient through argmin is 0
        z = z - lr * grad
        traj.append(z.copy())
    return np.array(traj), np.array(loss_hist)


traj_ste, loss_ste = run(True)
traj_raw, loss_raw = run(False)
print(f'STE final loss: {loss_ste[-1]:.3f}   |   raw-argmin final loss: {loss_raw[-1]:.3f}')

np.random.seed(2) # fixed codebook on a line so we can watch z_e hop between cells codebook2 = np.array([[-2.0, 0.0], [-0.7, 0.0], [0.7, 0.0], [2.0, 0.0]]) target = np.array([1.9, 0.0]) # we want the quantized output to reach the rightmost code lr = 0.15 steps = 60 def run(use_ste): z = np.array([-1.8, 1.2]) # start near the leftmost code traj, loss_hist = [z.copy()], [] for _ in range(steps): _, zq = quantize(z[None, :], codebook2) zq = zq[0] loss = ((zq - target) ** 2).sum() loss_hist.append(loss) if use_ste: grad = 2 * (zq - target) # STE: dL/dz_e := dL/dz_q else: grad = np.zeros(2) # true gradient through argmin is 0 z = z - lr * grad traj.append(z.copy()) return np.array(traj), np.array(loss_hist) traj_ste, loss_ste = run(True) traj_raw, loss_raw = run(False) print(f'STE final loss: {loss_ste[-1]:.3f} | raw-argmin final loss: {loss_raw[-1]:.3f}')

STE final loss: 0.010   |   raw-argmin final loss: 15.210

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('STE lets gradient flow through a non-differentiable lookup', fontsize=13, fontweight='bold')

ax = axes[0]
ax.scatter(codebook2[:, 0], codebook2[:, 1], c='black', marker='*', s=240,
           edgecolor='white', linewidth=1.2, label='codes', zorder=5)
ax.scatter(*target, c='crimson', marker='X', s=160, label='target', zorder=6)
ax.plot(traj_ste[:, 0], traj_ste[:, 1], '-o', ms=3, color='seagreen', label='STE: $z_e$ path')
ax.plot(traj_raw[:, 0], traj_raw[:, 1], '-o', ms=3, color='gray', label='raw argmin: stuck')
for x in (-1.35, 0.0, 1.35):
    ax.axvline(x, color='lightgray', ls='--', lw=0.8)
ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1')
ax.set_title('Encoder output trajectory (dashed = Voronoi boundaries)')
ax.legend(loc='upper center', fontsize=9)

ax2 = axes[1]
ax2.plot(loss_ste, color='seagreen', lw=2, label='STE')
ax2.plot(loss_raw, color='gray', lw=2, label='raw argmin')
ax2.set_xlabel('gradient step'); ax2.set_ylabel('quantized loss $\\|z_q - t\\|^2$')
ax2.set_title('Loss over training')
ax2.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('STE lets gradient flow through a non-differentiable lookup', fontsize=13, fontweight='bold') ax = axes[0] ax.scatter(codebook2[:, 0], codebook2[:, 1], c='black', marker='*', s=240, edgecolor='white', linewidth=1.2, label='codes', zorder=5) ax.scatter(*target, c='crimson', marker='X', s=160, label='target', zorder=6) ax.plot(traj_ste[:, 0], traj_ste[:, 1], '-o', ms=3, color='seagreen', label='STE: $z_e$ path') ax.plot(traj_raw[:, 0], traj_raw[:, 1], '-o', ms=3, color='gray', label='raw argmin: stuck') for x in (-1.35, 0.0, 1.35): ax.axvline(x, color='lightgray', ls='--', lw=0.8) ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1') ax.set_title('Encoder output trajectory (dashed = Voronoi boundaries)') ax.legend(loc='upper center', fontsize=9) ax2 = axes[1] ax2.plot(loss_ste, color='seagreen', lw=2, label='STE') ax2.plot(loss_raw, color='gray', lw=2, label='raw argmin') ax2.set_xlabel('gradient step'); ax2.set_ylabel('quantized loss $\\|z_q - t\\|^2$') ax2.set_title('Loss over training') ax2.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The gray run never moves — its loss is a flat line — because the true gradient through argmin is zero. The green STE run marches $z_e$ to the right, crossing Voronoi boundaries (dashed lines); each time it crosses, $z_q$ jumps to the next code and the loss drops in a staircase, finally settling at the code nearest the target.
2. Why it looks this way: STE pretends quantization is the identity on the backward pass, so $z_e$ feels the decoder's gradient $2(z_q - t)$ and keeps moving even though the forward output is piecewise constant. The staircase is the tell-tale sign of a discrete bottleneck: loss can only change when $z_e$ actually switches cells.
3. Key takeaway: Without STE the encoder is untrainable; with it, learning works — but notice the gradient is computed at $z_q$ and applied at $z_e$, the deliberate bias that makes VQ train at all.

---

Experiment 3: Rate–distortion¶

Question: A bigger codebook can describe data more finely. How fast does the average quantization error fall as we grow the codebook from K codes, and what is the cost?

Each token carries $\log_2 K$ bits (the rate). The mean squared distance $\lVert z_e - e_{k^\star}\rVert^2$ is the distortion — and a lower bound on reconstruction error. We fit codebooks of increasing size with k-means (a proxy for an optimally-placed codebook) and trace the trade-off.

In [6]:

np.random.seed(3)

# data: a curved 2D manifold (a noisy spiral) so structure matters
t = np.random.uniform(0, 3 * np.pi, size=4000)
r = 0.4 * t
data = np.stack([r * np.cos(t), r * np.sin(t)], axis=1) + 0.15 * np.random.randn(4000, 2)

Ks = [2, 4, 8, 16, 32, 64, 128, 256]
distortions = []
fitted = {}
for K in Ks:
    km = KMeans(n_clusters=K, n_init=3, random_state=0).fit(data)
    _, dq = quantize(data, km.cluster_centers_)
    distortions.append(((data - dq) ** 2).sum(1).mean())
    fitted[K] = km.cluster_centers_
distortions = np.array(distortions)
print('distortion per K:', np.round(distortions, 4))

np.random.seed(3) # data: a curved 2D manifold (a noisy spiral) so structure matters t = np.random.uniform(0, 3 * np.pi, size=4000) r = 0.4 * t data = np.stack([r * np.cos(t), r * np.sin(t)], axis=1) + 0.15 * np.random.randn(4000, 2) Ks = [2, 4, 8, 16, 32, 64, 128, 256] distortions = [] fitted = {} for K in Ks: km = KMeans(n_clusters=K, n_init=3, random_state=0).fit(data) _, dq = quantize(data, km.cluster_centers_) distortions.append(((data - dq) ** 2).sum(1).mean()) fitted[K] = km.cluster_centers_ distortions = np.array(distortions) print('distortion per K:', np.round(distortions, 4))

distortion per K: [2.9798 1.2622 0.3934 0.1213 0.0471 0.0253 0.0127 0.006 ]

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('More codes = less distortion, with diminishing returns', fontsize=13, fontweight='bold')

ax = axes[0]
bits = np.log2(Ks)
ax.plot(bits, distortions, '-o', color='darkorange', lw=2)
ax.set_xlabel('rate = $\\log_2 K$ (bits per token)')
ax.set_ylabel('distortion (mean squared quantization error)')
ax.set_title('Rate–distortion curve')
ax.set_yscale('log')
for K, b, d in zip(Ks, bits, distortions):
    ax.annotate(f'K={K}', (b, d), textcoords='offset points', xytext=(5, 6), fontsize=8)

ax2 = axes[1]
ax2.scatter(data[:, 0], data[:, 1], s=6, alpha=0.25, color='lightgray', label='data')
for K, color in [(8, 'royalblue'), (64, 'crimson')]:
    c = fitted[K]
    ax2.scatter(c[:, 0], c[:, 1], s=40, color=color, edgecolor='white',
                linewidth=0.6, label=f'codebook K={K}')
ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1')
ax2.set_title('Codebooks tile the data manifold')
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)) fig.suptitle('More codes = less distortion, with diminishing returns', fontsize=13, fontweight='bold') ax = axes[0] bits = np.log2(Ks) ax.plot(bits, distortions, '-o', color='darkorange', lw=2) ax.set_xlabel('rate = $\\log_2 K$ (bits per token)') ax.set_ylabel('distortion (mean squared quantization error)') ax.set_title('Rate–distortion curve') ax.set_yscale('log') for K, b, d in zip(Ks, bits, distortions): ax.annotate(f'K={K}', (b, d), textcoords='offset points', xytext=(5, 6), fontsize=8) ax2 = axes[1] ax2.scatter(data[:, 0], data[:, 1], s=6, alpha=0.25, color='lightgray', label='data') for K, color in [(8, 'royalblue'), (64, 'crimson')]: c = fitted[K] ax2.scatter(c[:, 0], c[:, 1], s=40, color=color, edgecolor='white', linewidth=0.6, label=f'codebook K={K}') ax2.set_xlabel('dim 0'); ax2.set_ylabel('dim 1') ax2.set_title('Codebooks tile the data manifold') 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 steeply at first then flattens — going from K=2 to K=16 slashes distortion, but doubling again buys ever less. The right panel shows why: with K=8 the codes coarsely cover the spiral; with K=64 they tile it densely, so every data point sits close to a code.
2. Why it looks this way: Each extra bit of rate roughly halves the volume each cell must cover, but error shrinks only with the cell diameter, so distortion drops sub-linearly. This is the classic rate–distortion trade-off, and it explains why you cannot get arbitrarily fine reconstruction from a single VQ layer without an exponentially large codebook.
3. Key takeaway: A single codebook hits diminishing returns — the motivation for stacking quantizers (Residual VQ) to add precision cheaply instead of paying $K = 2^{\text{bits}}$ all at once.

---

Experiment 4: Image → token grid¶

Question: What does it concretely mean to say VQ "turns an image into a sequence of tokens" that a Transformer can read like text?

We treat each pixel's RGB vector as a $z_e \in \mathbb{R}^3$, fit a small colour codebook, and quantize. The image becomes a grid of integer indices — exactly the discrete representation that the tokenized world model (later in this tier) feeds to a sequence model.

In [8]:

np.random.seed(4)

# build a synthetic 48x48 RGB image: smooth gradient + a couple of coloured blobs
H = W = 48
yy, xx = np.mgrid[0:H, 0:W] / H
img = np.stack([
    0.5 + 0.5 * xx,                       # red ramps left->right
    0.5 + 0.5 * yy,                       # green ramps top->bottom
    0.5 + 0.4 * np.sin(6 * (xx + yy)),    # blue ripples
], axis=-1)
# a bright blob
blob = np.exp(-(((xx - 0.7) ** 2 + (yy - 0.3) ** 2) / 0.02))
img[..., 0] = np.clip(img[..., 0] + 0.5 * blob, 0, 1)
img = np.clip(img, 0, 1)

K = 8
pixels = img.reshape(-1, 3)
km = KMeans(n_clusters=K, n_init=4, random_state=0).fit(pixels)
idx_map, zq_pixels = quantize(pixels, km.cluster_centers_)
token_grid = idx_map.reshape(H, W)
recon = zq_pixels.reshape(H, W, 3)
print(f'Image {H}x{W} -> token grid {token_grid.shape}, vocabulary size K={K}')
print('Top-left 6x6 block of token ids:')
print(token_grid[:6, :6])

np.random.seed(4) # build a synthetic 48x48 RGB image: smooth gradient + a couple of coloured blobs H = W = 48 yy, xx = np.mgrid[0:H, 0:W] / H img = np.stack([ 0.5 + 0.5 * xx, # red ramps left->right 0.5 + 0.5 * yy, # green ramps top->bottom 0.5 + 0.4 * np.sin(6 * (xx + yy)), # blue ripples ], axis=-1) # a bright blob blob = np.exp(-(((xx - 0.7) ** 2 + (yy - 0.3) ** 2) / 0.02)) img[..., 0] = np.clip(img[..., 0] + 0.5 * blob, 0, 1) img = np.clip(img, 0, 1) K = 8 pixels = img.reshape(-1, 3) km = KMeans(n_clusters=K, n_init=4, random_state=0).fit(pixels) idx_map, zq_pixels = quantize(pixels, km.cluster_centers_) token_grid = idx_map.reshape(H, W) recon = zq_pixels.reshape(H, W, 3) print(f'Image {H}x{W} -> token grid {token_grid.shape}, vocabulary size K={K}') print('Top-left 6x6 block of token ids:') print(token_grid[:6, :6])

Image 48x48 -> token grid (48, 48), vocabulary size K=8
Top-left 6x6 block of token ids:
[[4 5 5 5 5 5]
 [5 5 5 5 5 5]
 [5 5 5 5 5 5]
 [5 5 5 5 5 5]
 [5 5 5 5 5 5]
 [5 5 5 5 5 5]]

In [9]:

fig, axes = plt.subplots(1, 3, figsize=(14, 4.8))
fig.suptitle('A continuous image becomes a grid of discrete tokens', fontsize=13, fontweight='bold')

axes[0].imshow(img)
axes[0].set_title('Original (continuous RGB)')
axes[0].set_xlabel('x'); axes[0].set_ylabel('y')

im = axes[1].imshow(token_grid, cmap='tab10', vmin=0, vmax=K - 1)
axes[1].set_title(f'Token id map (K={K} vocabulary)')
axes[1].set_xlabel('x'); axes[1].set_ylabel('y')
cbar = fig.colorbar(im, ax=axes[1], fraction=0.046, ticks=range(K))
cbar.set_label('token id')

axes[2].imshow(recon)
axes[2].set_title('Reconstruction from codes')
axes[2].set_xlabel('x'); axes[2].set_ylabel('y')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 3, figsize=(14, 4.8)) fig.suptitle('A continuous image becomes a grid of discrete tokens', fontsize=13, fontweight='bold') axes[0].imshow(img) axes[0].set_title('Original (continuous RGB)') axes[0].set_xlabel('x'); axes[0].set_ylabel('y') im = axes[1].imshow(token_grid, cmap='tab10', vmin=0, vmax=K - 1) axes[1].set_title(f'Token id map (K={K} vocabulary)') axes[1].set_xlabel('x'); axes[1].set_ylabel('y') cbar = fig.colorbar(im, ax=axes[1], fraction=0.046, ticks=range(K)) cbar.set_label('token id') axes[2].imshow(recon) axes[2].set_title('Reconstruction from codes') axes[2].set_xlabel('x'); axes[2].set_ylabel('y') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The left image is smooth continuous colour. The middle panel is the same image after quantization — now every pixel is one of 8 integer ids (flat colour patches), and the printed 6×6 block shows the actual integers. The right panel reconstructs colour by replacing each id with its code; it is a posterized version of the original.
2. Why it looks this way: With only 8 colour codes, smooth gradients collapse into banded regions — each band is one Voronoi cell in RGB space. The middle panel is no longer an "image" of intensities but a map of symbols, which is precisely the form a language-model-style Transformer consumes.
3. Key takeaway: This integer grid is the bridge of Tier 9: once an observation is a sequence of tokens, autoregressive sequence modelling ("predict the next token") applies to images and world states just as it does to text.

---

Summary¶

Experiment	Key insight
1. Voronoi lookup	argmin partitions space into Voronoi cells; quantization snaps each vector to its cell's code, paying the snap distance as error.
2. Straight-through estimator	The raw argmin gradient is zero, so the encoder cannot learn; STE copies the gradient from $z_q$ to $z_e$, trading a deliberate bias for trainability.
3. Rate–distortion	Distortion falls as $\log_2 K$ grows but with diminishing returns — one codebook cannot be both small and precise.
4. Image → token grid	Quantizing every position turns a continuous signal into a grid of integer tokens — the form a sequence model can read.

Connection to the broader theory¶

These experiments isolate VQ as a standalone operator: a lossy, non-differentiable map from vectors to symbols, made trainable by STE. The staircase loss (Exp 2) and the diminishing returns (Exp 3) are exactly the pain points the rest of Tier 9 addresses — commitment loss and EMA codebook updates stabilise training, codebook collapse explains why many codes die, and Residual VQ / FSQ rework the quantizer itself. Experiment 4 is the payoff: discrete tokens are what let a latent space be modelled with the full machinery of language models.