Latent-Variable Representations › Discrete Latents

Many useful world-state variables can be modeled categorically, and discrete codes connect world models to the machinery of token-based sequence models. This section covers why discretization helps, how to train through non-differentiable quantization with straight-through gradients and Gumbel-softmax, and how codebook collapse arises and is mitigated.

Continuous Gaussian latents are a convenient default, but many distinctions useful to a world model can be treated categorically: an object may be a chair or a table, a light red or green, a door open or closed, a grid cell occupied or empty. Forcing such structure through a continuous Gaussian costs representational effort: the model must find a continuous embedding of a discrete distinction and hope that the dynamics over that embedding behave like the underlying discrete process. Discrete latents provide a natural representational form for categorical structure, but unsupervised code indices do not inherit the domain's semantic ontology. Their identities still need objectives and evaluation that align them with the distinctions the task uses.

Discrete latents let world models reuse autoregressive transformer machinery and next-token objectives. The analogy to language is useful but limited: visual tokens live on spatial and temporal lattices, their serialization changes dependencies, and their semantics and rates differ from text. Video tokenizers therefore require modality-specific architecture and training choices rather than inheriting every language-model recipe or scaling law automatically.

The problem is mechanical: you cannot backpropagate through argmin or argmax. Both operations are piecewise constant; a vanishingly small change to the input almost never changes the output, and when it does, the change is a discontinuous jump. Either way, the derivative is zero or undefined, and gradient descent has no purchase. Two families of workarounds dominate, and understanding both is worthwhile because they embody different philosophies: one quantizes after the fact and repairs the gradient, the other replaces the discrete sample with a differentiable relaxation.

The vector-quantized VAE (van den Oord, Vinyals, and Kavukcuoglu, 2017) maintains a codebook, a set of $K$ learned embedding vectors $e_1, \dots, e_K \in \mathbb{R}^{d}$. The encoder outputs a continuous vector $z_e(x)$, which is quantized to the nearest code:

$$z_q(x) = e_k, \qquad k = \underset{j}{\mathrm{argmin}} \; \| z_e(x) - e_j \|_2^2.$$

The decoder receives the discrete code's embedding. The quantization itself is cleanly non-differentiable, and the gradient problem is handled by the straight-through estimator: in the forward pass the decoder sees $z_q$, but in the backward pass the gradient at the decoder input is copied directly onto $z_e$, as if quantization were the identity:

$$z_{\text{st}} = z_e + \mathrm{sg}(z_q - z_e),$$

where $\mathrm{sg}(\cdot)$ denotes the stop-gradient operation. Read this expression carefully, because it is doing two jobs at once. In the forward pass, the two terms sum to exactly $z_q$, so the decoder sees the true quantized value. In the backward pass, the stop-gradient blocks all gradient flow through $z_q - z_e$, so the derivative of $z_{\text{st}}$ with respect to anything upstream is the derivative of $z_e$ alone, exactly as if the quantization step had been the identity map. This is a biased gradient estimate, since the true derivative of the quantization is essentially zero or undefined, but a simple and surprisingly effective one, on the principle that an approximate gradient direction is usually better than none.

The codebook itself is trained with two auxiliary losses:

$$\mathcal{L}_{\mathrm{VQ}} = \underbrace{\| x - \hat{x} \|^2}_{\text{reconstruction}} + \underbrace{\| \mathrm{sg}(z_e) - e_k \|^2}_{\text{codebook loss}} + \underbrace{\lambda_{\mathrm{commit}} \| z_e - \mathrm{sg}(e_k) \|^2}_{\text{commitment loss}}.$$

The codebook loss pulls embeddings toward encoder outputs. The stop-gradient on $z_e$ makes this term train only the assigned code. The commitment loss routes the opposite pressure to the encoder, with its independent weight $\lambda_{\mathrm{commit}}$. This is not the VAE's KL multiplier $\beta$. Without commitment pressure, encoder and codebook can chase one another or suffer scale instability, so convergence is less reliable. The original VQ-VAE used $\lambda_{\mathrm{commit}}=0.25$ as a starting value; treat it as a hyperparameter, not a universal constant.

The second workaround constructs a differentiable relaxation of a categorical draw. Add Gumbel noise to class logits and apply a temperature-controlled softmax:

$$\tilde{y}_k = \frac{\exp\big((\ell_k + g_k) / \tau\big)}{\sum_{k'} \exp\big((\ell_{k'} + g_{k'}) / \tau\big)}, \qquad g_k \stackrel{\mathrm{iid}}{\sim} \mathrm{Gumbel}(0, 1),$$

where $\ell_k$ are logits and $\tau > 0$ is a temperature. Adding Gumbel noise and taking the argmax samples exactly from the categorical distribution implied by the logits. As $\tau \to 0$, each relaxed sample approaches a one-hot vector. Low temperatures can saturate gradients away from category boundaries and make them sharp near ties. A straight-through variant uses a hard forward sample and the soft relaxation for backward computation, introducing a biased surrogate gradient. Temperature annealing is a common heuristic, but its endpoints and schedule depend on logit scale and observed gradient behavior.

The role of $\tau$ is easiest to see at the level of individual relaxed samples. The next cell fixes the logits, uses the same 2000 seeded Gumbel draws at every temperature, and records the mean sample entropy and mean maximum category mass. Entropy falls and maximum mass rises as individual samples sharpen.

A local engineering heuristic is to try Gumbel-softmax for a small categorical variable and vector quantization for high-rate perceptual tokenization. This is not a general regime boundary: both methods appear outside those settings, and the choice should be validated against estimator variance, required rate, and architecture.

Discrete latents bring their own silent failure: codebook collapse, where only a small subset of codes is used. Codes receive no direct codebook update in batches where they are unassigned, so imbalance can become self-reinforcing and persistent, though later encoder movement can revive a code. Data-based initialization, dead-code restarts, and exponential-moving-average updates can help; EMA stabilizes assigned-code updates but does not eliminate the need to monitor dead codes. A flat histogram is neither required nor sufficient evidence of a healthy codebook; interpret it against the expected state density and downstream behavior. Track perplexity, persistence, reconstruction, and downstream performance; prolonged zero use is a strong warning.