PaperlayerView exact revision on arXiv

Source-linked AI summary

Don't Blame the ELBO! A Linear VAE Perspective on Posterior Collapse

James Lucas, George Tucker, Roger Grosse, Mohammad Norouzi

arXiv:1911.02469v1cs.LGstat.ML

TL;DR

Posterior collapse is usually associated with the KL term, but this paper asks whether it can instead arise from the probabilistic model’s optimization landscape. Using linear VAEs and their correspondence with pPCA, it shows that collapsed local maxima can occur in log marginal likelihood while ELBO adds no such maxima, with related patterns observed in nonlinear VAEs.

  • Problem

    Posterior collapse reduces the latent information used by VAEs, while its mathematical mechanism remains poorly understood and is commonly attributed to the ELBO KL term.

  • Method

    The paper analyzes linear VAEs through their exact correspondence with pPCA, characterizing stationary points and ELBO optima before testing the insights in nonlinear VAEs.

  • Results

    The ELBO adds no spurious local maxima beyond pPCA log marginal likelihood; collapsed maxima are linked to observation noise and related patterns appear in deep nonlinear VAEs.

  • Takeaways & Limitations

    Posterior collapse can arise from ill-conditioned stationary points in the probabilistic model, so linear VAEs provide a useful test case for VAE hypotheses.

Abstract

from arXiv · show

Posterior collapse in Variational Autoencoders (VAEs) arises when the variational posterior distribution closely matches the prior for a subset of latent variables. This paper presents a simple and intuitive explanation for posterior collapse through the analysis of linear VAEs and their direct correspondence with Probabilistic PCA (pPCA). We explain how posterior collapse may occur in pPCA due to local maxima in the log marginal likelihood. Unexpectedly, we prove that the ELBO objective for the linear VAE does not introduce additional spurious local maxima relative to log marginal likelihood. We show further that training a linear VAE with exact variational inference recovers an identifiable global maximum corresponding to the principal component directions. Empirically, we find that our linear analysis is predictive even for high-capacity, non-linear VAEs and helps explain the relationship between the observation noise, local maxima, and posterior collapse in deep Gaussian VAEs.

1 Introduction

The paper studies posterior collapse through linear VAEs and their correspondence with pPCA, showing that collapse can arise from local maxima in log marginal likelihood rather than solely from the ELBO KL term.

  • Posterior collapse occurs when the generative model ignores a subset of latent variables, reducing the information used by the decoder.
  • The paper analyzes linear VAEs to investigate the connection between posterior collapse and spurious local maxima in the ELBO objective.
  • A linear VAE with diagonal variational covariance recovers the exact principal components at the global ELBO maximum.
  • The ELBO for a linear VAE introduces no additional local maxima beyond those already present in pPCA log marginal likelihood.
  • Observation noise influences the stability of collapsed stationary points, and the analysis extends empirically to high-capacity nonlinear VAEs.

2 Preliminaries

The preliminaries introduce VAEs, pPCA, the ELBO, and posterior collapse, emphasizing that the paper tests whether log marginal likelihood itself can encourage collapse.

  • pPCA uses Gaussian latent variables and spherical Gaussian observation noise, making its marginal distribution and posterior analytically tractable.
  • VAEs optimize a tractable ELBO because marginal likelihood is intractable for deep latent variable models.
  • The ELBO combines expected conditional log-likelihood with a KL divergence between the variational distribution and the prior.
  • The paper hypothesizes that log marginal likelihood itself can encourage posterior collapse, alongside the commonly implicated KL term.
  • Posterior collapse is defined here as a variational distribution close to the prior, preventing the decoder from using all latent dimensions.

3 Related Work

Related work attributes posterior collapse to KL regularization, inference dynamics, approximate inference, or decoder capacity; this paper instead emphasizes optimization in the probabilistic model.

  • Other explanations link collapse to inference-network lag or mismatch between the variational and true posteriors.
  • The paper reports that posterior collapse can occur with weak linear decoders and attributes it entirely to log marginal likelihood in the linear case.
  • Prior work commonly reduces the KL term through annealing, free-bits, or alternative prior and variational distributions.
  • Unlike approaches that modify the objective by ignoring collapsed dimensions, this work studies unmodified linear VAEs with diagonal covariance.
  • Linear VAEs recover principal components exactly, whereas related linear autoencoder results recover them only up to rotations.

4 Analysis of linear VAE

The analysis establishes a direct correspondence between linear VAEs and pPCA, characterizes collapsed stationary points, and shows that ELBO optimization adds no spurious local maxima.

  • 4.1 Probabilistic PCA Revisited: pPCA’s global maximum recovers the principal-component subspace, while other stationary points correspond to zeroed decoder columns and posterior collapse.
  • 4.1 Probabilistic PCA Revisited: Increasing observation noise makes collapsed stationary points stable local maxima, preventing learning of additional components.
  • 4.1 Probabilistic PCA Revisited: When observation noise is learned simultaneously, decreasing it increases likelihood at collapsed stationary points, making them saddle points.
  • 4.2 Linear VAEs recover pPCA: A linear VAE with globally shared diagonal covariance recovers the pPCA global optimum and the true posterior at the optimum.
  • 4.2 Linear VAEs recover pPCA: Distinct data covariance eigenvalues allow the linear VAE to identify individual principal components rather than only their spanning subspace.
  • 4.2 Linear VAEs recover pPCA: The ELBO introduces no additional local maxima relative to pPCA because non-orthogonal decoder configurations can be improved by infinitesimal rotations.

5 Deep Gaussian VAEs

The deep Gaussian VAE ELBO combines a prior-matching KL term with a noise-scaled expected reconstruction term. Its observation-noise parameter creates a natural comparison with β-VAE objectives.

  • The deep Gaussian VAE uses neural networks for qφ(z|x) and pθ(x|z), with the encoder parameterizing the variational distribution and the decoder reconstructing inputs.
  • The ELBO contains a KL divergence between qφ(z|x) and the prior plus an expected reconstruction term scaled by 1/(2σ2).
  • Changing σ2 can produce gradient updates equivalent to changing β in β-VAE, with an appropriate learning rate.
  • Using small β values can encourage powerful decoders to use latent variables, but may reduce ELBO and require downstream-task model selection.

6 Experiments

Experiments test whether linear-VAE theory matches practice and whether its optimization insights extend to deep nonlinear VAEs. Results connect analytic training, encoder structure, observation-noise initialization, and posterior collapse.

  • Experiments first verify the linear-VAE theory and then test whether its insights apply to deep nonlinear VAEs.
  • 6.1 Linear VAEs: Linear VAEs with learnable σ2 achieved final ELBOs essentially indistinguishable from pPCA maximum likelihood across hidden dimensions.
  • 6.1 Linear VAEs: Analytic ELBO gradients converged faster and reached a better final ELBO than stochastic gradients, 950.7 versus 939.3.
  • 6.1 Linear VAEs: With linear decoders and varying encoders, the linear encoder was not outperformed, while nonlinear encoders can introduce additional stationary points in general.
  • 6.2 Investigating posterior collapse in deep nonlinear VAEs: Posterior collapse is measured as the percentage of latent dimensions whose KL divergence remains below ε for at least 99% of inputs, using δ = 0.01.
  • 6.2 Investigating posterior collapse in deep nonlinear VAEs: For nonlinear VAEs, overly large fixed σ2 values yielded worse learned representations and lower final ELBOs, as predicted by the linear model.
  • 6.2 Investigating posterior collapse in deep nonlinear VAEs: Learning σ2 reduced but did not eliminate performance gaps across initializations, indicating that correctly learning noise is harder in nonlinear models.
  • 6.2 Investigating posterior collapse in deep nonlinear VAEs: Large σ2 initializations produced close prior matching even when σ2 was learned, suggesting local optima contribute to posterior collapse in deep VAEs.

7 Discussion

The discussion argues that linear VAEs provide a tractable theoretical test case for posterior collapse and optimization. Their global solution recovers principal components, while ELBO training adds no spurious local maxima in the linear setting.

  • Linear VAEs recover the pPCA posterior, and their global optimum identifies the principal-component directions rather than only their spanning subspace.
  • ELBO training for linear VAEs introduces no additional spurious local maxima relative to log marginal likelihood training.
  • Posterior collapse can occur with weak linear decoders because it arises from local maxima of the log marginal likelihood, not necessarily decoder capacity.
  • The stationary-point analysis represents decoder weights through singular vectors aligned with eigenvectors of the sample covariance matrix.
  • At the global optimum, decoder-column norms recover the ordering of principal components, allowing the permutation matrix to be computed exactly.

A.1 Stability of stationary point solutions

Fixed observation noise controls the stability of pPCA stationary points. When noise differs from its maximum-likelihood value, zero-column local optima can correspond to posterior collapse.

  • The stability analysis perturbs decoder columns and determines local-maxima behavior from the sign of the likelihood-gradient response.
  • When σ2 equals its maximum-likelihood estimate, non-leading-component solutions are saddle points because some perturbation direction is unstable.
  • For other fixed σ2 values, stationary points may have no unstable directions and become local optima with zero decoder columns.
  • These zero-column local optima correspond to posterior-collapse solutions observed in nonlinear VAEs.
  • When σ2 is learned in non-degenerate linear cases, the corresponding local maxima become saddle points because the noise gradient can reduce σ2.

B Identifiability of the linear VAE

Linear VAEs resolve the identifiability ambiguity of linear autoencoders while retaining a direct connection to pPCA. Their ELBO formulation selects a unique global optimum up to eigenvalue ordering and recovers the principal components themselves.

  • Identifiability: Linear VAEs fully identify the principal components, whereas linear autoencoders recover only their spanning subspace.The variational formulation removes the broader identifiability ambiguity described for linear autoencoders.
  • Identifiability: A diagonal rescaling leaves the VAE output invariant but changes the variational distribution and therefore the ELBO through its KL term.This breaks the corresponding transformation symmetry of the output model.
  • Identifiability: At the global optimum, the decoder-column norms recover eigenvalue ordering, so the remaining permutation can be computed exactly.The ordering is obtained by sorting the squared Euclidean norms of the columns of W.
  • Model formulation: The linear VAE uses a linear encoder and decoder, Gaussian prior, Gaussian observation model, and globally shared diagonal variational covariance.This restricted covariance structure is sufficient for analyzing the model's global optimum.
  • Variational structure: The ELBO's KL term is minimized when the variational distribution equals the true posterior, which is possible when decoder columns are orthogonal.The KL, reconstruction, and marginal-likelihood terms can be expressed in closed form for the linear VAE.

C.2 Finding stationary points

The stationary-point analysis explains posterior collapse through pPCA optimization and observation noise, while showing why the Bernoulli extension is analytically harder. Non-orthogonal stationary points can be shown to be saddles, but Bernoulli models lack closed-form likelihood and posterior calculations.

  • Finding stationary points: The ELBO introduces no additional local maxima beyond those of the pPCA log marginal likelihood.The variational mean recovers the true posterior at stationary points, and orthogonal decoder columns recover the marginal likelihood exactly.
  • Finding stationary points: A stationary point with non-orthogonal decoder columns is a saddle because a small rotation preserves log marginal likelihood while increasing the ELBO.The rotation reduces the posterior KL through the covariance-dependent log-determinant term.
  • Bernoulli Probabilistic PCA: For Bernoulli observations, neither the marginal likelihood nor the posterior has a closed-form solution, making stationary-point analysis difficult.The logit-normal expectation also lacks a closed form, and the ELBO reconstruction term contains another intractable integral.

E Experiment details

Experiments evaluate linear and deep VAEs on MNIST and CelebA, focusing on whether KL annealing can reduce posterior collapse and how learning σ2 affects the outcome.

  • Experiment setup: Linear and deep VAE experiments used TensorFlow and single-GPU training, with MNIST and CelebA serving as evaluation settings.MNIST models used the architecture 784-1024-512-k-512-1024-784, while CelebA models used a convolutional architecture.
  • Experiment setup: Linear VAEs were trained with 200 hidden dimensions on 1000 MNIST samples using full-batch Adam optimization.A grid search over four learning rates found 0.001 gave the best final ELBO for both stochastic and analytic models after 12000 steps.
  • KL annealing: KL annealing initially reduced posterior collapse but ultimately returned to collapsed solutions when σ2 was fixed.This behavior was observed as β increased from 0 to 1 over the first 100 epochs.
  • Learning σ2: When σ2 was learned, KL annealing enabled a much smaller σ2 and ultimately reduced posterior collapse, whereas standard ELBO training retained substantial collapse.The comparison was presented in Figures 8 and 9 and was interpreted as evidence that nonlinear VAE dynamics can resemble the linear case under suitable conditioning.
  • Collapse measurement: Deep VAE evaluations measured collapse as the percentage of latent dimensions within 0.01 KL of the prior for at least 99% of encoder inputs.This definition was used for averaged evaluations over 5 trials on real-valued MNIST and CelebA.

E.1.3 Qualitative Results

Qualitative evaluation examines CelebA reconstructions and latent interpolations from the KL-annealed convolutional VAE, alongside collapse measurements under fixed and learned σ2.

  • Latent interpolations: Latent interpolations use the variational means of three input points and linear interpolation on their plane, extrapolating to a fourth point.The fourth point lies on the plane defined by the other three inputs.
  • Collapse evaluation: Posterior collapse is plotted against the ε-threshold for deep CelebA VAEs with both fixed and learned σ2.Collapse is defined as the proportion of latent dimensions staying within ε KL divergence of the prior for at least 1 − δ of training data points, with δ = 0.01.
  • Evaluation coverage: The reported evaluation tables cover real-valued CelebA and MNIST with and without nonlinear preprocessing, averaging results over 5 trials.Collapse percent uses the 0.01 KL and 99% encoder-input thresholds.
Loading 1911.02469v1…