Title: ERGODIC MEASURE PRESERVING FLOWS

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. Flow direction (the order of HMPTs) Training iterations

1. INTRODUCTION:

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.

2. BACKGROUND 2.1 BAYESIAN NEURAL NETWORKS:

Given data D = {x n , y n } N n=1 formed by feature vectors x n and corresponding scalar targets y n , we can assume that each y n is obtained as y n = f (x n ; θ) + 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(y 1 , . . . , y n |x 1 , . . . , x n , θ) 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.

2.2. DEEP GENERATIVE MODELS:

Generative models extract intrinsic structure from data by making use of latent variables. Let D be a dataset with n data points {x n } N n=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.

2.3. VARIATIONAL INFERENCE AND NORMALIZING FLOWS:

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

EQUATION

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 f 1 , . . . , f L to a random variable z 0 with tractable density q 0 (z 0 ) and exact simulation. The resulting output is a random

variable z L = f L • • • • • f 1 (z 0 ) with density q φ (z L |x) = q 0 (z 0 |x) L l=1 |∂f l (z l-1 )/∂z l-1 | -1

and with φ being the parameters of f 1 , . . . , f L .

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 f 1 , . . . , f L have to be rather simple to ensure invertibility and to reduce computational costs.

2.4. MARKOV CHAIN MONTE CARLO:

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 z 0 , z 1 , . . . in which the transition from state z l to z l+1 is defined by the conditional probability distribution of z l+1 given z l , denoted by K(z l , z l+1 ). Markov chains in MCMC methods have the following strong stationary property: if one state z l of the chain follows the stationary distribution π, so does the next state z l+1 . In particular,

π(z l+1 ) = π(z l )K(z l , z l+1 ) dz l .

(2)

If z l follows a distribution π l that is different from the stationary one, then the distribution of z l+1

EQUATION

is guaranteed to be closer to π than π l . This property implies that, with sufficiently many transitions, the distribution π l of z l converges to π irrespectively of the distribution the initial state z 0 .

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.

3. ERGODIC MEASURE PRESERVING FLOWS:

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.

3.1. DEFINITIONS:

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 all A ∈ F . If T is a one-to-one mapping onto Ω, then T preserves µ: µ(A) = µ(T -1 T A) = µ(T A).

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(z 0 , z 1 , . . . , z L ) = q 0 (z 0 )

L l=1 K(z l-1 , z l )

, where q 0 is the distribution of the initial state z 0 . The density of the last state z L is then obtained by integrating out all the previous states:

q L (z L ) = q 0 (z 0 ) L l=1 K(z l-1 , z l ) dz 0 dz 1 . . . dz L-1 . (4)

If the Markov chain is ergodic, q L 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 kernel K becomes a deterministic transformation T r : Z → Z with stochastic auxiliary input r following distribution µ. We can then define z L as the result of applying these deterministic transformations to z 0 , that is,

z L = T r L • • • • • T r1 (z 0 )

. By following the rule of changing variables, it is then straightforward to derive the density of z L as

EQUATION

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, q L , 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 ∈ R d with unnormalised density function π * , the HMC kernel is essentially applying a deterministic transformation H to the previous state z i with an auxiliary random variable r ∈ R n with the density function µ(•). The transformation H is given by the solution of Hamiltonian dynamics

EQUATION

where ż 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 Ḣ(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(z L ) = q(z L , r 1:L )dr 1:L = q(z 0 )µ(r)δ(z L -H r L • • • • • H r1 (z 0 )) dz 0 dr 1:L . (7)

We call the MPF generated by Hamiltonian dynamics a Hamiltonian MPF (HMPF).

3.2. UNDERSTANDING MEASURE PRESERVING CONDITIONS:

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 T r 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 T r : 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 T r : 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, T r 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 of H r can be very complicated.

3.3. VARIATIONAL INFERENCE WITH MPFS:

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 (z l-1 , r l ) to (z l , r l ) = T φ l (z l-1 , r l ). Similarly, the composition

T φ L • • • • • T φ1 (z 0 , r 1:L

), which we denote by T φ , transforms (z 0 , r 1:L ) to (z L , 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

EQUATION

EQUATION

where q 0 (z 0 ) 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 z L , 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 q 0 (z 0 ) as

EQUATION

We call this expression the simple ELBO. We can then multiplying by the density of the auxiliary variables µ(r

EQUATION

)

We can then replace (z 0 , r 1:L ) with (z L , r 1:L ) in ( 11) by making use of using transformation T φ , (8) and ( 9). The result is

EQUATION

where we have omitted the dependence of (z L , r 1:L ) on φ, since these variables are determined by the hyper-parameters of the MPTs. We call (12) the reparameterised ELBO.

3.4. ERGODIC LOWER BOUND AND ERGODIC INFERENCE:

The reparameterised ELBO is of limited use, because it can only be as tight as the ELBO with initial proposal distribution q 0 . 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

EQUATION

)

where D L KL is the Kullback-Liebler divergence between q L (z L , r 1:L ; φ) and q L (z L ; φ) 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

EQUATION

)

This is a tighter lower bound than the simple ELBO because the difference between L(x; θ, φ) and L(x; θ) in ( 10) is D L KL ≥ 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 D L KL = 0.

The complete proof is included in appendix.

Recall that q L (z L ; φ) 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, q L (z L ; φ) 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 q L (z L ; φ), which we denote by H[q L (z L ; φ)] = -log q L (z L ; φ)q L (z L ; φ) dz L .

In particular,

EQUATION

When optimizing this quantity w.r.t. φ, the first term in the RHS will encourage q L (z L ; φ) to have high density in regions where p θ (x, z L ) is high, while the second term will favor high entropy solutions and will prevent q L (z L ; φ) from converging to a Dirac delta centered at the maximizer of log p θ (x, z L ). Note that the first term in the RHS of ( 15) can be easily approximated by Monte Carlo, while the second term is intractable because q L (z L ; φ) is not available. However, since q L (z L ; φ) converges to p θ (z|x) as L increases, we expect the effect of H[q L (z L ; φ)] on φ to be small and that most of the similarity of q L (z L ; φ) 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,

EQUATION

If with the initial flow parameter φ 0 , the objective F < E p θ (x,z) [log p θ (x, z)], we expect optimising F produces faster convergence of q L (z L ; φ) towards p θ (x, z). Importantly, the model parameters θ can also be adjusted by optimizing F because the omitted H[q L (z L ; φ)] does not depend on θ and, consequently, optimizing F(x; θ, φ) and L(x; θ, φ) w.r.t. θ are equivalent operations.

3.5. IMPLEMENTATION OF HMPFS:

The hyper-parameters to be tuned in a HMPF include the parameters of q 0 and the transformation parameters for the Hamiltonian simulation, that is, φ = (φ 1 , . . . , φ L ). A natural choice for q 0 is multivariate Gaussian with mean µ = (µ 1 . . . , µ d ) and diagonal covariance matrix with entries σ 2 = (σ 2 1 , . . . , σ 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,

EQUATION

)

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 = (∆t l,1 , . . . , ∆t l,d ). The pseudo code for ergodic inference with HMPFs is shown in Algorithm 1. 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 -∂ z U (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.

4. EXPERIMENTS:

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.

4.1. DEMONSTRATION OF CONVERGENCE:

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 HMPFs1 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

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.

Encoders:

Training hours Training Epochs Test log(x) ESS Conv VAE(n h =300) (Salimans et al., 2015) ---83.20 -HVI(1HMPF-16LF, n h =800) (Salimans et al., 2015) ---81.94 -HVAE(1HMPF-20LF, n h = 300) (Caterini et al., (Salimans et al., 2015) to test HVI. In (Salimans et al., 2015) , the test likelihood is estimated using importence-weighted samples from the encoder network. In our experiment, we use a more reliable estimation method based on Hamiltonian annealled importance sampling and report the effective sample size (ESS). For HMPF encoder, we use 30 HMPTs with 5 Leapfrog steps per HMPT. The initial distribution q 0 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 

4.3. BAYESIAN NEURAL NETWORKS:

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 q 0 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 

5. SUMMARY:

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.

A THE PROOF OF PROPOSITION 1

Proof. Because T preserves the probability measure π(x, ), then for any measurable set A in Borel set B x × B , we have π

(A) = π(T -1 A).

Given a set B ∈ B x , the set A x generated by B × E, is measurable under π with the measure

π(A x ) = π(B, )d ,

that is essentially a probability measure in the space of x, also known as the marginal probability

π(B) = π(A x ).

Because T preserves the joint measure, applying T on A x gives

T A x = T (B, E) = T B × E,

where T B denotes the projection of T A x in the space of x. Follow the definition of measure preserving transformations, we have

π(A x ) = π(B, )d = π(T -1 B, )d = π(T -1 A x ),

where T -1 denotes the preimage of B under T . Because π(A x ) is essentially the marginal probability π(B), B ∈ B x , we know the marginal distribution π(x) is preserved by T because

π(B) = π(A x ) = π(T -1 A x ) = π(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 T A x 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.

B THE PROOF OF PROPOSITION 2:

Proof. Let µ be the probability measure of auxiliary variable R. Following equation 14, the difference of ERLBO in each measure preserving transformation T L is given by dL z (x; θ, φ 1:L ) = L z (x; θ, φ 1:L ) -L z (x; θ, φ 1:L-1 )

EQUATION

where Q L denotes Q L (Z L , R 1:L ) and QL denotes Q L (Z L ) L l=1 µ(R l ). The KL divergence can be written as the integral D KL Q L QL = log q L (z L , r 1:L ; φ) q L (z L ; φ) L l=1 µ(r l (φ))

q L (z L , r 1:L ; φ)dz L dr 1:L .

Because (z L , r L ) is generated from (z L-1 , r L ) by the deterministic transformation T L that preserves Lebesgue measure in the phase space, we have the equality q L (z L , r 1:L ) = q L-1 (z L-1 , r 1:L-1 )µ(r L ).

Using the reparameterisation (z L , r L ) = T (z L-1 , r L ),

D EXPERIMENTAL RESULTS:



D.1 CONFIGURATION OF HMPFS:

The HMPFs in all the experiments share the following common configuration. The initial distribution of HMPF q 0 (x) is given by independent Gaussian N (µ, σ 2 ), where σ 2 = (σ 2 1 , . . . , σ 2 n ) is the vector of variance. We implement Hamiltonian measure preserving transformation (HMPT) f r 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 f l . 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. 

:

The code of HMPFs for all three experiments will be available at https://github.com/ firstauthor/hmpfs

https://github.com/yburda/iwae

Figures and Tables:

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; z 0 ∼ N (µ, diag(σ 2 )); / * Start of simulation of HMPF * / for l = 1, . . . , L do r ∼ N (0, 1); z l ←-H(z l-1 , r; U (z; x, θ), φ l ); / * Leapfrog simulation * / end / * End of simulation of HMPF * / obj ←-U (z L ; x, θ); / * one sample Monte Carlo Approx. of F(x; θ, φ) * / θ ←-AdamUpdate(θ, ∂ θ obj); φ ←-AdamUpdate(φ, ∂ φ obj); end

Figure 2: The histogram of perfect samples from the targets using rejection sampling.

Comparison in terms of compuational efficiency and approximate test log-likelihood. For fair comparison, we implemented the deconvolutional decoder network in

Table content: <table><tr><td>2018)</td><td>-</td><td>-</td><td>-84.78</td><td>-</td></tr></table>



Table content: <table><tr><td>Encoders</td><td colspan="4">Training hours Training Epochs Test log(x) ESS</td></tr><tr><td>Conv VAE(n h =500)</td><td>6.00</td><td>3000</td><td>-104.90</td><td>26.3</td></tr><tr><td>HMPF(30HMPT-5LF, n h =500, no encoder)</td><td>6.00</td><td>200</td><td>-103.087</td><td>16.2</td></tr></table>

The comparison of log marginal likelihood on fashion MNIST between convolutional VAE and HMPFs. We also evaluate HMPFs with different setting of HAIS that gives higher effective sample size (ESS), but the result of test log likelihood is roughly the same.

Table content: <table/>

Overall, HMPFs produce significantly better results than SGHMC.

Table content: <table><tr><td>Method/Dataset</td><td>Boston</td><td>Yacht</td><td>Concrete</td><td>Wine</td></tr><tr><td>SGHMC (best average) (Springenberg et al., 2016)</td><td colspan="4">-3.47±0.51 -13.58±0.98 -4.87±0.05 -1.82±0.75</td></tr><tr><td colspan="5">SGHMC (tuned per dataset) (Springenberg et al., 2016) -2.49±0.15 -1.75±0.19 -4.16±0.72 -1.29±0.28</td></tr><tr><td>SGHMC (scale-adapted) (Springenberg et al., 2016)</td><td colspan="4">-2.54±0.04 -1.11±0.08 -3.38±0.24 -1.04±0.17</td></tr><tr><td>HMPFs</td><td colspan="4">-2.17±0.07 -0.47±0.06 -2.71±0.03 -0.71±0.03</td></tr></table>

The test log-likelihood of Bayesian neural networks on UCI datasets averaged over 20 splits with 100 sampled network parameters from HMPFs.

Table content: <table/>

Peer review (iclr2019):
 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. The paper is written in an unnecessarily complicated and formal way. On first read it seems like the proposed method has much more formal justification than it really has. The discussion up to section 3.5 makes it seem as if there is some new kind of tractable variational bound (the ERLBO) that is optimized, but in practice the actual objective in equation 16 is simply the likelihood at the final sample of the MCMC chain, that is Monte Carlo EM as e.g. used by Kingma & Welling, 2013 as a baseline. The propositions and theorems seem to apply to an idealized setting, but not to the actual algorithm that is used. They could have been put in an appendix, or even a reference to the exisiting literature on HMC would have sufficed. The experiments do not clearly demonstrate that the method is much better than previous methods from the literature, although it is much more expensive. (The reported settings require 150 likelihood evaluations per example per minibatch update, versus 1 likelihood evaluation for a VAE). Also see my previous comments about evaluation in this paper*s thread. - Please explain why tuning the HMC algo by maximizing eq 16 should work. I don*t think it is a method that generally would work, e.g. if the initial sample z0 ~ q(z|x) is drawn from a data dependent encoder as in HVI (Salimans et al) then I would expect the step size of the HMC to simply go to zero as the encoder gets good. However in your case this does not happen as the initial sample is unconditional from x. Are there general guidelines or guarantees we can conclude from this? - The authors write "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." This is true for the idealized flow in continuous time, but HMC with finite step size does generally NOT converge to the correct distribution. This is why practical use of HMC includes a Metropolis-Hastings correction step. You omit this step in your algorithm, with the justification that we don*t care about asymptotic convergence in this case. Fair enough, but then you should also omit all statements in the paper that claim that your method converges to the correct posterior in the limit. E.g. the writing makes it seem like Proposition 2 and Theorem 1 apply to your algorithm, but it in fact they do not apply for finite step size. Maybe the statements are still correct if we take the limit with L->inf and the stepsize delta->0 at a certain rate? This is not obvious to me. In practice, you learn the stepsize delta. Do we have any guarantees this will make delta go to zero at the right rate as we increase the number of steps L? I.e. is this statement from your abstract true? -> "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". (convergence of uncorrected HMC only holds in the idealized case with step size -> 0)