Approximate Inference with Amortised MCMC
Yingzhen Li Richard E. Turner Qiang Liu



A Message Passing Approach for Decision Fusion in Adversarial Multi-Sensor Networks
    [
    December 30, 2023
===================================================================================



 
We propose a novel approximate inference framework that approximates a target distribution by amortising the dynamics of a user-selected Markov chain Monte Carlo (MCMC) sampler. 
The idea is to initialise MCMC using samples from an approximation network, apply the MCMC operator to improve these samples, and finally use the samples to update the approximation network thereby improving its quality.




This provides a new generic framework for approximate inference, allowing us to deploy highly complex, or implicitly defined approximation families with intractable densities,  
including approximations produced by warping a source of randomness through a deep neural network. Experiments consider Bayesian neural network classification and image modelling with deep generative models. Deep models trained using amortised MCMC are shown to generate realistic looking samples as well as producing diverse imputations for images with regions of missing pixels.





§ INTRODUCTION


Probabilistic modelling provides powerful tools for analysing and making future predictions from data. The Bayesian paradigm offers well-calibrated uncertainty estimates on unseen data, by representing the variability of model parameters given the current observations through the posterior distribution. 

However, Bayesian methods are typically computationally expensive, due to the intractability of 
evaluating posterior or marginal probabilities. 
This is especially true for complex models like neural networks, for which a Bayesian would treat all the weight matrices as random variables and integrate them out. Hence approximations have to be applied to overcome this computational intractability in order to make Bayesian methods practical for modern machine learning tasks.


This work considers the problem of amortised inference, 
in which we approximate a given intractable posterior distribution p with a sampler q, a distribution from which we can draw exact samples. Compared with the typical Monte Carlo (MC) which approximates p with a fixed set of samples, 
the amortised inference distributes the computation cost to training the sampler q. 
This allows us to quickly generate a large number of samples at the testing time, and 
can significantly save time and storage when inference is required repeatedly as inner loops of other algorithms, such as in training latent variable models and structured prediction. 

Variational inference <cit.> and its stochastic variants <cit.> provide a straightforward approach for amortised inference, 
in which we find an optimal q from a parametric family of distributions 𝒬 
by minimizing certain divergence measure (often KL divergence) D[q||p]. 
Unfortunately, except a few very recent attempts <cit.>, most existing variational approaches require the distributions in 𝒬 to have  
computationally tractable density functions in order to solve the optimization problem. 
This forms a major restriction on the choice of the approximation set 𝒬, since 
exact samplers, in the most general form, 



are random variables of form 
x = f(ϵ), where f is a (non-linear) transform function, 
and ϵ is some standard distribution such as Gaussian distribution. Except simple cases, e.g. when f is linear, 
it is difficult to explicitly calculate the density function of such random variables.  


Therefore, a key challenge is to develop efficient approximate inference algorithms using generic samplers, which we refer as wild approximations, without needing to calculate the density functions.
. 
Such algorithms would allow us to 
deploy more flexible families of approximate distributions to obtain better posterior approximations. 




In this paper we develop a new, simple principle for wild variational inference based on amortising MCMC dynamics. 
Our method deploys a student-teacher, or actor-critic framework 
to leverage existing MCMC dynamics to a supervisor for training samplers q, by iterating the following steps: 






  (1) the sampler q (student) generates initial samples which are shown to an MCMC sampler;

  (2) the MCMC sampler (teacher) improves the samples by running MCMC transitions;

  (3) the sampler q takes feedback from the teacher and adjust itself in order to generate the improved samples next time. 



This framework is highly generic, works for arbitrary sampler families, and can take the advantage of any existing MCMC dynamics for efficient amortised inference. 
Empirical results show that our method works efficiently on Bayesian neural networks and deep generative modelling. 











§ BACKGROUND





  
Bayesian inference 
Consider a probabilistic model p(x|z, θ) along with a prior distribution p_0(z),
where x denotes an observed variable, z an 
unknown latent variable, 
and θ a hyper-parameter that 
is assumed to be given, or will be learned by 
maximizing the marginal likelihood log p( x | θ). 

The key computational task of interest is to approximate the posterior distribution of z: 

    p(z|x, θ) = 1/p(x|θ) p_0(z) p(x|z, θ), 
           
    p(x|θ) = ∫ p_0(z) p(x|z, θ) d z.

This includes both drawing samples from p(z|x, θ) in order to estimate related average quantities, 
as well as estimating the normalization constant p(x|θ) for hyper-parameter optimisation or model selection.  
Since both the data x and θ are assumed to be fixed in inference, we may drop the dependency on them when it is clear from the context. 




  
MCMC basics
MCMC provides a powerful, flexible framework for 
drawing (approximate) samples from given distributions.  
An MCMC algorithm is typically specified by its transition kernel 𝒦(z' |  z)
whose unique stationary distribution equals the target distribution p of interest, that is,   

    q = p 
                        
    q(z) = ∫ q(z') 𝒦(z|z') dz',     ∀z.

This fixed point equation fully characterizes the target distribution p,
and hence the 
 inference regarding p can be framed as (approximately) solving equation (<ref>). 

 In particular, MCMC algorithms 
can be viewed as stochastic approximations for solving (<ref>), 
in which we start with drawing z_0 from an initial distribution q_0 
and iteratively draw sample z_t at the t-th iteration 
from the transition kernel conditioned on the previous state, 

i.e. 
z_t | z_t - 1∼𝒦(z_t | z_t-1). 
In this way, the distribution q_t of z_t can be viewed as obtained by a fixed point update of form 

    q_t(z) ←𝒦 q_t-1 ( z),       where   𝒦 q_t-1( z) :=  ∫ q_t-1(z') 𝒦(z|z') dz',

so that recursively, we have q_t = 𝒦_t q_0, where 𝒦_t denotes the t-step  transition kernel. 





The standard theory of Markov chains suggests that 
the Markov transition monotonically decreases the KL divergence (<cit.>, see also Lemma 1 in the appendix), that is, 

    D_KL[q_t ||  p] ≤D_KL [q_t-1 ||  p].

Therefore, q_t converges to the stationary distribution p as t→∞ under proper conditions.  





























§ AMORTISED MCMC


MCMC can be viewed as approximating 
the fixed point update (<ref>) 
in a non-parametric fashion, 
returning a set of fixed samples for approximating p. 
This motivates amortised MCMC which uses more general parametric approximations of the fixed point update (<ref>) to train parametric samplers for amortised inference. 
In the sequel, we introduce our generic framework in Section <ref>, 
discuss in Section <ref> some particular algorithmic choices, 

and apply amortised MCMC to approximate maximum likelihood estimation (MLE) in Section <ref>. 










 §.§ Main idea: learning to distil MCMC
 
Let 𝒬 = {q_ϕ}
be a set of candidate samplers parametrised by ϕ. 
Our goal is to find an optimal q_ϕ to closely approximate the posterior distribution p of interest. 
We achieve this by approximating the fixed point update (<ref>). 
Because of the parametrisation, 
an additional projection step is required to maintain q inside 𝒬, 
motivating the following update rule at the t-th iteration: 

    ϕ_t_ϕ D[ q_T ||    q_ϕ],         
    q_T := 𝒦_T q_ϕ_t-1.

where D[· || ·] is 
some divergence measure between distributions
whose choice is discussed in Section <ref>. 
Note that we extended (<ref>) to use the T-step transition kernel 𝒦_T.  
If 𝒬 is taken to be large enough so that 𝒦_T q_ϕ_t-1∈𝒬, 
then the projection update (<ref>) (with T=1) reduces to (<ref>). 
In practice, a gradient descent method can be used to solve (<ref>): 

    ϕ_t←ϕ_t-1 - η∇_ϕD[ 𝒦_T q_ϕ_t-1 ||    q_ϕ]|_ϕ = ϕ_t-1.

It is often intractable to evaluate ∇_ϕD[ q_T  ||    q_ϕ] thus an approximation is needed. 
This can be done by approximating 
 q_ϕ_t-1 with samples { z^k_0} drawn from it, 
 and approximating q_T = 𝒦_T q_ϕ_t-1
 with sample { z^k_T} drawn 
 by following the Markov transition 𝒦(· | ·) for T times starting at { z^k_0}. 
These samples are then used to 
estimate the gradient and update ϕ_t  by (<ref>) 
in order to “move” q_ϕ_t-1 towards q_T = 𝒦_T q_ϕ_t-1, 
which is closer to the target distribution according to (<ref>). 


















































To summarise, our generic framework requires three main ingredients: 




	
  (1) a parametric set 𝒬 = {q_ϕ} of the sampler distributions (the student);
	
  (2) an MCMC dynamics with kernel 𝒦(z_t | z_t-1) (the teacher);
	
  (3) a divergence D[·||·] and update rule for ϕ (the feedback). 



By selecting these components tailored to a specific approximate inference task, 

the method provides a highly generic framework, applicable to both continuous and discrete distribution cases, and as we shall see later, extensible to wild approximations without a tractable density.



  
Remark
MCMC can be viewed as a special case of our framework, in which the samplers 𝒬
are empirical distributions parametrised by the MCMC samples z_t, that is, q_ϕ_t( z)  = δ( z- ϕ_t),   
ϕ_t =  z_t, where the sample z_t is treated as the parameter ϕ_t in our framework, 
and is the only possible output of the sampler q_ϕ_t.  
Our framework allows more flexible parametrisations of samplers, which significantly saves running time and storage at test time. 




  
Remark
The same type of projected fixed point updates as (<ref>) 
have been widely used in reinforcement learning (RL), 

including deep Q learning (DQN) <cit.>, 
and temporal difference learning with function approximations in general <cit.>.
In this scenario the Q- or V- networks 
are iteratively adjusted by applying projected fixed point updates of the Bellman equation. 
This provides an opportunity to 
strengthen our method with the vast RL literature. 


For example, similar to the case of RL, 
the convergence of updates of form (<ref>) are not theoretically guaranteed in general, especially when the parametric set 𝒬 
 is complex or non-convex.
However, the practical stabilisation tricks developed in the DQN literature <cit.> can be potentially 
 applied to our case, 
 and theoretical analysis developed in RL <cit.>
 can be borrowed to establish convergence of our method  under simple assumptions (e.g., when 𝒬 is linear or convex). 
 











 §.§ The choice of update rule
 
The choice of the update signal and ways to estimate it plays a crucial role in our framework, 
and it should be carefully selected to ensure both strong discrimination power and 
computational tractability with respect to the parametric family of q that we use. 
We discuss three update rules in the following.




  
KL divergence minimisation A simple approach is to use the inclusive KL-divergence in (<ref>): 

    D_KL[q_T || q_ϕ] = 𝔼_q_T[ log q_T(z|x) - log q_ϕ(z|x) ],

and for the purpose of optimising ϕ, it only requires an MC estimate of -𝔼_q_T[log q_ϕ(z|x)]. 
This gives a simple algorithm that is a hybrid of MCMC and VI, and appears to be new to our knowledge. Similar algorithms include the so called cross entropy method <cit.> which replaces q_T with an importance weighted distribution, and methods for tuning proposal distributions for sequential Monte Carlo <cit.>.



  
Adversarially estimated divergences 
Unfortunately, the inclusive KL divergence requires the density of the sampler q_ϕ to be evaluated, 
and can not be used directly for wild approximations.
Instead, we need to estimate the divergences based on samples { z_0^k}∼ q_ϕ and { z_T^k}∼ q_T. 




To address this, we borrow the 

idea of generative adversarial networks (GAN) <cit.> to construct a sample-based estimator of the selected divergence. As an example, consider the Jensen-Shannon divergence:
D_JS[q_T||q] = 1/2D_KL[q_T || q̃] + 1/2D_KL[q || q̃],
with q̃ = 1/2 q + 1/2 q_T.
Since none of the three distributions have tractable density, a discriminator d_ψ(z|x) is trained to provide a stochastic lower-bound 

    D_JS[q_T||q] ≥D_adv[{z_T^k } || {z_0^k } ] = 1/K∑_k=1^K logσ(d_ψ(z_T^k | x))
    + 1/K∑_k=1^K log (1 - σ(d_ψ(z_0^k | x))),

with σ(·) the sigmoid function and z_0^k, z_T^k samples from q and q_T, respectively. 

Recent work <cit.> extends adversarial training to f-divergences where the two KL-divergences are special cases in that rich family. In such case D_adv also corresponds to the variational lower-bound to the selected f-divergence and the discriminator can also be defined accordingly. Furthermore, the density ratio estimation literature <cit.> suggests that the discriminator d_ψ in (<ref>) can be used to estimate log (q_T/q_ϕ), i.e. the objective function for q_ϕ could be decoupled from that for the discriminator <cit.>.



  
Energy matching 
An alternative approach to matching
q_ϕ with 
 q_T is to match their first M moments. 
In particular, matching the mean and variance is equivalent to minimising D_KL[q_T || q_ϕ] with q_T fixed and q_ϕ a Gaussian distribution. However it is difficult to know beforehand which moments are important for accurate approximations. Instead we propose energy matching, which matches the expectation of the log of the joint distribution p( x,  z) ∝ p(z| x) under q_ϕ and q_T:

    D_em[q_T  ||  q_ϕ]  = | 𝔼_q_T(z_T|x)[ log p(x, z_T) ] - 𝔼_q_(z|x)[ log p(x, z) ] |^β,   β > 0,
 

Although D_em[· || · ] is not a valid divergence measure since D_em[q_T  ||  q_ϕ]=0 does not imply q_ϕ = q_T, 
this construction puts emphasis on moments that are mostly important for accurate predictive distribution approximation, which is still good for inference. 
Furthermore, as (<ref>) can be approximated with Monte Carlo methods which only require samples from q, this energy matching objective can be applied to wild variational inference settings as it does not require density evaluation of q_ϕ and q_T. Another motivation from contrastive divergence <cit.> is also discussed in the appendix.



 §.§ Approximate MLE with amortised MCMC

Learning latent variable models have become an important topic with increasing interest. Our amortised inference method can be used to develop more flexible and accurate approximations for learning.  

Consider the variational auto-encoder (VAE) <cit.> which approximates maximum likelihood training (MLE) by maximising the following variational lower-bound over the model parameters θ and variational parameters ϕ:

    max_θ,  ϕ{𝔼_q_ϕ[ log p(x | z; θ) ] - D_KL[q_ϕ(z|x)||p_0(z)]
    = log p(x|θ) - D_KL[q_ϕ(z|x) || p(z| x, θ) ] }.

Our method can be directly used to update ϕ in (<ref>). 
This can be done using the inclusive KL divergence if the density  q_ϕ is tractable, 
and adversarially estimated divergence or energy matching for wild approximation when the density q_ϕ is intractable.  




Next we turn to the optimisation of the hyper-parameters θ where we decouple their objective function from that of ϕ. 

Because D_KL[q || p] ≥D_KL[q_T || p] when p is the stationary distribution of the MCMC, 

the following  objective  forms a tighter lower-bound to the marginal likelihood: 

    log p(x|θ) - D_KL[q_T(z|x) || p(z| x, θ) ] = 𝔼_q_T[ log p(x | z, θ) ] + const of θ.


Empirical evidences <cit.> suggested tighter lower-bounds often lead to better results. Monte Carlo estimation is applied to estimate the lower-bound (<ref>) with samples {z_T^k }∼ q_T.

The full method when using adversarially estimated divergences 
is presented in Algorithm <ref>, in which we train a discriminator d_ϕ to estimate the selected divergence, and propagate learning signals back through the samples from q_ϕ. 
 Note here the update step for the discriminator and the q_ψ could be executed for more iterations in order to achieve better approximations to the current posterior. This strategy turns the algorithm into a stochastic EM with MCMC methods approximating the E-step <cit.>. 
RKHS-based and energy-based moments <cit.> can also be applied as the discrepancy measure in step 2, but this is not explored here. 















§ RELATED WORK

Since <cit.>, generative adversarial networks (GANs) have attracted large attention from both academia and industry. 

We should distinguish the problem scope of our work with that of GANs: 
amortised MCMC aims to match an (implicitly defined) q to the posterior distribution p, 
while GANs aim to match a q with an observed sample, which we leverage as an inner loop for the divergence minimisation. Hence our framework could also benefit from the recent advances in this area <cit.>.

The amortisation framework is in similar spirit to <cit.> in that both approaches “distil” an MCMC sampler with a parametric model. Unlike the presented framework, <cit.> and <cit.> used a student model to approximate the predictive likelihood, and that student model is not used to initialise the MCMC transitions. We believe that initialising MCMC with the student model is important in amortising dynamics, as the teacher can “monitor” the student's progress and provide learning signals tailored to the student's need. Moreover, since the initialisation is improved after each student update, the quality of the teacher's samples also improves. Another related, but different approach <cit.> considered speeding-up Hybrid Monte Carlo <cit.> by approximating the transition kernel using a Gaussian process. Amortised MCMC could benefit from this line of work if the MCMC updates are too expensive.  

Perhaps the most related approaches to our framework (in the sense of using q of flexible forms) are operator variational inference (OPVI, <cit.>), amortised SVGD  <cit.>, and adversarial variational Bayes (AVB <cit.>, also concurrently proposed by <cit.>). These works assumed the q_ϕ distribution to be represented by a neural network warping input noise. OPVI minimises the Stein discrepancy <cit.> between the exact and approximate posterior, where the optimal test function is determined by optimising a discriminator. 
Though theoretically appealing, this method seems still impractical for large scale problems. 
Amortised SVGD can be viewed as a special case of our framework, which specifically uses a deterministic Stein variational gradient dynamic <cit.> and an l_2-norm as the divergence measure.  


AVB estimates the KL-divergence D_KL[q||p_0] in the variational lower-bound (<ref>) with GAN and density ratio estimation, making it closely related to the adversarial auto-encoder <cit.>. However we conjecture that the main learning signal of AVB comes from the “reconstruction error” term 𝔼_q[log p(x|z, θ)], and the regularisation power strongly depends on the adversarial estimation of D_KL[q||p_0], which can be weak as the discriminator is non-optimal in almost all cases. 






§ EXPERIMENTS

We evaluate amortised MCMC with both toy and real-world examples. For simplicity we refer the proposed framework as AMC. In the appendix experimental settings are further presented. Code will be released at <https://github.com/FirstAuthor/AmortisedMCMC>. 



 §.§ Synthetic example: fitting a mixture of Gaussians


We first consider fitting a Gaussian mixture 
p(z) = 1/2𝒩(z; -3, 1) + 1/2𝒩(z; 3, 1) with the variational program proposed by <cit.> as the following:
ϵ_1, ϵ_2, ϵ_3 ∼𝒩(ϵ; 0, 1), z = 1_ϵ_3 ≥ 0ReLU(w_1 ϵ_1 + b_1) - 1_ϵ_3 < 0ReLU(w_2 ϵ_2 + b_2). We further tested a small multi-layer perceptron (MLP) model of size [3, 20, 20, 1] which warps ϵ to generate the samples. 

The Jensen-Shannon divergence is adversarially estimated with an MLP of the same architecture. The MCMC sampler is Langevin dynamics with rejection (MALA <cit.>), and in training 10 parallel chains are used.

The fitted approximations are visualised in Figure <ref>. Both models cover both modes, however the variational program performs better in terms of estimating the variance of each Gaussian component. Thus an intelligent design of the q network can achieve better performance with much fewer number of parameters.

We empirically investigate the effect of the chain length T on the approximation quality using the MLP approximation. Time steps T=1, 5, 10 with step-sizes η = 0.1, 0.02, 0.01 are tested (each repeating 10 times), where by making Tη = 1.0 a constant, the particles approximately move equal distances during MCMC transitions. In Figure <ref> we shown the Kernel Stein Discrepancy (KSD <cit.>) as a metric of approximation error. With small chain length the student quickly learns the posterior shape, but running more MCMC transitions results in better approximation accuracy. A potential way to balance the time-accuracy trade-off is to initially use short Markov chains initially for AMC, but to lengthen them as AMC converges. This strategy has been widely applied to contrastive-divergence like methods <cit.>. We leave the exploration of this idea to future work.





 §.§ Bayesian neural network classification


Next we apply amortised MCMC to classification using Bayesian neural networks. Here the random variable z denotes the neural network weights which could be thousands of dimensions. In such high dimensions a large number of MCMC samples are typically required for good approximation. Instead we consider AMC as an alternative, which allows us to use much fewer samples during training as it decouples the samples from evaluation, leading to massive savings of time and storage. To validate this, we take from the UCI repository <cit.> 7 binary classification datasets and train a 50-unit Bayesian neural network. For comparison we test the mean-field Gaussian approximation trained by VI with 10 samples, and MALA with 100 particles simulated and stored. The approximated posterior for AMC is constructed by first taking a mean-field Gaussian approximation, then normalising the K=10 samples by their empirical mean and variance. This wild approximation is trained with the energy matching objective with β = 2, where we also use MALA as the dynamics with T=1.

The results are reported in Table <ref>. For test log-likelihood MALA is generally better than VI as expected. More importantly, AMC performs similarly, and for some datasets even better, than MALA. VI returns the best test error metrics on three of the datasets, with AMC and MALA topped for the rest. In order to demonstrate the speed-accuracy trade-off we visualise in Figure <ref> the (negative) test-LL and error as a function of CPU time. In this case we further test MALA with 10 samples, which is much faster than the 100 samples version, but it pays the price of slightly worse results. However AMC achieves better accuracy in a smaller amount of time, and in general out-performs the 10-sample MALA. These observations show that AMC with energy matching can be used to train Bayesian neural works and achieves a balance between computational complexity and performance.





 §.§ Deep generative models

The final experiment considers training deep generative models on the dynamically binarised MNIST dataset, containing 60,000 training datapoints and 10,000 test images <cit.>. For benchmark a convolutional VAE with dim(z) = 32 latent variables is tested. The Gaussian encoder consists of a convolutional network with 5 × 5 filters, stride 2 and [16, 32, 32] feature maps, followed by a fully connected network of size [500, 32 × 2]. The generative model has a symmetric architecture but with stride convolution replaced by deconvolution layers. This generative model architecture is fixed for all the tests. We also test AMC with inclusive KL divergence on Gaussian encoders, and compare to the naive approach which trains the encoder by maximising variational lower-bound (MCMC-VI). 

We construct two non-Gaussian encoders for AVB and AMC (see appendix). Both encoders start from a CNN followed by a reshaping operation. Then the first model (CNN-G) splits the CNN's output vector into [h(x), μ(x), logσ(x)], samples a Gaussian noise ϵ∼𝒩(ϵ; μ(x), diag[σ^2(x)]), and feeds [h(x), ϵ] to an MLP for generating z.
The second encoder (CNN-B) simply applies multiplicative Bernoulli noise with dropout rate 0.5 to the CNN output vector, and uses the same MLP architecture as CNN-G. The discriminator consists of a CNN acting on the input image x only, and an MLP acting on both z and the CNN output. Batch normalisation is applied to non-Gaussian encoders and the discriminator. The learning rate of Adam <cit.> is tuned on the last 5000 training images. Rejection steps are not used for Langevin dynamics as we found it slows down the learning.






  
Test Log-likelihood Results
We report the test log-likelihood (LL) results in Table <ref>. We first follow <cit.> to estimate the test log-likelihood with importance sampling using 5000 samples, and for the non-Gaussian encoders we train another Gaussian encoder with VI as the proposal distribution.

VAE appears to be the best method by this metric, and the best AMC model is about 2nats behind. However, effective sample size results (IW-ESS) show that the estimation results for AMC and AVB are unreliable.

Indeed approximate MLE using the variational lower-bound biases the generative network towards the model whose exact posterior is close to the inference network q <cit.>. As the MCMC-guided approximate MLE trains the generative model with q_T (which could be highly non-Gaussian), the VI-fitted Gaussian proposal, employed in the IWAE, can under-estimates the true test log-likelihood by a significant amount.

To verify the conjecture, we estimate the test-LL again but using Hamiltonian annealed importance sampling (HAIS) as suggested by <cit.>. We randomly select 1,000 test images for evaluation, and run HAIS with 10,000 intermediate steps and 100 parallel chains. Estimation results demonstrate that IW-LL significantly under-estimates the test LL for models trained by wild approximations. In this metric the CNN-G model with T=50 performs the best, which is significantly better than benchmark VAE. To demonstrate the improvement brought by the wild approximation we further train a generative model with “persistent MCMC”, by initialising the Markov chain with previous samples and ignoring the posterior changes. The HAIS-LL results shows that out best model is about 0.6nats better, which is a significant improvement on MNIST.

Although test log-likelihood is an important measures of model quality, <cit.> has shown that this metric is often largely orthogonal to one that tests visual fidelity when the data is high dimensional. We visualise the generated images in Figure <ref>, and we see that AMC trained models produce samples of similar quality to VAE samples. 



  
Missing Data Imputation
We also consider missing data imputation with pixels missing from contiguous sections of the image, i.e. not at random. 

We follow <cit.> using an approximate Gibbs sampling procedure for imputation. With observed and missing pixels denoted as x_o and x_m, the approximate sampling procedure iteratively applies the following transition steps: (1) sample z∼ q(z|x_o, x_m) given the imputation x_m, and (2) sample x^* ∼ p(x^*|z, θ) and set x_m ←x_m^*.

In other words, the encoder q(z|x) is used to approximately generate samples from the exact posterior. As ambiguity exists, the exact conditional distribution p(x_m|x_o, θ) is expected to be multi-modal.

Figure <ref> visualises the imputed images, where starting from the third column the remaining ones show samples for every 2 Gibbs steps. Clearly the approximate Gibbs sampling for VAE is trapped in local modes due to the uni-modal approximation q to the exact posterior. On the other hand, models trained with AMC return diverse imputations that explore the space of compatible images quickly, for instance, CNN-B returns imputations for digit “9” with answers 2, 0, 9, 3 and 8. To quantify this, we simulate the approximate Gibbs sampling for T=100 steps on the first 100 test images (10 for each class), find the nearest neighbour (in l_1-norm) of the imputations in the training dataset, and compute the entropy on the label distribution over these training images. The entropy values and the average l_1-distance to the nearest neighbours (divided by the number of pixels dim(x) = 784) are presented in Table <ref>. These metrics indicate that AMC trained models generate more diverse imputations compared to VAE, yet these imputed images are about the same distance from the training data.




§ CONCLUSION AND FUTURE WORK

We have proposed an MCMC amortisation algorithm which deploys a student-teacher framework to learn the approximate posterior. By using adversarially estimated divergences and energy matching, the algorithm allows approximations of arbitrary form to be learned. Experiments on Bayesian neural network classification showed that the amortisation method can be used as an alternative to MCMC when computational resources are limited. Application to training deep generative networks returned models that could generate high quality images, and the learned approximation captured multi-modality in generation.

Future work should cover both theoretical and practical directions. Convergence of the amortisation algorithm will be studied. Flexible approximations will be designed to capture multi-modality. Efficient MCMC samplers should be applied to speed-up the fitting process. Practical algorithms for approximating discrete distributions will be further developed. 












plain




§ EXAMPLES OF WILD APPROXIMATIONS

We provide several examples of wild approximations in the following.




(Deterministic transform) 

Sampling z∼ q(z|x) is defined by first sampling some random noise ϵ∼ p(ϵ), then transforming it with a deterministic mapping z = f(ϵ, x), which might be defined by a (deep) neural network. These distributions are also called variational programs in <cit.>, or implicit models in the generative model context <cit.>. An important note here is that f might not be invertible, which differs from the invertible transform techniques discussed in <cit.>.



(Truncated Markov chain)

Here the samples z∼ q(z|x) are defined by finite-step transitions of a Markov chain. Examples include Gibbs sampling in contrastive divergence <cit.>, or finite-step simulation of an SG-MCMC algorithm such as SGLD <cit.>. It has been shown in <cit.> that the trajectory of SGD can be viewed as a variational approximation to the exact posterior. In these examples the variational parameters are the parameters of the transition kernel, e.g. step-sizes and/or preconditioning matrices. Related work includes <cit.> which integrates MCMC into VI objective. These methods are more expensive as they require evaluations of ∇_zlog p(x, z| θ), but they can be much cheaper than sampling from the exact posterior.




(Stochastic regularisation techniques (SRT))

SRT for deep neural network training, e.g. dropout <cit.> and related variants <cit.>, have been re-interpreted as a variational inference method for network weights z = {W} <cit.>. The variational parameters ϕ = {M} are the weight matrices of a Bayesian neural network without SRT. The output is computed as h = σ((ϵ⊙x) M), with σ(·) the activation function and ϵ some randomness. This is equivalent to setting W = diag(ϵ) M, making SRT a special case of example <ref>. Fast evaluation of q(z|x) during training is intractable as different noise values ϵ are sampled for different inputs in a mini-batch. This means multiple sets of weights are processed if we were to evaluate the density, which has prohibitive costs especially when the network is wide and deep.






§ APPROXIMATE MLE WITH MCMC: MATHEMATICAL DETAILS

In the main text we stated that D_KL[q||p] ≥D_KL[q_T||p] if q is the stationary distribution of the kernel 𝒦. This is a direct result of the following lemma, and we provide a proof from <cit.> for completeness.


<cit.>
Let q and r be two distributions for z_0. Let q_t and r_t be the corresponded distributions of state z_t at time t, induced by the transition kernel 𝒦. Then D_KL[q_t||r_t] ≥D_KL[q_t+1||r_t+1] for all t ≥ 0.





    D_KL[q_t||r_t]    = 𝔼_q_t[ logq_t(z_t)/r_t(z_t)] 
       = 𝔼_q_t(z_t) 𝒦(z_t+1|z_t) [ logq_t(z_t) 𝒦(z_t+1|z_t) /r_t(z_t) 𝒦(z_t+1|z_t)] 
       = 𝔼_q_t+1(z_t+1) q_t+1(z_t|z_t+1)[ logq_t+1(z_t+1) q(z_t|z_t+1) /r_t+1(z_t+1) r(z_t|z_t+1)] 
       = D_KL[q_t+1||r_t+1] + 𝔼_q_t+1D_KL[q_t+1(z_t | z_t+1) || r_t+1(z_t | z_t+1)].










§ ENERGY MATCHING AND CONTRASTIVE DIVERGENCE

The energy matching method in Section <ref> can also be roughly motivated by contrastive divergence <cit.> as follows. 


First define 

    Δ_CD[q_ || p]  :=D_KL[q_ || p] - D_KL[ q_T ||p],

where q_T =  𝒦_T q_. 
From Lemma <ref>, 
Δ_CD[q_ || p]  ≥ 0 and can be used as a minimisation objective to fit an approximate posterior q. 
Expanding this contrastive divergence equation, we have:

    Δ_CD[q_ || p]  
    
    	   = 𝔼_q_T(z_T|x)[ log p(x, z_T)  ] - 𝔼_q_ (z|x)[ log p(x, z)  ]   +  R ,

where 

    R := 𝔼_q_T(z|x) 𝒦_T(z_T | z)[ log q_T(z | z_T, x)/𝒦_T(z_T | z)],

and q_T(z | z_T, x) = q_(z|x) 𝒦_T(z_T|z) / q_ T(z_T|x) is the “posterior” of z given the sample z_T after T-step MCMC. We ignore the residual term R in when T is small as now  and _T are highly correlated, which motivates the energy matching objective (<ref>). For large T one can also use density ratio estimations methods to estimate the third term, but we leave the investigation to future work.





§ UNCORRELATED VERSUS CORRELATED SIMULATIONS


In our algorithm, we generate { z_T^k} by simulating Markov transition for T steps starting from { z_0^k}∼ q, and use  these two samples 
{ z_0^k} and { z_T^k} to estimate the divergence D[q_T ||q].  
However, note that 
{ z_0^k} and { z_T^k} is corrected, and this may introduce bias in the divergence estimation. A way to eliminate this correction is to simulate another {  z_0^k} independently and use it to replace the original sample. 


In practice, we find that the correlated/uncorrelated samples exhibit different behaviour during training. We consider the extreme case K=1 and small T as an example. Using correlated samples would cause the teacher and the student's samples remaining in the same mode with high probability and thus easily confuse the discriminator and the student (generator) improves fast. On the other hand, if z_T is simulated from a Markov chain independent with z_0, then these samples might be far away from each other (especially when q is forced to be multi-modal), hence the discriminator can easily get saturated, providing no learning signal to the student. 

The above problem could potentially be solved using advanced techniques, e.g. Wasserstein GAN <cit.> which proposed minimising (an adversarial estimate of) Wasserstein-1 distance. In that case the gradient of q won't saturate even when the two sets of samples are separated. But minimising Wasserstein distance would fit the q distribution to the posterior in an “optimal transport” way, which presumably prefers moving the q samples to their nearest modes in the exact posterior. 



§ EXPERIMENTAL DETAILS



 §.§ Bayesian neural networks: settings

We use a one-hidden-layer neural network with 50 hidden units and ReLU activation. The mini-batch size is 32 and the learning rate is 0.001. The step-size of MALA is adaptively adjusted to achieve acceptance rate 0.99. The experiments are repeated on 20 random splits of the datasets. We time all the tests on a desktop machine with Intel(R) Core(TM) i7-4930K CPU @ 3.40GHz.



 §.§ Deep generative models: settings

We construct two non-Gaussian encoders for tests of AVB and AMC. Both encoders start from a CNN with 3 × 3 filters, stride 2 and [32, 64, 128, 256] feature maps, followed by a reshaping operation. Then the first model (CNN-G) splits the output vector of the CNN into [h(x), μ(x), logσ(x)], samples a Gaussian noise variable of 32 dimensions ϵ∼𝒩(ϵ; μ(x), diag[σ^2(x)]), and feeds [h(x), ϵ] to an MLP which has hidden layer sizes [500, 500, 32].
The second encoder (CNN-B) simply applies multiplicative Bernoulli noise with dropout rate 0.5 to the CNN output vector, and uses the same MLP architecture as CNN-G. A discriminator is trained for AVB and AMC methods, with a CNN (of the same architecture as the encoders) acting on the input image x only, and a MLP of [32+1024, 500, 500, 1] layers. All networks use leaky ReLU activation functions with slope parameter 0.2, except for the output of the deconvolution network which uses sigmoid activation. Batch normalisation is applied to non-Gaussian encoders and the discriminator. The Adam optimiser <cit.> is used with learning rates tuned on the last 5000 training images. Rejection steps are not used in the Langevin dynamics as we found this slows down the learning.




§ MORE VISUALISATION RESULTS











