Source-linked AI summary
InfoVAE: Information Maximizing Variational Autoencoders
Shengjia Zhao, Jiaming Song, Stefano Ermon
TL;DR
Variational autoencoder objectives can produce inaccurate amortized inference and can cause flexible decoders to ignore latent variables. InfoVAE makes the trade-offs between inference, data fitting, and latent-variable reliance explicit through efficiently optimizable objectives. Experiments report that MMD-VAE performs on-par or better across multiple metrics while improving inference and latent-variable use.
Problem
Existing VAE objectives can favor data fitting over accurate amortized inference, while expressive decoders may ignore latent variables and undermine meaningful representations.
Method
InfoVAE generalizes the ELBO with terms and parameters that control inference accuracy, data fitting, and preference for latent-variable use, using efficiently optimizable divergence replacements.
Results
MMD-VAEs perform on-par or better than other approaches across multiple performance metrics, while improving amortized inference and using latent variables with flexible decoders.
Takeaways & Limitations
The framework exposes ELBO trade-offs as explicit modeling choices and provides a model family that can address both inference quality and latent-feature use.
Takeaways & Limitations
When α > 0, the objective requires bounded mutual information because otherwise it can increase without bound; nonvanishing encoder variance regularizes experiments.
Abstract
from arXiv · showhide
A key advance in learning generative models is the use of amortized inference distributions that are jointly trained with the models. We find that existing training objectives for variational autoencoders can lead to inaccurate amortized inference distributions and, in some cases, improving the objective provably degrades the inference quality. In addition, it has been observed that variational autoencoders tend to ignore the latent variables when combined with a decoding distribution that is too flexible. We again identify the cause in existing training criteria and propose a new class of objectives (InfoVAE) that mitigate these problems. We show that our model can significantly improve the quality of the variational posterior and can make effective use of the latent features regardless of the flexibility of the decoding distribution. Through extensive qualitative and quantitative analyses, we demonstrate that our models outperform competing approaches on multiple performance metrics.
1. Introduction
Variational autoencoders address intractable learning and inference with amortized approximate inference, but their ELBO objective can sacrifice inference accuracy and latent-feature use. InfoVAE makes these preferences explicit and supports efficient optimization while improving performance across multiple metrics.
- 1. Introduction: Amortized approximate inference makes it possible to optimize a lower bound to marginal likelihood and inference quality efficiently.This approach underlies variational autoencoders.
- 1. Introduction: The ELBO can favor fitting the data distribution over correct amortized inference, especially when model capacity is limited.Improving the objective can therefore worsen inference quality relative to the true posterior.
- 1. Introduction: Expressive conditional decoders can ignore latent variables, preventing the model from learning meaningful latent representations.The decoder may use a single conditional component and fail to exploit the VAE's mixture modeling capability.
- 1. Introduction: InfoVAE introduces objectives that separately weight correct inference, data fitting, and reliance on latent variables.The added terms have an equivalent form that remains efficiently optimizable.
- 1. Introduction: MMD-VAE performs better or comparably on all considered metrics, while improving amortized inference and using latent variables with flexible decoders.The evaluation includes log-likelihood, sampling quality, and semi-supervised performance.
2. Variational Autoencoders
Variational autoencoders replace intractable likelihood optimization with an amortized inference distribution and an ELBO objective. Equivalent ELBO formulations expose the relationships among the generative and inference distributions for subsequent analysis.
- 2. Variational Autoencoders: A latent-variable model combines a prior p(z) with a conditional distribution pθ(x|z) to model the data distribution.The conditional is often parameterized by a neural network, while the prior is typically simple.
- 2. Variational Autoencoders: Direct likelihood optimization is intractable because computing pθ(x) requires integration over z.The likelihood marginalizes the conditional pθ(x|z) against the prior p(z).
- 2. Variational Autoencoders: An amortized inference distribution qφ(z|x) enables joint optimization of a lower bound to the log likelihood.This provides the variational-autoencoder training approach.
- 2. Variational Autoencoders: The ELBO can be equivalently rewritten in several forms to support analysis of the objective.The paper defines corresponding joint, conditional, and marginal distributions for the generative and inference viewpoints.
- 2. Variational Autoencoders: The equivalent ELBO forms differ by additive constants or applications of the additive property of KL divergence.All three forms are used in the paper's analysis.
3. Two Problems of Variational Autoencoders
The paper identifies two failures in VAE training: ELBO optimization can produce inaccurate amortized inference, and expressive decoders can ignore latent variables. These failures reflect both objective properties and modeling biases, motivating objectives that separately control inference, data fit, and latent-variable use.
- Inaccurate amortized inference: ELBO optimization can favor fitting the data distribution over accurate amortized inference when model capacity is limited.The objective may maximize despite an inaccurate variational posterior, especially when fitting the observed space conflicts with inference quality.
- Inaccurate amortized inference: In a two-sample Gaussian example, the ELBO can become arbitrarily large while D_KL(q_φ(z|x)||p_θ(z|x)) tends to +∞.The learned inference distribution pushes probability mass toward infinity and becomes infinitely far from the true posterior.
- Modeling bias: Higher dimensionality of X magnifies modeling errors relative to Z, causing ELBO optimization to sacrifice latent-space divergence when the objectives conflict.For fixed ε, Euclidean mean distance scales as Θ(√n), whereas the corresponding KL divergence scales as Θ(n).
- Modeling bias: These failures matter because accurate inference supports downstream tasks such as semi-supervised learning, while matching finite empirical data too closely can cause overfitting.The paper identifies accurate inference and generalization as reasons not to optimize data fit alone.
- Information preference: Expressive decoders can ignore latent variables, making mutual information between z and x vanishingly small and undermining meaningful latent representations.The decoder may use the same conditional distribution for every z, so x becomes independent of z.
4. The InfoVAE Model Family
InfoVAE generalizes the ELBO by explicitly controlling the trade-offs between fitting the data, matching the latent prior, and preserving mutual information. Its formulation supports tractable optimization with alternative divergences and recovers several existing models as special cases.
- InfoVAE objective: InfoVAE adds a scaling parameter λ for the divergence between qφ(z) and p(z), plus a mutual-information term favoring use of the latent code.The mutual-information term is Iq(x; z) under qφ(x, z).
- Optimization: An equivalent objective enables efficient optimization: its first two terms use the reparameterization trick, while the latent divergence can use likelihood-free methods.The divergence may be replaced by any strict divergence, defined as zero exactly when qφ(z) = p(z).
- Theoretical guarantee: For α < 1 and λ > 0, fixing Iq(x; z), global optimization recovers pθ(x) = pD(x) and qφ(z|x) = pθ(z|x).The proposition applies to continuous X and Z and requires bounded mutual information.
- Connections to prior models: The model family includes the original ELBO, β-VAE, and adversarial autoencoders through particular parameter and divergence choices.The original ELBO uses α = 0 and λ = 1; AAE uses α = 1, λ = 1, and Jensen-Shannon divergence.
- Hyper-parameter choices: The authors recommend balancing losses on X and Z with λ, setting α = 0 for simple decoders, and α = 1 for complex decoders when information preference matters.They also report a divergence choice that performs consistently better across almost all considered metrics.
- Divergence choices: The paper compares Jensen-Shannon divergence, Stein variational gradient, and maximum-mean discrepancy as divergence choices.Adversarial training may be unstable and slow, Stein methods are difficult to apply efficiently in high dimensions, and MMD uses the kernel trick.
5. Experiments
Experiments on toy data and binarized MNIST show that InfoVAE improves posterior accuracy, prior matching, sampling consistency, and performance across several evaluation metrics. MMD regularization is generally stable and performs well across latent dimensions and tasks.
- Toy data and MNIST: ELBO training leads to poor inference and significantly over-estimated qφ(z), whereas InfoVAE demonstrates more stable behavior on the toy setting.
- Toy data and MNIST: ELBO significantly over-estimates the variance of qφ(z), especially with small training sets, while large λ lets InfoVAE avoid this problem.The metric is log det(Cov[qφ(z)]), whose target is 0 for the standard Gaussian prior.
- Toy data and MNIST: InfoVAE matches the prior qφ(z) significantly better than ELBO when trained on 500 examples.
- Toy data and MNIST: InfoVAE produces mostly accurate approximate posteriors, while ELBO consistently produces poor approximate posteriors relative to the true posterior.The comparison uses testing data after training on 500 samples and computes the true posterior by rejection sampling.
- Toy data and MNIST: InfoVAE generates samples of consistent quality, unlike ELBO’s sharp training reconstructions but poor ancestral samples; this behavior remains reasonable after training on 500 examples.The ELBO discrepancy indicates overfitting and mismatch between qφ(z) and p(z).
- Comprehensive comparison: MMD regularization performs well across latent dimensions, while InfoVAE with MMD performs better in almost all metrics and has the best stability and training behavior.The evaluated metrics include MMD, covariance log determinant, class-distribution cross entropy, semi-supervised performance, and log likelihood.
6. Conclusion
The paper traces failures in amortized inference and latent-feature learning to the ELBO criterion, proposes a modified model family, and reports strong MMD-VAE performance across multiple metrics.
- The paper identifies ELBO training as the source of failures in amortized inference and meaningful latent-feature learning.These failures occur despite the success of variational autoencoders.
- The proposed model family modifies the ELBO objective to address both problems.The approach is evaluated through extensive experiments.
- MMD-VAEs perform on-par or better than other approaches on multiple performance metrics.
A. Proofs
The proofs show that ELBO over-fitting can receive unbounded reward while inference quality diverges, whereas the modified InfoVAE objective can attain both data and posterior-matching optima under its stated conditions.
- For a symmetric two-sample dataset with Gaussian conditionals, ELBO over-fitting has unbounded reward as c →∞ and λ →0.The ELBO growth rate exceeds that of the regularization term, and their sum is maximized to +∞.
- As c →∞, the variational gap DKL(q(z|x)∥p(z|x)) →∞ because the true posterior has bounded second-order moments.
- Under α < 1 and λ > 0, the modified InfoVAE objective separates into two terms whose maxima can be optimized under fixed mutual information I0.The proof considers the terms independently when β > 0 and γ > 0.
- A sufficiently flexible variational family can partition latent and data supports into equally probable subsets to construct qφ(z|x) with qφ(z) = p(z).The construction maps each data subset to a corresponding latent subset.
- The constructed solution maximizes both objectives and implies qφ(x,z) = pθ(x,z), hence pθ(z|x) = qφ(z|x) and pθ(x) = pD(x).
B. Stein Variational Gradient
The Stein variational gradient framework transforms particles toward a target distribution and is used to regularize variational autoencoders through minibatch feature samples and an implementable surrogate loss.
- Stein variational gradients match a distribution q to p by descending the variational gradient of DKL(q(z)||p(z)).
- The transformation T(z) = z + ϵφ(z) induces a new distribution from samples z ∼q, with φ∗q,p giving the steepest direction toward p.In practice, q(z) can be represented by minibatch particles.
- The proposed VAE regularizer computes Stein gradients from minibatch features and uses a surrogate loss whose gradient equals the Stein variational gradient of the true KL divergence.The surrogate loss is compatible with standard automatic differentiation.
C. Experimental Setup
Experiments use Gaussian priors, conditional PixelCNN decoders, and CNN encoders, with dataset-specific PixelCNN variants and end-to-end Adam training.
- All experiments use a zero-mean identity-covariance Gaussian prior, a latent-conditioned PixelCNN decoder, and a CNN encoder outputting factored-Gaussian parameters.
- MNIST uses a simplified conditional PixelCNN, while CIFAR uses the public PixelCNN++ implementation.
- Both models use a convolutional encoder based on Radford et al. and are trained end to end with Adam.
D. CIFAR Samples
On CIFAR, Stein and MMD regularization produce samples that are more globally coherent than ELBO regularization.
- Stein and MMD regularization yield CIFAR samples that are more coherent globally than samples generated with ELBO regularization.