Latent-Variable Representations

Imagine a robot with a camera pointed at a pendulum. The camera delivers 16 by 16 grayscale images, 256 numbers per frame, yet each snapshot lies on a one-dimensional manifold parameterized by the pendulum's angle. Modeling the raw frames exposes a dynamics model to a much higher-dimensional observation representation than that underlying snapshot manifold. It must learn or preserve the renderer, including the bob's glow, the fixed background, and the correlations among pixels, in addition to the pendulum's dynamics. What the robot needs is a compact internal description, a latent state, that retains the information required for prediction and control. When observation generation is part of the model, a decoder maps that state to a distribution over observations; the state need not preserve every detail of the particular input frame.

Chapter 2 of this part asked what makes a state representation sufficient for control. This chapter asks a harder, more practical question: when nobody hands you the state, how do you learn it from raw observations? Latent-variable models are one influential family of approaches in modern world modeling. You will learn how autoencoders compress observations, why plain reconstruction is a trap, how variational inference turns compression into probabilistic modeling, why bottlenecks can collapse, how to work with discrete codes, and what makes one latent geometry better than another.

The throughline is a single idea: a good latent representation keeps the information the world model needs and discards the rest on purpose. A reconstruction-only autoencoder selects information by reconstruction utility, not by downstream decision relevance, so useful and ornamental factors may remain mixed. The machinery we build in this chapter, from the ELBO's regularization term to the bottleneck coefficient $\beta$ to the discrete codes of a VQ-VAE, gives us ways to control and audit what the latent retains. None of these tools identifies task relevance automatically; that still has to be measured downstream.

Reconstruction is the simplest training signal for a learned representation, and this section shows why it falls short for world modeling. We introduce the autoencoder, identify three structural defects of a purely reconstructive objective, and explain what a probabilistic latent-variable model adds beyond a compression map.

A simple learned compressor is the autoencoder. You train two networks jointly:

• Encoder $q_\phi(z \mid x)$: in a probabilistic autoencoder, maps $x$ to the parameters of a low-dimensional latent distribution; in a deterministic autoencoder, it maps $x$ directly to a code $z$.
• Decoder $p_\theta(x \mid z)$: maps the code back to a reconstruction of the observation.

The training objective is reconstruction error, typically squared error or cross-entropy, minimized over a training set of observations. A narrow bottleneck removes the direct coordinatewise identity route and creates pressure to compress, so the encoder must find a lower-dimensional route through the network.

This is the compression half of what a world model needs. If the decoder can rebuild the image from $z$, then $z$ is, in a reconstruction sense, a sufficient description of the observation. But dimensionality alone is not a strict information budget: continuous codes with enough precision and expressive networks can assign distinct values to a finite training set and memorize the inverse map. Finite precision, noise, quantization, regularization, limited model capacity, and the requirement to generalize are what turn the architectural bottleneck into a practical information constraint. Whatever survives the round trip is the information the trained code makes available to its decoder, not necessarily the information a downstream agent needs.

Notice what is not part of the objective. There is no term asking the code to predict the future, no term asking it to respect the actions the agent will take, and no term asking it to express uncertainty. The autoencoder is a purely static, purely descriptive device: it learns to describe the present, in whatever encoding happens to fall out of gradient descent. That will turn out to be both its strength, because it needs no labels or interactions, and its fundamental limitation, because a world model is not a describer of the present but a predictor of futures.

Here is the trap. Reconstruction quality tells you the code retains information, but it says almost nothing about the structure of the code. Three problems recur.

Arbitrary geometry. Nothing in the reconstruction objective cares how codes are arranged in latent space. Two nearby observations can land far apart, and two very different observations can land next to each other. If your planning algorithm later interpolates between latents, or treats distance in latent space as similarity in the world, it is trusting structure that was never trained to exist. This failure is invisible at training time: reconstruction error is perfectly happy with a code assignment that scatters smoothly varying observations chaotically across latent space, as long as the decoder can memorize the mapping. But the moment a second consumer, such as a dynamics model, a planner, or a controller, starts using the latent as a space rather than as a list, the missing geometry becomes a live defect.

Holes in the space. The encoder only assigns codes to points near the training data. The rest of latent space is unsupported territory. Unconstrained sampling or planning can enter these regions and decode arbitrary outputs unless an additional latent prior or support constraint keeps proposals near encoded data. A common MLP decoder still returns a continuous output there, but continuity does not make the extrapolation reliable.

No notion of probability by itself. A deterministic autoencoder gives one point code per observation; it does not say how surprising the observation is or express ambiguity among possible states. A belief-state system can still maintain a distribution over deterministic codes with particles, ensembles, or another probabilistic layer, but that uncertainty is external to the autoencoder. A probabilistic encoder supplies a more direct starting point, though its uncertainty must still be calibrated and embedded in a temporal model.

The absence of probability and the danger of unsupported sampling motivate an explicit generative model: a probability model over observations and latents. Arbitrary geometry still requires separate objectives and evaluation. Probability models assign mass or density and support sampling; joint models additionally support conditioning. Variational inference turns the encoder into an explicit approximate posterior and defines a probability measure over latent codes. That measure makes likelihood-based reasoning possible, but it does not by itself give latent coordinates semantic or task-aligned geometry. Predictive or control-sufficient representations need not always model a normalized observation density, especially when they carry state through another predictive pathway rather than regenerate observations.

This section builds the variational machinery step by step: the latent-variable model and its intractable marginal, the evidence lower bound with its two-term decomposition, and the two techniques, reparameterization and amortization, that make the objective trainable by gradient descent.

This is the modeling direction we want: generate an observation $x$ from a latent $z$, and infer a distribution over $z$ after observing $x$. The factorization is statistical, not automatically causal. Calling $z$ a physical cause requires additional structural assumptions and validation through time or interventions. Even without that causal interpretation, the direction is useful: it supports sampling, conditioning, and later models of action-conditioned latent transitions.

Training means adjusting $\theta$ so that $p_\theta(x)$ is high on real observations. For a generic nonlinear neural-network decoder, the marginal usually has no tractable closed form, even with a low-dimensional continuous $z$; analytic special cases such as the linear-Gaussian model do exist. The ELBO gives us a differentiable surrogate for the generic case.

The classical move, which you met in the filtering chapter, is to approximate the posterior with variational inference. Introduce a family of distributions $q_\phi(z \mid x)$, the approximate posterior or encoder, and optimize it to stay close to the true posterior $p_\theta(z \mid x)$. There is a pleasing symmetry in the setup: the generative model runs in the direction $z \to x$, from state to observation, while the inference model runs in the reverse direction $x \to z$, from observation to belief. In a Bayes filter, the direct analogue of $p_\theta(x \mid z)$ is the observation likelihood, and inference combines that likelihood with a prior belief in the measurement update. A separate transition model supplies the prior belief when temporal dynamics are introduced. Here the observation and inference directions are learned together, and the next section shows how a single objective trains them.

We cannot minimize $\mathrm{KL}(q_\phi(z \mid x) \,\|\, p_\theta(z \mid x))$ directly because the true posterior contains the intractable marginal: by Bayes' rule, $p_\theta(z \mid x) = p_\theta(x, z) / p_\theta(x)$, and it is exactly $p_\theta(x)$ that we cannot compute. The trick is to rewrite the log marginal density in a form that sidesteps it. For the following identity, assume $q_\phi(z \mid x)$ is positive wherever $p_\theta(x,z)$ is positive for the fixed $x$, and that the displayed expectations exist. Start from the definition of the marginal and multiply and divide by $q_\phi(z \mid x)$:

This manipulation puts the chosen approximation $q$ inside the integral for $\log p_\theta(x)$, creating an expectation over $q$ of the ratio $p_\theta(x, z)/q_\phi(z \mid x)$. The logarithm of an expectation is hard, so apply Jensen's inequality, which holds because $\log$ is concave:

The right-hand side is the evidence lower bound, or ELBO. Two facts make it the workhorse of representation learning.

First, the gap is exactly the quantity we could not compute:

For fixed generative parameters $\theta$, maximizing the ELBO over $\phi$ reduces the KL gap to the true posterior within the chosen variational family. Joint optimization over $\theta$ and $\phi$ also raises the bound used to fit the generative model, although neither quantity is guaranteed to improve monotonically at every stochastic update. When the gap is zero, the encoder has recovered the exact posterior for that model.

Second, the ELBO splits into two interpretable terms. Expand $\log p_\theta(x, z) = \log p_\theta(x \mid z) + \log p_\theta(z)$:

The first term rewards reconstruction; the second penalizes each approximate posterior for deviating from the prior. This defines a reference measure for sampling and indirectly pressures the aggregated posterior toward that prior. It does not guarantee hole-free coverage, plausible interpolation, or semantic distance. Those geometric properties require direct evaluation or additional structure. What the KL does provide is an explicit price for information: every posterior deviation from the prior has to earn its reconstruction benefit.

It is worth being precise about what the regularization term is not. It does not require encoded training data to fill the prior, and prior density is not semantic geometry. Interpolation and metric alignment remain empirical questions.

Two engineering problems stand between the ELBO and gradient descent, and both have elegant, now-standard solutions.

Problem 1: a black-box draw blocks ordinary pathwise backpropagation to the distribution parameters. Score-function estimators are possible, but often have high variance. The reparameterization trick instead writes a draw as a differentiable transform of parameter-free noise. For a Gaussian encoder,

where $\mu_\phi$ and $\sigma_\phi$ are network outputs and $\odot$ is elementwise multiplication. All randomness now lives in $\varepsilon$, which does not depend on any parameter, and gradients flow through $\mu_\phi$ and $\sigma_\phi$. The conceptual shift is subtle but important: instead of sampling a value of $z$, we sample a source of randomness and then compute $z$ deterministically from it. The function from parameters to samples is now smooth, so the expectation over $\varepsilon$ can be estimated by a handful of Monte Carlo draws, often just one, and each estimate carries useful gradients. This trick was one key ingredient that made the variational autoencoder of Kingma and Welling (2013) and Rezende, Mohamed, and Wierstra (2014) practical; amortized inference, introduced next, supplies the other ingredient needed to encode new observations efficiently. Reparameterization recurs throughout this book whenever a stochastic layer must be threaded through a differentiable computation.

Problem 2: a separate optimization per data point. Traditional local variational inference often optimizes a dedicated set of variational parameters for each observation. That is poorly suited to a world model that must encode new observations under a tight latency budget, because a long inner optimization loop may not fit between camera frames. Instead we amortize inference: one encoder network $q_\phi(z \mid x)$ learns to output the variational parameters for any $x$. The term emphasizes that the cost of learning an inference procedure is shared across the data distribution rather than paid through a fresh optimization for every observation. Amortization trades per-observation optimization cost for a single upfront training cost, and it is what lets a learned world model ingest a fresh camera frame in one forward pass. The trade has a known downside worth naming: the encoder can be systematically imprecise on inputs it has rarely seen, a form of amortization error that dedicated per-observation optimization can reduce, though not necessarily eliminate. Local optimization may remain incomplete, and the chosen variational family can still be misspecified. For agents with tight latency budgets, amortized inference is often the practical default; whether to add bounded per-observation refinement depends on the available compute, the approximation error, and the cost of an inference mistake.

The closed form of the Gaussian KL shows what the regularizer asks of the encoder. Each dimension is minimized when the encoder's mean $\mu_j$ is zero and its variance $\sigma_j^2$ is one. The KL pressure pulls each encoding distribution toward zero mean and unit variance. The reconstruction term pushes the other way by rewarding informative codes.

The ELBO can be read as an information-theoretic tradeoff between compression and reconstruction, which makes the strength of the bottleneck an explicit knob. This section develops that reading, introduces the beta-weighted objective, and studies posterior collapse, the silent failure in which the model stops using its latent state.

The ELBO also has a rate-distortion reading. Define the joint distribution induced by the data and encoder as $q(x,z) = p_{\text{data}}(x)q(z\mid x)$, and define its marginal $q(z) = \int p_{\text{data}}(x)q(z\mid x)\,dx$, called the aggregated posterior. Let $I_q(X;Z)$ denote mutual information under that induced joint. Averaged over data, the KL term then decomposes as

The expected KL is therefore an upper bound on the mutual information $I_q(X;Z)$ and serves as an operational rate surrogate; reconstruction error serves as distortion. A $\beta$ sweep produces achieved operating points, but those points need not lie on the optimal rate-distortion frontier.

This reading matters because a world model need not preserve every pixel equally: required fidelity depends on its likelihood, downstream task, and observation model. A driving simulator may need the positions of cars, the road layout, and the traffic light state without regenerating every leaf on every tree. A bottleneck can squeeze out detail that is irrelevant to those requirements. Bigger, sharper reconstructions are therefore not sufficient evidence of better representations. Detail that the dynamics model must then carry forward is not free; it expands the state the model and planner must process. Task-aware compression can be a contribution to a world model rather than merely a compromise.

For a scalar memoryless Gaussian source with variance $\sigma^2$, mean-squared-error distortion, and $0 < D \leq \sigma^2$, the optimal curve is $R(D) = \frac{1}{2}\log(\sigma^2 / D)$. It is an analytic benchmark for optimal codes, not a guarantee about a trained VAE's sweep.

The knob that controls the squeeze is $\beta$:

With $\beta = 1$ you have the standard VAE. A $\beta$-VAE with $\beta > 1$ charges more for posterior deviation from the prior and can empirically encourage factorized representations in some settings. It does not guarantee disentanglement or semantic factor selection; optimization, decoder capacity, data statistics, and inductive biases determine what survives. Lower $\beta$ permits a higher achieved rate and often improves reconstruction, but the added information may be useful structure or nuisance detail. Treat $\beta$ as a swept hyperparameter selected with downstream sufficiency and uncertainty criteria.

A particularly consequential VAE failure mode is not bad reconstruction, but the model quietly learning not to use its latents at all.

Collapse can be a stable stationary point, and in some decoder and data regimes an optimum, because the KL is perfectly satisfied while a capable decoder remains competitive without $z$. Aggregate loss may look normal. An instrumented log does distinguish collapse: near-zero KL, few active units, and weak decoder sensitivity to latent traversals are direct warnings.

Collapse is common in sequence VAEs with powerful autoregressive decoders because local context can reduce the marginal value of a global latent. It can become self-reinforcing: as the decoder ignores $z$, reconstruction gradients reaching the encoder weaken, making recovery difficult, though not mathematically impossible. Optimization schedule and model mismatch matter alongside decoder capacity.

The symptoms are worth memorizing, because they are easy to miss:

• Per-dimension KL values near zero for most or all latent dimensions.
• Samples from the prior that look as good as reconstructions, used as a supporting symptom rather than a diagnosis by itself.
• Latent traversals that change nothing in the decoded output.

Collapse arises from objective incentives interacting with optimization and decoder capacity. Mitigations target one or more of those factors.

• KL annealing: start with $\beta = 0$ so the decoder learns to use $z$, then ramp $\beta$ up. Once dependence on $z$ improves reconstruction, later removal may be less attractive because it incurs a reconstruction penalty.
• Free bits: replace each KL penalty by $\max(\mathrm{KL}_j, \lambda)$. Below $\lambda$, the penalty is constant and exerts no gradient pressure to shrink further; above it, the usual KL gradient applies. This offers up to $\lambda$ nats of penalty-free capacity per dimension, but it does not guarantee that the model uses that capacity.

The effect of the clamp is easiest to see on a concrete set of per-dimension KL values. The next cell uses a fixed, hand-chosen array of hypothetical per-dimension KLs, spanning dead, weakly used, and heavily used dimensions, and applies the free-bits clamp $\max(\mathrm{KL}_j, \lambda)$ with $\lambda = 0.5$ nats. No model or randomness is involved.

• Weaken the decoder: remove skip connections that bypass the bottleneck, reduce decoder capacity, or drop out decoder inputs. This raises the reconstruction cost of collapse and can make informative latents favorable over a practical range of $\beta$.

For architectures whose downstream dynamics and planner rely on $z$ as their observation-dependent state, posterior collapse is an existential risk rather than a nuisance. Collapse empties that state of information about $x$. The consequences can compound: the dynamics model learns transitions over an uninformative variable, and the planner optimizes over a state that omits observation-dependent distinctions. Aggregate reconstruction loss may rise, as it does in the high-$\beta$ sweep below, or it may remain competitive while failing to reveal which downstream information was lost. Architectures with another recurrent or autoregressive state path can retain predictive information even when one stochastic latent collapses, so diagnose every state pathway explicitly. This is why later chapters on world-model architectures treat per-dimension KL and latent-sensitivity probes as standard operating instrumentation.

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:

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:

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:

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:

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.

Two latent spaces can match reconstruction error while differing enormously in usefulness for an agent. This section defines the geometric properties that matter, works through the linear case where the answer is exact, and states the identifiability limits of unsupervised representation learning.

Suppose two VAEs achieve identical reconstruction error. Their latent spaces can still differ enormously in usefulness for an agent. What distinguishes a good latent geometry?

• Supported coverage: avoid unintended gaps along the encoded data manifold, and interpolate along valid manifold paths. Topological holes, such as the interior of a circle, should be preserved rather than filled.
• Topological fidelity: the latent geometry should mirror the topology of the true state space. A pendulum's angle lives on a circle; a latent that flattens it onto a line creates a coordinate seam between the "ends" of the swing. An agent can plan discontinuously around a full rotation if its planner or dynamics model uses that nonperiodic metric without explicit wrap-around features, kernels, or boundary handling.
• Metric alignment: distances in latent space should be predictive of distances in world state, at least for the aspects the agent must control. Local differentiability or sufficient smoothness of the encoded transition or observation map enables local linearization. Metric alignment instead makes the Euclidean neighborhoods and step sizes used by that approximation meaningful for retrieval, planning, and control.
• Factor separation: when the downstream task requires independent intervention, editing, or interpretation, distinct world factors should be separable or independently traversable in latent space. That structure lets a planner vary one aspect of an imagined future without accidentally changing the others. Axis-aligned separation is not a universal requirement; an entangled representation can still be control-sufficient.

None of these properties is enforced by reconstruction alone, which is precisely why they deserve their own attention. The next two subsections make the claim concrete, first with an example where everything works out analytically, then with the theorem that bounds how far the good news extends.

To see that these claims are not hand-waving, consider the simplest possible latent-variable model: a Gaussian prior $z \sim \mathcal{N}(0, I_d)$ and a linear-Gaussian likelihood $p_\theta(x \mid z) = \mathcal{N}(W z + \mu, \sigma^2 I)$, where $x \in \mathbb{R}^n$. For this model the true posterior is analytically tractable, because the product of two Gaussians remains Gaussian and the integral in the marginal can be carried out exactly. Completing the square in the exponent of $p(z) \, p(x \mid z)$ gives

So the exact posterior mean is a fixed linear map and its covariance is constant. For centered data with $d<n$, let $\lambda_1 \ge \cdots \ge \lambda_n$ be the sample-covariance eigenvalues and suppose the retained eigenvalues satisfy $\lambda_d > \sigma^2$, where $\sigma^2$ is the fitted isotropic noise floor. At a global probabilistic-PCA optimum,

where $U_d$ contains leading sample-covariance eigenvectors, $\Lambda_d$ their eigenvalues, and $R$ is orthogonal. Thus the column space of $W$ spans a leading $d$-dimensional principal subspace. If an eigenvalue tie crosses the cutoff at $d$, the choice within that tied eigenspace is not unique; the theorem identifies the corresponding leading invariant subspace only up to that ambiguity.

A one-component PCA example illustrates the target subspace in that theorem. The next cell draws 300 points from an anisotropic Gaussian (seeded with np.random.default_rng(3)), computes the principal direction from the sample covariance, and records it together with the data for plotting. The scatter and the recovered axis are drawn in the figure cell below; the code computes PCA directly rather than fitting a VAE.

This PPCA result is due to Tipping and Bishop (1999); Lucas et al. (2019) analyze its connection to linear VAEs. It teaches two lessons.

First, the deep VAE extends the classical methods from the system identification chapter by replacing linear maps with nonlinear encoders and decoders. Gaussian factor analysis and PCA provide a rigorous anchor for the nonlinear case. When a nonlinear VAE behaves strangely, asking what the linear analogue would do can be a useful diagnostic.

Second, distinguish PCA eigenvectors from the latent loading basis. With distinct eigenvalues, PCA eigenvectors are fixed up to sign. In an isotropic-prior linear latent model, however, the loading matrix can be right-multiplied by an orthogonal rotation without changing the generative marginal likelihood. The principal subspace is identified, while equivalent latent coordinates inside it are not. Nonlinear latent models admit broader prior-preserving reparameterizations, so semantic coordinates still require additional constraints.

This arbitrariness has a formal version. Locatello et al. (2019) showed that unsupervised disentanglement is not identifiable without inductive biases. A prior-measure-preserving invertible transformation, or a consistent transformation of both prior and decoder with the appropriate change of variables, can leave the observation distribution unchanged while scrambling latent semantics. A latent space can look organized without factoring the world as hoped, and reconstruction quality alone cannot resolve the ambiguity.

The practical consequences for world model builders:

• Do not trust organization you did not measure. Probe latents against known generating factors when you have them. In simulation you can regress latent coordinates on true state variables and report how much variance each latent direction explains of each factor; if you cannot run this check, treat claims of disentanglement with suspicion.
• Exploit what you have. Actions, temporal consistency, and multi-view observations are inductive biases. Encouraging latents to be predictable from actions, or to change slowly, can reduce the family of equivalent representations left by pure reconstruction. Which transformations remain depends on the objective, data and action coverage, and model class. Each bias is an additional constraint that a reparametrization must preserve, and the more such constraints you impose, the smaller the family of equivalent solutions can become. This is precisely where the next chapter, on temporal state and memory, picks up.
• Judge latents by downstream control, not by aesthetics. A representation that is entangled but control-sufficient can be perfectly usable; a disentangled-looking one that hides the reward-relevant factor is not. The state abstraction criterion from the previous chapter, not the visual appeal of latent traversals, is the standard that matters.

Let us make all of this concrete. We will build a static image dataset where the ground-truth factor for each snapshot is a single number, the pendulum angle $\theta$, but the agent only sees 16 by 16 images of a glowing bob. This is a minimal laboratory for factor retention, geometry, topology, bottleneck strength, and collapse because the data has a known, one-dimensional, circular factor. It does not test temporal prediction or control sufficiency: there are no transitions, actions, rewards, or control objective, and a dynamical pendulum would generally also need angular velocity in its state. Then we will:

1. Train a small VAE on the images and inspect its 2-dimensional latent space.
2. Sweep the bottleneck strength $\beta$ and watch reconstruction, KL, and latent geometry respond.
3. Train a tiny VQ-VAE on the same data, and check codebook usage.

The example uses CPU-compatible PyTorch operations. First, the hidden setup cell that the book's plotting conventions require.

Each observation is a Gaussian blob placed at angle $\theta$ on a ring of radius 5 pixels around the image center. The only thing that varies across the dataset is $\theta$, so the data manifold is intrinsically one-dimensional and topologically a circle. That makes it a sharp test of latent geometry: a topology-faithful 2-dimensional latent space should trace a ring. It is also a deliberately friendly test because the single factor of variation is visually dominant. Reduced data complexity makes failures easier to diagnose here, though optimization, architecture, sampling, and loss-scaling confounds remain.

A glowing dot on a ring. No explicit angle label is supplied to the VAE; angle is only implicit in the blob's pixel location. The VAE must exploit the one-dimensional nonlinear correlation structure generated by motion around a closed curve in pixel space.

The model is deliberately tiny: an encoder and decoder with two hidden layers of 128 units and a 2-dimensional Gaussian latent. The physical snapshot has one intrinsic degree of freedom, while two Euclidean coordinates are the minimum for a seam-free embedding of a circle. The decoder outputs unconstrained reconstruction means. Under a unit-variance Gaussian likelihood, the negative log likelihood is one-half the summed squared error plus a parameter-independent constant, so the training loop uses that exact scale.

The moment of truth. Because this controlled dataset retains the true angle as evaluation metadata, we can use hue ordering as a supervised probe of whether the VAE retains that factor. A topology-faithful two-dimensional encoding would form a closed curve without a seam, but angle retention alone does not imply that geometry. We compare an informative run at $\beta = 0.01$ with the near-collapsed endpoint at $\beta = 0.5$.

At $\beta=0.01$, hue follows an ordered one-dimensional curve, so the means retain angle. The curve is spiral-like rather than a seam-free circle, directly illustrating that reconstruction and a Gaussian prior do not guarantee topological fidelity. At $\beta=0.5$, the means retain a tiny residual trace, but their mean radius is about 42 times smaller and total KL is about 0.003 nats. Samples from that near-prior posterior carry essentially no usable angle information, while reconstruction error rises sharply. Use measured KL, posterior scale, and latent sensitivity together; posterior-mean radius alone is not an information measure.

The per-$\beta$ summary quantities tell the quantitative side of the same story.

These are achieved $\beta$ operating points, not an empirical estimate of the optimal rate-distortion curve. At low $\beta$, both coordinates carry substantial KL. That establishes posterior deviation from the prior, but by itself establishes neither observation-dependent usage nor topology because expected KL also contains aggregated-posterior mismatch. A seam-free Euclidean embedding of a circular factor needs at least two coordinates; this learned curve is not seam-free. At $\beta=0.5$, both KL values approach zero while reconstruction error jumps: this seed crosses into near-complete posterior collapse rather than pruning an unused coordinate. Per-dimension KL is a primary warning signal, not a complete discriminator; pair it with reconstruction, latent-sensitivity or supervised probes, and downstream evaluation to learn whether and what each dimension encodes.

Now the discrete variant. The encoder is deterministic and its output is snapped to the nearest of 8 learned codes, with straight-through gradients. We initialize the codebook from encoder outputs of real data, placing every code initially in a region visited by the untrained encoder. Whether those codes remain active still has to be measured.

The VQ-VAE emits one token from an eight-code vocabulary, at most three bits before accounting for the learned usage distribution, so it can preserve the angle only coarsely. The seeded counts are [80, 49, 18, 44, 20, 35, 24, 30]: uneven, but with no dead code. Unequal partition sizes and decoder preferences can produce this pattern without collapse. Prolonged zero use is a strong dead-code warning, not a mathematical point of no return; track perplexity and stability over training and reinitialize codes that remain unused.

The continuous VAE exposes posterior parameters that can be incorporated into a temporal belief model after calibration and consistency checks. The VQ-VAE exposes a discrete vocabulary suited to token dynamics. Neither static representation is automatically a belief state or decision-sufficient. The experiment shows both the benefit of a useful bottleneck and the failure boundary: the low-$\beta$ VAE and VQ-VAE retain the angle at different resolutions, while the high-$\beta$ VAE nearly abandons it.

The key parameters for the VAE and VQ-VAE are:

• dim_z: Latent dimensionality. Two dimensions are enough to embed this dataset's one-angle snapshot manifold and make its geometry directly plottable; a dynamical pendulum state generally also includes angular velocity. In general, dim_z is a design choice for latent representational capacity. Intrinsic dimension informs it, but topology, uncertainty parameterization, and the chosen coordinate family can require more coordinates. Too small a bottleneck destroys necessary information; too large a bottleneck can leave redundant or collapsed dimensions and may expose downstream models to nuisance variation unless they learn to ignore it.
• beta: Bottleneck strength multiplying the KL term. Larger values penalize posterior rate more strongly and can reduce active dimensions or cause full collapse. Sweep it and measure reconstruction, KL, geometry, and downstream sufficiency.
• n_codes: Number of codebook entries in the VQ-VAE. Eight codes quantize the ring into arcs; too many invite dead codes, too few lose fidelity. The right number depends on the desired rate-distortion tradeoff, task, data distribution, architecture, and optimization, so usage diagnostics and sweeps are essential.
• lambda_commit: VQ-VAE commitment-loss weight, separate from the VAE KL multiplier $\beta$. It keeps encoder outputs near assigned codes. The original VQ-VAE used 0.25; treat it as a historical baseline and sweep it for the application.
• dim_h: Hidden layer width shared by encoder and decoder. A width of 128 reduces the likelihood of an obvious capacity bottleneck here, but a failure still requires checks across optimization, likelihood, seed, and architecture.

Latent-variable representations are foundational components of many influential modern world models. A compressed, decodable latent lets a model express dynamics over state-like variables rather than reproduce raw sensor coordinates at every step. Raw-frame rollout instead operates in a much higher-dimensional representation and must preserve rendering structure, which can substantially increase modeling and planning cost. Dimensionality does not translate into a simple “256 times larger” search space; complexity depends on model structure and how candidates are represented.

Several important limitations remain active research problems, while others have partial mitigations.

Non-identifiability is fundamental, not incidental. Reconstruction alone cannot pin down intended latent semantics. Temporal and action-conditioned objectives are especially relevant biases for world models, alongside supervision, geometry, equivariance, multiple views, and interventions.

The bottleneck is a fragile dial. $\beta$ couples how much detail is kept, the posterior's uncertainty, and whether the latent is used at all. Total loss can look acceptable during collapse, but decomposed KL, active-unit counts, code usage, and latent-sensitivity probes expose it. These diagnostics are permanent instrumentation, not a one-time debugging step.

Reconstruction is a proxy, and agents exploit the gap. A latent space optimized for pixel fidelity may discard exactly the feature a downstream policy needs, like a small but fast-approaching object. Under an unweighted pixelwise loss, that small object may contribute little total reconstruction error despite its decision importance; perceptual, focal, object-weighted, or likelihood-based objectives can rank it differently. This is the open-loop versus closed-loop distinction in miniature: a model can be an excellent predictor and a poor basis for decisions. Chapters on evaluation return to this gap repeatedly, and the state abstraction lens from the previous chapter is the right tool for diagnosing it.

Observation-only latents without temporal inference are insufficient in hidden-state environments. They need a transition model and a belief update that integrate history, separate hidden state from observation noise, and expose uncertainty over time. That is the temporal machinery of the next chapter.

• Latent-variable models posit unobserved variables $z$ and a likelihood $p(x \mid z)$. For generic nonlinear decoders the marginal likelihood is usually intractable, so we optimize the ELBO; analytic special cases such as linear-Gaussian models exist.
• The reparameterization trick makes sampling differentiable, and amortization lets one encoder network perform inference for any new observation in a single forward pass.
• The expected KL is a rate upper bound. Scaling it by $\beta$ trades reconstruction fidelity for compression; task-aware compression can be valuable, but its strength must be selected with downstream sufficiency and uncertainty criteria.
• Posterior collapse (the encoder ignores $z$) and codebook collapse (most discrete codes go unused) are two important silent failures of latent-variable training. Watch per-dimension KL and code usage the way you watch loss curves.
• Discrete latents, via VQ-VAE straight-through estimation or Gumbel-softmax, connect world modeling to token-based sequence models and provide a compact form for categorical structure; learned code semantics still require validation.
• Latent geometry matters independently of reconstruction quality. Useful latents preserve supported manifold paths, relevant topology, and task-aligned distances. Reconstruction alone cannot guarantee these properties.
• The next step, and the subject of the next chapter, is to add temporal inference and action-conditioned prediction. These objectives constrain the representation, but task sufficiency, uncertainty, and semantic alignment still require direct evaluation.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about latent-variable representations, variational autoencoders, and discrete latents.