Learning to Draw Samples with Amortized Stein Variational Gradient Descent

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 there is a need to perform 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 personalized prediction, or need fast inference as inner loops for other algorithms such as learning latent variable models (such as variational autoencoder (Kingma & Welling, 2013)) or unnormalized distributions. Therefore, it is highly desirable to develop intelligent probabilistic inference systems that can adaptively improve their own performance to fully optimize the computational efficiency, and generalize to new tasks with similar structures. Developing such systems requires solving the following learning-to-sample problem:

Given a distribution with density $p(z)$ on set $\mathcal{Z}$ and a simulator $z=f(\xi;~{}\eta)$, such as a neural network, which takes a parameter $\eta$ and a random seed $\xi$ drawn from $q_{0}$ and outputs a value $z$ in $\mathcal{Z}$, we want to find an optimal parameter $\eta$ so that the density of the random output $z=f(\xi;~{}\eta)$ with $\xi\sim q_{0}$ closely matches the target $p$.

Here, we assume that we do not know the analytical form of the simulator $f(\cdot)$ (which we call inference network), and we can only query it through the output value $f(\xi;~{}\eta)$ and derivative $\partial_{\eta}f(\xi;~{}\eta)$ for given $\eta$ and $\xi$. We also assume that the random seed distribution $q_{0}$ is unknown and we can only access it through the draws of the random input $\xi$; that is, $q_{0}$ can be arbitrarily complex, and can be discrete, continuous or hybrid.

Because of the above assumption, we cannot directly calculate the density $q_{\eta}(z)$ of the output variable $z=f(\xi;~{}\eta)$; this makes it difficult to solve Problem 1 using the typical variational inference (VI) methods. Recall that VI finds an optimal parameter $\eta$ to approximate the target $p$ with $q_{\eta}$, by minimizing the KL divergence:

However equation (1) requires calculating the density $q_{\eta}(z)$ or its derivative, which is intractable by our assumption (and is called implicit models in Mohamed & Lakshminarayanan (2016)). This holds true even when Monte Carlo gradient estimates (Hoffman et al., 2013) and the reparametrization trick (Kingma & Welling, 2013) are applied.

This requirement of calculating $q_{\eta}(z)$ makes it difficult for practitioners to use variational inference in an automatic way, especially in domains where it is critical to use expressive inference networks to achieve good approximation quality. Methods that do not require to explicitly calculate $q_{\eta}(z)$, referred to as wild variational inference, or variational programming (Ranganath et al., 2016), can significantly simplify the design and expand the applications of VI methods, allowing practioners to focus more on choosing proposals that work best with their specific tasks.

A similar problem also appears in importance sampling (IS), where it requires calculating the IS proposal density $q(z)$ in order to calculate the importance weight $w(z)=p(z)/q(z)$. However, there exist methods that use no explicit information of the proposal $q(z)$, and, seemingly counter-intuitively, give better asymptotic variance or converge rate 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 date back to O’Hagan (1987), who argued that “Monte Carlo (that uses the proposal information) is fundamentally unsound”, and developed Bayesian Monte Carlo (O’Hagan, 1991) as an instance that uses no information of proposal $q(z)$, yet gives better convergence rate than the typical Monte Carlo $O(n^{-1/2})$ convergence 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 using the information of density function $q(z)$ explicitly.

which motivates us to develop a project gradient like algorithm that iteratively calculates the Stein variational gradient, and projects it to the finite dimensional $\eta$-space to update parameter $\eta$. At the convergence, the samples drawn from $q_{\eta}$ reach the equilibrium state of SVGD, and hence form a good approximation of $p$. We can view this method as “amortizing” or distilling the SVGD dynamics using parametric family $q_{\eta}$ or its inference network $f(\xi;~{}\eta)$ and call it amortized SVGD. See Section 3 for the detailed description of amortized SVGD.

Our algorithm provides a simple approach for the wild inference problem in Problem 1, enabling wide application in approximate learning and inference. We explore two examples in this paper. In Section 4.1, we apply amortized SVGD to learn complex encoder functions in variational autoencoder (VAE), allowing it to model complex latent variable space. In Section 4.2 we use amortized SVGD to learn hyper-parameters of MCMC samplers, which allows us to adaptively improve the efficiency of Bayesian computation when performing a large number of similar tasks.

Related Work

There has been a number of very recent work that studies advanced variational inference methods that do not require explicitly calculating $q_{\eta}(z)$ (Problem 1). This includes adversarial variational Bayesian (Mescheder et al., 2017) which approximates the KL divergence with a density ratio estimator, operator variational inference (Ranganath et al., 2016) which replaces the KL divergence with an alternative operator variational objective that is equivalent to Stein discrepancy (see Appendix A for more discussion), and amortized MCMC (Li et al., 2017) which proposes to amortize arbitrary MCMC dynamics to make it applicable to discrete models and gradient-free settings.

The key advantage of our method is its simplicity. both Mescheder et al. (2017) and Ranganath et al. (2016) require training some type of auxiliary networks, while our main algorithm (Algorithm 1 with update (10)) is very simple, and is essentially a generalization of the typical gradient descent rule that replaces the typical gradient with the stein variational gradient.

It is also possible to adopt the auxiliary variational inference methods (e.g., Agakov & Barber, 2004; Salimans et al., 2015) to solve Problem 1 by treating $\xi$ as a hidden variable. however, in order to frame a tractable joint distribution $p(z,\xi)$, one would need to add additional noise on $z$, e.g., assume $z=f(\xi;~{}\eta)+\mathcal{N}(0,\sigma^{2})$ with a careful choice of noise variance $\sigma$, and also need to introduce an additional auxiliary network to approximate $q(\xi~{}|~{}z)$, which makes the algorithm more complex than ours.

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, 2015) and date-driven designs of Monte Carlo based methods (e.g., Paige & Wood, 2016), to name only a few. Balan et al. (2015) also explored the idea of amortizing or distil MCMC samplers to obtain compact fast posterior representation. Most of these methods uses typical variational inference methods and hence need to use simple $q_{\eta}$ to ensure tractability.

There is a large literature on traditional adaptive MCMC methods (e.g., Andrieu & Thoms, 2008; Roberts & Rosenthal, 2009) which can be used to adaptively adjust the proposal distribution of MCMC by exploiting the special theoretical properties of MCMC (e.g., by minimizing the autocorrelation). Our method is simpler, more generic, and works efficiently in practice thanks to the use of gradient-based back-propagation. Finally, connections between stochastic gradient descent and variational inference have been discussed and exploited in Mandt et al. (2016); Maclaurin et al. (2015).

STEIN VARIATIONAL GRADIENT DESCENT

Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a nonparametric variational inference algorithm that iteratively transports a set of particles $\{z_{i}\}_{i=1}^{n}$ to approximate the target distribution $p$ by performing a type of functional gradient descent on the KL divergence. We give a quick overview of the SVGD in this section.

that is, ${\boldsymbol{\phi}}$ should yield a maximum decreasing rate on the KL divergence between the particle distribution and the target distribution. Here, ${\mathcal{F}}$ is a function set that includes the possible velocity fields and is chosen to be the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS) $\mathcal{H}=\mathcal{H}_{0}\times\cdots\times\mathcal{H}_{0}$, where $\mathcal{H}_{0}$ is a RKHS formed by scalar-valued functions associated with a positive definite kernel $k(z,z^{\prime})$, that is, ${\mathcal{F}}=\{{\boldsymbol{\phi}}\in\mathcal{H}\colon||{\boldsymbol{\phi}}||_{\mathcal{H}}\leq 1\}$. This choice of ${\mathcal{F}}$ allows us to consider velocity fields in infinite dimensional function spaces while still obtaining a closed form solution. Liu & Wang (2016) showed that the objective function in (4) equals a simple linear functional of ${\boldsymbol{\phi}}$:

where ${\mathcal{T}}_{p}$ is a linear operator acting on a velocity field ${\boldsymbol{\phi}}$ and returns a scalar-valued function; ${\mathcal{T}}_{p}$ is called the Stein operator in connection with Stein’s identity which shows that the RHS of (5) equals zero if $p=q$:

This is a result of integration by parts assuming the values of $p(z){\boldsymbol{\phi}}(z)$ vanish on the boundary of the integration domain. Therefore, the optimization in (4) reduces to

Observe that (8) is “simple” in that it is a linear functional optimization on a unit ball of a Hilbert space. Therefore, it is not surprise to derive a closed form solution:

where $k(z,z^{\prime})$ is the positive definite kernel associated with RKHS $\mathcal{H}_{0}$. See Liu et al. (2016) for the derivation. We call ${\boldsymbol{\phi}}^{*}$ the Stein variational gradient direction since it provides the optimal direction for pushing the particles towards the target distribution $p$.

In the practical SVGD algorithm, we start with a set of initial particles, calculate its corresponding ${\boldsymbol{\phi}}^{*}$ by replacing the expectation under $q$ with the empirical average of particles, and use it to update the particles:

The two terms in ${\boldsymbol{\phi}}^{*}(z_{i})$ play two different roles: the term with the gradient $\nabla_{z}\log p(z)$ drives the particles toward the high probability regions of $p(z)$, while the term with $\nabla_{z}k(z,z_{i})$ serves as a repulsive force to encourage different particles to be different from each other as shown in Liu & Wang (2016). Overall, this procedure provides diverse points for approximating distribution $p$ when it converges.

It is easy to see from (10) that ${\boldsymbol{\phi}}^{*}(z_{i})$ reduces to the typical gradient $\nabla_{z}\log p(z_{i})$ when there is only a single particle ($n=1$) and $\nabla_{z}k(z,z_{i})=0$ when $z=z_{i}$, in which case SVGD reduces to the standard gradient ascent for maximizing $\log p(z)$ (i.e., maximum a posteriori (MAP)).

By substituting the ${\boldsymbol{\phi}}^{*}$ in (9) into (8), one can show that (Liu et al., 2016; Chwialkowski et al., 2016; Oates et al., 2017)

where $\kappa_{p}(z,z^{\prime})$ is a positive definite kernel obtained by applying Stein operator on $k(z,z^{\prime})$ twice, as a function of $z$ and $z^{\prime}$, respectively. It has the following computationally tractable form:

where ${\boldsymbol{s}}_{p}(z)=\nabla_{z}\log p(z)$. The form of KSD in (11) provides a computationally tractable way for estimating the Stein discrepancy between a set of samples $\{z_{i}\}$ (e.g., drawn from an unknown $q$) and a distribution $p$ specified by its score function $\nabla_{z}\log p(z)$ (which is independent of its normalization constant),

AMORTIZED SVGD: TOWARDS AN AUTOMATIC NEURAL SAMPLER

SVGD and other particle-based methods become inefficient when we need to apply them repeatedly on a large number of different, but similar target distributions for multiple tasks, because they can not leverage the similarity between the different distributions and may require a large memory to restore a large number of particles. This problem can be addressed by training a neural network $f(\xi;~{}\eta)$ to output particles that would have been produced by SVGD; this amounts to “amortizing” or compressing the nonparametric SVGD into a parametric network, yielding a solution for wild variational inference in Problem 1 of Section 1.

One straightforward way to achieve this is to run SVGD until convergence and train $f(\xi;~{}\eta)$ to fit the resulting SVGD particles (e.g., by using generative adversarial networks (GAN) (Goodfellow et al., 2014)). This, however, requires running 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 $z=f(\xi;~{}\eta)$ improves by moving along the Stein variational gradient direction in (10), in order to move towards the target distribution.

Specifically, denote by $\eta^{t}$ the parameter estimated 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 $z_{i}=f(\xi_{i};~{}\eta^{t})$ based on $\eta^{t}$, where $m$ is a mini-batch size (e.g., $m=100$). The Stein variational gradient ${\boldsymbol{\phi}}^{*}(z_{i})$ in (10) would then ensure that $z^{\prime}_{i}=z_{i}+\epsilon{\boldsymbol{\phi}}^{*}(z_{i})$ forms a better approximation of the target distribution $p$. Therefore, we should adjust $\eta$ to make it output $\{z^{\prime}_{i}\}$ instead of $\{z_{i}\}$, that is, we want to update $\eta$ by

where $z_{i}^{\prime}=z_{i}+\epsilon{\boldsymbol{\phi}}^{*}(z_{i})$. This process is repeated until convergence, in which case the outputs of network $f$ can no longer be improved by SVGD and hence should form a good approximation of the target $p$. See Algorithm 1.

If we assume $\epsilon$ is very small, then (13) can be approximated by a least square optimization. To see this, note that $f(\xi_{i};~{}\eta)\approx f(\xi_{i};~{}\eta^{t})+\partial_{\eta}f(\xi_{i};~{}\eta^{t})(\eta-\eta^{t})$ by Taylor expansion. Since $z_{i}=f(\xi_{i};~{}\eta^{t})$, we have

As a result, (13) reduces to a least square optimization:

It may be still slow to solve a least square problem at each iteration. We can derive a more computationally efficient approximation by performing only one step of gradient descent of (13) starting at $\eta^{t}$ (or equivalently (14) starting at $\delta=0$), which gives

Although update (15) is derived as an approximation of (13) or (14), it is computationally faster and it works effectively in practice; this is because when $\epsilon$ is small, one step of gradient update can be sufficiently close to the optimum.

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

APPLICATIONS OF AMORTIZED SVGD

With amortized SVGD, we can use expressive inference networks to obtain better approximation and explore new applications where traditional VI methods cannot be applied. In this section, we introduce two different applications of amortized SVGD. One is training variational autoencoders (Kingma & Welling, 2013) with complex, non-Gaussian encoders, and the other is training “smart” MCMC samplers that adaptively improve their own hyper-parameters from past experience.

Variational autoencoders (VAEs) (Kingma & Welling, 2013) are latent variable models of form $p_{\theta}(x)=\int_{z}p_{\theta}(x|z)p_{\theta}(z)dz$ where $x$ is an observed variable and $z$ is an un-observed latent variable. Assume the empirical distribution of the observed variable is $\hat{p}(x)$, VAE learns the parameter $\theta$ using a variational EM algorithm which approximates the posterior distribution $p_{\theta}(z|x)$ with a simple encoder $q_{\eta}(z|x)$, and updates $\theta$ and $\eta$ alternatively by

which alternates between updating $\theta$ by maximizing the joint likelihood (M-step (18)) and approximating the posterior distribution $p_{\theta}(z|x)$ given fixed $\theta$ with variational inference (VI) (E-step (19)). In standard VAE, (18) is performed using standard VI with the reparameterization trick (16), which requires $q_{\eta}$ to be tractable. Therefore, $q_{\eta}(z\mid x)$ is often defined as a Gaussian distribution with mean and diagonal covariance parameterized by neural networks with $x$ as input. This Gaussian assumption potentially limits the quality of the resulting generative models, and more expressive encoders can improve the performance as shown in recent works (e.g, Kingma et al., 2016; Mescheder et al., 2017, to name a few).

By applying amortized SVGD to solve the posterior inference problem in (19), we obtain simple algorithms that work with more complex inference networks. Specifically, we assume that $z\sim q_{\eta}(z|x)$ is generated by $z=f(\xi,x;~{}\eta)$, and optimize $\eta$ using update (13)-(15). See Algorithm 2. In our experiment, we take $z=f(\xi,x;~{}\eta)$ to be a deep neural network with binary Bernoulli dropout noise at the input layer of the network which is more effective in approximating multi-modal posteriors than the simple Gaussian encoders.

Similar idea has also been explored in Pu et al. (2017). Besides, we propose an entropy regularized VAE to improve the standard VAE and get more diverse images by adding an entropy regularization term on the encoder networks. Here, we replace the $\eta$-update in (19) with

where the temperature parameter $(1+\alpha)$ becomes a weight coefficient of the repulsive force; a high temperature (or equivalent a large entropy regularization) yields a strong repulsive force and push the particles to be further away from each other.

2 Training Langevin Samplers

By viewing typical MCMC procedures as a simulator $f(\cdot)$, we can apply amortized SVGD to adaptively improve hyperparameters in MCMC inference. This is useful when we need to perform Bayesian inference on a large number of different, but similar datasets or posteriors, where we can adaptively improve the MCMC sampler for future tasks by leveraging the information of the previous tasks. An example of this, which we consider in this work, is adaptively learning optimal step sizes for Langevin dynamics.

To specify the general framework, we assume that we are interested in drawing samples from a set of distributions

indexed by parameter $\vartheta$. We are interested in learning a network $f(\xi,p_{\vartheta};~{}\eta)$ which maps the distribution $p_{\vartheta}$ to stochastic posterior samples. Note that this is a generalization of Problem 1 which focuses on approximating an individual distribution. In practice, $p_{\vartheta}$ could be the posterior distributions of unknown parameters of interest conditioning on different observed data, or models of different individuals in hierarchical models.

In order to learn the network $f(\xi,p_{\vartheta};~{}\eta)$, we modify Algorithm 1, to perform amortized SVGD on a randomly selected $p_{\vartheta}$ from ${\mathcal{Q}}$ at each iteration. See Algorithm 3. In this way, we expect that the trained network $f(\xi,p_{\vartheta};~{}\eta)$ can perform well on similar $p_{\vartheta}$ drawn from the same distribution, but never seen by the training algorithm.

A useful perspective is that typical MCMC methods can be viewed as neural networks $f(\xi,p_{\vartheta};~{}\eta)$ handcrafted by researchers, with nice theoretical properties. We can leverage the structure of existing MCMC to design the architecture of $f(\cdot)$, and use amortized SVGD to adaptively improve its parameters across different tasks.

As an example, Langevin dynamics draws samples from $p_{\vartheta}$ by starting with an initial sample $z^{0}$ and performing iterative random updates of form $z^{t+1}\leftarrow f_{t}(z^{t})$ with

which performs gradient ascent with a Gaussian perturbation. Here $\eta^{t}$ denotes a vector-valued step size at the $t$-th iteration and “$\odot$” denotes element-wise product, and $\xi^{t}$ is a standard Gaussian random vector of the same size as $z^{t}$.

We can view $T$ iterations of Langevin dynamics as a $T$-layer neural network:

in which the initial samples and the Gaussian noise form the random seeds of the network, that is, $\xi=\{\xi^{t}\}_{t=0}^{T-1}\cup\{z^{0}\}$, and the step sizes $\eta=\{\eta^{t}\}_{t=0}^{T-1}\vspace{0.1 cm}$ form the parameters which we can estimate using Algorithm 3. In cases when it is difficult to calculate $\nabla_{z}\log p_{\vartheta}$ exactly, such as the case of Bayesian inference with big datasets, we can use stochastic gradient Langevin dynamics (Welling & Teh, 2011) to approximate $\nabla_{z}\log p_{\vartheta}$ with subsampling, which introduces another source of randomness. This, however, does not influence the application of Algorithm 3, since our method does not need to know the structure of the random seed distribution.

In practice, a large value of $T$ would result in a deep network and cause a vanishing gradient problem. We address this problem by partitioning the $T$ layers into small blocks of size $5$ or $10$, and evaluate the gradient of the parameters in each block by back-propagating the Stein variational gradient from the output of its own block.

EXPERIMENTS

In this section, we use amortized SVGD (Algorithm 3) to learn the step size parameters in the Langevin sampler in (23). We test a number of distribution families, including Gaussian Mixture, Gaussian Bernoulli RBM, Bayesian logistic regression and Bayesian neural networks. In all the cases, we train the sampler with a set of “training distributions” and evaluate the sampler on “test distributions” that are not seen by the algorithm during training. We compare our method to the typical Langevin sampler with power decay step size, selected to be the best from $\eta^{t}=10^{a}/(t+b)^{\gamma}$ where $\gamma=0.55$, $a\in\{-6,...,2\}$, $b\in\{0,...,9\}$.

We first train the Langevin samplers to learn to sample from simple Gaussian mixtures. We consider a family of Gaussian Mixtures $q_{\vartheta}(z)=\frac{1}{10}\sum_{i=1}^{10}\mathcal{N}(z;\vartheta_{i},0.1^{2})$, where $\vartheta$ is the mean parameter.

Restricted Boltzmann Machine (RBM)

Bayesian Classification

We test our method on Bayesian Logistic Regression and Bayesian neural networks for binary classification on real world datasets. In this case, the distribution of interest has a form of $p_{\vartheta}(z)=p(z|D)$, where $z$ is the network weights in logistic regression and neural networks, and $D$ is the dataset for binary classification, which we view as the parameter $\vartheta$, that is, different dataset $D$ yields different posterior $p(z|D)$, and we are interested in training the Langevin sampler on a set of available datasets, and hope it performs well on future datasets that have similar structures. This setting can be useful, for example, in the streaming setting where we use existing datasets to adaptively improve the Langevin sampler. In our experiment, we take nine similar datasets (a1a-a9) from the libsvm repositoryhttps://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/; we train our Langevin sampler on a9a, and evaluate the sampler on the remaining 8 datasets (a1a-a8a). Our training and evaluation steps are as follows:

1. Estimate the step sizes $\{\eta^{t}\}_{t=0}^{T-1}$ of the Langevin sampler using amortized SVGD based on dataset a9a.

2. Apply the Langevin sampler with the estimated step size to the training subsets of a$k$a, $k=1,\ldots,8$, to obtain posterior samples $\{z_{i}^{k}\}$ of the classification weights.

3. Calculate the test likelihood of $\{z_{i}^{k}\}$ on the testing subsets of a$k$a, $k=1,\ldots,8$. Report the averaged testing likelihood averaged on the 8 datasets in Figure 4.

Because each dataset is relatively large, we use the stochastic gradient approximation as suggested by Welling & Teh (2011) (with a minibatch size of 100) in Langevin samplers. We find that the $T=10$ Langevin samplers trained by amortized SVGD on a9a is comparable with the $T=10^{3}$ Langevin sampler with the best power decay step size.

2 Training VAE With Amortized SVGD

We compare the entropy regularized VAE trained with amortized SVGD (denoted by ESteinVAE), which is Algorithm 2 with update (21), with the standard VAE and entropy regularized standard VAE (denoted by EVAE) on the dynamically binarized MNIST dataset (Burda et al., 2015). We tested the following settings:

1. A standard VAE (VAE-f) with a fully connected encoder consisting of one hidden layer with 400 hidden units, and a Gaussian output hidden variable with diagonal covariance.

2. A standard convolutional Gaussian VAE (VAE-CNN) with a convolutional encoder consisting of 2 convolution layers with $5\times 5$ filters, stride 2 and $$ features maps, followed by a fully connected layer with 512 hidden units.

3. A convolutional Gaussian entropy regularized (EVAE-CNN) with the same encoder and decoder structures as the standard convolutional VAE (VAE-CNN).

4. Entropy regularized VAEs trained by our amortized SVGD in Algorithm 2, with the same encoder architecture as VAE-f and VAE-CNN, respectively, but removing the Gaussian noise on the top layer and adding a multiplicative Bernoulli noise with a dropout rate of 0.3 to the each layer of the encoder. These two cases are denoted by ESteinVAE-f, and ESteinVAE-CNN, respectively.

Table 1 reports the test log-likelihood of all the methods estimated using Hamiltonian annealed importance sampling (HAIS) (Wu et al., 2016) with 100 independent AIS chains and 10,000 intermediate transitions, averaged on 5000 test images. We find that ESteinVAE-f significantly outperforms VAE-f, and ESteinVAE-CNN slightly outperforms VAE-CNN and EVAE-CNN. Table 1 also reports the effective sample size (ESS) of the HAIS estimates. The fact that the effective sample sizes of all the methods are close suggests that accuracy of the different NLL estimates are comparable.

Missing data imputation

Each column (starting from the third column) of Figure 5 shows an independent run of this procedure, where we can see that ESteinVAE-CNN is able to generate diverse construction when ambiguity exists, while the EVAE-CNN and VAE-CNN tend to be trapped in a local mode. This suggests that the diagonal variance of the latent variables in the Gaussian encoder of VAE-CNN tends to be small, underestimating the posterior uncertain, while ESteinVAE can capture the multi-modal posterior due to the dropout noise.

Table 2 is the quantitative result of the imputation experiment, in which the “accuracy” column denotes the number of original images whose digit is in its reconstructed images. and the “entropy” column denotes the entropy of the probability of the reconstructed images belonging to different digit classes. EsteinVAE obtains more diverse images and slightly more accurate reconstructed images.

CONCLUSION

We propose a new method to train neural samplers for given distributions, together with various applications to learning to draw samples using neural samplers. Future directions include exploring more efficient neural architectures and theoretical understanding of our method.

Acknowledgments This work is supported in part by NSF CRII 1565796. We thank Yingzhen Li from University of Cambridge for her valuable comments and feedbacks.

References

Appendix A KSD Variational Inference

Kernelized Stein discrepancy (KSD) provides a discrepancy measure between distributions and can be in principle used as a variational objective function in replace of KL divergence. In fact, thanks to the special form of KSD ((11)-(12)), one can derive a standard stochastic gradient descent for minimizing KSD without needing to estimate $q_{\eta}(z)$ explicitly, which provides a conceptually simple wild variational inference algorithm. Although this work mainly focuses on amortized SVGD which we find to be easier to implement and tend to perform superior to KSD variational inference in practice (see Figure 6), we think the KSD approach is of theoretical interest and hence give a brief discussion here.

where $z_{i}=f(\xi_{i};~{}\eta)$. This enables a wild variational inference method based on directly minimizing $\eta$ with standard (stochastic) gradient descent. We call this algorithm amortized KSD. Note that (24) is similar to (15) in form, but replaces ${\boldsymbol{\phi}}^{*}(z_{i})$ with

Here $\bar{\boldsymbol{\phi}}^{*}$ depends on the second order derivative of $\log p$ because $\kappa_{p}(z,z^{\prime})$ depends on $\nabla\log p$, which makes it more difficult to implement amortized KSD than amortized SVGD.

Intuitively, minimizing KSD can be viewed as seeking a stationary point of KL divergence under SVGD updates. To see this, recall that $q_{[\epsilon{\boldsymbol{\phi}}]}$ denotes the density of $z^{\prime}=z+\epsilon{\boldsymbol{\phi}}(z)$ when $z\sim q$. From (4), we have for small $\epsilon$,

where $\nabla_{{\boldsymbol{\phi}}}F({\boldsymbol{\phi}})$ denotes the functional gradient of the function $F({\boldsymbol{\phi}})$ w.r.t. ${\boldsymbol{\phi}}$ defined in RKHS $\mathcal{H}^{d}$, and $\nabla_{{\boldsymbol{\phi}}}F({\boldsymbol{\phi}})$ is also an element in $\mathcal{H}^{d}$. Therefore, in contrast to amortized SVGD which attends to minimize the KL objective $F({\boldsymbol{\phi}})$, KSD variational inference minimizes the gradient magnitude $||\nabla_{\boldsymbol{\phi}}F(0)||_{\mathcal{H}^{d}}$ of KL divergence.

This idea is closely related to the operator variational inference (Ranganath et al., 2016), which directly minimizes the variational form of Stein discrepancy in (4) and (8) with ${\mathcal{F}}$ replaced by sets of parametric neural networks. Specifically, Ranganath et al. (2016) assumes ${\mathcal{F}}$ consists of a neural network ${\boldsymbol{\phi}}_{\tau}(z)$ with parameter $\tau$, and find $\tau$ jointly with $\eta$ by solving a min-max game:

This yields a more challenging computation problem, although it is possible that the neural networks provide stronger discrimination than RKHS in practice. The main advantage of the KSD based approach is that it leverages the closed form solution in RKHS, yields a simpler optimization formulation based on standard gradient descent.

Figure 6 shows results of Langevin samplers trained by amortized SVGD and amortized KSD, respectively, for learning simple Gaussian mixtures under the same setting as that in Section 5.1. We find that amortized KSD tends to perform worse (Figure 6(c)) than amortized SVGD (Figure 6(b)); given that it is also less straightforward to implement amortized KSD (for requiring calculating $\nabla_{z}\kappa_{p}(z,z^{\prime})$ in (24)), we did not test it in our other experiments.

Appendix B Solving the Projection Step Using Different Numbers Gradient Steps

Amortized SVGD requires us to solve the projection step using either (13) or (14) at each iteration. In practice, we approximately solve it using only one step of gradient descent starting from the old values of $\eta$ for the sake of computational efficiency.

In order to study the trade-off of accuracy and computational cost here, we plot in Figure 7 the results when we solve Eq (14) using different numbers of gradient descent steps (the result is almost identical when we solve Eq (13) instead). We can see that when using more gradient steps, although the training time per iteration increases, the overall convergence speed may still improve, because it may take less iterations to converge. Figure 7 seems to suggest that using 5, 10, 20 steps gives better convergence than using a single step, but this may vary in different cases. We suggest to search for the best gradient step if the convergence speed is a primary concern. On the other hand, the number of gradient steps seems to have minor influence on the final result at the convergence as shown in Figure 7.

Appendix C Images Generated by Different VAEs

Figure 8 shows the images generated by the standard VAE-CNN, the entropy regularized VAE-CNN and ESteinVAE-CNN. We can see that both EVAE-CNN and ESteinVAE-CNN can generate images of good quality.