Year: 2018

Abstract:
Training probabilistic models with neural network components is intractable in most cases and requires to use approximations such as Markov chain Monte Carlo (MCMC), which is not scalable and requires significant hyper-parameter tuning, or mean-field variational inference (VI), which is biased. While there has been attempts at combining both approaches, the resulting methods have some important limitations in theory and in practice. As an alternative, we propose a novel method which is scalable, like mean-field VI, and, due to its theoretical foundation in ergodic theory, is also asymptotically accurate, like MCMC. We test our method on popular benchmark problems with deep generative models and Bayesian neural networks. Our results show that we can outperform existing approximate inference methods.

text: Training probabilistic models with neural network components is intractable in most cases and requires to use approximations such as Markov chain Monte Carlo (MCMC), which is not scalable and requires significant hyper-parameter tuning, or mean-field variational inference (VI), which is biased. While there has been attempts at combining both approaches, the resulting methods have some important limitations in theory and in practice. As an alternative, we propose a novel method which is scalable, like mean-field VI, and, due to its theoretical foundation in ergodic theory, is also asymptotically accurate, like MCMC. We test our method on popular benchmark problems with deep generative models and Bayesian neural networks. Our results show that we can outperform existing approximate inference methods.

heading: 1 INTRODUCTION
text: Approximate statistical inference with unnormalised density functions is fundamentally important problem both Bayesian and frequentist inference. In particular, the successes of many sophisticated generative models in machine learning rely on power inference algorithms. Markov chain Monte Carlo (MCMC) methods and variational inference (VI), originally developed in statistics and physics, are two most important approximate inference methods in machine learning, which has been widely used in all kinds of probabilistic models, like latent topic models (Blei et al., 2003), Boltzmann machines (Hinton, 2002; Salakhutdinov & Larochelle, 2010), Bayesian non-parametric models (Neal, 2000; Kurihara et al., 2007). However, they are facing great challenges in the recent research on probabilistic modelling with deep neural networks (NNs). In particular, Bayesian deep neural networks become popular in recent works, because it exploits the Bayesian framework to overcome the overfitting and data demanding problems in deep learning (Neal, 2012). Another interesting research direction is to use deep neural networks in latent variable models to transform the simple latent random variables into complex distribution, which is also known as to Deep Generative Models (DGM). DGMs have been proved to be very powerful generative models.
Inspired by DGMs, many recent works on variational inference adopted NN to construct flexible approximate distributions. In particular, variational autoencoders (Kingma & Welling, 2014) and normalising flows (Rezende & Mohamed, 2015) are two most influential works in this direction. However, due to lack of understanding of the convergence of specific NNs, the research of NNbased inference is focused on engineering the architecture of inference NNs based on experiments and heuristics.
In this work, we propose a novel approximate inference method based on the classic inference theory of MCMC. Our method is inspired by the idea of parallel simulations of MCMC and the recent advances in variational inference with NNs. Like these variational methods, it is straightforward to accelerate the computation of our method using parallelised simulations on Graphical Processing Units. More importantly, with solid theoretical foundations in the theory of MCMC, the proposed method guarantees asymptotic convergence to arbitrary distributions of interest. It is a great advantage over variational inference, because of the approximation bias in variational methods. Our method is also attractive to a wide range of probabilistic models without NNs and Bayesian NNs.

heading: 2 BACKGROUND

heading: 2.1 BAYESIAN NEURAL NETWORKS
text: Given data D = {xn, yn}Nn=1 formed by feature vectors xn and corresponding scalar targets yn, we can assume that each yn is obtained as yn = f(xn;θ) + n, where f(·;θ) is the output of a deep neural network with weights θ and the n are independent noise variables with n ∼ N (0, σ2). This model specifies a likelihood function p(y1, . . . , yn|x1, . . . ,xn,θ) which can be combined with a Gaussian prior on θ to obtain a posterior distribution p(θ|D). Predictions for the y? corresponding to a new feature vector x? are then obtained by using the predictive distribution p(y?|x?,D) =∫ p(y?|x?,θ)p(θ|D) dθ. However, integrating with respect to the posterior distribution p(θ|D) is intractable and approximations have to be performed in practice, with the most popular methods for this being VI and MCMC.

heading: 2.2 DEEP GENERATIVE MODELS
text: Generative models extract intrinsic structure from data by making use of latent variables. LetD be a dataset with n data points {xn}Nn=1. Given a latent representation z, the data point x is assumed to be sampled from the conditional distribution pθ(x|z), which is specified in terms of some parameters θ. This conditional distribution is often refered to as the decoder. Given a prior distribution p(z) over the latent variables, the joint distribution of data and latent variables is pθ(x, z) = p(z)pθ(x|z). The marginal probability of data x under the model is then pθ(x) = ∫ p(z)pθ(x|z)dz.
Until recenlty, pθ(x|z) was typically specified using a simple distributional family, e.g. generalized linear models (Murphy, 2012; Bishop, 2006). However, more recently, deep generative models (DGMs) use deep neural networks with weights θ to specify the decoder (Kingma & Welling, 2014; Goodfellow et al., 2014).
Maximum likelihood is a straightforward way to train probabilistic models with latent variables. However, the marginal likelihood pθ(x) is intractable to compute in DGMs and approximations are needed in practice. As before, the most popular methods for this are based on VI and MCMC. We briefly describe these methods in the following sections.

heading: 2.3 VARIATIONAL INFERENCE AND NORMALIZING FLOWS
text: VI approximates a complex posterior distribution with another simpler parametric distribution which is found by optimizing a lower bound on the marginal likelihood. Let the complex posterior be pθ(z|x), with associated marginal likelihood p(x), and let qφ(z), parameterized by φ, be a simpler tractable distribution. The lower bound of the marginal likelihood is then defined as
log pθ(x) ≥ Eqφ [log pθ(x, z)− log qφ(z)] , (1)
which is often known as the evidence lower bound (ELBO). The more flexible the parametric family qφ, the better the approximation quality to the true posterior and the tighter the value of the ELBO.
Mean-field VI uses a form for qφ(z) which assumes independence between random variables. This reduces the computational cost of the optimization problem but often leads to poor performance with complex posterior distributions, such as the ones arising in DGMs or in Bayesian neural networks.
Amortization can be used to accelerate convergence and reduce computational cost when multiple inference problems have to be solved simultaneously. The optimization of (1) can be amortized by making qφ depend explicitly on the data x. In this case, qφ(z) is replaced with qφ(z|x), where φ are now the weights of a neural network which computes the parameteres of a tractable parametric distribuion on z from x. In this manner, for any new value of x, we can readily obtain a correspoding variational approximation given by qφ(z|x). Variational auto-encoders (VAEs) (Kingma & Welling, 2014) are DGMs trained by using mean-field VI with a Gaussian parametric distribuion and amortization. Rezende & Mohamed (2015) improve over this method by using a more flexible variational family called normalizing flows (NFs). The NF family is obtained by applying L invertible non-linear transformations f1, . . . , fL to a random variable z0 with tractable density q0(z0) and exact simulation. The resulting output is a random
variable zL = fL ◦ · · · ◦ f1(z0) with density qφ(zL|x) = q0(z0|x) ∏L l=1 |∂fl(zl−1)/∂zl−1|
−1 and with φ being the parameters of f1, . . . , fL.
Stochastic gradient descent (SGD), in combination with the reparameterisation trick (Kingma & Welling, 2014), can be used for the scalable optimization of the ELBO in VAEs and NFs. However, the main limitation of VAEs and NFs is the bias present in their variational approximations. This bias can be quite high, even in the case of NFs, since the transformations f1, . . . , fL have to be rather simple to ensure invertibility and to reduce computational costs.

heading: 2.4 MARKOV CHAIN MONTE CARLO
text: Markov chain Monte Carlo (MCMC) is an approximate inference method which does not have the aforementioned bias problem. MCMC works by simulating a stationary Markov chain that generates asymptotically unbiased samples from the distributions of interest. Formally, a Markov chain is a sequence of random variables z0, z1, . . . in which the transition from state zl to zl+1 is defined by the conditional probability distribution of zl+1 given zl, denoted by K(zl, zl+1). Markov chains in MCMC methods have the following strong stationary property: if one state zl of the chain follows the stationary distribution π, so does the next state zl+1. In particular,
π(zl+1) = ∫ π(zl)K(zl, zl+1) dzl . (2)
If zl follows a distribution πl that is different from the stationary one, then the distribution of zl+1
πl+1(zl+1) = ∫ πl(zl)K(zl, zl+1) dzl , (3)
is guaranteed to be closer to π than πl. This property implies that, with sufficiently many transitions, the distribution πl of zl converges to π irrespectively of the distribution the initial state z0.
Despite beign asymptotically unbiased, MCMC methods are less popular than VI for two reasons. First, they are computationally more expensive and second, they typically include hyper-parameters in the Markov kernel K which are highly problem dependent and are hard to tune in practice.

heading: 3 ERGODIC MEASURE PRESERVING FLOWS
text: In this section, we describe an inference method that combines the strengths of MCMC and VI and avoids their drawbacks. The idea is to use the output distribution of a MCMC chain, given by (3), as the variational distribution and optimize a simple to evaluate objective function for tuning MCMC parameters. Since MCMC converges to the target asymptotically, our variational approximation can be arbitrarily accurate. The result is a computationally efficient method which avoids the bias of parametric approximations and which can do automatic tuning of hyper-parameters.

heading: 3.1 DEFINITIONS
text: Given the target distribution π with unnormalised density function π∗, we define an approximate distribution q by a mixture of sequential deterministic transformations that preserves the measure π∗. We call such approximate distributions measure preserving flows (MPFs). The transformations preserving a given measure π∗ are formally defined as follows (Billingsley, 1986). Definition 3.1. Measure Preserving Transformation (MPT). Let (Ω,F , P ) be a probability space and µ be a consistent measure with P . A mapping T : Ω→ Ω is a measure preserving transformation if T is measurable in both the input filed F and the output field F and µ(A) = µ(T−1(A)) for allA ∈ F . If T is a one-to-one mapping onto Ω, then T preserves µ: µ(A) = µ(T−1TA) = µ(TA).
An example of MPT is any transformation with Jacobian determinant equal to 1, which preserves the Lebesgue measure. In practice, it is straightforward to verify if a transformation T preserves the measure with density function π∗ with the following conditions:
(i) Bijection: T is invertible, (ii) Preservation of density function: π∗(z) = π∗(T (z)) for all z.
(iii) Preservation of base measure: the Jacobian determinant is one if the Lebesgue measure is the base measure.
In probability theory, MPTs are often used within the area of ergodic stochastic processes since many of these processes can be reformulated as a composition of MPTs. In particular, MCMC kernels are MPTs and stationary Markov chains in MCMC are ergodic processes (Robert & Casella, 2005).
The joint probability of states in a MCMC chain is q(z0, z1, . . . , zL) = q0(z0) ∏L l=1K(zl−1, zl), where q0 is the distribution of the initial state z0. The density of the last state zL is then obtained by integrating out all the previous states:
qL(zL) = ∫ q0(z0) L∏ l=1 K(zl−1, zl) dz0 dz1 . . . dzL−1 . (4)
If the Markov chain is ergodic, qL converges to the stationary distribution π in total variation distance as the length L of the chain increases (Robert & Casella, 2005).
We define a measure preserving flow (MPF) as a representation of (4) in which the kernelK becomes a deterministic transformation Tr : Z → Z with stochastic auxiliary input r following distribution µ. We can then define zL as the result of applying these deterministic transformations to z0, that is, zL = TrL ◦ · · · ◦ Tr1(z0). By following the rule of changing variables, it is then straightforward to derive the density of zL as
qL(zL) = ∫ q(zL, r1:L)dr1:L = ∫ q(z0)µ(r1:L)δ(zL − TrL ◦ · · · ◦ Tr1(z0)) dz0 dr1:L , (5)
where δ denotes the Dirac delta function. Note that there is no Jacobian term in (5) because of the preservation of the Lebesgue measure. Because MPFs are equivalent to ergodic Markov chains, the density obtained at the output of an MPF, that is, qL, will converge to the stationary distribution π as L increases.
Hamiltonian Monte Carlo (HMC) is one of the most successful MCMC methods, which also can be interpreted as an MPF. Given the target random variable z ∈ Rd with unnormalised density function π∗, the HMC kernel is essentially applying a deterministic transformationH to the previous state zi with an auxiliary random variable r ∈ Rn with the density function µ(·). The transformation H is given by the solution of Hamiltonian dynamics
ż(t) = ∂rK(r), ṙ(t) = −∂zU(z) , (6) where ż denotes the derivative of z w.r.t. the time of the dynamics t, U(z) = − log π∗(z) and K(r) = − logµ(r). The preservation of total Hamiltonian energy, H(z, r) = U(z) + K(r), is the characteristic property of Hamiltonian dynamics, which can be easily verified by Ḣ(z, r) = 0 using (6). It is straightforward to see that H preserves the joint measure π(z)µ(r). In particular, H is a bijective transformation because the dynamics are deterministic and time reversible (Neal, 2010) and the preservation of the Hamiltonian energy implies the preservation of density. Finally, the volume preservation of H in the space of (z, r) is a well known property of Hamiltonian dynamics, which can be proved by Liouville’s theorem (Leimkuhler & Reich, 2004; Neal, 2010).
We can write the marginal distribution of the last sample generated by HMC as an MPF:
q(zL) = ∫ q(zL, r1:L)dr1:L = ∫ q(z0)µ(r)δ(zL −HrL ◦ · · · ◦ Hr1(z0)) dz0 dr1:L . (7)
We call the MPF generated by Hamiltonian dynamics a Hamiltonian MPF (HMPF).

heading: 3.2 UNDERSTANDING MEASURE PRESERVING CONDITIONS
text: We would like to address a common misunderstanding on the preservation of volume condition stated by (iii) in Section 3.1. Note that we are interested in sampling the random variable z, but the Hamiltonian dynamics preserve the joint measure π(z, r) rather than π(z). Following the conditions of MPTs in Section 3.1, it seems necessary to show that, for any specific value of r, any Tr used within an MPF should preserve volume in z space. However, this is not the case since the measure preservation conditions in the augmented space (z, r) are enough to guarantee the preservation of the marginal distribution in z space. Formally, we have the following proposition:
Proposition 1. Let T : Z×E → Z×E preserve the distribution π(z, r). Then, if r is sampled from π(r) = ∫ π(z, r) dr, the marginal distribution
π(z) = ∫ π(z, r) dr
is also preserved by the projection of T in the space of z, that is, by Tr : Z → Z.
Proposition 1 gives us some insights on the difference between MPFs and normalising flows (NFs). As mentioned earlier, NFs also use a sequence of transformations Tr : Z → Z. However, these do not preserve the distribution of z and, consequently, they require the computation of Jacobian determinants by the rule of changing variables. By contrast, Proposition 1 implies that, in MPFs, Tr preserves the marginal π(z) if T preserves the joint, which means that there is no need to include any Jacobian computations. For this reason, the transformations used in MPFs can be much more complicated than those used in NFs. For example, in Hamiltonian MPFs, for a given r, the Jacobian ofHr can be very complicated.

heading: 3.3 VARIATIONAL INFERENCE WITH MPFS
text: Given an unormalized posterior distribution pθ(x, z) for z, we can construct an MPF that preserves pθ(x, z)µ(r), where µ(r) is a simple distribution with tractable density and sampling algorithm. Let Tφl be the l-th transformation in the flow, where φl are hyper-parameters. This transformation maps the state of the flow from (zl−1, rl) to (zl, r′l) = T
φl(zl−1, rl). Similarly, the composition TφL ◦ · · · ◦ Tφ1(z0, r1:L), which we denote by Tφ, transforms (z0, r1:L) to (zL, r′1:L). By the preservation of density and the preservation of Lebesgue measure of Tφ, as given by conditions (ii) and (iii) in Section 3.1), we have the following equalities
pθ(x, z0) L∏ l=1 µ(rl) = pθ(x, z1)µ(r ′ 1) L∏ l=2 µ(rl) = · · · = pθ(x, zL) L∏ l=1 µ(r′l), (8)
q0(z0) L∏ l=1 µ(rl) = qL(z1, r ′ 1;φ1) L∏ l=2 µ(rl) = · · · = qL(zL, r′1, r′2 . . . , r′L;φ) , (9)
where q0(z0) is an initial proposal distribution and φ = (φ1, . . . ,φL) are the transformation hyperparameters. It is important to clarify that, according to (9), the joint density of zL, r′1, r ′ 2, . . . , r ′ L is known, but the marginal density for these variables is intractable to compute in general.
Following (1), we can obtain the ELBO for the initial proposal distribution q0(z0) as L(x;θ) = ∫ log pθ(x, z0)
q0(z0) q0(z0) dz0 . (10)
We call this expression the simple ELBO. We can then multiplying by the density of the auxiliary variables µ(r1:L) = ∏L l=1 µ(rl) to obtain
L(x;θ) = ∫ log pθ(x, z0) ∏L l=1 µ(rl)
q0(z0) ∏L l=1 µ(rl) q0(z0) L∏ l=1 µ(rl) dz0 dr1:L . (11)
We can then replace (z0, r1:L) with (zL, r′1:L) in (11) by making use of using transformation T φ, (8) and (9). The result is
L(x;θ,φ) = ∫ log pθ(x, zL) ∏L l=1 µ(r ′ l)
qL(zL, r′1:L;φ) qL(zL, r
′ 1:L;φ) dzL dr ′ 1:L , (12)
where we have omitted the dependence of (zL, r′1:L) on φ, since these variables are determined by the hyper-parameters of the MPTs. We call (12) the reparameterised ELBO.

heading: 3.4 ERGODIC LOWER BOUND AND ERGODIC INFERENCE
text: The reparameterised ELBO is of limited use, because it can only be as tight as the ELBO with initial proposal distribution q0. This seems to erase the benefits of using an ergodic MPF, which we know
will converge to the target posterior distribution given a sufficiently long flow. To overcome the drawback of the reparameterised ELBO, we propose another ELBO tailored to the MPF framework, which becomes arbitrarily tight as the length of the flow grows. We call such an ELBO ergodic lower bound (ERLBO).
To derive ERLBO, we first rewrite (12) as L(x;θ,φ) = ∫ log pθ(x, zL)
qL(zL;φ) qL(zL;φ) dzL +
∫ log ∏L l=1 µ(r ′ l)
qL(r′1:L|zL;φ) qL(zL, r
′ 1:L;φ) dzL dr ′ 1:L
= ∫ log pθ(x, zL)
qL(zL;φ) qL(zL;φ) dzL −DLKL . (13)
where DLKL is the Kullback-Liebler divergence between qL(zL, r ′ 1:L;φ) and qL(zL;φ) ∏L l=1 µ(r ′ l). It is straightforward to show that the first term on the RHS in (13) is a lower bound of the marginal likelihood by Jensen’s inequality. This leads to the ERLBO given by
L̃(x;θ,φ) = ∫ log pθ(x, zL)
qL(zL;φ) qL(zL;φ) dzL = L(x;θ,φ) +DLKL . (14)
This is a tighter lower bound than the simple ELBO because the difference between L̃(x;θ,φ) and L(x;θ) in (10) isDLKL ≥ 0. Moreover, the ERLBO can be shown to monotonically increase w.r.t. L. Proposition 2. The lower bound L̃(x;θ,φ) in (14) becomes tighter and tighter as L increases, that is, L̃(x;θ,φ1:L) ≥ L̃(x;θ,φ1:L−1) and the equality holds if and only if DLKL = 0.
The complete proof is included in appendix.
Recall that qL(zL;φ) is obtained by a sequence of transformations that preserve the probability measure p(z|x). It is well-known that MCMC chains have a unique invariant distribution and so do the MCMC-equivalent MPFs. Therefore, we know that if L̃(x;θ,φ1:L) stops growing, qL(zL;φ) must converge to p(z|x). This is formally described by the following theorem. Theorem 1. Given an ergodic measure preserving flow with invariant measure π, the ergodic lower bound L̃(x;θ,φ) increases in the length of the flow L and becomes an unbiased estimator of the marginal log p(x) as L increases to infinity. We could tune φ by optimizing the ERLBO. To better understand the values of φ that would be favored by this optimisation process, we can rewrite the ERLBO by making explicit its dependence on the entropy of qL(zL;φ), which we denote by H[qL(zL;φ)] = − ∫ log qL(zL;φ)qL(zL;φ) dzL. In particular,
L̃(x;θ,φ) = EqL(zL;φ) [log pθ(x, zL)] + H[qL(zL;φ)] . (15) When optimizing this quantity w.r.t. φ, the first term in the RHS will encourage qL(zL;φ) to have high density in regions where pθ(x, zL) is high, while the second term will favor high entropy solutions and will prevent qL(zL;φ) from converging to a Dirac delta centered at the maximizer of log pθ(x, zL). Note that the first term in the RHS of (15) can be easily approximated by Monte Carlo, while the second term is intractable because qL(zL;φ) is not available. However, since qL(zL;φ) converges to pθ(z|x) as L increases, we expect the effect of H[qL(zL;φ)] on φ to be small and that most of the similarity of qL(zL;φ) to pθ(z|x) will be captured by the first term in the RHS of (15). Therefore, we propose to tune φ by optimizing the tractable objective given by the first term, that is,
F̃(x;θ,φ) = EqL(zL;φ) [log pθ(x, zL)] . (16)
If with the initial flow parameter φ0, the objective F̃ < Epθ(x,z) [log pθ(x, z)], we expect optimising F̃ produces faster convergence of qL(zL;φ) towards pθ(x, z). Importantly, the model parameters θ can also be adjusted by optimizing F̃ because the omitted H[qL(zL;φ)] does not depend on θ and, consequently, optimizing F̃(x;θ,φ) and L̃(x;θ,φ) w.r.t. θ are equivalent operations.

heading: 3.5 IMPLEMENTATION OF HMPFS
text: The hyper-parameters to be tuned in a HMPF include the parameters of q0 and the transformation parameters for the Hamiltonian simulation, that is, φ = (φ1, . . . ,φL). A natural choice for q0
is multivariate Gaussian with mean µ = (µ1 . . . , µd) and diagonal covariance matrix with entries σ2 = (σ21 , . . . , σ 2 d), where d is the dimensionality of the sample space. The most popular algorithm for simulating Hamiltonian dynamics is the vanilla Leapfrog integrator. We refer to the tutorial of Neal (2010) for more detailed description of the implementation of this algorithm. Leapfrog is a numeric integrator that approximates the Hamiltonian dynamics (6) by an iterative procedure with discretized time ∆t, that is,
x(t+ ∆t) = x(t) + ∆t∂rK(r(t)) , r(t+ ∆t) = r(t)−∆t∂xU(x(t)) . (17) For the flow parameters, we consider the total simulation time T . Given a fixed number of Leapfrog iterations m, the simulation time T can be reparameterized as the time step size ∆t = T/m. Neal (2010) shows that it is possible to use different ∆t for each dimension of the sample space to improve the quality of Leapfrog. Therefore, we consider the parameters of the l-th Hamiltonian simulation in the flow to be φl = (∆tl,1, . . . ,∆tl,d). The pseudo code for ergodic inference with HMPFs is shown in Algorithm 1.
Algorithm 1: Ergodic Inference on Hamiltonian Measure Preserving Flow. input : potential function U(z;x,θ), dataset D and large L output: optimal decoder and flow parameters θ∗ and φ∗
initialize θ and φ; while not converged do
x←− sample one data point from D; z0 ∼ N (µ, diag(σ2)); /* Start of simulation of HMPF */ for l = 1, . . . , L do
r ∼ N (0, 1); zl ←− H(zl−1, r;U(z;x,θ),φl); /* Leapfrog simulation */ end /* End of simulation of HMPF */ obj←− U(zL;x,θ); /* one sample Monte Carlo Approx. of F̃(x;θ,φ) */ θ ←− AdamUpdate(θ, ∂θobj); φ←− AdamUpdate(φ, ∂φobj);
end
We do not include any Metropolis-Hastings (MH) correction steps in our method since this is not necessary. The MH steps are included in MCMC methods to ensure asymptotic convergence to the correct target with an unlimited number of transitions. By contrast, MPFs are in the finite-length regime and the main concern is to accelerate the convergence of MPFs to be as close as possible to the target measure. In this setting, it is not immediately clear that MH steps would be helpful. In particular, in Section 4.1 we provide empirical evidence of how HMPFs can converge to the correct target without MH steps.
Due to the composition of the MPTs, computing the gradient can be expensive when the flow is long. To speed up training, we stop the gradient computations when evaluating −∂zU(z) in the Leapfrog steps. This trick leads to incorrect gradients. However, we noticed that the optimization was not significantlly affected by this and still worked very well in practice. Finally, note that working with incorrect gradients does not affect the convergence of the flow to the correct target distribution because that is guaranteed by the convergence of the ERLBO as we mentioned in previous section.

heading: 4 EXPERIMENTS
text: We provide empirical evidence of HMPFs in three inference tasks. Our goal is to show that HMPFs can provide better approximations than other approximate inference methods.

heading: 4.1 DEMONSTRATION OF CONVERGENCE
text: To verify the theoretical results on the convergence of MPFs in Section 3.4, we test HMPFs on 8 bivariate distributions. The full list of benchmark distributions and results are included in the
appendix. Here, we focus on two multimodal benchmarks. The first testing target distribution is a bimodal moon shaped distribution as shown in Figure 2a. We call this target dual moon. Dual moon is one of the benchmarks in normalising flows (Rezende et al., 2014). The second testing target is a mixture of 6 Gaussian distributions placed in a circle. We use 15 Hamiltonian transformations with 5 Leapfrog steps each. The architecture detail of HMPFs 1 can be found in Section D.1. The initial state of MPFs is sampled from a standard Gaussian distribution. The gradient of the objective function is estimated using 1000 samples from HMPFs.
To illustrate the convergence of HMPFs to the target distribution, figures 1c and 1d show histograms of samples as a function of the flow length and the training iterations. To confirm the convergence numerically, we compute the ERLBO using a numeric method for the estimation of the entropy. Plots for the ERLBO and the ground truth log normalization constants are show in figures 1a and 1b

heading: 4.2 DEEP GENERATIVE MODELS
text: MNIST is a standard benchmark for testing approximate inference algorithms for training deep generative models. This dataset contains 60,000 grey level 28×28 images of handwritten digits. For fair comparison with previous work, we use the 10,000 prebinarised MNIST test images from (Burda et al., 2015)2. Our benchmark deep generative model is based on the deconvolutional network used by Salimans et al. (2015) for testing Hamiltonian variational inference (HVI). In particular, the decoder pθ(x, z) consists of 32 dimensional latent variables z with isotropic Gaussian prior p(z) = N (0, I) and a deconvolutional network with the architecture from top to bottom including a single fully-connected layer with 500 RELU hidden units, then three deconvolutional layers with 5 × 5 filters, [16, 32, 32] feature maps and RELU activations and the final output layer is simply element-wise logistic activation functions. In the convolutional VAE, the encoder network mirrors the architecture of the decoder.
The code for HVI (Salimans et al., 2015) is not publicly available. Nevertheless, we reimplemented their convolutional VAE and were able to reproduce the marginal likelihood reported by Salimans et al. (2015), as shown in Table 1. This verifies that our implementation of the convolutional VAE is correct and that our results are comparable to the ones reported originally by Salimans et al. (2015). We also implemented HVI following Salimans et al. (2015). We used a single hidden layer network with 640 hidden units and RELU activations as the reverse model for the HMC transitions. We also implemented another VI method similar to HVI and called the Hamiltonian variational encoder (HVAE) (Caterini et al., 2018). Unlike HVI, HVAE does not use a reverse model. This method optimizes instead a bound derived from the stationary distribution of reverse momenta. Futhermore, HVAE uses tempering Hamiltonian dynamics that requires additional Jacobian corrections. In our implementation of HVAE, we simply ignore the temperature for computational efficiency.
1The code of HMPFs for all three experiments will be available at https://github.com/ firstauthor/hmpfs
2https://github.com/yburda/iwae
For HMPF encoder, we use 30 HMPTs with 5 Leapfrog steps per HMPT. The initial distribution q0 is 32 dimensional independent Gaussian. More detailed description of the architecture of HMPFs is in the appendix Section D.1. We optimise the HMPF encoder and the decoder jointly using Adam. The the initial state of the flow is sampled from independent Gaussian. The mean and variance of the initial Gaussian is also optimised jointly with flow and model parameters and their gradients are computed by back propagation given the value of momenta. However, we noticed that with sufficient number of transformations, the effect of optimising initial Gaussian distribution is not significant. Table 1 shows the marginal likelihood of HMPFs and other methods estimated using 100 Hamiltonian annealled importance samples (HAIS) (Wu et al., 2017). We also report the effective sample size (ESS). Overall, HMPFs produce better results and are faster than the baselines.
We also tested the same decoder with HMPFs and convolutional encoder on dynamically binarised Fashion-MNIST (Xiao et al., 2017). The results of test marginal likelihood can be found Table 2.

heading: 4.3 BAYESIAN NEURAL NETWORKS
text: In our final experiment we approximate the posterior distribution of Bayesian neural networks. We use four UCI datasets and compare HMPFs with relevant stochastic gradient Hamilton Monte Carlo (SGHMC) methods from (Springenberg et al., 2016). The networks used in these experiments have 50 hidden layers and 1 real valued output unit, as stated in (Springenberg et al., 2016). The HMPFs contain 50 HMC transformation with 3 Leapfrog steps each. The distribution the initial state of the flow is independent Gaussian with mean and variance parameters obtained by fitting a variational Gaussian proposal q0 with Adam optimiser for 200 iterations. To reduce the cost of the Leapfrog iterations, we split training data into 19 mini-batches and only use one random sampled mini-batch for computing the gradient of the potential energy. We train our HMPFs for 10 epochs and the stationary distribution of the flow is chosen as approximate posterior on a random sampled minibatch. The resulting test log-likelihoods are shown in Table 3. Overall, HMPFs produce significantly better results than SGHMC.

heading: 5 SUMMARY
text: We have proposed a novel method for approximate inference that combines advantages of variational inference and MCMC methods. We call this method ergodic measure preserving flows (EMPFs). Different from most previous works combining HMC and variational inference, EMPFs enjoy the same asymptotic convergence as HMC and can tune sampling hyper-parameters by optimizing a tractable objective function at a low computational cost. We have shown that EMPFs achieve better results than existing baselines on standard benchmarks. For future work, it will be interesting to study the convergence rate of EMPFs to the target distribution with increasing flow length. Finally, the proposed method can be easily extended to consider recent Riemannian-manifold HMC methods (Zhang & Sutton, 2014; Girolami & Calderhead, 2011) for the construction of the flow.

heading: A THE PROOF OF PROPOSITION 1
text: Proof. Because T preserves the probability measure π(x, ), then for any measurable set A in Borel set Bx × B , we have
π(A) = π(T−1A).
Given a set B ∈ Bx, the set Ax generated by B × E , is measurable under π with the measure
π(Ax) =
∫ π(B, )d ,
that is essentially a probability measure in the space of x, also known as the marginal probability
π(B) = π(Ax).
Because T preserves the joint measure, applying T on Ax gives
TAx = T (B, E) = T B × E ,
where T B denotes the projection of TAx in the space of x. Follow the definition of measure preserving transformations, we have
π(Ax) =
∫ π(B, )d = ∫ π(T−1 B, )d = π(T −1Ax),
where T−1 denotes the preimage of B under T . Because π(Ax) is essentially the marginal probability π(B), B ∈ Bx, we know the marginal distribution π(x) is preserved by T because
π(B) = π(Ax) = π(T −1Ax) = π(T −1 B).
This implies that the marginal distribution π(x) is preserved by the stochastic mapping
T : X → X, ∼ π( ), where π( ) = ∫ π(x, )dx. If T is invertible for any , we have
π(B) = ∫ π(B, )d = ∫ π(T B, )d = π(T B),
where T B denotes the projection of TAx in the space of x. Therefore, we know that if we sample from π(x, ) and apply T on x, then the marginal distribution π(x) is preserved.

heading: B THE PROOF OF PROPOSITION 2
text: Proof. Let µ be the probability measure of auxiliary variable R. Following equation 14, the difference of ERLBO in each measure preserving transformation TL is given by
dLz(x; θ, φ1:L) = Lz(x; θ, φ1:L)− Lz(x; θ, φ1:L−1)
= DKL ( QL ∣∣∣∣∣∣Q̂L)−DKL (QL−1∣∣∣∣∣∣Q̂L−1) , (18) where QL denotes QL(ZL,R′1:L) and Q̂L denotes QL(ZL) ∏L l=1 µ(R ′ l). The KL divergence can be written as the integral
DKL ( QL ∣∣∣∣∣∣Q̂L) = ∫ log qL(zL, r′1:L;φ) qL(zL;φ) ∏L l=1 µ(r ′ l(φ)) qL(zL, r ′ 1:L;φ)dzLdr ′ 1:L. (19)
Because (zL, r′L) is generated from (zL−1, rL) by the deterministic transformation TL that preserves Lebesgue measure in the phase space, we have the equality
qL(zL, r ′ 1:L) = qL−1(zL−1, r ′ 1:L−1)µ(rL).
Using the reparameterisation (zL, r ′ L) = T (zL−1, rL),
we can rewrite the second KL term in equation 18 as
DKL ( QL−1 ∣∣∣∣∣∣Q̂L−1) = ∫ log qL−1(zL−1, r′1:L−1;φ) qL−1(zL−1;φ) ∏L−1 l=1 µ(r ′ l(φ)) qL−1(zL−1, r ′ 1:L−1;φ)dzL−1dr ′ 1:L−1
= ∫ log
qL−1(zL−1, r ′ 1:L−1;φ)µ(rL) qL−1(zL−1;φ)µ(rL) ∏L−1 l=1 µ(r ′ l(φ)) qL−1(zL−1, r ′ 1:L−1;φ)µ(rL)dzL−1dr ′ 1:L−1drL
= ∫ log
qL(zL, r ′ 1:L;φ) qL(zL, r′L;φ) ∏L−1 l=1 µ(r ′ l(φ)) qL(zL, r ′ 1:L;φ)dzLdr ′ 1:L, (20)
where qL(zL, r′L;φ) comes from qL−1(zL−1;φ)µ(rL) by the rule of changing variables
qL−1(zL−1;φ)µ(rL) = ∫ qL−1(zL−1, r ′ 1:L−1;φ)µ(rL)dr ′ 1:L−1
= ∫ qL(zL, r ′ 1:L;φ)dr ′ 1:L−1
= qL(zL, r ′ L;φ).
Subtract equation 20 from equation 19, then we have the difference in KL as dLz(x; θ, φ1:L) = ∫ log hhhhhhqL(zL, r ′ 1:L;φ)
qL(zL;φ) ∏L l=1 µ(r ′ l(φ)) qL(zL, r ′ 1:L;φ)dzLdr ′ 1:L
− ∫ log
hhhhhhqL(zL, r ′ 1:L;φ) qL(zL, r′L;φ) ∏L−1 l=1 µ(r ′ l(φ)) qL(zL, r ′ 1:L;φ)dzLdr ′ 1:L
= ∫ log qL(zL, r ′ L;φ) XXXXXXX ∏L−1 l=1 µ(r ′ l(φ))
qL(zL;φ)µ(r′L(φ)) XXXXXXX ∏L−1 l=1 µ(r ′ l(φ)) qL(zL, r ′ 1:L;φ)dzLdr ′ 1:L
= ∫ log qL(zL, r ′ L;φ)
qL(zL;φ)µ(r′L(φ)) qL(zL, r
′ 1:L;φ)dzLdr ′ 1:L
= ∫ log qL(zL, r ′ L;φ)
qL(zL;φ)µ(r′L(φ)) qL(zL, r
′ L;φ)dzLdr ′ L. (21)
From equation 21, it is easy to see that dLz(x; θ, φ1:L) is essentially KL divergence between qL(zL, r ′ L;φ) and qL(zL;φ)µ(r ′ L(φ)). This verifies that dLz(x; θ, φ1:L) = Lz(x; θ, φ1:L)− Lz(x; θ, φ1:L−1) ≥ 0. This implies that the ergodic lower bound stops increasing if and only if the KL divergence between qL(zL, r ′ L;φ) and qL(zL;φ)µ(r ′ L(φ)) is equal to 0.

heading: C PROOF OF THEOREM 1
text: Proof. Because the flow is ergodic with invariant distribution p(z|x), with sufficient many transformations, qL(z) can converge to p(z|x) in total variation distance, that is
lim L→∞
||qL(z)− p(z|x)||TV = 0.
This is implied by the monotonic convergence of flow marginal qL(x) to π(x) in total variation simply follows monotonic convergence of MCMC chains in total variation distance by Theorem 6.51 and Proposition 6.52 in (Robert & Casella, 2005).
Then, we have
lim L→∞ Lz(x; θ, φL) = lim L→∞
∫ log pθ(x, z)
qL(z;φ) qL(z;φ)dz
= ∫ log pθ(x, z)
q∞(z∞;φ) q∞(z;φ)dz
= ∫ log pθ(x, z)
pθ(z|x) pθ(z|x)dz
= log p(x).

heading: D EXPERIMENTAL RESULTS
text: D.1 CONFIGURATION OF HMPFS
The HMPFs in all the experiments share the following common configuration. The initial distribution of HMPF q0(x) is given by independent Gaussian N (µ,σ2), where σ2 = (σ21 , . . . , σ2n) is the vector of variance. We implement Hamiltonian measure preserving transformation (HMPT) fr using vanilla Leapfrog integrator for simulating Hamiltonian dynamics H and independent Gaussian momentum variable r. The implementation of Leapfrog algorithm follows the tutorial of Neal (Neal, 2010). The momentum variables in each HMPT are independent and the each momentum variable has different variance. We consider separate step size = ( 1, . . . , n) for each dimension of x. Neal (Neal, 2010) shows that tuning leapfrog step size per dimension in HMC is equivalent to tuning the variance vector of momentum variables. So, we generate momentum variables from standard normal and assign an independent Leapfrog step size l for each HMPT fl. The number of iterations in Leapfrog integrator is a fixed parameters based manual tuning. We found that 5 to 10 Leapfrog iterations are often good enough. More Leapfrog steps than that do not give better results. This is consistent with the practice of tuning HMC (Neal, 2010). The intuitive explanation to this is that because Hamiltonian dynamics often have strong oscillation, the longer simulation does not lead to further exploration in sample space.
D.2 THE TRUE HISTOGRAM OF 2D BENCHMARKS
D.3 PLOTS OF MNIST AND FASHION MNIST

Review summaries:
This paper proposes training latent variable models (as in VAE decoders) by running HMC to approximate the posterior of the latents, and then estimating model parameters by maximizing the complete data log-likelihood. This is not a new idea by itself and is used e.g. as a baseline in Kingma and Welling's original VAE paper. The novelty in this paper is that it proposes tuning the parameters of the HMC inference algo by maximizing the likelihood achieved by the final sample in the MCMC chain. This seems to work well in practice and might be a useful method, but it is not clear under what conditions it should work.