Recognizing Posterior Collapse in Count VAEs

Why Posterior Collapse Needs More Than One Diagnostic

A count VAE is designed to represent cell-to-cell variation through latent variables while modeling observed gene counts with a count likelihood. Posterior collapse occurs when the fitted encoder returns a distribution close to the prior for most cells, so the latent variables carry little cell-specific information. The decoder can then produce its predictions with little or no useful input from the inferred latents. This is a failure of the intended representation, not simply a low reconstruction quality or a particular value of the training objective.

A single small KL term is a warning, not a complete diagnosis. A flexible decoder may fit counts well while ignoring its latent input; alternatively, a small average KL may reflect a representation in which only a few dimensions carry useful information. A sound assessment therefore combines KL contributions, evidence that encoded latents vary with cells, sensitivity of reconstructions to those latents, and the trajectory of these quantities during training.

The usual count-VAE objective maximizes the evidence lower bound (ELBO), whose per-cell form is

The first term rewards the decoder for assigning probability to observed counts, and the KL term penalizes departure of the encoder from the prior. Training traces often report the negative of these terms as losses; check the sign convention before comparing curves. This tutorial's calculations use a standard normal prior \(p(z)=N(0,I)\) and a diagonal Gaussian encoder, \(q_\phi(z\mid x)=N(\mu(x),\operatorname{diag}(\sigma^2(x)))\). The Gaussian formulae below rely on these assumptions. With other priors or encoder families, compute the corresponding KL rather than applying the diagonal-Gaussian expression.

Reading KL Contributions by Dimension

For a diagonal Gaussian encoder and standard normal prior, the KL contribution of dimension \(j\) for one cell is

It is nonnegative and equals zero exactly when that cell's posterior coordinate has mean zero and variance one, matching the prior coordinate. Averaging \(\operatorname{KL}_j(x)\) across cells gives a useful per-dimension summary. Very small averages across nearly all dimensions suggest that the encoder is staying close to the prior. A few appreciable dimensions among many near-zero ones indicate selective use, not necessarily collapse of the entire latent representation.

The average KL has a useful interpretation. Under the empirical distribution of cells, let \(q(z)\) denote the aggregate encoder distribution, the mixture of \(q_\phi(z\mid x)\) over cells. Then

Here \(I_q(X;Z)\) is the mutual information between the observed cell and its encoded latent under the joint distribution \(q(x,z)=q(x)q_\phi(z\mid x)\). To justify the identity, write \(\log(q_\phi(z\mid x)/p(z))=\log(q_\phi(z\mid x)/q(z))+\log(q(z)/p(z))\), then average over \(q(x,z)\). The first term averages to \(I_q(X;Z)\); the second averages over \(q(z)\) and is \(\operatorname{KL}(q(z)\|p(z))\). Thus average KL includes both cell-specific information and aggregate mismatch to the prior. It is not itself a direct estimate of mutual information.

Check Whether the Decoder Uses the Latents

KL summaries describe how the encoder differs from the prior; they do not by themselves show whether the decoder relies on the latent input. For a fixed observed cell, examine the decoder's conditional count distribution under its encoded latent and compare it with distributions under altered latents. In a negative binomial model, for example, compare the decoder's predicted mean counts (and, where relevant, dispersion) while holding library-size inputs and other covariates fixed. A substantial, reproducible change indicates sensitivity; negligible change across sensible perturbations is evidence that the decoder may be ignoring \(z\).

Useful comparisons include the encoded latent, a latent drawn from the prior, and an encoded latent from a different cell. Shuffling encoded latents across cells preserves the collection of latent vectors but breaks their cell assignments. If cell-specific decoder predictions barely change under these substitutions, latent usage is weak. If predictions change, that alone does not prove the representation is biologically meaningful: changes may be unstable, driven by nuisance variation, or concentrated in a few dimensions. Treat these as diagnostics of model behavior, not biological validation.

Use Training Traces to Locate the Failure

Track the reconstruction term and KL term separately, preferably also by latent dimension, over training. A common warning pattern is that KL contributions fall rapidly toward zero early in optimization while reconstruction improves little or remains adequate because the decoder can fit broad count patterns without cell-specific latents. The timing matters: a late decline in KL may reflect a trade-off under the objective, whereas near-zero KL from the start may indicate that the encoder never established useful latent variation.

A different pattern is a small total KL but several dimensions with sustained, nontrivial contributions. That is evidence of selective latent use, not a basis for declaring collapse. Conversely, an appreciable total KL is not a guarantee of useful representation: KL can reflect aggregate mismatch to the prior, as the decomposition above shows. Plot per-dimension values and inspect latent-dependent predictions as well as the total.

Minibatch estimates can be noisy, especially with sparse counts and small batches. Compare traces smoothed over a stated window, and check that the KL scaling is consistent with the ELBO implementation: summing over cells or dimensions in one term while averaging in another can change the effective penalty. Compare runs using the same data split, objective scaling, and evaluation procedure. A downward KL curve alone cannot diagnose why a model changed.

Interventions and Their Limits

KL warm-up changes the optimization path by gradually increasing the KL weight. During warm-up, the objective has the form

Starting with a smaller \(\beta\) temporarily reduces pressure to match the prior, which can give the encoder and decoder time to establish a useful latent pathway. This is a training intervention, not a diagnostic and not a change to the final target if \(\beta\) is eventually set to one. Evaluate the converged model under the intended objective; a model that uses latents only while the KL penalty is reduced has not resolved collapse for the final objective.

Decoder changes can also alter the balance. A decoder that can explain all count variation without \(z\) may have little incentive to use it. Reducing excessive decoder capacity or providing a suitable pathway from \(z\) to decoder parameters can encourage reliance on latents. Such changes involve trade-offs: too restrictive a decoder can worsen count fit, and a reconstruction gain does not by itself establish that latent variation is reliable. Compare alternatives using the same diagnostics and objective, rather than treating a particular architecture as a guaranteed remedy.

Other causes should be checked before adjusting the model: incorrect KL scaling, a mismatch between the assumed count likelihood and data behavior, optimization instability, or a coding error in which the decoder input does not depend on \(z\). These issues can produce misleading traces or mimic a training failure. The likelihood model's assumptions and the count-data setup should be reviewed before interpreting an intervention.

Common Mistakes and Practical Decision Rule

• Calling every low KL a collapse. A low average can coexist with a few actively used dimensions. Inspect contributions by dimension and cell, then test decoder sensitivity.
• Calling any positive KL useful information. Average KL also includes aggregate posterior mismatch to the prior. Check whether encoded latents carry cell-specific differences and affect predictions.
• Using reconstruction quality alone. A decoder can fit counts while ignoring \(z\). Compare encoded, shuffled, and prior-latent predictions with other inputs fixed.
• Claiming warm-up solved the problem from its early trace. Warm-up temporarily changes the objective. Reassess after the KL weight reaches its intended final value.
• Interpreting a latent perturbation as biology. Decoder sensitivity shows model dependence, not a causal or biological meaning for the coordinate.

A practical diagnosis is strongest when several findings agree: KL is near zero across most dimensions and cells; encoded latents show little cell-specific variation; decoder predictions change little when encoded latents are shuffled or replaced; and these patterns persist near the end of training. If only one finding is present, describe it as a warning sign and investigate the others before concluding collapse.