InfoVAE: Information Maximizing Variational Autoencoders

Introduction

Generative models have shown great promise in modeling complex distributions such as natural images and text (Radford et al., 2015; Zhu et al., 2017; Yang et al., 2017; Li et al., 2017). These are directed graphical models which represent the joint distribution between the data and a set of hidden variables (features) capturing latent factors of variation. The joint is factored as the product of a prior over the latent variables and a conditional distribution of the visible variables given the latent ones. Usually a simple prior distribution is provided for the latent variables, while the distribution of the input conditioned on latent variables is complex and modeled with a deep network.

Both learning and inference are generally intractable. However, using an amortized approximate inference distribution it is possible to use the evidence lower bound (ELBO) to efficiently optimize both (a lower bound to) the marginal likelihood of the data and the quality of the approximate inference distribution. This leads to a very successful class of models called variational autoencoders (Kingma & Welling, 2013; Jimenez Rezende et al., 2014; Kingma et al., 2016; Burda et al., 2015).

However, variational autoencoders have several problems. First, the approximate inference distribution is often significantly different from the true posterior. Previous methods have resorted to using more flexible variational families to better approximate the true posterior distribution (Kingma et al., 2016). However we find that the problem is rooted in the ELBO training objective itself. In fact, we show that the ELBO objective favors fitting the data distribution over performing correct amortized inference. When the two goals are conflicting (e.g., because of limited capacity), the ELBO objective tends to sacrifice correct inference to better fit (or worse overfit) the training data.

Another problem that has been observed is that when the conditional distribution is sufficiently expressive, the latent variables are often ignored (Chen et al., 2016). That is, the model only uses a single conditional distribution component to model the data, effectively ignoring the latent variables and fail to take advantage of the mixture modeling capability of the VAE. In addition, one goal of unsupervised learning is to learn meaningful latent representations but this fails because the latent variables are ignored. Some solutions have been proposed in (Chen et al., 2016) by limiting the capacity of the conditional distribution, but this requires manual and problem-specific design of the features we would like to extract.

In this paper we propose a novel solution by framing both problems as explicit modeling choices: we introduce new training objectives where it is possible to weight the preference between correct inference and fitting data distribution, and specify a preference on how much the model should rely on the latent variables. This choice is implicitly made in the ELBO objective. We make this choice explicit and generalize the ELBO objective by adding additional terms that allow users to select their preference on both choices. Despite of the addition of seemingly intractable terms, we find an equivalent form that can still be efficiently optimized.

Our new family also generalizes known models including the $\beta$-VAE (Higgins et al., 2016) and Adversarial Autoencoders (Makhzani et al., 2015). In addition to deriving these models as special cases, we provide generic principles for hyper-parameter selection that work well in all the experimental settings we considered. Finally we perform extensive experiments to evaluate our newly introduced model family, and compare with existing models on multiple metrics of performance such as log-likelihood, sampling quality, and semi-supervised performance. An instantiation of our general framework called MMD-VAE achieves better or on-par performance on all metrics we considered. We further observe that our model can lead to better amortized inference, and utilize the latent variables even in the presence of a very flexible decoder.

Variational Autoencoders

A latent variable generative model defines a joint distribution between a feature space $z\in\mathcal{Z}$, and the input space $x\in\mathcal{X}$. Usually we assume a simple prior distribution $p(z)$ over features, such as Gaussian or uniform, and model the data distribution with a complex conditional distribution $p_{\theta}(x|z)$, where $p_{\theta}(x|z)$ is often parameterized with a neural network. Suppose the true underlying distribution is $p_{\mathcal{D}}(x)$ (that is approximated by a training set), then a natural training objective is maximum (marginal) likelihood

We further average this over the data distribution $p_{\mathcal{D}}(x)$ to obtain the final optimization objective

There are several ways to equivalently rewrite the ELBO objective that will become useful in our following analysis. We define the joint generative distribution as

In fact we can correspondingly define a joint “inference distribution”

Note that the two definitions are symmetrical. In the former case we start from a known distribution $p(z)$ and learn the conditional distribution on $\mathcal{X}$, in the latter we start from a known (empirical) distribution $p_{\mathcal{D}}(x)$ and learn the conditional distribution on $\mathcal{Z}$. We also correspondingly define any conditional and marginal distributions as follows:

For the purposes of optimization, the ELBO objective can be written equivalently (up to an additive constant) as

We prove the first equivalence in the appendix. The second and third equivalence are simple applications of the additive property of KL divergence. All three forms of ELBO in Eqns. (2),(3),(4) are useful in our analysis.

Two Problems of Variational Autoencoders

Under ideal conditions, optimizing the ELBO objective using sufficiently flexible model families for $p_{\theta}(x|z)$ and $q_{\phi}(z|x)$ over $\theta,\phi$ will achieve both goals of correctly capturing $p_{\mathcal{D}}(x)$ and performing correct amortized inference. This can be seen by examining Eq. (3). This form indicates that the ELBO objective is minimizing the KL divergence between the data distribution $p_{\mathcal{D}}(x)$ and the (marginal) model distribution $p_{\theta}(x)$, as well as the KL divergence between the variational posterior $q_{\phi}(z|x)$ and the true posterior $p_{\theta}(z|x)$. However, with finite model capacity the two goals can be conflicting and subtle tradeoffs and failure modes can emerge from optimizing the ELBO objective.

In particular, one limitation of the ELBO objective is that it might fail to learn an amortized inference distribution $q_{\phi}(z|x)$ that approximates the true posterior $p_{\theta}(z|x)$. This can happen for two different reasons:

Inherent properties of the ELBO objective: the ELBO objective can be maximized (even to $+\infty$ in pathological cases) even with a very inaccurate variational posterior $q_{\phi}(z|x)$.

Implicit modeling bias: common modeling choices (such as the high dimensionality of $\mathcal{X}$ compared to $\mathcal{Z}$) tend to sacrifice variational inference vs. data fit when modeling capacity is not sufficient to achieve both.

We will explain in turn why these failures happen.

We first provide some intuition to this phenomena, then formally prove the result for a pathological case of continuous spaces and Gaussian distributions. Finally we justify in the experiments section that this happens in realistic settings on real datasets (in both continuous and discrete spaces).

Let $\mathcal{D}$ be a dataset with two samples $\{-1,1\}$, and $p_{\theta}(x|z)$ be selected from the family of all functions $\mu^{p}_{\theta},\sigma^{p}_{\theta}$ that map $z\in\mathcal{X}$ to a Gaussian $\mathcal{N}(\mu^{p}_{\theta}(z),\sigma^{p}_{\theta}(z))$ on $\mathcal{X}$, and $q_{\phi}(z|x)$ be selected from the family of all functions $\mu^{q}_{\phi},\sigma^{q}_{\phi}$ that map $x\in\mathcal{X}$ to a Gaussian $\mathcal{N}(\mu^{q}_{\phi}(z),\sigma^{q}_{\phi}(z))$ on $\mathcal{Z}$. Then $\mathcal{L}_{ELBO}$ can be maximized to $+\infty$ when

1.2 Modeling Bias

In the above example we indicated a potential problem with the ELBO objective where the model tends to push the probability mass of $q_{\phi}(z|x)$ too far from . This tendency is a property of the ELBO objective and true for any $\mathcal{X}$ and $\mathcal{Z}$. However this is made worse by the fact that $\mathcal{X}$ is often higher dimensional compared to $\mathcal{Z}$, so any error in fitting $\mathcal{X}$ will be magnified compared to $\mathcal{Z}$.

For example, consider fitting an $n$ dimensional distribution $\mathcal{N}(0,I)$ with $\mathcal{N}(\epsilon,I)$ using KL divergence, then

Regardless of the cause (properties of ELBO or modeling choices), this is generally an undesirable phenomenon for two reasons:

1) One may care about accurate inference more than generating sharp samples. For example, generative models are often used for down stream tasks such as semi supervised learning.

2) Overfitting: Because $p_{\mathcal{D}}$ is an empirical (finite) distribution in practice, matching it too closely can lead to poor generalization.

Both issues are observed in the experiments section.

2 The Information Preference Property

Using a complex decoding distribution $p_{\theta}(x|z)$ such as PixelRNN/PixelCNN (van den Oord et al., 2016b; Gulrajani et al., 2016) has been shown to significantly improve sample quality on complex natural image datasets. However, this approach suffers from a new problem: it tends to neglect the latent variables $z$ altogether, that is, the mutual information between $z$ and $x$ becomes vanishingly small. Intuitively, the reason is that the learned $p_{\theta}(x|z)$ is the same for all $z\in\mathcal{Z}$, implying that the $z$ is not dependent on the input $x$. This is undesirable because a major goal of unsupervised learning is to learn meaningful latent features which should depend on the inputs.

This effect, which we shall refer to as the information preference problem, was studied in (Chen et al., 2016) with a coding efficiency argument. Here we provide an alternative interpretation, which sheds light on a novel solution to solve this problem.

We inspect the ELBO in the form of Eq.(3), and consider the two terms respectively. We show that both can be optimized to without utilizing the latent variables.

The InfoVAE Model Family

In order to remedy these two problems we define a new training objective that will learn both the correct model and amortized inference distributions. We begin with the form of ELBO in Eq. (4)

First we add a scaling parameter $\lambda$ to the divergence between $q_{\phi}(z)$ and $p(z)$ to increase its weight and counter-act the imbalance between $\mathcal{X}$ and $\mathcal{Z}$ (cf. discussion in section 3.1.1). Next we add a mutual information maximization term that prefers high mutual information between $x$ and $z$. This encourages the model to use the latent code and avoids the information preference problem. We arrive at the following objective

where $I_{q}(x;z)$ is the mutual information between $x$ and $z$ under the distribution $q_{\phi}(x,z)$.

Even though we cannot directly optimize this objective, we can rewrite this into an equivalent form that we can optimize more efficiently (We prove this in the Appendix)

Note that in the proposition we have the additional requirement that the mutual information $I_{q}(x;z)$ is bounded. This is inevitable because if $\alpha>0$ the objective can be optimized to $+\infty$ by simply increasing the mutual information infinitely. In our experiments simply ensuring that $q_{\phi}(z|x)$ does not have vanishing variance is sufficient to regularize the behavior of the model.

Relation to Adversarial Autoencoders (AAE): When $\alpha=1$, $\lambda=1$ and $D$ is chosen to be the Jensen Shannon divergence we get the adversarial autoencoders in (Makhzani et al., 2015). This paper generalizes AAE, but more importantly we provide a deeper understanding of the correctness and desirable properties of AAE. Furthermore, we characterize settings when AAE is preferable compared to VAE (i.e. when we would like to have $\alpha=1$).

Our generalization introduces new parameters, but the meaning and effect of the various parameter choices is clear. We recommend setting $\lambda$ to a value that makes the loss on $\mathcal{X}$ approximately the same as the loss on $\mathcal{Z}$. We also recommend setting $\alpha=0$ when $p_{\theta}(x|z)$ is a simple distribution, and $\alpha=1$ when $p_{\theta}(x|z)$ is a complex distribution and information preference is a concern. The final degree of freedom is the divergence $D(q_{\phi}(z)\|p(z))$ to use. We will explore this topic in the next section. We will also show in the experiments that this design choice is also easy to choose: there is a choice that we find to be consistently better in almost all metrics of performance.

We consider and compare three divergences in this paper.

Adversarial Training: Adversarial autoencoders (AAE) proposed in (Makhzani et al., 2015) use an adversarial discriminator to approximately minimize the Jensen-Shannon divergence (Goodfellow et al., 2014) between $q_{\phi}(z)$ and $p(z)$. However, when $p(z)$ is a simple distribution such as Gaussian, there are preferable alternatives. In fact, adversarial training can be unstable and slow even when we apply recent techniques for stabilizing GAN training (Arjovsky et al., 2017; Gulrajani et al., 2017).

Maximum-Mean Discrepancy: Maximum-Mean Discrepancy (MMD) (Gretton et al., 2007; Li et al., 2015; Dziugaite et al., 2015) is a framework to quantify the distance between two distributions by comparing all of their moments. It can be efficiently implemented using the kernel trick. Letting $k(\cdot,\cdot)$ be any positive definite kernel, the MMD between $p$ and $q$ is

Experiments

We first perform some simple experiments on toy data and MNIST to demonstrate that ELBO suffers from inaccurate inference in practice, and adding the scaling term $\lambda$ in Eq.(5) can correct for this. Next, we will perform a comprehensive set of experiments to carefully compare different models on multiple performance metrics.

We verify the conclusions in Proposition 1 by using the same setting in that proposition. We use a three layer deep network with 200 hidden units in each layer to simulate the highly flexible function family. For InfoVAE we choose the scaling coefficient $\lambda=500$, information preference $\alpha=0$, and divergence optimized by MMD.

The results are shown in Figure 1. It can be observed that the predictions of the theory are reflected in the experiments: ELBO training leads to poor inference and a significantly over-estimated $q_{\phi}(z)$, while InfoVAE demonstrates a more stable behavior.

1.2 MNIST

To make this more intuitive we plot in Figure 3 the contour plot of $q_{\phi}(z)$ when training on 500 examples. It can be seen that with ELBO $q_{\phi}(z)$ matches $p(z)$ very poorly, while InfoVAE matches significantly better.

To verify that ELBO trains inaccurate amortized inference we plot in Figure 4 the comparison between samples from the approximate posterior $q_{\phi}(z|x)$ and samples from the true posterior $p_{\theta}(z|x)$ computed by rejection sampling. The same trend can be observed. ELBO consistently gives very poor approximate posterior, while the InfoVAE posterior is mostly accurate.

Finally show the samples generated by the two models in Figure 5. ELBO generates very sharp reconstructions, but very poor samples when sampled ancestrally $x\sim p(z)p_{\theta}(x|z)$. InfoVAE, on the other hand, generates samples of consistent quality, and in fact, produces samples of reasonable quality after only training on a dataset of 500 examples. This reflect InfoVAE’s ability to control overfitting and demonstrate consistent training time and testing time behavior.

2 Comprehensive Comparison

In this section, we perform extensive qualitative and quantitative experiments on the binarized MNIST dataset to evaluate the performance of different models. We would like to answer these questions:

1) Compare the models on a comprehensive set of numerical metrics of performance. Also compare the stability and training speed of different models.

2) Evaluate and compare the possible types of divergences (Adversarial, Stein, MMD).

For all two questions, we find InfoVAE with MMD regularization to perform better in almost all metrics of performance and demonstrate the best stability and training speed. The details are presented in the following sections.

For models we use ELBO, Adversarial autoencoders, InfoVAE with Stein variational gradient, and InfoVAE with MMD ($\alpha=1$ because information preference is a concern, $\lambda=1000$ which can put the loss on $\mathcal{X}$ and $\mathcal{Z}$ on the same order of magnitude). In this setting we also use a highly flexible PixelCNN as the decoder $p_{\theta}(x|z)$ so that information preference is also a concern. Detailed experimental setup is explained in the Appendix.

We consider multiple quantitative evaluations, including the quality of the samples generated, the training speed and stability, the use of latent features for semi-supervised learning, and log-likelihoods on samples from a separate test set.

Distance between $q_{\phi}(z)$ and $p(z)$: To measure how well $q_{\phi}(z)$ approximates $p(z)$, we use two numerical metrics. The first is the full batch MMD statistic over the full data. Even though MMD is also used during training of MMD-VAE, it is too expensive to train using the full dataset, so we only use mini-batches for training. However during evaluation we can use the full dataset to obtain more accurate estimates. The second is the log determinant of the covariance matrix of $q_{\phi}(z)$. Ideally when $p(z)$ is the standard Gaussian $\Sigma_{q_{\phi}}$ should be the identity matrix, so $\log\det(\Sigma_{q_{\phi}})=0$. In our experiments we plot the log determinant divided by the dimensionality of the latent space. This measures the average under/over estimation per dimension of the learned covariance.

The results are plotted in Figure 6 (A,B). This is different from the experiments in Figure 2 because in this case the decoder is a highly complex pixel recurrent model and the concern that we highlight is failure to use latent features rather than inaccurate posterior. MMD achieves the best performance except for ELBO. Even though ELBO achieves extremely low error, this is trivial because for this experimental setting of flexible decoders, ELBO learns a latent code $z$ that does not contain any information about $x$, and $q_{\phi}(z|x)\approx p(z)$ for all $z$.

Sample distribution: If the generative model $p(z)p_{\theta}(x|z)$ has true marginal $p_{data}(x)$, then the distribution of different object classes should also follow the distribution of classes in the original dataset. On the other hand, an incorrect generative model is unlikely to generate a class distribution that is identical to the ground truth. We let $\bm{c}$ denote the class distribution in the real dataset, and $\hat{\bm{c}}$ denote the class distribution of the generated images, computed by a pretrained classifier. We use cross entropy loss $l_{ce}(\bm{c},\hat{\bm{c}})=-\bm{c}^{T}(\log\hat{\bm{c}}-\log\bm{c})$ to measure the deviation from the true class distribution.

The results for this metric are plotted in Figure 6 (C). In general, Stein regularization performs well only with a small latent space with 2 dimensions, whereas the adversarial regularization performs better with larger dimensions; MMD regularization generally performs well in all the cases and the performance is stable with respect to the latent code dimensions.

Training Speed and Stability: In general we would prefer a model that is stable, trains quickly and requires little hyperparameter tuning. In Figure 6 (D) we plot the change of MMD statistic vs. the number of iterations. In this respect, adversarial autoencoder becomes less desirable because it takes much longer to converge, and sometimes converges to poor results even if we consider more power GAN training techniques such as Wasserstein GAN with gradient penalty (Gulrajani et al., 2017).

Log likelihood: To be consistent with previous results, we use the stochastically binarized version first used in (Salakhutdinov & Murray, 2008). Estimation of log likelihood is achieved by importance sampling. We use 5-dimensional latent features in our log likelihood experiments. The values are shown in Table 1. Our results are slightly worse than reported in PixelRNN (van den Oord et al., 2016b), which achieves a log likelihood of 79.2. However, all the regularizations perform on-par or superior compared to our ELBO baseline. This is somewhat surprising because we do not explicitly optimize a lower bound to the true log likelihood, unless we are using the ELBO objective.

Conclusion

Despite the recent success of variational autoencoders, they can fail to perform amortized inference, or learn meaningful latent features. We trace both issues back to the ELBO learning criterion, and modify the ELBO objective to propose a new model family that can fix both problems. We perform extensive experiments to verify the effectiveness of our approach. Our experiments show that a particular subset of our model family, MMD-VAEs perform on-par or better than all other approaches on multiple metrics of performance.

Acknowledgements

We thank Daniel Levy, Rui Shu, Neal Jean, Maria Skoularidou for providing constructive feedback and discussions.

References

Appendix A Proofs

where the unique valid solution is $\sigma=2\sqrt{q(z<0|x=1)}$, therefore optimally

Let the dataset contain two samples $\{-1,1\}$, and $p(x|z),q(z|x)$ be arbitrary one dimensional Gaussians, then by symmetry of the problem, at optimality for $\mathcal{L}_{ELBO}$, we have

The optimal solution $\sigma^{*}$ is achieved where

where the unique valid solution is $\sigma=2\sqrt{q(z<0|x=1)}$, therefore optimally

Where $q(z<0|x=1)$ is the tail probability of a Gaussian. In the limit $\lambda\to 0,c\to\infty$, we have

Then in the limit of $\lambda\to 0$, $c\to\infty$

Notice that by our condition of $\alpha<1,\lambda>0$, we have $1-\alpha>0,\alpha+\lambda-1>0$. For convenience we will rename

In addition we consider our objective in two separate terms

We will prove that whenever $\beta>0,\gamma>0$, the two terms are maximized under the condition in the proposition respectively. First consider $\mathcal{L}_{1}$, because $I_{q_{\phi}}(x,z)=I_{0}$:

Therefore for any $q_{\phi}(z|x)$, $p_{\theta^{*}}(x|z)$ that optimizes $\mathcal{L}_{1}$ satisfies $\forall z$, $p_{\theta^{*}}(x|z)=q_{\phi}(x|z)$, and we have for any given $q_{\phi}$, the optimal $\mathcal{L}_{1}$ is

where we use $H_{p}(x)$ to denote the entropy of $p(x)$. Notice that $\mathcal{L}_{1}$ is dependent on $q_{\phi}$ only by $I_{q_{\phi}}(x;z)$ therefore when $I_{q_{\phi}}(x;z)=I_{0}$, $\mathcal{L}_{1}$ is maximized regardless of the choice of $q_{\phi}$. So we only have to independently maximize $\mathcal{L}_{2}$ subject to fixed $I_{0}$. Notice that $\mathcal{L}_{1}$ is maximized when $q_{\phi}(z)=p(z)$, and we show that this can be achieved. When $\{q_{\phi}\}$ is sufficiently flexible we simply have to partition the support set $\mathcal{A}$ of $p(z)$ into $N=\lceil e^{I_{0}}\rceil$ subsets $\{\mathcal{A}_{1},\cdots,\mathcal{A}_{N}\}$, so that each subset satisfies $\int_{\mathcal{A}_{i}}p(z)dz=1/N$. Similarly we partition the support set $\mathcal{B}$ of $p_{\mathcal{D}}(x)$ into $N$ subsets $\{\mathcal{B}_{1},\cdots,\mathcal{B}_{N}\}$, so that each subset satisfies $\int_{\mathcal{B}_{i}}p_{\mathcal{D}}(x)dx=1/N$. Then we construct $q_{\phi}(z|x)$ mapping each $B_{i}$ to $A_{i}$ as follows

such that for any $x_{i}\in B_{i}$. It is easy to see that this distribution is normalized because

Also it is easy to see that $p(z)=q_{\phi}(z)$. In addition

Therefore we have constructed a $q_{\phi}(z|x),p_{\theta}(x|z)$ so that we have reached the maximum for both objectives

so there sum must also be maximized. Under this optimal solution we have that $q_{\phi}(x|z)=p_{\theta}(x|z)$ and $q_{\phi}(z)=p(z)$, this implies $q_{\phi}(x,z)=p_{\theta}(x,z)$, which implies that both $p_{\theta}(z|x)=q_{\phi}(z|x)$ and $p_{\theta}(x)=p_{\mathcal{D}}(x)$. ∎

Appendix B Stein Variational Gradient

as shown in Lemma 3.2 in (Liu & Wang, 2016). Intuitively $\phi^{*}_{q,p}$ is the steepest direction that transforms $q(z)$ towards $p(z)$. In practice, $q(z)$ can be represented by the particles in a mini-batch.

We propose to use Stein variational gradient to regularize variational autoencoders using the following process. For a mini-batch $\bm{x}$, we compute the corresponding mini-batch features $\bm{z}=q_{\phi}(\bm{x})$. Based on this mini-batch we compute the Stein gradients by empirical samples

The gradients wrt. the model parameters are

In practice we can define a surrogate loss

where $\mbox{stop\_gradient}(\cdot)$ indicates that this term is treated as a constant during back propagation. Note that this is not really a divergence, but simply a convenient loss function that we can implement using standard automatic differentiation software, whose gradient is the Stein variational gradient of the true KL divergence.

Appendix C Experimental Setup

In all our experiments in Section 5.2, we choose the prior $p(z)$ to be a Gaussian with zero mean and identity covariance, $p_{\theta}(x|z)$ to be a PixelCNN conditioned on the latent code (van den Oord et al., 2016a), and $q_{\phi}(z|x)$ to be a CNN where the final layer outputs the mean and standard deviation of a factored Gaussian distribution.

For MNIST we use a simplified version of the conditional PixelCNN architecture (van den Oord et al., 2016a). For CIFAR we use the public implementation of PixelCNN++ (Salimans et al., 2017). In either case we use a convolutional encoder with the same architecture as (Radford et al., 2015) to generate the latent code, and plug this into the conditional input for both models. The entire model is trained end to end by Adam (Kingma & Ba, 2014). PixelCNN on ImageNet take about 10h to train to convergence on a single Titan X, and CIFAR take about two days to train to convergence. We will make the code public upon publication.

Appendix D CIFAR Samples

We additional perform experiments on the CIFAR (Krizhevsky & Hinton, 2009) dataset. We show samples from models with different regularization methods - ELBO, MMD, Stein and Adversarial in Figure 7. In all cases, the model accurately matches $q_{\phi}(z)$ with $p(z)$, and samples generated with Stein regularization and MMD regularization are more coherent globally compared to the samples generated with ELBO regularization.