$\mathbfλ$-VAE: Variance Equalization for Posterior Collapse

2026-07-06

AI Summary

This work addresses the pervasive issue of posterior collapse in variational autoencoders (VAEs), which renders latent codes uninformative due to two intertwined mechanisms: gradient imbalance and the information gap. To mitigate both simultaneously, the authors propose λ-VAE, a novel approach that applies dimension-wise exponential scaling to the noise in the reparameterization trick while preserving the original KL divergence penalty. This design achieves variance balancing across latent dimensions, effectively suppressing collapse. Notably, λ-VAE provides the first unified framework that concurrently alleviates both collapse mechanisms and admits a closed-form optimal solution for the scaling exponents. Empirical evaluations demonstrate that λ-VAE substantially reduces the number of collapsed dimensions across multiple benchmark datasets, yielding up to a 2.8-nat increase in information capacity and up to a 0.33 bits-per-dim improvement in reconstruction quality.

Abstract

Variational Autoencoders (VAEs) frequently suffer from posterior collapse, a failure mode in which the approximate posterior converges to the prior, rendering the latent code uninformative. Despite extensive research, a unified account of why collapse occurs has remained an open question. We identify and formalize two logically independent but coupled causes. \emph{Gradient imbalance} occurs when the decoder's reconstruction signal vanishes faster than the $\mathbb{KL}$ regularization pressure as the posterior widens. \emph{Information gap} occurs when the stochastic sampling step discards a substantial fraction of the encoder's computed representation, attenuating decoder sensitivity and making collapse inexpensive. Both causes share the same collapse trajectory, and we show that the information gap is algebraically equivalent to mismatch between the aggregate posterior and the prior, unifying two pathologies. Subsequently, we introduce $λ$-VAE, which resolves both causes through a single modification to the reparameterization step: the sampling noise is scaled by per-dimension exponent, while the $\mathbb{KL}$ penalty retains the original posterior variance. This asymmetry shifts the stable training attractor away from the degenerate collapsed state, driving all latent dimensions toward the same equilibrium -- a mechanism we term \emph{variance equalization}. A closed-form optimal exponent per dimension follows from a net information gain objective, with a single hyperparameter controlling the reconstruction-generation tradeoff. We validate on standard benchmarks (Binary MNIST, Binary Omniglot, CIFAR-10, CelebA-64), showing consistent reductions in collapsed dimensions, information capacity gains of up to $2.8\times$ nats, and reconstruction quality improvements of up to $+0.33$ BPD.

Problem

• posterior collapse
• variational autoencoders
• gradient imbalance
• information gap
• latent representation

Innovation

• posterior collapse
• variance equalization
• gradient imbalance
• information gap
• λ-VAE