A-NICE-MC: Adversarial Training for MCMC

Jiaming Song, Shengjia Zhao, Stefano Ermon

Introduction

Variational inference (VI) and Monte Carlo (MC) methods are two key approaches to deal with complex probability distributions in machine learning. The former approximates an intractable distribution by solving a variational optimization problem to minimize a divergence measure with respect to some tractable family. The latter approximates a complex distribution using a small number of typical states, obtained by sampling ancestrally from a proposal distribution or iteratively using a suitable Markov chain (Markov Chain Monte Carlo, or MCMC).

Recent progress in deep learning has vastly advanced the field of variational inference. Notable examples include black-box variational inference and variational autoencoders , which enabled variational methods to benefit from the expressive power of deep neural networks, and adversarial training , which allowed the training of new families of implicit generative models with efficient ancestral sampling. MCMC methods, on the other hand, have not benefited as much from these recent advancements. Unlike variational approaches, MCMC methods are iterative in nature and do not naturally lend themselves to the use of expressive function approximators . Even evaluating an existing MCMC technique is often challenging, and natural performance metrics are intractable to compute . Defining an objective to improve the performance of MCMC that can be easily optimized in practice over a large parameter space is itself a difficult problem .

To address these limitations, we introduce A-NICE-MC, a new method for training flexible MCMC kernels, e.g., parameterized using (deep) neural networks. Given a kernel, we view the resulting Markov Chain as an implicit generative model, i.e., one where sampling is efficient but evaluating the (marginal) likelihood is intractable. We then propose adversarial training as an effective, likelihood-free method for training a Markov chain to match a target distribution. First, we show it can be used in a learning setting to directly approximate an (empirical) data distribution. We then use the approach to train a Markov Chain to sample efficiently from a model prescribed by an analytic expression (e.g., a Bayesian posterior distribution), the classic use case for MCMC techniques. We leverage flexible volume preserving flow models and a “bootstrap” technique to automatically design powerful domain-specific proposals that combine the guarantees of MCMC and the expressiveness of neural networks. Finally, we propose a method that decreases autocorrelation and increases the effective sample size of the chain as training proceeds. We demonstrate that these trained operators are able to significantly outperform traditional ones, such as Hamiltonian Monte Carlo, in various domains.

Notations and Problem Setup

where Tθ(x)T_{\theta}(\cdot|x) is a time-homogeneous stochastic transition kernel parametrized by θΘ\theta\in\Theta and π0\pi^{0} is some initial distribution for x0x_{0}. In particular, we assume that TθT_{\theta} is defined through an implicit generative model fθ(x,v)f_{\theta}(\cdot|x,v), where vp(v)v\sim p(v) is an auxiliary random variable, and fθf_{\theta} is a deterministic transformation (e.g., a neural network). Let πθt\pi^{t}_{\theta} denote the distribution for xtx_{t}. If the Markov chain is both irreducible and positive recurrent, then it has an unique stationary distribution \pi_{\theta}=\raisebox{2.15277pt}{\scalebox{0.8}{\displaystyle\lim_{t\rightarrow\infty}\;}}\pi^{t}_{\theta}. We assume that this is the case for all the parameters θΘ\theta\in\Theta.

Low bias: The stationary distribution is close to the target distribution (minimize πθpd)|\pi_{\theta}-p_{d}|).

Efficiency: {πθt}t=0\{\pi_{\theta}^{t}\}_{t=0}^{\infty} converges quickly (minimize tt such that πθtpd<δ|\pi_{\theta}^{t}-p_{d}|<\delta).

Low variance: Samples from one chain {xt}t=0\{x_{t}\}_{t=0}^{\infty} should be as uncorrelated as possible (minimize autocorrelation of {xt}t=0\{x_{t}\}_{t=0}^{\infty}).

We think of πθ\pi_{\theta} as a stochastic generative model, which can be used to efficiently produce samples with certain characteristics (specified by pdp_{d}), allowing for efficient Monte Carlo estimates. We consider two settings for specifying the target distribution. The first is a learning setting where we do not have an analytic expression for pd(x)p_{d}(x) but we have access to typical samples {si}i=1mpd\{s_{i}\}_{i=1}^{m}\sim p_{d}; in the second case we have an analytic expression for pd(x)p_{d}(x), possibly up to a normalization constant, but no access to samples. The two cases are discussed in Sections 3 and 4 respectively.

Adversarial Training for Markov Chains

Consider the setting where we have direct access to samples from pd(x)p_{d}(x). Assume that the transition kernel Tθ(xt+1xt)T_{\theta}(x_{t+1}|x_{t}) is the following implicit generative model:

Assuming a stationary distribution πθ(x)\pi_{\theta}(x) exists, the value of πθ(x)\pi_{\theta}(x) is typically intractable to compute. The marginal distribution πθt(x)\pi^{t}_{\theta}(x) at time tt is also intractable, since it involves integration over all the possible paths (of length tt) to xx. However, we can directly obtain samples from πθt\pi_{\theta}^{t}, which will be close to πθ\pi_{\theta} if tt is large enough (assuming ergodicity). This aligns well with the idea of generative adversarial networks (GANs), a likelihood free method which only requires samples from the model.

Generative Adversarial Network (GAN) is a framework for training deep generative models using a two player minimax game. A generator network GG generates samples by transforming a noise variable zp(z)z\sim p(z) into G(z)G(z). A discriminator network D(x)D({{\bf x}}) is trained to distinguish between “fake” samples from the generator and “real” samples from a given data distribution pdp_{d}. Formally, this defines the following objective (Wasserstein GAN, from )

In our setting, we could assume pd(x)p_{d}(x) is the empirical distribution from the samples, and choose zπ0z\sim\pi^{0} and let Gθ(z)G_{\theta}(z) be the state of the Markov Chain after tt steps, which is a good approximation of πθ\pi_{\theta} if tt is large enough. However, optimization is difficult because we do not know a reasonable tt in advance, and the gradient updates are expensive due to backpropagation through the entire chain.

Therefore, we propose a more efficient approximation, called Markov GAN (MGAN):

We use two types of samples from the generator for training, optimizing θ\theta such that the samples will fool the discriminator:

Samples obtained after bb transitions xˉπθb\bar{x}\sim\pi_{\theta}^{b}, starting from x0π0x_{0}\sim\pi^{0}.

Samples obtained after mm transitions, starting from a data sample xdpdx_{d}\sim p_{d}.

Intuitively, the first condition encourages the Markov Chain to converge towards pdp_{d} over relatively short runs (of length bb). The second condition enforces that pdp_{d} is a fixed point for the transition operator. We provide a more rigorous justification in Appendix B. Instead of simulating the chain until convergence, which will be especially time-consuming if the initial Markov chain takes many steps to mix, the generator would run only (b+m)/2(b+m)/2 steps on average. Empirically, we observe better training times by uniformly sampling bb from [1,B][1,B] and mm from [1,M][1,M] respectively in each iteration, so we use BB and MM as the hyperparameters for our experiments.

We experiment with a distribution pdp_{d} over images, such as digits (MNIST) and faces (CelebA). In the experiments, we parametrize fθf_{\theta} to have an autoencoding structure, where the auxiliary variable vN(0,I)v\sim\mathcal{N}(0,I) is directly added to the latent code of the network serving as a source of randomness:

where β\beta is a hyperparameter we set to 0.10.1. While sampling is inexpensive, evaluating probabilities according to Tθ(xt)T_{\theta}(\cdot|x_{t}) is generally intractable as it would require integration over vv. The starting distribution π0\pi_{0} is a factored Gaussian distribution with mean and standard deviation being the mean and standard deviation of the training set. We include all the details, which ares based on the DCGAN architecture, in Appendix E.1. All the models are trained with the gradient penalty objective for Wasserstein GANs , where λ=1/3\lambda=1/3, B=4B=4 and M=3M=3.

We visualize the samples generated from our trained Markov chain in Figures 1 and 3, where each row shows consecutive samples from the same chain (we include more images in Appendix F) From Figure 1 it is clear that xt+1x_{t+1} is related to xtx_{t} in terms of high-level properties such as digit identity (label). Our model learns to find and “move between the modes” of the dataset, instead of generating a single sample ancestrally. This is drastically different from other iterative generative models trained with maximum likelihood, such as Generative Stochastic Networks (GSN, ) and Infusion Training (IF, ), because when we train Tθ(xt+1xt)T_{\theta}(x_{t+1}|x_{t}) we are not specifying a particular target for xt+1x_{t+1}. In fact, to maximize the discriminator score the model (generator) may choose to generate some xt+1x_{t+1} near a different mode.

To further investigate the frequency of various modes in the stationary distribution, we consider the class-to-class transition probabilities for MNIST. We run one step of the transition operator starting from real samples where we have class labels y{0,,9}y\in\{0,\ldots,9\}, and classify the generated samples with a CNN. We are thus able to quantify the transition matrix for labels in Figure 3. Results show that class probabilities are fairly uniform and range between 0.090.09 and 0.110.11.

Although it seems that the MGAN objective encourages rapid transitions between different modes, it is not always the case. In particular, as shown in Figure 3, adding residual connections and highway connections to an existing model can significantly increase the time needed to transition between modes. This suggests that the time needed to transition between modes can be affected by the architecture we choose for fθ(xt,v)f_{\theta}(x_{t},v). If the architecture introduces an information bottleneck which forces the model to “forget” xtx_{t}, then xt+1x_{t+1} will have higher chance to occur in another mode; on the other hand, if the model has shortcut connections, it tends to generate xt+1x_{t+1} that are close to xtx_{t}. The increase in autocorrelation will hinder performance if samples are used for Monte Carlo estimates.

Adversarial Training for Markov Chain Monte Carlo

We now consider the setting where the target distribution pdp_{d} is specified by an analytic expression:

where U(x)U(x) is a known “energy function” and the normalization constant in Equation (5) might be intractable to compute. This form is very common in Bayesian statistics , computational physics and graphics . Compared to the setting in Section 3, there are two additional challenges:

We want to train a Markov chain such that the stationary distribution πθ\pi_{\theta} is exactly pdp_{d};

We do not have direct access to samples from pdp_{d} during training.

We use ideas from the Markov Chain Monte Carlo (MCMC) literature to satisfy the first condition and guarantee that {πθt}t=0\{\pi^{t}_{\theta}\}_{t=0}^{\infty} will asymptotically converge to pdp_{d}. Specifically, we require the transition operator Tθ(x)T_{\theta}(\cdot|x) to satisfy the detailed balance condition:

for all xx and xx^{\prime}. This condition can be satisfied using Metropolis-Hastings (MH), where a sample xx^{\prime} is first obtained from a proposal distribution gθ(xx)g_{\theta}(x^{\prime}|x) and accepted with the following probability:

Therefore, the resulting MH transition kernel can be expressed as Tθ(xx)=gθ(xx)Aθ(xx)T_{\theta}(x^{\prime}|x)=g_{\theta}(x^{\prime}|x)A_{\theta}(x^{\prime}|x) (if xxx\neq x^{\prime}), and it can be shown that pdp_{d} is stationary for Tθ(x)T_{\theta}(\cdot|x) .

The idea is then to optimize for a good proposal gθ(xx)g_{\theta}(x^{\prime}|x). We can set gθg_{\theta} directly as in Equation (1) (if fθf_{\theta} takes a form where the probability gθg_{\theta} can be computed efficiently), and attempt to optimize the MGAN objective in Eq. (3) (assuming we have access to samples from pdp_{d}, a challenge we will address later). Unfortunately, Eq. (7) is not differentiable - the setting is similar to policy gradient optimization in reinforcement learning. In principle, score function gradient estimators (such as REINFORCE ) could be used in this case; in our experiments, however, this approach leads to extremely low acceptance rates. This is because during initialization, the ratio gθ(xx)/gθ(xx)g_{\theta}(x|x^{\prime})/g_{\theta}(x^{\prime}|x) can be extremely low, which leads to low acceptance rates and trajectories that are not informative for training. While it might be possible to optimize directly using more sophisticated techniques from the RL literature, we introduce an alternative approach based on volume preserving dynamics.

2 Hamiltonian Monte Carlo and Volume Preserving Flow

To gain some intuition to our method, we introduce Hamiltonian Monte Carlo (HMC) and volume preserving flow models . HMC is a widely applicable MCMC method that introduces an auxiliary “velocity” variable vv to gθ(xx)g_{\theta}(x^{\prime}|x). The proposal first draws vv from p(v)p(v) (typically a factored Gaussian distribution) and then obtains (x,v)(x^{\prime},v^{\prime}) by simulating the dynamics (and inverting vv at the end of the simulation) corresponding to the Hamiltonian

where xx and vv are iteratively updated using the leapfrog integrator (see ). The transition from (x,v)(x,v) to (x,v)(x^{\prime},v^{\prime}) is deterministic, invertible and volume preserving, which means that

MH acceptance (7) is computed using the distribution p(x,v)=pd(x)p(v)p(x,v)=p_{d}(x)p(v), where the acceptance probability is p(x,v)/p(x,v)p(x^{\prime},v^{\prime})/p(x,v) since gθ(x,vx,v)/gθ(x,vx,v)=1g_{\theta}(x^{\prime},v^{\prime}|x,v)/g_{\theta}(x,v|x^{\prime},v^{\prime})=1. We can safely discard vv^{\prime} after the transition since xx and vv are independent.

Let us return to the case where the proposal is parametrized by a neural network; if we could satisfy Equation 9 then we could significantly improve the acceptance rate compared to the “REINFORCE” setting. Fortunately, we can design such an proposal by using a volume preserving flow model .

and can be optimized by maximum likelihood.

In the case of a volume preserving flow model ff, the determinant of the Jacobian f(h)h\frac{\partial f(h)}{\partial h} is one. Such models can be constructed using additive coupling layers, which first partition the input into two parts, yy and zz, and then define a mapping from (y,z)(y,z) to (y,z)(y^{\prime},z^{\prime}) as:

where m()m(\cdot) can be an expressive function, such as a neural network. By stacking multiple coupling layers the model becomes highly expressive. Moreover, once we have the forward transformation ff, the backward transformation f1f^{-1} can be easily derived. This family of models are called Non-linear Independent Components Estimation (NICE).

3 A NICE Proposal

HMC has two crucial components. One is the introduction of the auxiliary variable vv, which prevents random walk behavior; the other is the symmetric proposal in Equation (9), which allows the MH step to only consider p(x,v)p(x,v). In particular, if we simulate the Hamiltonian dynamics (the deterministic part of the proposal) twice starting from any (x,v)(x,v) (without MH or resampling vv), we will always return to (x,v)(x,v).

Auxiliary variables can be easily integrated into neural network proposals. However, it is hard to obtain symmetric behavior. If our proposal is deterministic, then fθ(fθ(x,v))=(x,v)f_{\theta}(f_{\theta}(x,v))=(x,v) should hold for all (x,v)(x,v), a condition which is difficult to achieve The cycle consistency loss (as in CycleGAN ) introduces a regularization term for this condition; we added this to the REINFORCE objective but were not able to achieve satisfactory results.. Therefore, we introduce a proposal which satisfies Equation (9) for any θ\theta, while preventing random walk in practice by resampling vv after every MH step.

Our proposal considers a NICE model fθ(x,v)f_{\theta}(x,v) with its inverse fθ1f_{\theta}^{-1}, where vp(v)v\sim p(v) is the auxiliary variable. We draw a sample xx^{\prime} from the proposal gθ(x,vx,v)g_{\theta}(x^{\prime},v^{\prime}|x,v) using the following procedure:

Randomly sample vp(v)v\sim p(v) and uUniformu\sim\text{Uniform};

If u>0.5u>0.5, then (x,v)=fθ(x,v)(x^{\prime},v^{\prime})=f_{\theta}(x,v);

If u0.5u\leq 0.5, then (x,v)=fθ1(x,v)(x^{\prime},v^{\prime})=f_{\theta}^{-1}(x,v).

We call this proposal a NICE proposal and introduce the following theorem.

For any (x,v)(x,v) and (x,v)(x^{\prime},v^{\prime}) in their domain, a NICE proposal gθg_{\theta} satisfies

4 Training A NICE Proposal

Given any NICE proposal with fθf_{\theta}, the MH acceptance step guarantees that pdp_{d} is a stationary distribution, yet the ratio p(x,v)/p(x,v)p(x^{\prime},v^{\prime})/p(x,v) can still lead to low acceptance rates unless θ\theta is carefully chosen. Intuitively, we would like to train our proposal gθg_{\theta} to produce samples that are likely under p(x,v)p(x,v).

Although the proposal itself is non-differentiable w.r.t. xx and vv, we do not require score function gradient estimators to train it. In fact, if fθf_{\theta} is a bijection between samples in high probability regions, then fθ1f_{\theta}^{-1} is automatically also such a bijection. Therefore, we ignore fθ1f_{\theta}^{-1} during training and only train fθ(x,v)f_{\theta}(x,v) to reach the target distribution p(x,v)=pd(x)p(v)p(x,v)=p_{d}(x)p(v). For pd(x)p_{d}(x), we use the MGAN objective in Equation (3); for p(v)p(v), we minimize the distance between the distribution for the generated vv^{\prime} (tractable through Equation (10)) and the prior distribution p(v)p(v) (which is a factored Gaussian):

where LL is the MGAN objective, LdL_{d} is an objective that measures the divergence between two distributions and γ\gamma is a parameter to balance between the two factors; in our experiments, we use KL divergence for LdL_{d} and γ=1\gamma=1 The results are not very sensitive to changes in γ\gamma; we also tried Maximum Mean Discrepancy (MMD, see for details) and achieved similar results..

Our transition operator includes a trained NICE proposal followed by a Metropolis-Hastings step, and we call the resulting Markov chain Adversarial NICE Monte Carlo (A-NICE-MC). The sampling process is illustrated in Figure 4. Intuitively, if (x,v)(x,v) lies in a high probability region, then both fθf_{\theta} and fθ1f_{\theta}^{-1} should propose a state in another high probability region. If (x,v)(x,v) is in a low-probability probability region, then fθf_{\theta} would move it closer to the target, while fθ1f_{\theta}^{-1} does the opposite. However, the MH step will bias the process towards high probability regions, thereby suppressing the random-walk behavior.

5 Bootstrap

The main remaining challenge is that we do not have direct access to samples from pdp_{d} in order to train fθf_{\theta} according to the adversarial objective in Equation (12), whereas in the case of Section 3, we have a dataset to get samples from the data distribution.

In order to retrieve samples from pdp_{d} and train our model, we use a bootstrap process where the quality of samples used for adversarial training should increase over time. We obtain initial samples by running a (possibly) slow mixing operator Tθ0T_{\theta_{0}} with stationary distribution pdp_{d} starting from an arbitrary initial distribution π0\pi_{0}. We use these samples to train our model fθif_{\theta_{i}}, and then use it to obtain new samples from our trained transition operator TθiT_{\theta_{i}}; by repeating the process we can obtain samples of better quality which should in turn lead to a better model.

6 Reducing Autocorrelation by Pairwise Discriminator

An important metric for evaluating MCMC algorithms is the effective sample size (ESS), which measures the number of “effective samples” we obtain from running the chain. As samples from MCMC methods are not i.i.d., to have higher ESS we would like the samples to be as independent as possible (low autocorrelation). In the case of training a NICE proposal, the objective in Equation (3) may lead to high autocorrelation even though the acceptance rate is reasonably high. This is because the coupling layer contains residual connections from the input to the output; as shown in Section 3.1, such models tend to learn an identity mapping and empirically they have high autocorrelation.

We propose to use a pairwise discriminator to reduce autocorrelation and improve ESS. Instead of scoring one sample at a time, the discriminator scores two samples (x1,x2)(x_{1},x_{2}) at a time. For “real data” we draw two independent samples from our bootstrapped samples; for “fake data” we draw x2Tθm(x1)x_{2}\sim T_{\theta}^{m}(\cdot|x_{1}) such that x1x_{1} is either drawn from the data distribution or from samples after running the chain for bb steps, and x2x_{2} is the sample after running the chain for mm steps, which is similar to the samples drawn in the original MGAN objective.

The optimal solution would be match both distributions of x1x_{1} and x2x_{2} to the target distribution. Moreover, if x1x_{1} and x2x_{2} are correlated, then the discriminator should be able distinguish the “real” and “fake” pairs, so the model is forced to generate samples with little autocorrelation. More details are included in Appendix D. The pairwise discriminator is conceptually similar to the minibatch discrimination layer ; the difference is that we provide correlated samples as “fake” data, while provides independent samples that might be similar.

To demonstrate the effectiveness of the pairwise discriminator, we show an example for the image domain in Figure 5, where the same model with shortcut connections is trained with and without pairwise discrimination (details in Appendix E.1); it is clear from the variety in the samples that the pairwise discriminator significantly reduces autocorrelation.

Experiments

Code for reproducing the experiments is available at https://github.com/ermongroup/a-nice-mc.

To demonstrate the effectiveness of A-NICE-MC, we first compare its performance with HMC on several synthetic 2D energy functions: ring (a ring-shaped density), mog2 (a mixture of 2 Gaussians) mog6 (a mixture of 6 Gaussians), ring5 (a mixture of 5 distinct rings). The densities are illustrated in Figure 6 (Appendix E.2 has the analytic expressions). ring has a single connected component of high-probability regions and HMC performs well; mog2, mog6 and ring5 are selected to demonstrate cases where HMC fails to move across modes using gradient information. A-NICE-MC performs well in all the cases.

We use the same hyperparameters for all the experiments (see Appendix E.4 for details). In particular, we consider fθ(x,v)f_{\theta}(x,v) with three coupling layers, which update vv, xx and vv respectively. This is to ensure that both xx and vv could affect the updates to xx^{\prime} and vv^{\prime}.

We evaluate and compare ESS and ESS per second (ESS/s) for both methods in Table 1. For ring, mog2, mog6, we report the smallest ESS of all the dimensions (as in ); for ring5, we report the ESS of the distance between the sample and the origin, which indicates mixing across different rings. In the four scenarios, HMC performed well only in ring; in cases where modes are distant from each other, there is little gradient information for HMC to move between modes. On the other hand, A-NICE-MC is able to freely move between the modes since the NICE proposal is parametrized by a flexible neural network.

We use ring5 as an example to demonstrate the results. We assume π0(x)=N(0,σ2I)\pi_{0}(x)=\mathcal{N}(0,\sigma^{2}I) as the initial distribution, and optimize σ\sigma through maximum likelihood. Then we run both methods, and use the resulting particles to estimate pdp_{d}. As shown in Figures 7(a) and 7(b), HMC fails and there is a large gap between true and estimated statistics. This also explains why the ESS is lower than 1 for HMC for ring5 in Table 1.

Another reasonable measurement to consider is Gelman’s R hat diagnostic , which evaluates performance across multiple sampled chains. We evaluate this over the rings5 domain (where the statistics is the distance to the origin), using 32 chains with 5000 samples and 1000 burn-in steps for each sample. HMC gives a R hat value of 1.26, whereas A-NICE-MC gives a R hat value of 1.002 For R hat values, the perfect value is 1, and 1.1-1.2 would be regarded as too high.. This suggest that even with 32 chains, HMC does not succeed at estimating the distribution reasonably well.

Does training increase ESS?

We show in Figure 8 that in all cases ESS increases with more training iterations and bootstrap rounds, which also indicates that using the pairwise discriminator is effective at reducing autocorrelation.

Admittedly, training introduces an additional computational cost which HMC could utilize to obtain more samples initially (not taking parameter tuning into account), yet the initial cost can be amortized thanks to the improved ESS. For example, in the ring5 domain, we can reach an ESS of 121.54121.54 in approximately 550550 seconds (2500 iterations on 1 thread CPU, bootstrap included). If we then sample from the trained A-NICE-MC, it will catch up with HMC in less than 2 seconds.

Next, we demonstrate the effectiveness of A-NICE-MC on Bayesian logistic regression, where the posterior has a single mode in a higher dimensional space, making HMC a strong candidate for the task. However, in order to achieve high ESS, HMC samplers typically use many leap frog steps and require gradients at every step, which is inefficient when xU(x)\nabla_{x}U(x) is computationally expensive. A-NICE-MC only requires running fθf_{\theta} or fθ1f_{\theta}^{-1} once to obtain a proposal, which is much cheaper computationally. We consider three datasets - german (25 covariates, 1000 data points), heart (14 covariates, 532 data points) and australian (15 covariates, 690 data points) - and evaluate the lowest ESS across all covariates (following the settings in ), where we obtain 5000 samples after 1000 burn-in samples. For HMC we use 40 leap frog steps and tune the step size for the best ESS possible. For A-NICE-MC we use the same hyperparameters for all experiments (details in Appendix E.5). Although HMC outperforms A-NICE-MC in terms of ESS, the NICE proposal is less expensive to compute than the HMC proposal by almost an order of magnitude, which leads to higher ESS per second (see Table 2).

Discussion

To the best of our knowledge, this paper presents the first likelihood-free method to train a parametric MCMC operator with good mixing properties. The resulting Markov Chains can be used to target both empirical and analytic distributions. We showed that using our novel training objective we can leverage flexible neural networks and volume preserving flow models to obtain domain-specific transition kernels. These kernels significantly outperform traditional ones which are based on elegant yet very simple and general-purpose analytic formulas. Our hope is that these ideas will allow us to bridge the gap between MCMC and neural network function approximators, similarly to what “black-box techniques” did in the context of variational inference .

Combining the guarantees of MCMC and the expressiveness of neural networks unlocks the potential to perform fast and accurate inference in high-dimensional domains, such as Bayesian neural networks. This would likely require us to gather the initial samples through other methods, such as variational inference, since the chances for untrained proposals to “stumble upon” low energy regions is diminished by the curse of dimensionality. Therefore, it would be interesting to see whether we could bypass the bootstrap process and directly train on U(x)U(x) by leveraging the properties of flow models. Another promising future direction is to investigate proposals that can rapidly adapt to changes in the data. One use case is to infer the latent variable of a particular data point, as in variational autoencoders. We believe it should be possible to utilize meta-learning algorithms with data-dependent parametrized proposals.

Acknowledgements

This research was funded by Intel Corporation, TRI, FLI and NSF grants 1651565, 1522054, 1733686. The authors would like to thank Daniel Lévy for discussions on the NICE proposal proof, Yingzhen Li for suggestions on the training procedure and Aditya Grover for suggestions on the implementation.

References

Appendix A Estimating Effective Sample Size

Assume a target distribution p(x)p(x), and a Markov chain Monte Carlo (MCMC) sampler that produces a set of N correlated samples {xi}1N\{x_{i}\}_{1}^{N} from some distribution q({xi}1N)q(\{x_{i}\}_{1}^{N}) such that q(xi)=p(xi)q(x_{i})=p(x_{i}). Suppose we are estimating the mean of p(x)p(x) through sampling; we assume that increasing the number of samples will reduce the variance of that estimate.

Let V=Varq[i=1Nxi/N]V=\operatorname{Var}_{q}[\sum_{i=1}^{N}x_{i}/N] be the variance of the mean estimate through the MCMC samples. The effective sample size (ESS) of {xi}1N\{x_{i}\}_{1}^{N}, which we denote as M=ESS({xi}1N)M=ESS(\{x_{i}\}_{1}^{N}), is the number of independent samples from p(x)p(x) needed in order to achieve the same variance, i.e. Varp[j=1Mxj/M]=V\operatorname{Var}_{p}[\sum_{j=1}^{M}x_{j}/M]=V. A practical algorithm to compute the ESS given {xi}1N\{x_{i}\}_{1}^{N} is provided by:

where ρs\rho_{s} denotes the autocorrelation under qq of xx at lag ss. We compute the following empirical estimate ρ^s\hat{\rho}_{s} for ρs\rho_{s}:

where μ^\hat{\mu} and σ^\hat{\sigma} are the empirical mean and variance obtained by an independent sampler.

Due to the noise in large lags ss, we adopt the approach of where we truncate the sum over the autocorrelations when the autocorrelation goes below 0.05.

Appendix B Justifications for Objective in Equation 3

We consider two necessary conditions for pdp_{d} to be the stationary distribution of the Markov chain, which can be translated into a new algorithm with better optimization properties, described in Equation 3.

Consider a sequence of ergodic Markov chains over state space S\mathcal{S}. Define πn\pi_{n} as the stationary distribution for the nn-th Markov chain, and πnt\pi_{n}^{t} as the probability distribution at time step tt for the nn-th chain. If the following two conditions hold:

b>0\exists b>0 such that the sequence {πnb}n=1\{\pi^{b}_{n}\}_{n=1}^{\infty} converges to pdp_{d} in total variation;

ϵ>0\exists\epsilon>0, ρ<1\rho<1 such that M>0,n>M\exists M>0,\forall n>M, if πntpdTV<ϵ\lVert\pi_{n}^{t}-p_{d}\lVert_{\text{TV}}<\epsilon, then \lVert\pi_{n}^{t+1}-p_{d}\lVert_{\text{TV}}<\rho\|\pi_{n}^{t}-p_{d}\lVert_{\text{TV}}\;

then the sequence of stationary distributions {πn}n=1\{\pi_{n}\}_{n=1}^{\infty} converges to pdp_{d} in total variation.

The goal is to prove that δ>0\forall\delta>0, K>0,T>0\exists K>0,T>0, such that n>N,t>T\forall n>N,t>T, πntpdTV<δ\lVert\pi_{n}^{t}-p_{d}\lVert_{\text{TV}}<\delta.

According to the first assumption, N>0\exists N>0, such that n>N\forall n>N, πnbpdTV<ϵ\lVert\pi_{n}^{b}-p_{d}\lVert_{\text{TV}}<\epsilon.

Therefore, n>K=max(N,M)\forall n>K=\max(N,M), δ>0\forall\delta>0, T=b+max(0,logρδlogρϵ)+1\exists T=b+\max(0,\lceil\log_{\rho}\delta-\log_{\rho}\epsilon\rceil)+1, such that t>T\forall t>T,

The first inequality uses the fact that πnbpdTV<ϵ\lVert\pi_{n}^{b}-p_{d}\lVert_{\text{TV}}<\epsilon (from Assumption 1), and πnt+1pdTV/πntpdTV<ρ\lVert\pi_{n}^{t+1}-p_{d}\lVert_{\text{TV}}/\|\pi_{n}^{t}-p_{d}\lVert_{\text{TV}}<\rho (from Assumption 2). The second inequality is true because ρ<1\rho<1 by Assumption 2. The third inequality uses the fact that Tb>logρδlogρϵT-b>\lceil\log_{\rho}\delta-\log_{\rho}\epsilon\rceil (from definition of TT), so ρTb<δ/ϵ\rho^{T-b}<\delta/\epsilon. Hence the sequence {πn}n=1\{\pi_{n}\}_{n=1}^{\infty} converges to pdp_{d} in total variation. ∎

Moreover, convergence in total variation distance is equivalent to convergence in Jensen-Shannon (JS) divergence, which is what GANs attempt to minimize . This motivates the use of GANs to achieve the two conditions in Proposition 1. This suggests a new optimization criterion, where we look for a θ\theta that satisfies both conditions in Proposition 1, which translates to Equation 3.

Appendix C Proof of Theorem 1

For any (x,v)(x,v) and (x,v)(x^{\prime},v^{\prime}), gg satisfies:

Theorem 1 allows us to use the ration p(x,v)/p(x,v)p(x^{\prime},v^{\prime})/p(x,v) when performing the MH step.

Appendix D Details on the Pairwise Discriminator

Similar to the settings in MGAN objective, we consider two chains to obtain samples:

Starting from a data point xx, sample z1z_{1} in BB steps.

Starting from some noise zz, sample z2z_{2} in BB steps; and from z2z_{2} sample z3z_{3} in MM steps.

For the “generated” (fake) data, we use two type of pairs (x,z1)(x,z_{1}), and (z2,z3)(z_{2},z_{3}). This is illustrated in Figure 9. We assume equal weights between the two types of pairs.

Appendix E Additional Experimental Details

Code is available at https://github.com/ermongroup/markov-chain-gan.

Let ‘fc nn, (activation)’ denote a fully connected layer with nn neurons. Let ‘conv2d nn, kk, ss, (activation)’ denote a convolutional layer with nn filters of size kk and stride ss. Let ‘deconv2d nn, kk, ss, (activation)’ denote a transposed convolutional layer with nn filters of size kk and stride ss.

We use the following model to generate Figure 1 (MNIST).

We use the following model to generate Figure 3 (CelebA, top)

For the bottom figure in Figure 3, we add a residual connection such that the input to the second layer of the decoder is the sum of the outputs from the first layers of the decoder and encoder (both have shape 16×16×6416\times 16\times 64); we add a highway connection from input image to the output of the decoder:

where xˉ\bar{x} is the output of the function, x^\hat{x} is the output of the decoder, and α\alpha is an additional transposed convolutional output layer with sigmoid activation that has the same dimension as x^\hat{x}.

We use the following model to generate Figure 5 (CelebA, pairwise):

For the pairwise discriminator, we double the number of filters in each convolutional layer. According to , we only use batch normalization in the generator for all experiments.

E.2 Analytic Forms of Energy Functions

Let f(xμ,σ)f(x|\mu,\sigma) denote the log pdf of N(μ,σ2)\mathcal{N}(\mu,\sigma^{2}).

where μ1=\mu_{1}=, μ2=\mu_{2}=, σ1=σ2=[0.5,0.5]\sigma_{1}=\sigma_{2}=[0.5,0.5].

where μi=[siniπ3,cosiπ3]\mu_{i}=[\sin\frac{i\pi}{3},\cos\frac{i\pi}{3}] and σi=[0.5,0.5]\sigma_{i}=[0.5,0.5].

The analytic form of U(x)U(x) for ring5 is:

where ui=(x12+x22i)2/0.04u_{i}=(\sqrt{x_{1}^{2}+x_{2}^{2}}-i)^{2}/0.04.

E.3 Benchmarking Running Time

Since the runtime results depends on the type of machine, language, and low-level optimizations, we try to make a fair comparison between HMC and A-NICE-MC on TensorFlow .

Our code is written and executed in TensorFlow 1.0. Due to the optimization of the computation graphs in TensorFlow, the wall-clock time does not seem to be exactly linear in some cases, even when we force the program to use only 1 thread on the CPU. The wall-clock time is affected by 2 aspects, batch size and number of steps. We find that the wall-clock time is relatively linear with respect to the number of steps, and not exactly linear with respect to the batch size.

Given a fixed number of steps, the wall-clock time is constant when the batch size is lower than a threshold, and then increases approximately linearly. To perform speed benchmarking on the methods, we select the batch size to be the value around the threshold, in order to prevent significant under-estimates of the efficiency.

We found that the graph is much more optimized if the batch size is determined before execution. Therefore, we perform all the benchmarks on the optimized graph where we specify a batch size prior to running the graph. For the energy functions, we use a batch size of 2000; for Bayesian logistic regression we use a batch size of 64.

E.4 Hyperparameters for the Energy Function Experiments

For all the experiments, we use same hyperparameters for both A-NICE-MC and HMC. We sample x0N(0,I)x_{0}\sim\mathcal{N}(0,I) and run the chain for 1000 burn-in steps and evaluate the samples from the next 1000 steps.

For HMC we use 40 leapfrog steps and a step size of 0.1. For A-NICE-MC we consider fθ(x,v)f_{\theta}(x,v) with three coupling layers, which updates vv, xx and vv respectively. The motivation behind this particular architecture is to ensure that both xx and vv could affect the updates to xx^{\prime} and vv^{\prime}. In each coupling layer, we select the function m()m(\cdot) to be a one-layer NN with 400 neurons. The discriminator is a three layer MLP with 400 neurons each. Similar to the settings in Section 3.1, we use the gradient penalty method in to train our model.

For bootstrapping, we first collect samples by running the NICE proposal over the untrained fθf_{\theta}, and for every 500 iterations we replace half of the samples with samples from the latest trained model. All the models are trained with AdaM for 20000 iterations with B=4B=4, M=2M=2, batch size of 32 and learning rate of 10410^{-4}.

E.5 Hyperparameters for the Bayesian Logistic Regression Experiments

For HMC we tuned the step size parameter to achieve the best ESS possible on each dataset, which is 0.0050.005 for german, 0.010.01 for heart and 0.01150.0115 for australian (HMC performance on australian is extremely sensitive to the step size). For A-NICE-MC we consider f(x,v)f(x,v) with three coupling layers, which updates vv, xx and vv respectively; we set vv to have 50 dimensions in all the experiments. m()m(\cdot) is a one-layer NN with 400 neurons for the top and bottom coupling layer, and a two-layer NN with 400 neurons each for the middle layer. The discriminator is a three layer MLP with 800 neurons each. We use the same training and bootstrapping strategy as in Appendix E.4. All the models are trained with AdaM for 20000 iterations with B=16B=16, M=2M=2, batch size of 32 and learning rate of 5×1045\times 10^{-4}.

E.6 Architecture Details

The following figure illustrates the architecture details of fθ(x,v)f_{\theta}(x,v) for A-NICE-MC experiments. We do not use batch normalization (or other normalization techniques), since it slows the execution of the network and does not provide much ESS improvement.

Appendix F Extended Images

We only displayed a small number of images in the main text due to limited space. Here we include the extended version of images for our image generation experiments.

F.2 Extended Images for Figure 3

The following models are trained with the original MGAN objective (without pairwise discriminator).

F.3 Extended Images for Figure 5

The following images are trained on the same model with shortcut connections.