Commitment Loss — Visual Experiments¶

Goal: See why the VQ loss splits into two stop-gradient terms, and what the commitment term actually buys: a stable, committed encoder.

What we'll explore¶

#	Experiment	Question answered
1	Two forces	How does stop-gradient split one distance into a code-pulling and an encoder-pulling force?
2	Stabilising assignment	Without commitment, why does the chosen token flicker — and how does commitment stop it?
3	The β trade-off	What does turning β up and down do to commitment strength vs reconstruction?

Linked theory: research/02-commitment-loss.md

In [1]:

Copied!





import numpy as np
import matplotlib.pyplot as plt
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 nearest(z, codebook):
    """Return (index, code) of the nearest codebook entry to a single vector z."""
    d2 = ((codebook - z) ** 2).sum(1)
    k = int(np.argmin(d2))
    return k, codebook[k]


print('Setup complete.')
import numpy as np
import matplotlib.pyplot as plt
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 nearest(z, codebook):
    """Return (index, code) of the nearest codebook entry to a single vector z."""
    d2 = ((codebook - z) ** 2).sum(1)
    k = int(np.argmin(d2))
    return k, codebook[k]


print('Setup complete.')

Setup complete.

Experiment 1: Two forces¶

Question: The full VQ loss has two squared-distance terms that look almost identical. What does stop-gradient ($\mathrm{sg}$) make each one do?

Codebook loss $\lVert \mathrm{sg}[z_e] - e\rVert^2$ freezes $z_e$, so its gradient moves only the code $e$ toward the encoder. Commitment loss $\beta\lVert z_e - \mathrm{sg}[e]\rVert^2$ freezes $e$, so its gradient moves only the encoder output $z_e$ toward the code. Same distance, opposite ownership.

In [2]:

Copied!





np.random.seed(1)

codebook = np.array([[-1.5, -1.0], [1.8, 0.4], [0.2, 1.9], [-0.4, -1.8]])
z_e = np.array([0.9, 1.3])
k, e = nearest(z_e, codebook)

# gradient directions (negative gradient = descent direction)
codebook_force = z_e - e      # moves e toward z_e
commit_force = e - z_e        # moves z_e toward e
print(f'z_e chooses code #{k} at {e}')
np.random.seed(1)

codebook = np.array([[-1.5, -1.0], [1.8, 0.4], [0.2, 1.9], [-0.4, -1.8]])
z_e = np.array([0.9, 1.3])
k, e = nearest(z_e, codebook)

# gradient directions (negative gradient = descent direction)
codebook_force = z_e - e      # moves e toward z_e
commit_force = e - z_e        # moves z_e toward e
print(f'z_e chooses code #{k} at {e}')

z_e chooses code #2 at [0.2 1.9]

In [3]:

Copied!





fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))
fig.suptitle('Stop-gradient splits one distance into two one-way forces', fontsize=13, fontweight='bold')

for ax, title, moving, arrow_from, arrow_vec, color in [
    (axes[0], 'Codebook loss: $\\|\\mathrm{sg}[z_e]-e\\|^2$  (moves the CODE)', 'e', e, codebook_force, 'crimson'),
    (axes[1], 'Commitment loss: $\\beta\\|z_e-\\mathrm{sg}[e]\\|^2$  (moves the ENCODER)', 'z_e', z_e, commit_force, 'seagreen'),
]:
    ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=220,
               edgecolor='white', linewidth=1.0, zorder=4)
    ax.scatter(*z_e, c='royalblue', s=120, zorder=5, label='$z_e$ (encoder output)')
    ax.scatter(*e, c='darkorange', marker='*', s=320, edgecolor='white',
               linewidth=1.2, zorder=6, label='$e$ (chosen code)')
    ax.annotate('', xy=arrow_from + 0.85 * arrow_vec, xytext=arrow_from,
                arrowprops=dict(arrowstyle='-|>', color=color, lw=2.5))
    ax.plot([z_e[0], e[0]], [z_e[1], e[1]], ':', color='gray', lw=1)
    ax.set_xlim(-3, 3); ax.set_ylim(-3, 3)
    ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1')
    ax.set_title(title, fontsize=10.5)
    ax.legend(loc='lower left', fontsize=9)

plt.tight_layout()
plt.show()
fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))
fig.suptitle('Stop-gradient splits one distance into two one-way forces', fontsize=13, fontweight='bold')

for ax, title, moving, arrow_from, arrow_vec, color in [
    (axes[0], 'Codebook loss: $\\|\\mathrm{sg}[z_e]-e\\|^2$  (moves the CODE)', 'e', e, codebook_force, 'crimson'),
    (axes[1], 'Commitment loss: $\\beta\\|z_e-\\mathrm{sg}[e]\\|^2$  (moves the ENCODER)', 'z_e', z_e, commit_force, 'seagreen'),
]:
    ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=220,
               edgecolor='white', linewidth=1.0, zorder=4)
    ax.scatter(*z_e, c='royalblue', s=120, zorder=5, label='$z_e$ (encoder output)')
    ax.scatter(*e, c='darkorange', marker='*', s=320, edgecolor='white',
               linewidth=1.2, zorder=6, label='$e$ (chosen code)')
    ax.annotate('', xy=arrow_from + 0.85 * arrow_vec, xytext=arrow_from,
                arrowprops=dict(arrowstyle='-|>', color=color, lw=2.5))
    ax.plot([z_e[0], e[0]], [z_e[1], e[1]], ':', color='gray', lw=1)
    ax.set_xlim(-3, 3); ax.set_ylim(-3, 3)
    ax.set_xlabel('dim 0'); ax.set_ylabel('dim 1')
    ax.set_title(title, fontsize=10.5)
    ax.legend(loc='lower left', fontsize=9)

plt.tight_layout()
plt.show()

No description has been provided for this image

What you see¶

What the plot shows: Both panels have the same $z_e$ (blue) and the same chosen code (orange star) at the same distance apart. The red arrow (left) points the code toward the encoder; the green arrow (right) points the encoder toward the code. Each panel moves only one of the two.
Why it looks this way: The $\mathrm{sg}$ wrapper zeroes the gradient on whatever it wraps, so codebook loss can only update $e$ and commitment loss can only update $z_e$. This is the deliberate decoupling that prevents the two endpoints from chasing each other.
Key takeaway: Commitment loss is the encoder-side half of closing the quantization gap — it exists precisely because STE never sends gradient to $z_e$ from the codebook side.

Experiment 2: Stabilising assignment¶

Question: What concretely goes wrong when commitment loss is missing, and how does adding it fix the symptom?

We put two codes either side of the origin and let a noisy reconstruction gradient pull the encoder output toward the Voronoi boundary (the worst place — ambiguous between both codes). Without commitment, the chosen token flickers every step. Commitment pulls $z_e$ off the boundary into one cell, freezing the token.

In [4]:

Copied!





np.random.seed(2)

codebook2 = np.array([[-1.0, 0.0], [1.0, 0.0]])
target = np.array([0.0, 0.0])   # reconstruction pull sits ON the boundary -> ambiguous
steps = 120
lr = 0.1
noise = 0.25


def simulate(beta, seed=0):
    rng = np.random.default_rng(seed)
    z = np.array([0.05, 0.6])
    tokens, xs, gaps = [], [], []
    for _ in range(steps):
        k, e = nearest(z, codebook2)
        tokens.append(k)
        xs.append(z[0])
        gaps.append(np.linalg.norm(z - e))
        # STE reconstruction gradient pulls z_e toward target, plus noise
        recon_grad = 2 * (z - target) + noise * rng.standard_normal(2)
        commit_grad = 2 * beta * (z - e)   # commitment pulls z_e toward chosen code
        z = z - lr * (recon_grad + commit_grad)
    return np.array(tokens), np.array(xs), np.array(gaps)


tok_off, x_off, gap_off = simulate(beta=0.0)
tok_on, x_on, gap_on = simulate(beta=1.0)
flips_off = int(np.abs(np.diff(tok_off)).sum())
flips_on = int(np.abs(np.diff(tok_on)).sum())
print(f'token flips  --  no commitment: {flips_off}   |   with commitment: {flips_on}')
np.random.seed(2)

codebook2 = np.array([[-1.0, 0.0], [1.0, 0.0]])
target = np.array([0.0, 0.0])   # reconstruction pull sits ON the boundary -> ambiguous
steps = 120
lr = 0.1
noise = 0.25


def simulate(beta, seed=0):
    rng = np.random.default_rng(seed)
    z = np.array([0.05, 0.6])
    tokens, xs, gaps = [], [], []
    for _ in range(steps):
        k, e = nearest(z, codebook2)
        tokens.append(k)
        xs.append(z[0])
        gaps.append(np.linalg.norm(z - e))
        # STE reconstruction gradient pulls z_e toward target, plus noise
        recon_grad = 2 * (z - target) + noise * rng.standard_normal(2)
        commit_grad = 2 * beta * (z - e)   # commitment pulls z_e toward chosen code
        z = z - lr * (recon_grad + commit_grad)
    return np.array(tokens), np.array(xs), np.array(gaps)


tok_off, x_off, gap_off = simulate(beta=0.0)
tok_on, x_on, gap_on = simulate(beta=1.0)
flips_off = int(np.abs(np.diff(tok_off)).sum())
flips_on = int(np.abs(np.diff(tok_on)).sum())
print(f'token flips  --  no commitment: {flips_off}   |   with commitment: {flips_on}')

token flips  --  no commitment: 18   |   with commitment: 0

In [5]:

Copied!





fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Commitment loss stops the assigned token from flickering', fontsize=13, fontweight='bold')

ax = axes[0]
ax.step(range(steps), tok_off, where='mid', color='gray', lw=1.5, label=f'$\\beta=0$ ({flips_off} flips)')
ax.step(range(steps), tok_on, where='mid', color='seagreen', lw=1.8, label=f'$\\beta=1$ ({flips_on} flips)')
ax.set_yticks([0, 1]); ax.set_yticklabels(['code 0', 'code 1'])
ax.set_xlabel('training step'); ax.set_ylabel('assigned token')
ax.set_title('Chosen token over training')
ax.legend(loc='center right', fontsize=9)

ax2 = axes[1]
ax2.plot(x_off, color='gray', lw=1.5, label='$\\beta=0$')
ax2.plot(x_on, color='seagreen', lw=1.8, label='$\\beta=1$')
ax2.axhline(0.0, color='crimson', ls='--', lw=1, label='Voronoi boundary')
ax2.set_xlabel('training step'); ax2.set_ylabel('$z_e$ x-coordinate')
ax2.set_title('Encoder output position (boundary = ambiguous)')
ax2.legend(loc='upper right', fontsize=9)

plt.tight_layout()
plt.show()
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Commitment loss stops the assigned token from flickering', fontsize=13, fontweight='bold')

ax = axes[0]
ax.step(range(steps), tok_off, where='mid', color='gray', lw=1.5, label=f'$\\beta=0$ ({flips_off} flips)')
ax.step(range(steps), tok_on, where='mid', color='seagreen', lw=1.8, label=f'$\\beta=1$ ({flips_on} flips)')
ax.set_yticks([0, 1]); ax.set_yticklabels(['code 0', 'code 1'])
ax.set_xlabel('training step'); ax.set_ylabel('assigned token')
ax.set_title('Chosen token over training')
ax.legend(loc='center right', fontsize=9)

ax2 = axes[1]
ax2.plot(x_off, color='gray', lw=1.5, label='$\\beta=0$')
ax2.plot(x_on, color='seagreen', lw=1.8, label='$\\beta=1$')
ax2.axhline(0.0, color='crimson', ls='--', lw=1, label='Voronoi boundary')
ax2.set_xlabel('training step'); ax2.set_ylabel('$z_e$ x-coordinate')
ax2.set_title('Encoder output position (boundary = ambiguous)')
ax2.legend(loc='upper right', fontsize=9)

plt.tight_layout()
plt.show()

What you see¶

What the plot shows: With $\beta=0$ (gray) the token jumps back and forth between code 0 and code 1 dozens of times, and the right panel shows $z_e$ hovering right on the red boundary line. With $\beta=1$ (green) the token locks onto a single code after a few steps, because $z_e$ is pulled clearly to one side of the boundary.
Why it looks this way: The reconstruction pull alone parks $z_e$ on the boundary, where tiny noise flips the argmin every step. Commitment adds a force toward the currently chosen code, which is self-reinforcing: once slightly on one side, it gets pulled further that way and commits.
Key takeaway: Commitment loss converts an ambiguous, oscillating assignment into a stable one — this is the "commit" in its name, and why training is far more stable with it.

Experiment 3: The β trade-off¶

Question: β controls how hard the encoder is pulled toward its code. What does sweeping β do to the commitment gap and to reconstruction quality?

We sweep β and, for each value, measure two things at the end of training: the commitment gap $\lVert z_e - e\rVert$ (how tightly the encoder hugs its code) and the reconstruction distance $\lVert z_e - \text{target}\rVert$ (how well the encoder still serves the decoder's objective). The original paper's $\beta = 0.25$ sits in the sweet spot.

In [6]:

Copied!





np.random.seed(3)

# target is offset from the code so recon and commitment genuinely compete
codebook3 = np.array([[-1.0, 0.0], [1.0, 0.0]])
target3 = np.array([1.8, 0.0])   # pulls z_e to the right, past code 1
betas = np.array([0.0, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0])


def final_state(beta, steps=400, lr=0.05, seed=0):
    rng = np.random.default_rng(seed)
    z = np.array([0.9, 0.4])
    for _ in range(steps):
        k, e = nearest(z, codebook3)
        recon_grad = 2 * (z - target3) + 0.1 * rng.standard_normal(2)
        commit_grad = 2 * beta * (z - e)
        z = z - lr * (recon_grad + commit_grad)
    k, e = nearest(z, codebook3)
    return np.linalg.norm(z - e), np.linalg.norm(z - target3)


gaps, recons = [], []
for b in betas:
    g, r = final_state(b)
    gaps.append(g); recons.append(r)
gaps, recons = np.array(gaps), np.array(recons)
print('beta :', betas)
print('gap  :', np.round(gaps, 3))
print('recon:', np.round(recons, 3))
np.random.seed(3)

# target is offset from the code so recon and commitment genuinely compete
codebook3 = np.array([[-1.0, 0.0], [1.0, 0.0]])
target3 = np.array([1.8, 0.0])   # pulls z_e to the right, past code 1
betas = np.array([0.0, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0])


def final_state(beta, steps=400, lr=0.05, seed=0):
    rng = np.random.default_rng(seed)
    z = np.array([0.9, 0.4])
    for _ in range(steps):
        k, e = nearest(z, codebook3)
        recon_grad = 2 * (z - target3) + 0.1 * rng.standard_normal(2)
        commit_grad = 2 * beta * (z - e)
        z = z - lr * (recon_grad + commit_grad)
    k, e = nearest(z, codebook3)
    return np.linalg.norm(z - e), np.linalg.norm(z - target3)


gaps, recons = [], []
for b in betas:
    g, r = final_state(b)
    gaps.append(g); recons.append(r)
gaps, recons = np.array(gaps), np.array(recons)
print('beta :', betas)
print('gap  :', np.round(gaps, 3))
print('recon:', np.round(recons, 3))

beta : [0.   0.05 0.1  0.25 0.5  1.   2.   4.  ]
gap  : [0.816 0.778 0.743 0.655 0.547 0.411 0.275 0.167]
recon: [0.016 0.022 0.057 0.145 0.253 0.389 0.525 0.633]

In [7]:

Copied!





fig, ax = plt.subplots(figsize=(8.5, 5.5))
fig.suptitle('β balances commitment against reconstruction', fontsize=13, fontweight='bold')

ax.plot(betas, gaps, '-o', color='seagreen', lw=2, label='commitment gap $\\|z_e-e\\|$')
ax.plot(betas, recons, '-s', color='darkorange', lw=2, label='reconstruction dist $\\|z_e-t\\|$')
ax.axvline(0.25, color='crimson', ls='--', lw=1.2, label='paper default $\\beta=0.25$')
ax.set_xlabel('commitment weight $\\beta$ (log scale)')
ax.set_ylabel('distance at end of training')
ax.set_xscale('symlog', linthresh=0.05)
ax.set_title('Larger β tightens the code, looser on the decoder target')
ax.legend()

plt.tight_layout()
plt.show()
fig, ax = plt.subplots(figsize=(8.5, 5.5))
fig.suptitle('β balances commitment against reconstruction', fontsize=13, fontweight='bold')

ax.plot(betas, gaps, '-o', color='seagreen', lw=2, label='commitment gap $\\|z_e-e\\|$')
ax.plot(betas, recons, '-s', color='darkorange', lw=2, label='reconstruction dist $\\|z_e-t\\|$')
ax.axvline(0.25, color='crimson', ls='--', lw=1.2, label='paper default $\\beta=0.25$')
ax.set_xlabel('commitment weight $\\beta$ (log scale)')
ax.set_ylabel('distance at end of training')
ax.set_xscale('symlog', linthresh=0.05)
ax.set_title('Larger β tightens the code, looser on the decoder target')
ax.legend()

plt.tight_layout()
plt.show()

What you see¶

What the plot shows: As β grows, the green commitment gap shrinks monotonically — the encoder hugs its code ever tighter. But the orange reconstruction distance rises: the encoder is increasingly dragged away from the decoder's preferred target toward the code. At β=0 the encoder ignores the code entirely (large gap); at very large β it pins to the code at the cost of reconstruction.
Why it looks this way: β sets the tug-of-war between two pulls on the same $z_e$ — one toward the reconstruction target, one toward the chosen code. The end position is a weighted compromise, sliding from "all target" to "all code" as β increases.
Key takeaway: β is a stability/fidelity knob, not a free win. The paper's 0.25 keeps the encoder committed enough to be stable while still letting reconstruction dominate — which is why it generalises across many setups.

Summary¶

Experiment	Key insight
1. Two forces	Stop-gradient turns one distance into two one-way pulls: codebook loss moves the code, commitment loss moves the encoder.
2. Stabilising assignment	Without commitment, $z_e$ lingers on the Voronoi boundary and the token flickers; commitment pulls it into one cell and the token locks.
3. The β trade-off	Bigger β tightens the encoder-to-code gap but loosens reconstruction; β=0.25 is the practical sweet spot.

Connection to the broader theory¶

Commitment loss is the encoder-side cure for the gap that the straight-through estimator leaves open. It stabilises which code each input commits to, but it does nothing for codes that are never chosen — those still die. The codebook side of the gap is handled either by the codebook-loss term shown in Experiment 1 or, more stably, by EMA codebook updates, the next topic in this tier.