Source-linked AI summary
Preventing Posterior Collapse with delta-VAEs
Ali Razavi, Aäron van den Oord, Ben Poole, Oriol Vinyals
TL;DR
Posterior collapse lets expressive decoders ignore latent variables, creating a challenge for combining strong density models with useful representations. δ-VAEs enforce a minimum posterior–prior KL divergence without changing the ELBO or weakening the decoder, and experiments show useful representations with strong image and text modeling results. The sequential variant uses anti-causal encoding and correlated priors to support future-oriented latent information.
Problem
Expressive autoregressive decoders can ignore latent variables, causing the posterior to collapse to the prior and undermining informative representations.
Method
δ-VAEs constrain the variational posterior family so its KL divergence from the prior is bounded below by a committed rate δ, while retaining the ELBO and decoder capacity.
Results
δ-VAEs learn useful representations with powerful decoders on CIFAR-10, ImageNet 32 × 32, and LM1B, while achieving state-of-the-art image likelihoods.
Takeaways & Limitations
The approach enables powerful decoders and informative latent variables without sacrificing density-modeling performance.
Takeaways & Limitations
The systematic prior–posterior mismatch may create a large gap between the prior and aggregate posterior, exacerbating posterior holes.
Abstract
from arXiv · showhide
Due to the phenomenon of "posterior collapse," current latent variable generative models pose a challenging design choice that either weakens the capacity of the decoder or requires augmenting the objective so it does not only maximize the likelihood of the data. In this paper, we propose an alternative that utilizes the most powerful generative models as decoders, whilst optimising the variational lower bound all while ensuring that the latent variables preserve and encode useful information. Our proposed $δ$-VAEs achieve this by constraining the variational family for the posterior to have a minimum distance to the prior. For sequential latent variable models, our approach resembles the classic representation learning approach of slow feature analysis. We demonstrate the efficacy of our approach at modeling text on LM1B and modeling images: learning representations, improving sample quality, and achieving state of the art log-likelihood on CIFAR-10 and ImageNet $32\times 32$.
1 INTRODUCTION
Latent-variable models with expressive autoregressive decoders risk ignoring their latents, while simpler decoders limit representation quality. δ-VAEs address this tension by constraining the posterior family and demonstrate useful representations with strong image and text modeling.
- Expressive autoregressive decoders can model the data density while ignoring latent variables, causing the posterior to collapse to the prior.
- δ-VAEs prevent posterior collapse by selecting variational posterior families with a minimum KL divergence from the prior.
- The framework preserves the ELBO objective and decoder capacity instead of modifying the training objective or weakening the decoder.
- δ-VAEs model images on CIFAR-10 and ImageNet 32 × 32 and text on LM1B using powerful decoders.
- The experiments report useful latent representations, state-of-the-art image log-likelihoods, and no density-modeling performance sacrifice.
2 MITIGATING POSTERIOR COLLAPSE WITH δ-VAES
The method makes latent variables informative by enforcing a positive posterior–prior KL gap, then extends this principle to sequential latents with temporally correlated priors and anti-causal encoders. These design choices connect committed rate to temporal structure and future-oriented information encoding.
- VAEs optimize an ELBO comprising expected reconstruction quality minus the KL rate between posterior and prior.
- Posterior collapse occurs when the KL rate reaches zero, making the posterior equal the prior and leaving latents uninformed about the input.
- δ-VAEs impose minθ,φ DKL(qφ(z|x)∥pθ(z)) ≥ δ, where δ is the model’s committed rate.
- A fixed variance difference between Gaussian prior and posterior is a simple nonzero-rate construction explored in the experiments.
- 2.1 δ-VAE WITH SEQUENTIAL LATENT VARIABLES: Sequential δ-VAEs use a mean-field posterior and an AR(1) prior to model temporally correlated latent variables.
- 2.1 δ-VAE WITH SEQUENTIAL LATENT VARIABLES: The prior–posterior correlation mismatch creates a positive KL lower bound, whose value can be controlled by choosing α for a target δ.
- 2.1 δ-VAE WITH SEQUENTIAL LATENT VARIABLES: As α approaches one, AR(1) prior trajectories become smoother; as α approaches zero, the prior becomes temporally independent.
- 2.2 ANTI-CAUSAL ENCODER NETWORK: The anti-causal encoder prevents a timestep’s posterior parameters from depending on past observations, encouraging latents to encode information about the future.
3 RELATED WORK
Related work addresses posterior collapse through powerful-decoder architectures, modified objectives, constrained variational families, and dataset-level compression. δ-VAEs differ by allowing data-dependent KL rates while retaining flexible decoders.
- Powerful-decoder architectures: Prior work combines powerful autoregressive decoders with latent variables, including sequential image-generation architectures.The paper positions its image models among autoregressive approaches and distinguishes its row-by-row latent-variable generation from prior sequential designs.
- Objective modifications: Most prior methods prevent posterior collapse by modifying the training objective, often through rate annealing, fixed rate coefficients, auxiliary terms, or free bits.These methods can depart from direct lower-bound optimization or create tuning and optimization difficulties.
- Variational-family constraints: Other approaches constrain the variational family, including VQ-VAE and von Mises–Fisher posteriors with constant KL divergence.VQ-VAE obtains a fixed KL of log K under a uniform prior, while vMF methods use a fixed concentration parameter and uniform prior.
- Variational-family constraints: Unlike VQ-VAE and vMF approaches, δ-VAE permits higher KL for some inputs, allocating more bits to more complicated examples.The distinction is between a constant divergence for every data point and data-dependent rates.
- Dataset-level compression: Associative Compression Networks amortize code description length across similar training examples, treating compression as a procedure over whole datasets rather than individual examples.This approach departs from the classical i.i.d. training regime.
4 EXPERIMENTS
Experiments show that δ-VAEs combine powerful autoregressive decoders with informative latent variables across image and text modeling. On images, they achieve competitive density estimation, useful representations, and improved heldout rate-distortion behavior, while text models use latent variables despite slightly worse log-likelihood.
- 4.1 NATURAL IMAGES: δ-VAE image models nearly match powerful latent-free autoregressive baselines, reaching 2.83 bits per dimension on CIFAR-10 and on-par performance with Image Transformer on ImageNet 32 × 32.The CIFAR-10 result is reported as a new state of the art, while the ImageNet comparison is described as on par.
- 4.1 NATURAL IMAGES: The auxiliary prior reduces the latent coding rate by more than 50% on CIFAR-10 while preserving the model's density-estimation performance.The passage describes an average reduction of about 30 bits per image.
- 4.2 UTILIZATION OF LATENT VARIABLES: Fixed latent samples produce images with similar global structure but different details, indicating that latents encode global structure while the autoregressive decoder models local patterns.The ImageNet samples vary across conditional generations and auxiliary-prior samples.
- 4.2 UTILIZATION OF LATENT VARIABLES: 68% classification accuracy is achieved by the highest-rate CIFAR-10 model, which encodes 92 bits per image.Higher rate generally improves classification accuracy, although better log-likelihood does not necessarily improve linear classification.
- 4.3 ABLATION STUDIES: δ-VAEs dominate the heldout rate-distortion frontier and overfit less than β-VAEs across evaluated rates.β-VAEs reach the training-set rate-distortion frontier but overfit more than δ-VAEs.
- 4.3 ABLATION STUDIES: Without a committed information rate or another mitigation strategy, both anti-causal and non-causal encoders collapse, with KL divergence falling below 10^-8 bits per dimension and never recovering.The collapse occurs after only a few thousand training steps.
- 4.4 TEXT: On LM1B, the δ-VAE has slightly worse log-likelihood than its autoregressive counterpart but makes considerable use of latent variables.Samples and interpolations provide evidence of latent-variable use.
5 DISCUSSION
δ-VAEs address posterior collapse while retaining powerful decoders, preserving useful representations and state-of-the-art likelihoods. The paper identifies stronger autoregressive baselines and downstream representation gains as open challenges.
- δ-VAEs provide a simple solution to posterior collapse that enables pairing latent-variable models with powerful decoders.The approach does not require changing the objective or weakening the decoder.
- δ-VAEs learn useful representations while achieving state-of-the-art likelihoods.
- The paper leaves open whether latent-variable models can exceed strong autoregressive baselines.
- The paper leaves open whether learned representations can improve downstream applications such as classification.
A DERIVATION OF THE KL DIVERGENCE FOR SEQUENTIAL LATENT
The sequential KL-divergence derivation starts from the analytic Gaussian KL and simplifies the expectation under an equal-coordinate mean assumption. It then obtains an analytic form for the KL of the sequential latent variable.
- The derivation begins with the analytic KL-divergence form for two univariate Gaussian distributions.
- The derivation assumes, without loss of generality, that the mean vector has equal values in each dimension.
- The expectation inside the summation is simplified before substitution into the sequential KL expression.
- Substitution yields an analytic form for the KL-divergence of the sequential latent variable.
C DERIVATION OF THE LOWER-BOUND
The lower-bound derivation simplifies the sequential KL by removing non-negative quadratic terms and applying convexity. For multiple latent dimensions, the committed rate is additive across dimensions.
- Equation 3 is simplified by removing non-negative quadratic terms involving µ_i.
- Convexity of f_a(x) supplies a minimum used to derive the KL lower bound.The function reaches its minimum at x = a^-1.
- The resulting lower bound is obtained after substituting the minimizing variance expression.
- For multidimensional z_i at each timestep, the committed rate equals the sum of the KL terms for individual dimensions.
D INDEPENDENT δ-VAES
Independent Gaussian δ-VAEs constrain posterior parameters to guarantee a committed information rate, while temporal δ-VAEs achieve better density modeling at the same target rate. Both approaches prevent posterior collapse.
- A diagonal-covariance Gaussian posterior can guarantee a committed information rate δ by constraining its mean and variance against a standard Gaussian prior.
- The posterior parameterization allows the data-dependent rate to exceed the designated lower bound δ.
- The temporal δ-VAE with an AR(1) prior achieves significantly better marginal likelihood than the independent Gaussian δ-VAE.
- Both independent and temporal δ-VAEs achieve KL values around the target rate and perform similarly on downstream linear classification.
E ARCHITECTURE DETAILS
The image architecture uses an anti-causal encoder to construct row-wise latent variables that condition the decoder. Training uses specified Transformer-style components, optimization settings, and image-model hyperparameters.
- Image architecture: The encoder reverses and pads each input image, then crops and reverses its output to provide anti-causal context.Each pixel receives information from its own and future values.
- Image architecture: Average pooling converts the encoder representation into row-wise latent variables that condition the decoder.
- Optimization: Layer normalization after multi-head attention was essential for stabilizing training, alongside ReLU activations, dropout, and Adam optimization.
- Hyperparameters: Table 4 specifies encoder, decoder, auxiliary-prior, PixelCNN, and dropout hyperparameters for the image experiments.
- Text architecture: The text architecture adapts the Transformer by inverting causal attention masks into upper-triangular bias matrices for anti-causal encoding.
F ABLATION STUDIES
The ablations examine representation quality, rate-distortion behavior, and anti-causal encoder structure on CIFAR-10. Results indicate useful latent organization and reduced overfitting in the reported training-set comparison.
- Rate-distortion: δ-VAE is less prone to overfitting than β-VAE and free-bits in the CIFAR-10 training-set rate-distortion comparison.Its negative log-likelihood is relatively higher at larger rates.
- Encoder ablation: The anti-causal encoder is evaluated against a non-causal encoder across four capacity and information-rate configurations using temporal δ-VAE training.The low-rate and high-rate models use latent dimensions of 8 and 64, respectively.
- Latent-space visualization: CIFAR-10 latent means organize semantically related classes nearby, with cats and dogs and automobiles and trucks largely interspersed.Horse examples concentrate above the deer region in the t-SNE visualization.
H ADDITIONAL SAMPLES
The additional samples illustrate latent interpolation, reconstructions, text completion, and prior sampling for image and text models. The figures also compare auxiliary and AR(1) priors across rate settings.
- Image samples: CIFAR-10 and ImageNet figures include latent interpolations, day-dream sampling, occlusion reconstructions, and multiple posterior reconstructions.The CIFAR-10 examples use test-set originals, while the ImageNet examples use validation-set originals.
- Prior samples: The AR(1) prior can produce overly smooth high-rate CIFAR-10 samples because its gap from the marginal posterior is closed by the auxiliary prior.The high-rate and low-rate models have -ELBO values of 2.90 and 2.83, and KL values of 0.10 and 0.01 bits/dim, respectively.
- Prior samples: ImageNet prior samples contain four decodings of the same latent sample within each half-column of each block.The latent sample is drawn from the auxiliary prior.
- Text latent interpolation: Text figures vary one latent dimension at a time by linearly interpolating between µ ± 3σ while fixing the other dimension.The procedure is shown for both AR(1) and auxiliary priors.
- Text completion: Text completion conditions the decoder on interpolations between two latent samples after priming it with a validation-set fragment.