Learning to Draw Samples: With Application to Amortized MLE for Generative Adversarial Learning

Introduction

Modern machine learning increasingly relies on highly complex probabilistic models to reason about uncertainty. A key computational challenge is to develop efficient inference techniques to approximate, or draw samples from complex distributions. Currently, most inference methods, including MCMC and variational inference, are hand-designed by researchers or domain experts. This makes it difficult to fully optimize the choice of different methods and their parameters, and exploit the structures in the problems of interest in an automatic way. The hand-designed algorithm can also be inefficient when it requires to make fast inference repeatedly on a large number of different distributions with similar structures. This happens, for example, when we need to reason about a number of observed datasets in settings like online learning, or need fast inference as inner loops for other algorithms such as maximum likelihood training. Therefore, it is highly desirable to develop more intelligent probabilistic inference systems that can adaptively improve its own performance to fully the optimize computational efficiency, and generalize to new tasks with similar structures.

Specifically, denote by $p(x)$ a probability density of interest specified up to the normalization constant, which we want to draw sample from, or marginalize to estimate its normalization constant. We want to study the following problem:

Given a distribution with density $p(x)$ and a function $f(\eta;~{}\xi)$ with parameter $\eta$ and random input $\xi$, for which we only have assess to draws of the random input $\xi$ (without knowing its true distribution $q_{0}$), and the output values of $f(\eta;~{}\xi)$ and its derivative $\partial_{\eta}f(\eta;~{}\xi)$ given $\eta$ and $\xi$. We want to find an optimal parameter $\eta$ so that the density of the random output variable $x=f(\eta;~{}\xi)$ with $\xi\sim q_{0}$ closely matches the target density $p(x)$.

In fact, it is this requirement of calculating $q_{\eta}(x)$ that has been the major constraint for the designing of state-of-the-art variational inference methods with rich approximation families; the recent successful algorithms (e.g., Rezende & Mohamed, 2015b; Tran et al., 2015; Ranganath et al., 2015, to name only a few) have to handcraft special variational families to ensure the computational tractability of $q_{\eta}(x)$ and simultaneously obtain high approximation accuracy, which require substantial mathematical insights and research effects. Methods that do not require to explicitly calculate $q_{\eta}(x)$ can significantly simplify the design and applications of VI methods, allowing practical users to focus more on choosing proposals that work best with their specific tasks. We will use the term wild variational inference to refer to new variants of variational methods that require no tractability $q_{\eta}(x)$, to distinguish with the black-box variational inference (Ranganath et al., 2014) which refers to methods that work for generic target distributions $p(x)$ without significant model-by-model consideration (but still require to calculate the proposal density $q_{\eta}(x)$).

A similar problem also appears in importance sampling (IS), where it requires to calculate the IS proposal density $q(x)$ in order to calculate the importance weight $w(x)=p(x)/q(x)$. However, there exist methods that use no explicit information of $q(x)$, which, seemingly counter-intuitively, give better asymptotic variance or converge rates than the typical IS that uses the proposal information (e.g., Liu & Lee, 2016; Briol et al., 2015; Henmi et al., 2007; Delyon & Portier, 2014). Discussions on this phenomenon dates back to O’Hagan (1987), who argued that “Monte Carlo (that uses the proposal information) is fundamentally unsound” for violating the Likelihood Principle, and developed Bayesian Monte Carlo (O’Hagan, 1991) as an example that uses no information on $q(x)$, yet gives better convergence rate than the typical Monte Carlo $O(n^{-1/2})$ rate (Briol et al., 2015). Despite the substantial difference between IS and VI, these results intuitively suggest the possibility of developing efficient variational inference without calculating $q(x)$ explicitly.

In this work, we propose a simple algorithm for Problem 1 by iteratively adjusting the network parameter $\eta$ to make its output random variable changes along a Stein variational gradient direction (SVGD) (Liu & Wang, 2016) that optimally decreases its KL divergence with the target distribution. Critically, the SVGD gradient includes a repulsive term to ensure that the generated samples have the right amount of variability that matches $p(x).$ In this way, we “amortize SVGD” using a neural network, which makes it possible for our method to adaptively improve its own efficiency by leveraging fast experience, especially in cases when it needs to perform fast inference repeatedly on a large number of similar tasks. As an application, we use our method to amortize the MLE training of deep energy models, where a neural sampler is adaptively trained to approximate the likelihood function. Our method, which we call SteinGAN, mimics an adversarial game between the energy model and the neural sampler, and obtains realistic-looking images competitive with the state-of-the-art results produced by generative adversarial networks (GAN) (Goodfellow et al., 2014; Radford et al., 2015).

The idea of amortized inference (Gershman & Goodman, 2014) has been recently applied in various domains of probabilistic reasoning, including both amortized variational inference (e.g., Kingma & Welling, 2013; Rezende & Mohamed, 2015a), and data-driven proposals for (sequential) Monte Carlo methods (e.g., Paige & Wood, 2016), to name only a few. Most of these methods, however, require to explicitly calculate $q(x)$ (or its gradient). One exception is a very recent paper (Ranganath et al., 2016) that avoids calculating $q(x)$ using an idea related to Stein discrepancy (Gorham & Mackey, 2015; Liu et al., 2016; Oates et al., 2014; Chwialkowski et al., 2016). There is also a raising interest recently on a similar problem of “learning to optimize” (e.g., Andrychowicz et al., 2016; Daniel et al., 2016; Li & Malik, 2016), which is technically easier than the more general problem of “learning to sample”. In fact, we show that our algorithm reduces to “learning to optimize” when only one particle is used in SVGD.

Generative adversarial network (GAN) and its variants have recently gained remarkable success on generating realistic-looking images (Goodfellow et al., 2014; Salimans et al., 2016; Radford et al., 2015; Li et al., 2015; Dziugaite et al., 2015; Nowozin et al., 2016). All these methods are set up to train latent variable models (the generator) under the assistant of the discriminator. Our SteinGAN instead performs traditional MLE training for a deep energy model, with the help of a neural sampler that learns to draw samples from the energy model to approximate the likelihood function; this admits an adversarial interpretation: we can view the neural sampler as a generator that attends to fool the deep energy model, which in turn serves as a discriminator that distinguishes the real samples and the simulated samples given by the neural sampler. This idea of training MLE with neural samplers was first discussed by Kim & Bengio (2016); one of the key differences is that the neural sampler in Kim & Bengio (2016) is trained with the help of a heuristic diversity regularizer based on batch normalization, while SVGD enforces the diversity in a more principled way. Another method by Zhao et al. (2016) also trains an energy score to distinguish real and simulated samples, but within a non-probabilistic framework (see Section 5 for more discussion). Other more traditional approaches for training energy-based models (e.g., Ngiam et al., 2011; Xie et al., 2016) are often based on variants of MCMC-MLE or contrastive divergence (Geyer, 1991; Hinton, 2002; Tieleman, 2008), and have difficulty generating realistic-looking images from scratch.

Stein Variational Gradient Descent (SVGD)

Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a general purpose Bayesian inference algorithm motivated by Stein’s method (Stein, 1972; Barbour & Chen, 2005) and kernelized Stein discrepancy (Liu et al., 2016; Chwialkowski et al., 2016; Oates et al., 2014). It uses an efficient deterministic gradient-based update to iteratively evolve a set of particles $\{x_{i}\}_{i=1}^{n}$ to minimize the KL divergence with the target distribution. SVGD has a simple form that reduces to the typical gradient descent for maximizing $\log p$ when using only one particle $(n=1)$, and hence can be easily combined with the successful tricks for gradient optimization, including stochastic gradient, adaptive learning rates (such as adagrad), and momentum.

where $\epsilon$ is a step size, and ${\boldsymbol{\phi}}(x)$ is a “particle gradient direction” chosen to maximumly decrease the KL divergence between the distribution of particles and the target distribution, in the sense that

where $q_{[\epsilon{\boldsymbol{\phi}}]}$ denotes the density of the updated particle $x^{\prime}=x+\epsilon{\boldsymbol{\phi}}(x)$ when the density of the original particle $x$ is $q$, and ${\mathcal{F}}$ is the set of perturbation directions that we optimize over. We choose ${\mathcal{F}}$ to be the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS) $\mathcal{H}^{d}=\mathcal{H}\times\cdots\times\mathcal{H}$ with each $\mathcal{H}$ associating with a positive definite kernel $k(x,x^{\prime})$; note that $\mathcal{H}$ is dense in the space of continuous functions with universal kernels such as the Gaussian RBF kernel.

Critically, the gradient of KL divergence in (2) equals a simple linear functional of ${\boldsymbol{\phi}}$, allowing us to obtain a closed form solution for the optimal ${\boldsymbol{\phi}}$. Liu & Wang (2016) showed that

where ${\mathcal{T}}_{p}$ is considered as a linear operator acting on function ${\boldsymbol{\phi}}$ and is called the Stein operator in connection with Stein’s identity which shows that the RHS of (3) equals zero if $p=q$:

This is a result of integration by parts assuming the value of $p(x){\boldsymbol{\phi}}(x)$ vanishes on the boundary of the integration domain.

Therefore, the optimization in (2) reduces to

By approximating the expectation under $q$ with the empirical average of the current particles $\{x_{i}\}_{i=1}^{n}$, SVGD admits a simple form of update:

It is easy to see from (2) that $\Delta x_{i}$ reduces to the typical gradient $\nabla_{x}\log p(x_{i})$ when there is only a single particle ($n=1$) and $\nabla_{x}k(x,x_{i})$ when $x=x_{i}$, in which case SVGD reduces to the standard gradient ascent for maximizing $\log p(x)$ (i.e., maximum a posteriori (MAP)).

Amortized SVGD: Towards an Automatic Neural Sampler

SVGD and other particle-based methods become inefficient when we need to repeatedly infer a large number different target distributions for multiple tasks, including online learning or inner loops of other algorithms, because they can not improve based on the experience from the past tasks, and may require a large memory to restore a large number of particles. We propose to “amortize SVGD” by training a neural network $f(\eta;~{}\xi)$ to mimic the SVGD dynamics, yielding a solution for Problem 1.

One straightforward way to achieve this is to run SVGD to convergence and train $f(\eta;~{}\xi)$ to fit the SVGD results. This, however, requires to run many epochs of fully converged SVGD and can be slow in practice. We instead propose an incremental approach in which $\eta$ is iteratively adjusted so that the network outputs $x=f(\eta;~{}\xi)$ changes along the Stein variational gradient direction in (2) in order to decrease the KL divergence between the target and approximation distribution.

To be specific, denote by $\eta^{t}$ the estimated parameter at the $t$-th iteration of our method; each iteration of our method draws a batch of random inputs $\{\xi_{i}\}_{i=1}^{m}$ and calculate their corresponding output $x_{i}=f(\eta;~{}\xi_{i})$ based on $\eta^{t}$; here $m$ is a mini-batch size (e.g., $m=100$). The Stein variational gradient $\Delta x_{i}$ in (2) would then ensure that $x^{\prime}_{i}=x_{i}+\epsilon\Delta x_{i}$ forms a better approximation of the target distribution $p$. Therefore, we should adjust $\eta$ to make its output matches $\{x^{\prime}_{i}\}$, that is, we want to update $\eta$ by

See Algorithm 1 for the summary of this procedure. If we assume $\epsilon$ is very small, then (8) reduces to a least square optimization. To see this, note that $f(\eta;~{}\xi_{i})\approx f(\eta^{t};~{}\xi_{i})+\partial_{\eta}f(\eta^{t};~{}\xi_{i})(\eta-\eta^{t})$ by Taylor expansion. Since $x_{i}=f(\eta^{t};~{}\xi_{i})$, we have

As a result, (8) reduces to the following least square optimization:

Update (9) can still be computationally expensive because of the matrix inversion. We can derive a further approximation by performing only one step of gradient descent of (8) (or (9)), which gives

Although update (10) is derived as an approximation of (8)-(9), it is computationally faster and we find it works very effectively in practice; this is because when $\epsilon$ is small, one step of gradient update can be sufficiently close to the optimum.

Update (10) also has a simple and intuitive form: (10) can be thought as a “chain rule” that back-propagates the Stein variational gradient to the network parameter $\eta$. This can be justified by considering the special case when we use only a single particle $(n=1)$ in which case $\Delta x_{i}$ in (2) reduces to the typical gradient $\nabla_{x}\log p(x_{i})$ of $\log p(x)$, and update (10) reduces to the typical gradient ascent for maximizing

in which case $f(\eta;~{}\xi)$ is trained to maximize $\log p(x)$ (that is, learning to optimize), instead of learning to draw samples from $p$ for which it is crucial to use Stein variational gradient $\Delta x_{i}$ to diversify the network outputs.

With $\{\xi_{i}\}$ i.i.d. drawn from $q_{0}$ and $x_{i}=f(\eta;~{}\xi_{i}),~{}\forall i$, the standard stochastic gradient descent for minimizing the KL divergence is

Amortized MLE for Generative Adversarial Training

Our method allows us to design efficient approximate sampling methods adaptively and automatically, and enables a host of novel applications. In this paper, we apply it in an amortized MLE method for training deep generative models.

Maximum likelihood estimator (MLE) provides a fundamental approach for learning probabilistic models from data, but can be computationally prohibitive on distributions for which drawing samples or computing likelihood is intractable due to the normalization constant. Traditional methods such as MCMC-MLE use hand-designed methods (e.g., MCMC) to approximate the intractable likelihood function but do not work efficiently in practice. We propose to adaptively train a generative neural network to draw samples from the distribution during MLE training, which not only provides computational advantage, and also allows us to generate realistic-looking images competitive with, or better than the state-of-the-art generative adversarial networks (GAN) (Goodfellow et al., 2014; Radford et al., 2015) (see Figure 1-5).

To be specific, denote by $\{x_{i,obs}\}$ a set of observed data. We consider the maximum likelihood training of energy-based models of form

where $\phi(x;~{}\theta)$ is an energy function for $x$ indexed by parameter $\theta$ and $\Phi(\theta)$ is the log-normalization constant. The log-likelihood function of $\theta$ is

We call our method SteinGAN, because it can be intuitively interpreted as an adversarial game between the generative network $f(\eta;~{}\xi)$ and the energy model $p(x|\theta)$ which serves as a discriminator: The MLE gradient update of $p(x|\theta)$ effectively decreases the energy of the training data and increases the energy of the simulated data from $f(\eta;~{}\xi)$, while the SVGD update of $f(\eta;~{}\xi)$ decreases the energy of the simulated data to fit better with $p(x|\theta)$. Compared with the traditional methods based on MCMC-MLE or contrastive divergence, we amortize the sampler as we train, which gives much faster speed and simultaneously provides a high quality generative neural network that can generate realistic-looking images; see Kim & Bengio (2016) for a similar idea and discussions.

Empirical Results

We evaluated our SteinGAN on four datasets, MNIST, CIFAR-10, CelebA (Liu et al., 2015), and Large-scale Scene Understanding (LSUN) (Yu et al., 2015), on which we find our method tends to generate realistic-looking images competitive with, sometimes better than DCGAN (Radford et al., 2015) (see Figure 2 - Figure 3). Our code is available at https://github.com/DartML/SteinGAN.

In order to generate realistic-looking images, we define our energy model based on an autoencoder:

Some datasets include both images $x$ and their associated discrete labels $y$. In these cases, we train a joint energy model on $(x,y)$ to capture both the inner structure of the images and its predictive relation with the label, allowing us to simulate images with a control on which category it belongs to. Our joint energy model is defined to be

where $\sigma(\cdot,\cdot)$ is the cross entropy loss function of a fully connected output layer. In this case, our neural sampler first draws a label $y$ randomly according to the empirical counts in the dataset, and then passes $y$ into a neural network together with a $100\times 1$ random vector $\xi$ to generate image $x$. This allows us to generate images for particular categories by controlling the value of input $y$.

Stabilization

In practice, we find it is useful to modify (13) to be

where $\gamma$ is a discount factor (which we take to be $\gamma=0.7$). This is equivalent to maximizing a regularized likelihood:

where $\Phi(\theta)$ is the log-partition function; note that $\exp(\gamma\Phi(\theta))$ is a conjugate prior of $p(x|\theta)$.

We initialize the weights of both the generator and discriminator from Gaussian distribution $\mathcal{N}(0,0.02)$, and train them using Adam (Kingma & Ba, 2014) with a learning rate of $0.001$ for the generator and $0.0001$ for the energy model (the discriminator). In order to keep the generator and discriminator approximately aligned during training, we speed up the MLE update (16) of the discriminator (by increasing its learning rate to $0.0005$) when the energy of the real data batch is larger than the energy of the simulated images, while slow down it (by freezing the MLE update of $\theta$ in (16)) if the magnitude of the energy difference between the real images and the simulated images goes above a threshold of 0.5. We used the bag of architecture guidelines for stable training suggested in DCGAN (Radford et al., 2015).

Discussion

The MNIST dataset has a training set of $60,000$ examples. Both DCGAN and our model produce high quality images, both visually indistinguishable from real images; see figure 1.

CIFAR-10 is very diverse, and with only 50,000 training examples. Figure 2 shows examples of simulated images by DCGAN and SteinGAN generated conditional on each category, which look equally well visually. We also provide quantitively evaluation using a recently proposed inception score (Salimans et al., 2016), as well as the classification accuracy when training ResNet using $50,000$ simulated images as train sets, evaluated on a separate held-out testing set never seen by the GAN models. Besides DCGAN and SteinGAN, we also evaluate another simple baseline obtained by subsampling 500 real images from the training set and duplicating them 100 times. We observe that these scores capture rather different perspectives of image generation: The inception score favors images that look realistic individually and have uniformly distributed labels; as a result, the inception score of the duplicated 500 images is almost as high as the real training set. We find that the inception score of SteinGAN is comparable, or slightly lower than that of DCGAN. On the other hand, the classification accuracy measures the amount information captured in the simulated image sets; we find that SteinGAN achieves the highest classification accuracy, suggesting that it captures more information in the training set.

Figure 3 and 4 visualize the results on CelebA (with more than 200k face images) and LSUN (with nearly 3M bedroom images), respectively. We cropped and resized both dataset images into $64\times 64$.

Conclusion

We propose a new method to train neural samplers for given distributions, together with a new SteinGAN method for generative adversarial training. Future directions involve more applications and theoretical understandings for training neural samplers.

References