Latent-Variable Representations › The Information Bottleneck and Posterior Collapse

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.