Posterior Collapse and the KL Term That Wins

A latent variable model that ignores its latents still produces a decreasing loss. The decoder learns the marginal distribution of the data, the encoder outputs the prior for every input, and the objective settles at a value that is locally optimal given the architecture. The symptom is invisible in the training curve and only appears when the model is sampled from: every draw from the prior decodes to something near the dataset mean, and interpolating between two encodings does not move the output.

This behaviour is called posterior collapse, and it is the dominant failure mode of the variational autoencoder. It is worth understanding as an equilibrium of the objective rather than as a training bug, because the interventions that work are the ones that change the equilibrium, and the ones that do not are the ones that only change the optimiser’s path to it.

The equilibrium that has no latent information

Consider a decoder with enough capacity to model the data distribution without any conditioning. If the decoder can learn p(x) directly, then the reconstruction term is maximised by ignoring z. The KL term is minimised by making the encoder output exactly the prior. Both terms agree on the same solution: the encoder is constant, and the model is a plain autoregressive or convolutional density model with an extra input that is never used.

The gradient structure explains why this solution is reachable and why it is stable. Once the encoder’s mean and log variance are close to zero, the sample z is almost exactly zero, and the decoder sees a nearly constant input. The gradient of the reconstruction term with respect to the encoder parameters is then small, because the decoder has adapted to a constant input and its sensitivity to that input has been driven down. The KL gradient, by contrast, is still pulling the mean and variance toward the prior with a force proportional to their deviation. The two gradients therefore reinforce collapse near the collapsed point: the reconstruction gradient is weak and the KL gradient is strong, so any perturbation away from collapse is pushed back.

That argument assumes the encoder’s own gradient reaches the reconstruction term at all, which depends on how the sample is written as a function of the encoder’s outputs. The form used here, and the reason it keeps the reconstruction gradient alive even when the encoder’s outputs are near the prior, is developed in the gradient path through sampling.

The strength of the KL gradient scales with the latent dimension, because the KL is a sum over dimensions. A 512-dimensional latent contributes 512 independent restoring forces, while the reconstruction term contributes one signal spread across all of them. This is why collapse is more common in high-dimensional latent spaces and in models with high-capacity decoders, and why a small latent with a weak decoder often avoids it without any special handling.

Why the bound permits it

The objective does not distinguish between a model that uses its latent and one that does not, as long as both achieve the same bound value. If the decoder’s unconditional likelihood is as good as the conditional likelihood averaged over the posterior, the collapsed solution has the same bound as the informative one. The bound is a lower bound on log p(x); it says nothing about how the model represents the data internally.

There is a tighter statement available. The gap between the bound and the log evidence is the KL between the approximate and true posterior. In a collapsed model the true posterior is close to the prior, so the gap is small and the bound is nearly tight. The collapsed model is therefore not a poorly optimised model; it is a model for which the bound is a good approximation of the evidence, because the latent genuinely carries no information about the data under that decoder.

This is the sense in which collapse is a modelling outcome rather than an optimisation failure. A decoder with an autoregressive output distribution over pixels can represent the data exactly, so the latent is redundant. The same architecture with a factorised output distribution cannot represent pixel correlations, so the latent is forced to carry them, and collapse is much less likely. The architecture determines whether the latent is necessary, and the objective finds the answer.

Detection: what to measure and what the numbers mean

The KL term alone is not a reliable indicator, because its scale depends on the latent dimension and on the reconstruction scale. A per-dimension KL, averaged over a batch, is more useful: it shows whether a few dimensions carry all the information while the rest sit at the prior. A healthy model typically has a distribution of per-dimension values with a long tail rather than a flat near-zero profile.

The mutual information between x and z is the quantity of interest, and it can be estimated from the aggregate posterior. If the encoder’s output distribution, averaged over the dataset, equals the prior, the mutual information is zero by construction. A practical proxy is to compare the variance of the encoder means across the dataset against the average encoder variance: if the means barely vary, the latent is not encoding anything, regardless of what the KL term reports.

The reconstruction term is the second measurement. Compare the loss of the full model against the loss of the same decoder trained with the latent input removed or held at zero. If the two are indistinguishable, the latent is not being used. This comparison is cheap and unambiguous, and it is the check that distinguishes collapse from a model that has simply learned a smooth latent space.

import numpy as np

def collapse_report(mu, log_var, batch_axis=0):
    """mu, log_var: encoder outputs of shape (N, d)."""
    kl_per_dim = 0.5 * (mu ** 2 + np.exp(log_var) - log_var - 1.0)
    # spread of the encoder means across the dataset, per dimension
    mean_spread = mu.var(axis=batch_axis)
    # average encoder variance, per dimension
    mean_var = np.exp(log_var).mean(axis=batch_axis)
    active = kl_per_dim.mean(axis=batch_axis) > 0.01   # nats, a rough threshold
    return {
        "kl_total": kl_per_dim.sum(axis=1).mean(),
        "active_dims": int(active.sum()),
        "mean_spread_over_var": (mean_spread / mean_var).mean(),
    }

The ratio of mean spread to mean variance is the useful number. A value near zero means the encoder means do not move relative to the noise the encoder itself declares, so the latent code is dominated by sampling noise and the decoder cannot rely on it. A value well above one means the means move more than the declared noise, which is what an informative latent looks like.

Interventions and what each one actually changes

The standard intervention is a coefficient on the KL term, commonly written β, with values between 1 and 10 in the direction that increases the penalty and values below 1 in the direction that weakens it. Weakening the KL below 1 is the intervention that prevents collapse, because it reduces the restoring force while leaving the reconstruction signal intact. The cost is that the aggregate posterior no longer matches the prior, so samples drawn from the prior may fall outside the region the decoder was trained on. The model becomes a good autoencoder and a poor sampler, which is a real trade rather than a free fix.

Annealing the coefficient from zero upward is the same intervention applied on a schedule. Early in training the model is allowed to use the latent freely; the penalty is introduced gradually so that the encoder has already committed to an informative solution by the time the restoring force becomes significant. The schedule’s shape matters less than its endpoint: a schedule that ends at a coefficient of 1 has the same equilibrium as an unmodified model and can still collapse late if the decoder becomes expressive enough to make the latent redundant.

Architectural interventions change whether the latent is necessary. Restricting the decoder’s receptive field or its output distribution forces the latent to carry the information the decoder cannot represent. This is the mechanism behind models that use a latent to capture global structure and a local decoder for detail. The intervention is more reliable than loss weighting because it removes the collapsed equilibrium instead of making it less attractive.

Weakening the decoder directly has the same effect and is easier to reason about. Reducing its capacity, removing skip connections from the latent to the output, or constraining its receptive field all raise the reconstruction cost of ignoring z. The failure mode of overdoing it is a model whose samples are blurry because the decoder cannot represent fine detail even with a good latent.

Free-bits objectives replace the per-example KL with a hinge that only penalises the part above a threshold. Dimensions below the threshold contribute no gradient, so the model can keep a few dimensions informative without paying for the rest. The threshold is a target rate, and setting it too high produces a model with a latent that is informative but whose aggregate posterior is far from the prior.

The sequential case where collapse is structural

In sequence models the collapse is not a subtle equilibrium but an architectural inevitability. An autoregressive decoder predicts each token conditioned on all previous tokens. If the latent is also available, the decoder can predict the next token from the prefix alone, and the latent adds nothing. The reconstruction term’s gradient with respect to the latent vanishes because the decoder has already extracted the information from the prefix, and the KL term then drives the encoder to the prior.

The consequence is that a variational sequence model with a strong autoregressive decoder will collapse unless the decoder is prevented from seeing the full prefix. Training the decoder on a limited context, or dropping the latent connection for some fraction of steps, or using a decoder whose capacity is deliberately restricted, all change the balance. So does weakening the KL, but the equilibrium returns as soon as the decoder’s context is long enough to make the latent redundant.

The general lesson is that collapse is a statement about the division of labour between the latent and the decoder. The objective will assign information to whichever path is cheaper, and the KL term makes the latent path expensive. Any model where the decoder can do the job alone will end up doing it alone, and no amount of optimiser tuning changes that. The diagnostic in the previous section is what distinguishes this from a model that has simply not converged, and the interventions above are the ones that change the terms of the division rather than the path to it.

The related question of how the reconstruction term is scaled against the KL term is taken up in the closed-form Gaussian case, where the decoder variance sets the exchange rate between nats of reconstruction and nats of divergence. Getting that exchange rate wrong produces the same symptom as collapse, with the same flat latent and the same near-mean samples, for an entirely different reason.