Posterior Collapse and KL Annealing

Train a variational autoencoder with a strong decoder — say an autoregressive text or audio decoder — and something strange can happen: the latent code $z$ stops mattering. The encoder's posterior $q(z|x)$ collapses onto the prior $p(z)$ for every single input, meaning $z$ carries zero information about $x$. Reconstructions still look fine, because the decoder learned to reproduce plausible outputs on its own. But the latent space is useless: sample two different $z$'s and you get the same output, because the decoder never learned to listen to $z$ in the first place.

The analogy: a student who stops reading the prompt

Imagine a student graded only on essay quality, never on whether the essay actually answered the prompt. If the student is skilled enough to write a plausible essay from any one-word prompt, they will stop reading prompts carefully — there's no penalty for ignoring them, and ignoring them is easier. That is posterior collapse: the decoder is the skilled student, the latent code is the prompt, and the training objective never explicitly rewards the decoder for using it.

Why the ELBO invites this

The VAE objective, the evidence lower bound, is a sum of two competing terms:

$$\text{ELBO} = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{reconstruction}} - \underbrace{\text{KL}\big(q(z|x)\,\|\,p(z)\big)}_{\text{closeness to prior}}$$

In words: reward accurate reconstruction, but penalize the encoder for straying from the prior. Early in training, before $z$ has learned anything useful, the cheapest way to raise the ELBO is often to make $q(z|x)$ identical to the prior for every $x$ — the KL term drops to exactly zero, and if the decoder can still reconstruct reasonably well without $z$, reconstruction loss barely suffers. That local optimum is easy to fall into and hard to climb back out of, because once the decoder ignores $z$, the reconstruction gradient carries no signal telling the encoder to make $z$ informative again.

Worked example 1: the collapse is a real optimum, in numbers

Suppose, per sequence, reconstruction using an informative $z$ reaches log-likelihood $\log p(x|z) = -50$ nats, and the KL cost of that informative posterior is $8$ nats. A decoder good enough to ignore $z$ instead reaches $\log p(x) = -55$ nats (worse reconstruction, but no KL cost at all). Compare full ELBOs:

$$\text{informative: } -50 - 8 = -58, \qquad \text{collapsed: } -55 - 0 = -55$$

Collapse wins, $-55 > -58$, even though its reconstruction is 5 nats worse — the KL penalty outweighs the reconstruction gain. This is why plain ELBO training can prefer a useless latent space.

Worked example 2: annealing tips the balance

KL annealing multiplies the KL term by a weight $\beta(t)$ that ramps from $0$ up to $1$ over training, using $\text{recon} - \beta \cdot \text{KL}$. With a linear schedule $\beta(t) = \min(1, t/10{,}000)$:

• Step 2,000 ($\beta = 0.2$): informative $= -50 - 0.2(8) = -51.6$; collapsed $= -55$. Informative now wins.
• Step 5,000 ($\beta = 0.5$): informative $= -50 - 0.5(8) = -54$; collapsed $= -55$. Informative still wins, by a narrower margin.
• Step 10,000 ($\beta = 1$): back to the original numbers, informative loses on paper — but by then the decoder has already learned to depend on $z$, so abandoning it would hurt reconstruction, not just cost KL. The early low-$\beta$ window is what gets $z$ established before the full penalty arrives.

What this means in practice

Posterior collapse matters whenever a VAE's decoder is autoregressive or otherwise expressive enough to model $x$ on its own — language VAEs and audio VAEs are the classic victims, but any financial sequence model with a powerful decoder (a VAE meant to generate synthetic price paths, say) can suffer the same failure, silently producing a latent space that carries no interpretable structure for stress-testing or interpolation. Besides KL annealing, practitioners use "free bits" (only penalizing KL above a small floor per latent dimension) and weakening the decoder (dropout on its own inputs) to force it to lean on $z$.