Linear VAE and Posterior Collapse, Interactively

← text notes · Lucas, Tucker, Grosse and Norouzi (NeurIPS 2019). Three widgets: why a fixed noise level can trap a latent unit at zero, how many units that traps, and why a diagonal posterior forces the decoder's columns to be perpendicular.

1. One column: a valley or a hill at zero

Take a single latent unit whose decoder column points along one principal direction of the data, and ask how much the log-likelihood improves as the column's length s grows from zero. Here λ is the data's variance along that direction and σ² is the observation noise, held fixed. The curve is the gain per data point, measured against the collapsed column (s = 0).

Because the model's variance along the direction is σ² + s², the unit can only describe variance above the noise floor.

Things to try. With σ² below λ the curve dips at s = 0 and rises to two peaks (higher is better, so these are the likelihood's maxima): a zero column is not a resting place, because growing it helps. Raise σ² past λ and the two peaks merge into a single peak at zero: now the collapsed column is a local maximum, and no small push away from zero improves the likelihood. The switch happens exactly where the two sliders are equal. Nothing about the KL term, the encoder or the decoder's power has entered yet.

2. Fixed noise against learned noise

Now several directions at once. Each bar is a principal-component variance λ_i of the data. The model has q latent columns to spend on the top q directions. The horizontal line is σ². A column whose bar lies above the line can be used (filled, blue). One whose bar lies below it is stuck at zero (hollow): it is collapsed. The lower panel shows how much information each live unit carries, ½ ln(λ_i/σ²), which is the KL cost the unit pays.

Things to try. With σ² fixed high, only the few bars above the line survive and the rest of the q units collapse, however many you provide. Lower σ² and collapsed units come alive one by one, from the largest bar downward. Now switch to σ² learned: the line drops to the average of the directions you did not give a column to, which is never above the q-th bar, so every one of the q units is alive. That is the paper's reason a learned noise level turns the collapsed points into saddles. Notice in the lower panel that a unit's information is half the log of its signal-to-noise ratio, so it vanishes exactly as the bar reaches the line.

3. Rotating the decoder: same likelihood, different ELBO

With two latent units, the decoder is W = L Rᵀ: lengths L = diag(l₁, l₂) and a rotation R of the latent coordinates. The data distribution (left panel, its ellipse) does not depend on the rotation, so the likelihood is the same for every angle θ. The arrows are the two columns of W.

The middle panel is the latent space: the dark ellipse is the true posterior, and the coloured ellipse is what a variational posterior with diagonal covariance can offer. The right panel plots the gap KL(q ‖ posterior) against θ. The ELBO is the likelihood minus this gap.

Things to try. At θ = 0 (and 90°, 180°) the columns are perpendicular, the two ellipses coincide, and the gap is zero: the ELBO equals the likelihood. Between those angles the true posterior tilts while the variational one cannot, and the gap is positive. The ELBO therefore prefers perpendicular columns, and a small rotation toward them always raises it. Make l₁ equal l₂ and the true posterior is round at every angle: the gap is zero everywhere and the angle is no longer pinned down, which is the tie caveat in the notes.

Everything here is computed in the browser from the formulas in the notes; none of it is data from the paper.