Original: G ENERALIZING H AMILTONIAN M ONTE C ARLO WITH N EURAL N ETWORKS
Revised: G ENERALIZING H AMILTONIAN M ONTE C ARLO WITH N EURAL N ETWORKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We present a general-purpose method to train Markov chain Monte Carlo kernels, parameterized by deep neural networks, that converge and mix quickly to their target distribution.
Revised: We present a general-purpose method to train Markov chain Monte Carlo kernels, parameterized by deep neural networks, that converge and mix quickly to their target distribution.

Original: Our method generalizes Hamiltonian Monte Carlo and is trained to maximize expected squared jumped distance, a proxy for mixing speed.
Revised: Our method generalizes Hamiltonian Monte Carlo and is trained to maximize expected squared jumped distance, a proxy for mixing speed.

Original: We demonstrate large empirical gains on a collection of simple but challenging distributions, for instance achieving a 106 × improvement in effective sample size in one case, and mixing when standard HMC makes no measurable progress in a second.
Revised: We demonstrate large empirical gains on a collection of simple but challenging distributions, for instance achieving a 106 × improvement in effective sample size in one case, and mixing when standard HMC makes no measurable progress in a second.

Original: Finally, we show quantitative and qualitative gains on a real-world task: latent-variable generative modeling.
Revised: Finally, we show quantitative and qualitative gains on a real-world task: latent-variable generative modeling.

Original: Python source code will be open-sourced with the camera-ready paper.
Revised: We release an open source TensorFlow implementation of the algorithm

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: High-dimensional distributions that are only analytically tractable up to a normalizing constant are ubiquitous in many ﬁelds.
Revised: High-dimensional distributions that are only analytically tractable up to a normalizing constant are ubiquitous in many ﬁelds.

Original: For instance, they arise in protein folding (Sch ¨ utte et al., 1999), physics simulations (Olsson, 1995), and machine learning (Andrieu et al., 2003).
Revised: For instance, they arise in protein folding (Sch ¨ utte et al., 1999), physics simulations (Olsson, 1995), and machine learning (Andrieu et al., 2003).

Original: Sampling from such distributions is a critical task for learning and inference (MacKay, 2003), however it is an extremely hard problem in general.
Revised: Sampling from such distributions is a critical task for learning and inference (MacKay, 2003), however it is an extremely hard problem in general.

Original: Markov Chain Monte Carlo (MCMC) methods promise a solution to this problem.
Revised: Markov Chain Monte Carlo (MCMC) methods promise a solution to this problem.

Original: They operate by generating a sequence of correlated samples that converge in distribution to the target.
Revised: They operate by generating a sequence of correlated samples that converge in distribution to the target.

Original: This convergence is most often guaranteed through detailed balance, a sufﬁcient condition for the chain to have the target equilibrium distribution.
Revised: This convergence is most often guaranteed through detailed balance, a sufﬁcient condition for the chain to have the target equilibrium distribution.

Original: In practice, for any proposal distribution, one can ensure detailed balance through a Metropolis-Hastings (Hastings, 1970) accept/reject step.
Revised: In practice, for any proposal distribution, one can ensure detailed balance through a Metropolis-Hastings (Hastings, 1970) accept/reject step.

Original: Despite theoretical guarantees of eventual convergence, in practice convergence and mixing speed depend strongly on choosing a proposal that works well for the task at hand.
Revised: Despite theoretical guarantees of eventual convergence, in practice convergence and mixing speed depend strongly on choosing a proposal that works well for the task at hand.

Original: What’s more, it is often more art than science to know when an MCMC chain has converged (“burned-in”), and when the chain has produced a new uncorrelated sample (“mixed”).
Revised: What’s more, it is often more art than science to know when an MCMC chain has converged (“burned-in”), and when the chain has produced a new uncorrelated sample (“mixed”).

Original: Additionally, the reliance on detailed balance, which assigns equal probability to the forward and reverse transitions, often encourages random-walk behavior and thus slows exploration of the space (Ichiki & Ohzeki, 2013).
Revised: Additionally, the reliance on detailed balance, which assigns equal probability to the forward and reverse transitions, often encourages random-walk behavior and thus slows exploration of the space (Ichiki & Ohzeki, 2013).

Original: For densities over continuous spaces, Hamiltonian Monte Carlo (HMC; Duane et al., 1987; Neal, 2011) introduces independent, auxiliary momentum variables, and computes a new state by integrating Hamiltonian dynamics.
Revised: For densities over continuous spaces, Hamiltonian Monte Carlo (HMC; Duane et al., 1987; Neal, 2011) introduces independent, auxiliary momentum variables, and computes a new state by integrating Hamiltonian dynamics.

Original: This method can traverse long distances in state space with a single
Revised: This method can traverse long distances in state space with a single

Original: Metropolis-Hastings test.
Revised: Metropolis-Hastings test.

Original: This is the state-of-the-art method for sampling in many domains.
Revised: This is the state-of-the-art method for sampling in many domains.

Original: However, HMC can perform poorly in a number of settings.
Revised: However, HMC can perform poorly in a number of settings.

Original: While HMC mixes quickly spatially, it struggles at mixing across energy levels due to its volume-preserving dynamics.
Revised: While HMC mixes quickly spatially, it struggles at mixing across energy levels due to its volume-preserving dynamics.

Original: HMC also does not work well with multi-modal distributions, as the probability of sampling a large enough momentum to traverse a very low-density region is negligibly small.
Revised: HMC also does not work well with multi-modal distributions, as the probability of sampling a large enough momentum to traverse a very low-density region is negligibly small.

Original: Furthermore, HMC struggles with ill-conditioned energy landscapes (Girolami & Calderhead, 2011) and deals poorly with rapidly changing gradients (Sohl-Dickstein et al., 2014).
Revised: Furthermore, HMC struggles with ill-conditioned energy landscapes (Girolami & Calderhead, 2011) and deals poorly with rapidly changing gradients (Sohl-Dickstein et al., 2014).

Original: Recently, probabilistic models parameterized by deep neural networks have achieved great success at approximately sampling from highly complex, multi-modal empirical distributions (Kingma &
Revised: Recently, probabilistic models parameterized by deep neural networks have achieved great success at approximately sampling from highly complex, multi-modal empirical distributions (Kingma &

Original: Welling, 2013; Rezende et al., 2014; Goodfellow et al., 2014; Bengio et al., 2014; Sohl-Dickstein et al., 2015).
Revised: Welling, 2013; Rezende et al., 2014; Goodfellow et al., 2014; Bengio et al., 2014; Sohl-Dickstein et al., 2015).

Original: Building on these successes, we present a method that, given an analytically described distribution, automatically returns an exact sampler with good convergence and mixing properties, from a class of highly expressive parametric models.
Revised: Building on these successes, we present a method that, given an analytically described distribution, automatically returns an exact sampler with good convergence and mixing properties, from a class of highly expressive parametric models.

Original: The proposed family of samplers is a generalization of HMC; it transforms the HMC trajectory using parametric functions (deep networks in our experiments), while retaining theoretical guarantees with a tractable Metropolis-Hastings accept/reject step.
Revised: The proposed family of samplers is a generalization of HMC; it transforms the HMC trajectory using parametric functions (deep networks in our experiments), while retaining theoretical guarantees with a tractable Metropolis-Hastings accept/reject step.

Original: The sampler is trained to minimize a variation on expected squared jumped distance (similar in spirit to Pasarica & Gelman (2010)).
Revised: The sampler is trained to minimize a variation on expected squared jumped distance (similar in spirit to Pasarica & Gelman (2010)).

Original: Our parameterization reduces easily to standard HMC.
Revised: Our parameterization reduces easily to standard HMC.

Original: It is further capable of emulating several common extensions of HMC such as withintrajectory tempering (Neal, 1996) and diagonal mass matrices (Bennett, 1975).
Revised: It is further capable of emulating several common extensions of HMC such as withintrajectory tempering (Neal, 1996) and diagonal mass matrices (Bennett, 1975).

Original: We evaluate our method on distributions where HMC usually struggles, as well as on a the real-world task of training latent-variable generative models.
Revised: We evaluate our method on distributions where HMC usually struggles, as well as on a the real-world task of training latent-variable generative models.

Original: Our contributions are as follows:
Revised: Our contributions are as follows

Original:  • We introduce a generic training procedure which takes as input a distribution deﬁned by an energy function, and returns a fast-mixing MCMC kernel.
Revised: : • We introduce a generic training procedure which takes as input a distribution deﬁned by an energy function, and returns a fast-mixing MCMC kernel.

Original: • We show signiﬁcant empirical gains on various distributions where HMC performs poorly.
Revised: • We show signiﬁcant empirical gains on various distributions where HMC performs poorly.

Original: • We ﬁnally evaluate our method on the real-world task of training and sampling from a latent variable generative model, where we show improvement in the model’s log-likelihood, and greater complexity in the distribution of posterior samples.
Revised: • We ﬁnally evaluate our method on the real-world task of training and sampling from a latent variable generative model, where we show improvement in the model’s log-likelihood, and greater complexity in the distribution of posterior samples.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Adaptively modifying proposal distributions to improve convergence and mixing has been explored in the past (Andrieu & Thoms, 2008).
Revised: Adaptively modifying proposal distributions to improve convergence and mixing has been explored in the past (Andrieu & Thoms, 2008).

Original: In the case of HMC, prior work has reduced the need to choose step size (Neal, 2011) or number of leapfrog steps (Hoffman & Gelman, 2014) by adaptively tuning those parameters.
Revised: In the case of HMC, prior work has reduced the need to choose step size (Neal, 2011) or number of leapfrog steps (Hoffman & Gelman, 2014) by adaptively tuning those parameters.

Original: Salimans et al.
Revised: Salimans et al.

Original: (2015) proposed an alternate scheme based on variational inference.
Revised: (2015) proposed an alternate scheme based on variational inference.

Original: We adopt the much simpler approach of Pasarica & Gelman (2010), who show that choosing the hyperparameters of a proposal distributionto maximize expected squared jumped distance is both principled and effective in practice.
Revised: We adopt the much simpler approach of Pasarica & Gelman (2010), who show that choosing the hyperparameters of a proposal distribution maximize expected squared jumped Previous work has also explored applying models from machine learning to MCMC tasks.

Revised: Kernel methods have been used both for learning a proposal distribution (Sejdinovic et al., 2014) and for approximating the gradient of the energy (Strathmann et al., 2015).

Revised: In physics, Restricted and semiRestricted Boltzmann machines have been used both to build approximations of the energy function which allow more rapid sampling (Liu et al., 2017; Huang & Wang, 2017), and to motivate new hand-designed proposals (Wang, 2017).

Original: Most similar to ours is recent work from Song et al.
Revised: Most similar to our approach is recent work from Song et al.

Original: (2017), which uses adversarial training of a volume-preserving transformation, which is subsequently used as an MCMC proposal distribution.
Revised: (2017), which uses adversarial training of a volume-preserving transformation, which is subsequently used as an MCMC proposal distribution.

Original: While promising, this technique has several limitations.
Revised: While promising, this technique has several limitations.

Original: It does not use gradient information, which is often crucial to maintaining high acceptance rates, especially in high dimensions.
Revised: It does not use gradient information, which is often crucial to maintaining high acceptance rates, especially in high dimensions.

Original: It also can only indirectly measure the quality of the generated sample using adversarial training, which is notoriously unstable, suffers from “mode collapse” (where only a portion of a target distribution is covered), and often requires objective modiﬁcation to train in practice (Arjovsky et al., 2017).
Revised: It also can only indirectly measure the quality of the generated sample using adversarial training, which is notoriously unstable, suffers from “mode collapse” (where only a portion of a target distribution is covered), and often requires objective modiﬁcation to train in practice (Arjovsky et al., 2017).

Original: Finally, since the proposal transformation preserves volume, it can suffer from the same difﬁculties in mixing across energy levels as HMC, as we illustrate in Section 5.
Revised: Finally, since the proposal transformation preserves volume, it can suffer from the same difﬁculties in mixing across energy levels as HMC, as we illustrate in Section 5.

Original: To compute the Metropolis-Hastings acceptance probability for a deterministic transition, the operator must be invertible and have a tractable Jacobian.
Revised: To compute the Metropolis-Hastings acceptance probability for a deterministic transition, the operator must be invertible and have a tractable Jacobian.

Original: Recent work (Dinh et al., 2016), introduces
Revised: Recent work (Dinh et al., 2016), introduces

Original: RNVP, an invertible transformation that operates by, at each layer, modifying only a subset of the variables by a function that depends solely on the remaining variables.
Revised: RNVP, an invertible transformation that operates by, at each layer, modifying only a subset of the variables by a function that depends solely on the remaining variables.

Original: This is exactly invertible with an efﬁciently computable Jacobian.
Revised: This is exactly invertible with an efﬁciently computable Jacobian.

Original: Furthermore, by chaining enough of these layers, the model can be made arbitrarily expressive.
Revised: Furthermore, by chaining enough of these layers, the model can be made arbitrarily expressive.

Original: This parameterization will directly motivate our extension of the leapfrog integrator in HMC.
Revised: This parameterization will directly motivate our extension of the leapfrog integrator in HMC.

Original:  B ACKGROUND 3. MCMC METHODS AND M ETROPOLIS -H ASTINGS
Revised:  B ACKGROUND 3. MCMC METHODS AND M ETROPOLIS -H ASTINGS

Original: Let p be a target distribution, analytically known up to a constant, over a space X .
Revised: Let p be a target distribution, analytically known up to a constant, over a space X .

Original: Markov chain
Revised: Markov chain

Original: Monte Carlo (MCMC) methods (Neal, 1993) aim to provide samples from p .
Revised: Monte Carlo (MCMC) methods (Neal, 1993) aim to provide samples from p .

Original: To that end, MCMC methods construct a Markov Chain whose stationary distribution is the target distribution p .
Revised: To that end, MCMC methods construct a Markov Chain whose stationary distribution is the target distribution p .

Original: Obtaining samples then corresponds to simulating a Markov Chain, i.e., given an initial distribution π 0 and a transition kernel K , constructing the following sequence of random variables:
Revised: Obtaining samples then corresponds to simulating a Markov Chain, i.e., given an initial distribution π 0 and a transition kernel K , constructing the following sequence of random variables:

Original: [Equation]
Revised: [Equation]

Original: In order for p to be the stationary distribution of the chain, three conditions must be satisﬁed: K must be irreducible and aperiodic (these are usually mild technical conditions) and p has to be a ﬁxed point of K .
Revised: In order for p to be the stationary distribution of the chain, three conditions must be satisﬁed: K must be irreducible and aperiodic (these are usually mild technical conditions) and p has to be a ﬁxed point of K .

Original: This last condition can be expressed as: p ( x (cid:48) )
Revised: This last condition can be expressed as: p ( x (cid:48) )

Original: = (cid:82) K
Revised: = (cid:82) K

Original: ( x (cid:48) | x ) p ( x )d x .
Revised: ( x (cid:48) | x ) p ( x )d x .

Original: This condition is most often satisﬁed by satisfying the stronger detailed balance condition, which can be written as: p ( x (cid:48) )
Revised: This condition is most often satisﬁed by satisfying the stronger detailed balance condition, which can be written as: p ( x (cid:48) )

Original: K ( x | x (cid:48) )
Revised: K [Equation] x [Equation]

Original: =
Revised: =

Original: p ( x ) K ( x (cid:48) | x ) .
Revised: [Equation]

Original: Given any proposal distribution q , satisfying mild conditions, we can easily construct a transition kernel that respects detailed balance using Metropolis-Hastings (Hastings, 1970) accept/reject rules.
Revised: Given any proposal distribution q , satisfying mild conditions, we can easily construct a transition kernel that respects detailed balance using Metropolis-Hastings (Hastings, 1970) accept/reject rules.

Original: More formally, starting from x 0 ∼ π 0 , at each step t , we sample x (cid:48) ∼ q ( ·| X t ) , and with probability
Revised: More formally, starting from x 0 ∼ π 0 , at each step t , we sample x (cid:48) ∼ q ( ·| X t ) , and with probability

Original: A ( x (cid:48) | x t ) =
Revised: A ( x (cid:48) | x t ) =

Original: min (cid:16) 1 , p ( x (cid:48) )
Revised: min (cid:16) 1 , p ( x (cid:48) )

Original: q ( x t | x (cid:48) ) p ( x t ) q ( x (cid:48) | x t ) (cid:17) , accept x (cid:48) as the next sample x t +1 in the chain.
Revised: q ( x t | x (cid:48) ) p ( x t ) q ( x (cid:48) | x t ) (cid:17) , accept x (cid:48) as the next sample x t +1 in the chain.

Original: If we reject x (cid:48) , then we retain the previous state and x t +1 = x t .
Revised: If we reject x (cid:48) , then we retain the previous state and x t +1 = x t .

Original: For typical proposals this algorithm has strong asymptotic guarantees.
Revised: For typical proposals this algorithm has strong asymptotic guarantees.

Original: But in practice one must often choose between very low acceptance probabilities and very cautious proposals, both of which lead to slow mixing.
Revised: But in practice one must often choose between very low acceptance probabilities and very cautious proposals, both of which lead to slow mixing.

Original: For continuous state spaces, Hamiltonian Monte Carlo (HMC; Neal, 2011) tackles this problem by proposing updates that move far in state space while staying roughly on iso-probability contours of p .
Revised: For continuous state spaces, Hamiltonian Monte Carlo (HMC; Neal, 2011) tackles this problem by proposing updates that move far in state space while staying roughly on iso-probability contours of p .

Original: 3. H AMILTONIAN M ONTE C ARLO
Revised: 3. H AMILTONIAN M ONTE C ARLO

Original: Without loss of generality, we assume p ( x ) to be deﬁned by an energy function U ( x ) , s.t. p ( x ) ∝ exp( − U ( x )) , and where the state x ∈ R n .
Revised: Without loss of generality, we assume p ( x ) to be deﬁned by an energy function U ( x ) , s.t. p ( x ) ∝ exp( − U ( x )) , and where the state x ∈ R n .

Original: HMC extends the state space with an additional momentum vector v ∈ R n , where v is distributed independently from x , as p ( v ) ∝ exp( − 12 v T v ) (i.e., identity-covariance Gaussian).
Revised: HMC extends the state space with an additional momentum vector v ∈ R n , where v is distributed independently from x , as p ( v ) ∝ exp( − 12 v T v ) (i.e., identity-covariance Gaussian).

Original: From an augmented state ξ (cid:44) ( x, v ) , HMC produces a proposed state ξ (cid:48)
Revised: From an augmented state ξ (cid:44) ( x, v ) , HMC produces a proposed state ξ (cid:48)

Original: =
Revised: =

Original: ( x (cid:48) , v (cid:48) ) by approximately integrating Hamiltonian dynamics jointly on x and v , with U ( x ) taken to be the potential energy, and 12 v T v the kinetic energy.
Revised: ( x (cid:48) , v (cid:48) ) by approximately integrating Hamiltonian dynamics jointly on x and v , with U ( x ) taken to be the potential energy, and 12 v T v the kinetic energy.

Original: Since Hamiltonian dynamics conserve the total energy of a system, their approximate integration moves along approximate iso-probability contours of p ( x, v ) = p ( x ) p ( v ) .
Revised: Since Hamiltonian dynamics conserve the total energy of a system, their approximate integration moves along approximate iso-probability contours of p ( x, v ) = p ( x ) p ( v ) .

Original: The dynamics are typically simulated using the leapfrog integrator (Hairer et al., 2003; Leimkuhler & Reich, 2004), which for a single time step consists of:
Revised: The dynamics are typically simulated using the leapfrog integrator (Hairer et al., 2003; Leimkuhler & Reich, 2004), which for a single time step consists of:

Original: [Equation]
Revised: [Equation]

Original: Following Sohl-Dickstein et al.
Revised: Following Sohl-Dickstein et al.

Original: (2014), we write the action of the leapfrog integrator in terms of an operator L : L ξ (cid:44) L ( x, v ) (cid:44) ( x (cid:48) , v (cid:48) ) , and introduce a momentum ﬂip operator F : F ( x, v ) (cid:44) ( x, − v ) .
Revised: (2014), we write the action of the leapfrog integrator in terms of an operator L : L ξ (cid:44) L ( x, v ) (cid:44) ( x [Equation] , v (cid:48) ) , and introduce a momentum ﬂip operator F : F ( x, v ) (cid:44) ( x, − v ) .

Original: It is important to note two properties of these operators.
Revised: It is important to note two properties of these operators.

Original: First, the transformation FL is an involution, i.e. FLFL ( x, v ) =
Revised: First, the transformation FL is an involution, i.e. FLFL ( x, v ) =

Original: FL ( x (cid:48) , − v (cid:48) )
Revised: FL ( x (cid:48) , − v (cid:48) )

Original: =
Revised: =

Original: ( x, v ) .
Revised: ( x, v ) .

Original: Second, the transformations from ( x, v ) to ( x, v 12 ) , from ( x, v 12 ) to ( x (cid:48) , v 12 ) , and from ( x (cid:48) , v 12 ) to ( x (cid:48) , v (cid:48) ) are all volume-preserving shear transformations i.e., only one of the variables ( x or v ) changes, by an amount determined by the other one.
Revised: Second, the transformations from ( x, v ) to ( x, v 12 ) , from ( x, v 12 ) to ( x (cid:48) , v 12 ) , and from ( x (cid:48) , v 12 ) to ( x (cid:48) , v (cid:48) ) are all volume-preserving shear transformations i.e., only one of the variables ( x or v ) changes, by an amount determined by the other one.

Original: The determinant of the Jacobian, (cid:12)(cid:12)(cid:12) ∂
Revised: The determinant of the Jacobian, (cid:12)(cid:12) (cid:12) ∂

Original: [ FL ξ ] ∂ξ T (cid:12)(cid:12)(cid:12) , is thus easy to compute.
Revised: [ FL ξ ] ∂ξ T (cid:12)(cid:12) (cid:12) , is thus easy to compute.

Original: For vanilla HMC (cid:12)(cid:12)(cid:12) ∂ [ FL ξ ] ∂ξ T (cid:12)(cid:12)(cid:12) = 1 , but we will leave it in symbolic form for use in Section 4.
Revised: For vanilla HMC (cid:12) (cid:12)(cid:12) ∂ [ FL ξ ] ∂ξ T (cid:12) (cid:12)(cid:12) = 1 , but we will leave it in symbolic form for use in Section 4.

Original: The Metropolis-Hastings acceptance probability for the HMC proposal is made simple by these two properties, and is
Revised: The Metropolis-HastingsGreen (Hastings, 1970; Green, 1995) acceptance probability for the HMC proposal is made simple by these two properties, and is

Original: [Equation]
Revised: [Equation]

Original:  L2HMC: T RAINING MCMC S AMPLERS
Revised:  L2HMC: T RAINING MCMC S AMPLERS

Original: In this section, we describe our proposed method L2HMC 1 .
Revised: In this section, we describe our proposed method L2HMC (for ‘Learning To Hamiltonian Monte Carlo’).

Original: Given access to only an energy function U (and not samples), L2HMC learns a parametric leapfrog operator L θ over an augmented state space.
Revised: Given access to only an energy function U (and not samples), L2HMC learns a parametric leapfrog operator L θ over an augmented state space.

Original: We begin by describing what desiderata we have for L θ , then go into detail on how we parameterize our sampler.
Revised: We begin by describing what desiderata we have for L θ , then go into detail on how we parameterize our sampler.

Original: Finally, we conclude this section by describing our training procedure.
Revised: Finally, we conclude this section by describing our training procedure.

Original: 4. A UGMENTING HMC
Revised: 4. A UGMENTING HMC

Original: HMC is a powerful algorithm, but it can still struggle even on very simple problems.
Revised: HMC is a powerful algorithm, but it can still struggle even on very simple problems.

Original: For example, a two-dimensional multivariate Gaussian with an ill-conditioned covariance matrix can take arbitrarily long to traverse (even if the covariance is diagonal), whereas it is trivial to sample directly from it.
Revised: For example, a two-dimensional multivariate Gaussian with an ill-conditioned covariance matrix can take arbitrarily long to traverse (even if the covariance is diagonal), whereas it is trivial to sample directly from it.

Original: Another problem is that HMC can only move between energy levels via a random walk (Neal, 2011), which leads to slow mixing in some models.
Revised: Another problem is that HMC can only move between energy levels via a random walk (Neal, 2011), which leads to slow mixing in some models.

Original: Finally, HMC cannot easily traverse low-density zones.
Revised: Finally, HMC cannot easily traverse low-density zones.

Original: For example, given a simple Gaussian mixture model, HMC cannot mix between modes without recourse to additional tricks, as illustrated in Figure 1b.
Revised: For example, given a simple Gaussian mixture model, HMC cannot mix between modes without recourse to additional tricks, as illustrated in Figure 1b.

Original: These observations determine the list of desiderata for our learned MCMC kernel: fast mixing , fast burn-in , mixing across energy levels , and mixing between modes .
Revised: These observations determine the list of desiderata for our learned MCMC kernel: fast mixing , fast burn-in , mixing across energy levels , and mixing between modes .

Original: While pursuing these goals, we must take care to ensure that our proposal operator retains two key features of the leapfrog operator used in HMC: it must be invertible, and the determinant of its
Revised: While pursuing these goals, we must take care to ensure that our proposal operator retains two key features of the leapfrog operator used in HMC: it must be invertible, and the determinant of its

Original: Jacobian must be tractable.
Revised: Jacobian must be tractable.

Original: The leapfrog operator satisﬁes these properties by ensuring that each sub-update only affects a subset of the variables, and that no sub-update depends nonlinearly on any of the variables being updated.
Revised: The leapfrog operator satisﬁes these properties by ensuring that each sub-update only affects a subset of the variables, and that no sub-update depends nonlinearly on any of the variables being updated.

Original: We are free to generalize the leapfrog operator in any way that preserves these properties.
Revised: We are free to generalize the leapfrog operator in any way that preserves these properties.

Original: In particular, we are free to translate and rescale each sub-update of the leapfrog operator, so long as we are careful to ensure that these translation and scale terms do not depend on the variables being updated.
Revised: In particular, we are free to translate and rescale each sub-update of the leapfrog operator, so long as we are careful to ensure that these translation and scale terms do not depend on the variables being updated.

Original: 4.1. S TATE S PACE
Revised: 4.1. S TATE S PACE

Original: As in HMC, we begin by augmenting the current state x ∈ R n with a continuous momentum variable v ∈ R n drawn from a standard normal.
Revised: As in HMC, we begin by augmenting the current state x ∈ R n with a continuous momentum variable v ∈ R n drawn from a standard normal.

Original: We also introduce a binary direction variable d ∈ {− 1 , 1 } , drawn from a uniform distribution.
Revised: We also introduce a binary direction variable d ∈ {− 1 , 1 } , drawn from a uniform distribution.

Original: We will denote the complete augmented state as ξ (cid:44) ( x, v, d ) , with probability density p ( ξ )
Revised: We will denote the complete augmented state as ξ (cid:44) ( x, v, d ) , with probability density p ( ξ )

Original: = p ( x ) p ( v ) p ( d ) .
Revised: = p ( x ) p ( v ) p ( d ) .

Original: Finally, to each step t of the operator L θ we assign a ﬁxed random binary mask m t ∈ { 0 , 1 } n that will determine which variables are affected by each sub-update.
Revised: Finally, to each step t of the operator L θ we assign a ﬁxed random binary mask m t ∈ { 0 , 1 } n that will determine which variables are affected by each sub-update.

Original: We draw m t uniformly from the set of binary vectors satisfying (cid:80) ni
Revised: We draw m t uniformly from the set of binary vectors satisfying (cid:80) ni

Original: =1 m ti = (cid:98) n 2 (cid:99) , that is, half of the entries of m t are 0 and half are 1.
Revised: =1 m ti = (cid:98) n 2 (cid:99) , that is, half of the entries of m t are 0 and half are 1.

Original: For convenience, we write ¯ m t = 1 − m t and x m t = x (cid:12) m t ( (cid:12) denotes element-wise multiplication, and 1 the all ones vector).
Revised: For convenience, we write ¯ m t = 1 − m t and x m t = x (cid:12) m t ( (cid:12) denotes element-wise multiplication, and 1 the all ones vector).

Original: 4.1. U PDATE S TEPS
Revised: 4.1. U PDATE S TEPS

Original: We now describe the details of our augmented leapfrog integrator L θ , for a single time-step t , and for direction d = 1 .
Revised: We now describe the details of our augmented leapfrog integrator L θ , for a single time-step t , and for direction d = 1 .

Original: We ﬁrst update the momenta v .
Revised: We ﬁrst update the momenta v .

Original: This update can only depend on a subset ζ 1 (cid:44) ( x, ∂ x U ( x ) , t ) of the full state, which excludes v .
Revised: This update can only depend on a subset ζ 1 (cid:44) ( x, ∂ x U ( x ) , t ) of the full state, which excludes v .

Original: It takes the form
Revised: It takes the form

Original: [Equation]
Revised: [Equation]

Original: We have introduced three new functions of ζ 1 : T v , Q v , and S v .
Revised: We have introduced three new functions of ζ 1 : T v , Q v , and S v .

Original: T v is a translation, exp( Q v ) rescales the gradient, and exp( (cid:15) 2 S v ) rescales the momentum.
Revised: T v is a translation, exp( Q v ) rescales the gradient, and exp( (cid:15) 2 S v ) rescales the momentum.

Original: The determinant of the Jacobian of this transformation is exp (cid:0) (cid:15) 2 1 · S v ( ζ1 ) (cid:1) .
Revised: The determinant of the Jacobian of this transformation is exp (cid:0) [Equation] [Equation] ( ζ

Revised: [Equation]

Revised: ) (cid:1) .

Original: Note that if T v , Q v , and S v are all zero, then we recover the standard leapfrog momentum update.
Revised: Note that if T v , Q v , and S v are all zero, then we recover the standard leapfrog momentum update.

Original: We now update x .
Revised: We now update x .

Original: As hinted above, to make our transformation more expressive, we ﬁrst update a subset of the coordinates of x , followed by the complementary subset.
Revised: As hinted above, to make our transformation more expressive, we ﬁrst update a subset of the coordinates of x , followed by the complementary subset.

Original: The ﬁrst update, which yields x (cid:48) and affects only x m t , depends on the state subset ζ 2 (cid:44) ( x ¯ m t , v, t ) .
Revised: The ﬁrst update, which yields x (cid:48) and affects only x m t , depends on the state subset ζ 2 (cid:44) ( x ¯ m t , v, t ) .

Original: Conversely, with x (cid:48) deﬁned below, the second update only affects x (cid:48)
Revised: Conversely, with x (cid:48) deﬁned below, the second update only affects x (cid:48)

Original: ¯ m t and depends only on ζ 3 (cid:44) ( x (cid:48) m t , v, t ) :
Revised: ¯ m t and depends only on ζ 3 (cid:44) ( x (cid:48) m t , v, t ) :

Original: [Equation]
Revised: [Equation]

Original: [Equation] Again, T x is a translation, exp( Q x ) rescales the effect of the momenta, exp( (cid:15)S x ) rescales the positions x , and we recover the original leapfrog position update
Revised: Again, T x is a translation, exp( Q x ) rescales the effect of the momenta, exp( (cid:15)S x ) rescales the positions x , and we recover the original leapfrog position update

Original: if T x = Q x = S x = 0 .
Revised: if T x = Q x = S x = 0 .

Original: The determinant of the Jacobian of the ﬁrst transformation is exp ( (cid:15)m t · S x ( ζ 2 )) , and the determinant of the Jacobian of the second transformation is exp ( (cid:15) ¯ m t · S x ( ζ 3 )) .
Revised: The determinant of the Jacobian of the ﬁrst transformation is exp ( (cid:15)m t · S x ( ζ 2 )) , and the determinant of the Jacobian of the second transformation is exp ( (cid:15) ¯ m t · S x ( ζ 3 )) .

Original: Finally, we update v again, based on the subset ζ 4 (cid:44) ( x 
Revised: Finally, we update v again, based on the subset ζ 4 (cid:44) ( x (cid:48)(cid:48) , ∂ x U ( x

Original: [Equation]
Revised: [Equation]

Original:  , ∂
Revised: ) , t ) :

Original: [Equation]
Revised: [Equation]

Original: U

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: This update has the same form as the momentum update in equation 4.
Revised: This update has the same form as the momentum update in equation 4.

Original: 4.1. I NTUITION

Original: To give intuition into these terms, the scaling applied to the momentum can enable, among other things, acceleration in low-density zones, to facilitate mixing between modes.
Revised: To give intuition into these terms, the scaling applied to the momentum can enable, among other things, acceleration in low-density zones, to facilitate mixing between modes.

Original: The scaling term applied to the gradient of the energy may allow better conditioning of the energy landscape (e.g., by learning a diagonal inertia tensor), or partial ignoring of the energy gradient for rapidly oscillating energies.
Revised: The scaling term applied to the gradient of the energy may allow better conditioning of the energy landscape (e.g., by learning a diagonal inertia tensor), or partial ignoring of the energy gradient for rapidly oscillating energies.

Original: The corresponding integrator for d = − 1 is given in Appendix A; it essentially just inverts the updates in equations 4, 5 and 6.
Revised: The corresponding integrator for d = − 1 is given in Appendix A; it essentially just inverts the updates in equations 4, 5 and 6.

Original: For all experiments, the functions Q, S, T are implemented using multi-layer perceptrons, with shared weights.
Revised: For all experiments, the functions Q, S, T are implemented using multi-layer perceptrons, with shared weights.

Original: We encode the current time step in the MLP input.
Revised: We encode the current time step in the MLP input.

Original: Our leapfrog operator L θ corresponds to running M steps of this modiﬁed leapfrog, L θ ξ =
Revised: Our leapfrog operator L θ corresponds to running M steps of this modiﬁed leapfrog, L θ ξ =

Original: L θ ( x, v, d ) =
Revised: L θ ( x, v, d ) =

Original: ( x (cid:48)(cid:48)× M , v (cid:48)(cid:48)× M , d ) , and our ﬂip operator F reverses the direction variable d , F ξ =
Revised: ( x (cid:48)(cid:48)× M , v (cid:48)(cid:48)× M , d ) , and our ﬂip operator F reverses the direction variable d , F ξ =

Original: ( x, v, − d ) .
Revised: ( x, v, − d ) .

Original: Written in terms of these modiﬁed operators, our proposal and acceptance probability are identical to those for standard HMC.
Revised: Written in terms of these modiﬁed operators, our proposal and acceptance probability are identical to those for standard HMC.

Original: Note, however, that this parameterization enables learning non-volume-preserving transformations, as the determinant of the Jacobian is a function of S x and
Revised: Note, however, that this parameterization enables learning non-volume-preserving transformations, as the determinant of the Jacobian is a function of S x and

Original: S v that does not necessarily evaluate to 1 .
Revised: S v that does not necessarily evaluate to 1 .

Original: This quantity is derived in Appendix B.
Revised: This quantity is derived in Appendix B.

Original: 4.1. MCMC T RANSITIONS
Revised: 4.1. MCMC T RANSITIONS

Original: For convenience, we denote by R an operator that re-samples the momentum and direction.
Revised: For convenience, we denote by R an operator that re-samples the momentum and direction.

Original: I.e., given ξ = ( x, v, d ) , R ξ = ( x, v (cid:48) , d (cid:48) ) where v (cid:48) ∼ N (0 , I ) , d (cid:48) ∼ U ( {− 1 , 1 } ) .
Revised: I.e., given ξ = ( x, v, d ) , R ξ = ( x, v (cid:48) , d (cid:48) ) where v (cid:48) ∼ N (0 , I ) , d (cid:48) ∼ U ( {− 1 , 1 } ) .

Original: Sampling thus consists of alternating application of the FL θ and R , in the following two steps each of which is a
Revised: Sampling thus consists of alternating application of the FL θ and R , in the following two steps each of which is a

Original: Markov transition that satisﬁes detailed balance with respect to p :
Revised: Markov transition that satisﬁes detailed balance with respect to p :

Original: ξ (cid:48)

Original: =

Original: FL θ ξ with probability A

Original: ( FL θ ξ | ξ ) (Equation 3), otherwise ξ (cid:48)

Original: = ξ . 2.
Revised: [Equation]

Original: ξ (cid:48) = R ξ This parameterization is effectively a generalization of standard HMC as it is non-volume preserving, with learnable parameters, and easily reduces to standard HMC for Q, S, T = 0 .
Revised: This parameterization is effectively a generalization of standard HMC as it is non-volume preserving, with learnable parameters, and easily reduces to standard HMC for Q, S, T = 0 .

Original: 4. L OSS AND T RAINING P ROCEDURE
Revised: 4. L OSS AND T RAINING P ROCEDURE

Original: We need some criterion to train the parameters θ that control the functions Q , S , and T .
Revised: We need some criterion to train the parameters θ that control the functions Q , S , and T .

Original: We choose a loss designed to reduce mixing time.
Revised: We choose a loss designed to reduce mixing time.

Original: Speciﬁcally, we aim to minimize lag-one autocorrelation.
Revised: Speciﬁcally, we aim to minimize lag-one autocorrelation.

Original: This is equivalent to maximizing expected squared jumped distance (Pasarica & Gelman, 2010).
Revised: This is equivalent to maximizing expected squared jumped distance (Pasarica & Gelman, 2010).

Original: For ξ, ξ (cid:48) in the extended state space, we deﬁne [Equation] [Equation]
Revised: For ξ, ξ (cid:48) in the extended state space, we deﬁne δ ( ξ (cid:48) , ξ ) =

Original: [Equation] d )) =
Revised: δ (( x (cid:48) , v (cid:48)

Original: [Equation]
Revised: [Equation]

Original: Expected squared jumped distance is thus E ξ ∼ p ( ξ )[ δ [Equation] θ
Revised: Expected squared jumped distance is thus E ξ ∼ p ( ξ )

Original: [Equation]

Original: , ξ ) A

Original: [Equation] FL [Equation] ξ | ξ )] .
Revised: [ δ ( FL θ ξ, ξ ) A ( FL θ ξ | ξ )] .

Original: However, this loss need not encourage mixing across the entire state space.
Revised: However, this loss need not encourage mixing across the entire state space.

Original: Indeed, maximizing this objective can lead to regions of state space where almost no mixing occurs, so long as the average squared distance traversed remains high.
Revised: Indeed, maximizing this objective can lead to regions of state space where almost no mixing occurs, so long as the average squared distance traversed remains high.

Original: To optimize both for typical and worst case behavior, we include a reciprocal term in the loss,
Revised: To optimize both for typical and worst case behavior, we include a reciprocal term in the loss,

Original: [Equation]
Revised: [Equation]

Original: (7) where λ is a scale parameter, capturing the characteristic length scale of the problem.
Revised: where λ is a scale parameter, capturing the characteristic length scale of the problem.

Original: The second term encourages typical moves to be large, while the ﬁrst term strongly penalizes the sampler if it is ever in a state where it cannot move effectively – δ ( ξ, ξ (cid:48) ) being small resulting in a large loss value.
Revised: The second term encourages typical moves to be large, while the ﬁrst term strongly penalizes the sampler if it is ever in a state where it cannot move effectively – δ ( ξ, ξ (cid:48) ) being small resulting in a large loss value.

Original: We train our sampler by minimizing this loss over both the target distribution and initialization distribution.
Revised: We train our sampler by minimizing this loss over both the target distribution and initialization distribution.

Original: Formally, given an initial distribution π 0 over X , we deﬁne q ( ξ ) 
Revised: Formally, given an initial distribution π 0 over X , we deﬁne q ( ξ ) =

Revised: π 0 ( x ) N ( v ; 0 , I ) p ( d ) , and minimize

Original: [Equation]
Revised: [Equation]

Original: and minimize

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst term of this loss encourages mixing as it considers our operator applied on draws from the distribution; the second term rewards fast burn-in; λ b controls the strength of the ‘burn-in’ regularization.
Revised: The ﬁrst term of this loss encourages mixing as it considers our operator applied on draws from the distribution; the second term rewards fast burn-in; λ b controls the strength of the ‘burn-in’ regularization.

Original: Given this loss, we exactly describe our training procedure in Algorithm 1.
Revised: Given this loss, we exactly describe our training procedure in Algorithm 1.

Original: It is important to note that each training iteration can be done with only one pass through the network and can be efﬁciently batched.
Revised: It is important to note that each training iteration can be done with only one pass through the network and can be efﬁciently batched.

Original: We further emphasize that this training procedure can be applied to any learnable operator whose Jacobian’s determinant is tractable, making it a general framework for training
Revised: We further emphasize that this training procedure can be applied to any learnable operator whose Jacobian’s determinant is tractable, making it a general framework for training

Original: MCMC proposals.
Revised: MCMC proposals.

Original: Input: Energy function U : X → R and its gradient ∇ x U : X → X , initial distribution over the augmented state space q , number of iterations n iters , number of leapfrogs M , learning rate schedule ( α t ) t ≤ n iters , batch size N , scale parameter λ and regularization strength λ b .
Revised: Input: Energy function U : X → R and its gradient ∇ x U : X → X , initial distribution over the augmented state space q , number of iterations n iters , number of leapfrogs M , learning rate schedule ( α t ) t ≤ n iters , batch size N , scale parameter λ and regularization strength λ b .

Original: Initialize the parameters of the sampler θ .
Revised: Initialize the parameters of the sampler θ .

Original: Initialize { ξ ( i ) p } i ≤ N from q ( ξ ) .
Revised: Initialize { ξ ( i ) p } i ≤ N from q ( ξ ) .

Original: for t = 0 to n iters − 1 do
Revised: for t = 0 to n iters − 1 do

Original: Sample a minibatch { ξ ( i ) q } i ≤ N from q ( ξ ) .
Revised: Sample a minibatch { ξ ( i ) q } i ≤ N from q ( ξ ) .

Original: L ← 0 for i = 1 to N do ξ ( i ) p ← R ξ ( i ) p
Revised: L ← 0 for i = 1 to N do ξ ( i ) p ← R ξ ( i ) p

Original: L ← L + (cid:96) λ (cid:16) ξ ( i ) p , FL θ ξ ( i ) p , A
Revised: L ← L + (cid:96) λ (cid:16) ξ ( i ) p , FL θ ξ ( i ) p , A

Original: ( FL θ ξ ( i ) p | ξ ( i ) p ) (cid:17) + λ [Table] (cid:96) [Table]
Revised: ( FL θ ξ ( i ) p | ξ ( i ) p ) (cid:17) + λ b (cid:96) λ (cid:16) ξ

Original: (cid:16)

Original: [Table]

Original: ( i ) q , FL
Revised: ( i ) q , FL

Original: [Table] ( i ) q , A
Revised: θ ξ ( i ) q , A

Original: ( FL [Table] ( i [Table]
Revised: ( FL θ ξ ( i ) q | ξ ( i ) q ) (cid:17)

Original: |

Original: [Table] ( i ) [Table] ) (cid:17) ξ ( i ) p ← FL θ ξ ( i ) p with probability A
Revised: ξ ( i ) p ← FL θ ξ ( i ) p with probability A

Original: ( FL θ ξ ( i ) p | ξ ( i ) p ) .
Revised: ( FL θ ξ ( i ) p | ξ ( i ) p ) .

Original: end for θ ← θ − α t ∇ θ L end for
Revised: end for θ ← θ − α t ∇ θ L end for

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We present an empirical evaluation of our trained sampler on a diverse set of energy functions.
Revised: We present an empirical evaluation of our trained sampler on a diverse set of energy functions.

Original: We ﬁrst present results on a collection of toy distributions capturing common pathologies of energy landscapes, followed by results on a task from machine learning: maximum-likelihood training of deep generative models.
Revised: We ﬁrst present results on a collection of toy distributions capturing common pathologies of energy landscapes, followed by results on a task from machine learning: maximum-likelihood training of deep generative models.

Original: For each, we compare against HMC with well-tuned step length and show signiﬁcant gains in mixing time.
Revised: For each, we compare against HMC with well-tuned step length and show signiﬁcant gains in mixing time.

Revised: Code implementing our algorithm is available online 1 .

Original: 5. V ARIED C OLLECTION OF E NERGY F UNCTIONS
Revised: 5. V ARIED C OLLECTION OF E NERGY F UNCTIONS

Original: We evaluate our L2HMC sampler on a diverse collection of energy functions, each posing different challenges for standard HMC.
Revised: We evaluate our L2HMC sampler on a diverse collection of energy functions, each posing different challenges for standard HMC.

Original: Ill-Conditioned Gaussian (ICG): Gaussian distribution with diagonal covariance spaced loglinearly between 10 − 2 and 10 2 .
Revised: Ill-Conditioned Gaussian (ICG): Gaussian distribution with diagonal covariance spaced loglinearly between 10 − 2 and 10 2 .

Original: This demonstrates that L2HMC can learn a diagonal inertia tensor.
Revised: This demonstrates that L2HMC can learn a diagonal inertia tensor.

Original: Strongly correlated Gaussian (SCG): We rotate a diagonal Gaussian with variances [10 2 , 10 − 2 ] by π 4 .
Revised: Strongly correlated Gaussian (SCG): We rotate a diagonal Gaussian with variances [10 2 , 10 − 2 ] by π 4 .

Original: This is an extreme version of an example from Neal (2011).
Revised: This is an extreme version of an example from Neal (2011).

Original: This problem shows that, although our parametric sampler only applies element-wise transformations, it can adapt to structure which is not axis-aligned.
Revised: This problem shows that, although our parametric sampler only applies element-wise transformations, it can adapt to structure which is not axis-aligned.

Original: Mixture of Gaussians (MoG): Mixture of two isotropic Gaussians with σ 2 = 0 .
Revised: Mixture of Gaussians (MoG): Mixture of two isotropic Gaussians with σ 2 = 0 .

Original: 1 , and centroids separated by distance 4 .
Revised: 1 , and centroids separated by distance 4 .

Original: The means are thus about 12 standard deviations apart, making it almost impossible for HMC to mix between modes.
Revised: The means are thus about 12 standard deviations apart, making it almost impossible for HMC to mix between modes.

Original: Rough Well : Similar to an example from Sohl-Dickstein et al.
Revised: Rough Well : Similar to an example from Sohl-Dickstein et al.

Original: (2014), for a given η > 0 , U ( x ) = 12 x T x + η (cid:80) i cos( x i η ) .
Revised: (2014), for a given η > 0 , U ( x ) =2 x T x + η (cid:80) i cos( x i η ) .

Original: For small η the energy itself is altered negligibly, but its gradient is perturbed by a high frequency noise oscillating between − 1 and 1 .
Revised: For small η the energy itself is altered negligibly, but its gradient is perturbed by a high frequency noise oscillating between − 1 and 1 .

Original: In our experiments, we choose η = 10 − 2 .
Revised: In our experiments, we choose η = 10 − 2 .

Original: For each of these distributions, we compare against HMC with the same number of leapfrog steps and a well-tuned step-size.
Revised: For each of these distributions, we compare against HMC with the same number of leapfrog steps and a well-tuned step-size.

Original: To compare mixing time, we plot auto-correlation for each method and report effective sample size (ESS).
Revised: To compare mixing time, we plot auto-correlation for each method and [Figure] [Figure] [Figure]

Original: We compute those quantities in the same way as Sohl-Dickstein et al.
Revised: [Table] [Figure] −−−−   6 − Chain

Revised: [Figure]

Revised: with L2HMC −−−−   6 −

Revised: Chain samples with A-NICE

Revised: [Figure]

Revised: MC

Revised: [Figure]

Revised: report effective sample size (ESS).

Revised: We compute those quantities in the same way as Sohl-Dickstein et al.

Original: We observe that samplers trained with L2HMC show greatly improved autocorrelation and ESS on the presented tasks, providing more than 106 × improved ESS on the SCG task.
Revised: We observe that samplers trained with L2HMC show greatly improved autocorrelation and ESS on the presented tasks, providing more than 106 × improved ESS on the SCG task.

Original: In addition, for the MoG, we show that L2HMC can easily mix between modes while standard HMC gets stuck in a mode, unable to traverse the low density zone.
Revised: In addition, for the MoG, we show that L2HMC can easily mix between modes while standard HMC gets stuck in a mode, unable to traverse the low density zone.

Original: Experimental details, as well as a comparison with LAHMC (Sohl-Dickstein et al., 2014), are shown in Appendix C.
Revised: Experimental details, as well as a comparison with LAHMC (Sohl-Dickstein et al., 2014), are shown in Appendix C.

Original: Comparison to A-NICE-MC (Song et al., 2017)
Revised: Comparison to A-NICE-MC (Song et al., 2017)

Original: In addition to the well known challenges associated with adversarial training (Arjovsky et al., 2017), we note that parameterization using a [Figure] [Figure] [Figure]
Revised: In addition to the well known challenges associated with adversarial training (Arjovsky et al., 2017), we note that parameterization using a volume-preserving operator can dramatically fail on simple energy landscapes.

Original: [Figure]

Original: [Table]

Original: 2 . 61 × 10 −(cid:29)

Original: [Figure]

Original: −−−−   6 −

Original: Chain samples with L2HMC −−−−   6 −

Original: Chain samples with A-NICE-MC

Original: [Figure]

Original: volume-preserving operator can dramatically fail on simple energy landscapes.

Original: We build off of the mog2 experiment presented in (Song et al., 2017), which is a 2 -d mixture of isotropic Gaussians separated by a distance of 10 with variances 0 .
Revised: We build off of the mog2 experiment presented in (Song et al., 2017), which is a 2 -d mixture of isotropic Gaussians separated by a distance of 10 with variances 0 .

Original: We consider that setup but increase the ratio of variances: σ 21 = 3 , σ 22 = 0 .
Revised: We consider that setup but increase the ratio of variances: σ 21 = 3 , σ 22 = 0 .

Original: We show in Figure 1d sample chains trained with L2HMC and
Revised: We show in Figure 1d sample chains trained with L2HMC and

Original: A-NICE-MC; A-NICE-MC cannot effectively mix between the two modes as only a fraction of the volume of the large mode can be mapped to the small one, making it highly improbable to traverse.
Revised: A-NICE-MC; A-NICE-MC cannot effectively mix between the two modes as only a fraction of the volume of the large mode can be mapped to the small one, making it highly improbable to traverse.

Original: This is also an issue for HMC.
Revised: This is also an issue for HMC.

Original: On the other hand, L2HMC can both traverse the low-density region between modes, and map a larger volume in the left mode to a smaller volume in the right mode.
Revised: On the other hand, L2HMC can both traverse the low-density region between modes, and map a larger volume in the left mode to a smaller volume in the right mode.

Original: It is important to note that the distance between both clusters is less than in the mog2 case, and it is thus a good diagnostic of the shortcomings of volume-preserving transformations.
Revised: It is important to note that the distance between both clusters is less than in the mog2 case, and it is thus a good diagnostic of the shortcomings of volume-preserving transformations.

Original: 5. L ATENT -V ARIABLE G ENERATIVE M ODEL
Revised: 5. L ATENT -V ARIABLE G ENERATIVE M ODEL

Original: We apply our learned sampler to the task of training, and sampling from the posterior of, a latentvariable generative model.
Revised: We apply our learned sampler to the task of training, and sampling from the posterior of, a latentvariable generative model.

Original: The model consists of a latent variable
Revised: The model consists of a latent variable

Original: z ∼ p ( z ) , where we choose p ( z ) =
Revised: z ∼ p ( z ) , where we choose p ( z ) =

Original: N ( z ; 0 , I ) , and a conditional distribution p ( x | z ) which generates the image x .
Revised: N ( z ; 0 , I ) , and a conditional distribution p ( x | z ) which generates the image x .

Original: Given a family of parametric ‘decoders’ { z (cid:55)→ p ( x | z ; φ ) , φ ∈ Φ } , and a set of samples D
Revised: Given a family of parametric ‘decoders’ { z (cid:55)→ p ( x | z ; φ ) , φ ∈ Φ } , and a set of samples D

Original: = { x ( i ) } i ≤ N ,  training involves ﬁnding φ ∗ =
Revised: = { x ( i ) } i ≤ N , [Figure] training involves ﬁnding φ ∗ =

Original: arg max φ ∈
Revised: arg max φ ∈

Original: Φ p ( D ; φ ) .
Revised: Φ p ( D ; φ ) .

Original: However, the log-likelihood is intractable as p ( x ; φ ) =
Revised: However, the log-likelihood is intractable as p ( x ; φ ) =

Original: (cid:82) p ( x | z ; φ ) p ( z )
Revised: (cid:82) p ( x | z ; φ ) p ( z )

Original: d z .
Revised: d z .

Original: To remedy that problem, Kingma & Welling (2013) proposed jointly training an approximate posterior q ψ that maximizes a tractable lower-bound on the log-likelihood:
Revised: To remedy that problem, Kingma & Welling (2013) proposed jointly training an approximate posterior q ψ that maximizes a tractable lower-bound on the log-likelihood:

Original: [Equation]
Revised: [Equation]

Original: where q ψ ( z | x ) is a tractable conditional distribution with parameters ψ , typically parameterized by a neural network.
Revised: where q ψ ( z | x ) is a tractable conditional distribution with parameters ψ , typically parameterized by a neural network.

Original: Recently, to improve upon well-known pitfalls like over-pruning (Burda et al., 2015) of the VAE, Hoffman (2017) proposed HMC-DLGM.
Revised: Recently, to improve upon well-known pitfalls like over-pruning (Burda et al., 2015) of the VAE, Hoffman (2017) proposed HMC-DLGM.

Original: For a data sample x ( i ) , after obtaining a sample from the approximate posterior q ψ ( ·| x ( i ) ) , Hoffman (2017) runs a MCMC algorithm with[Figure] energy function U ( z, x ( i ) )
Revised: For a data sample x ( i ) , after obtaining a sample from the approximate posterior q ψ ( ·| x ( i ) ) , Hoffman (2017) runs a MCMC algorithm with energy function U ( z, x ( i ) )

Original: = − log p ( z ) − log p ( x ( i ) | z ; φ ) to obtain a more exact posterior sample from p ( z | x ( i ) ; φ ) .
Revised: = − log p ( z ) − log p ( x ( i ) | z ; φ ) to obtain a more exact posterior sample from p ( z | x ( i ) ; φ ) .

Original: Given that better posterior sample z (cid:48) , the algorithm maximizes log p ( x ( i ) | z (cid:48) ; φ ) .
Revised: Given that better posterior sample z (cid:48) , the algorithm maximizes log p ( x ( i ) | z (cid:48) ; φ ) .

Original: To show the beneﬁts of L2HMC, we borrow the method from Hoffman (2017), but replace
Revised: To show the beneﬁts of L2HMC, we borrow the method from Hoffman (2017), but replace

Original: HMC by jointly training an L2HMC sampler to improve the efﬁciency of the posterior sampling.
Revised: HMC by jointly training an L2HMC sampler to improve the efﬁciency of the posterior sampling.

Original: We call this model L2HMC-DLGM .
Revised: We call this model L2HMC-DLGM .

Original: A diagram of our model and a formal description of our training procedure are presented in Appendix D.We deﬁne, for ξ = { z, v, d } , r ( ξ | x ; ψ ) (cid:44) q ψ ( z | x ) N ( v ; 0 , I ) U ( d ; {− 1 , 1 } ) .
Revised: A diagram of our model and a formal description of our training procedure are presented in Appendix D. We deﬁne, for ξ = { z, v, d } , r ( ξ | x ; ψ ) (cid:44) q ψ ( z | x ) N ( v ; 0 , I ) U ( d ; {− 1 , 1 } ) .

Original: In the subsequent sections, we compare our method to the standard VAE model from Kingma &
Revised: In the subsequent sections, we compare our method to the standard VAE model from Kingma &

Original: Welling (2013) and HMC-DGLM from Hoffman (2017).
Revised: Welling (2013) and HMC-DGLM from Hoffman (2017).

Original: It is important to note that, since our sampler is trained jointly with p φ and q ψ , it performs exactly the same number of gradient computations of the energy function as HMC .
Revised: It is important to note that, since our sampler is trained jointly with p φ and q ψ , it performs exactly the same number of gradient computations of the energy function as HMC .

Original: We ﬁrst show that training a latent variable generative model with
Revised: We ﬁrst show that training a latent variable generative model with

Original: L2HMC results in better generative models both qualitatively and quantitatively.
Revised: L2HMC results in better generative models both qualitatively and quantitatively.

Original: We then show that our improved sampler enables a more expressive, non-Gaussian, posterior.
Revised: We then show that our improved sampler enables a more expressive, non-Gaussian, posterior.

Original: Implementation details: Our decoder ( p φ ) is a neural network with 2 fully connected layers, with 1024 units each and softplus non-linearities, and outputs Bernoulli activation probabilities for each pixel.
Revised: Implementation details: Our decoder ( p φ ) is a neural network with 2 fully connected layers, with 1024 units each and softplus non-linearities, and outputs Bernoulli activation probabilities for each pixel.

Original: The encoder ( q ψ ) has the same architecture, returning mean and variance for the approximate posterior.
Revised: The encoder ( q ψ ) has the same architecture, returning mean and variance for the approximate posterior.

Original: Our model was trained for 300 epochs with Adam (Kingma & Ba, 2014) and a learning rate α = 10 − 3 .
Revised: Our model was trained for 300 epochs with Adam (Kingma & Ba, 2014) and a learning rate α = 10 − 3 .

Original: All experiments were done on the dynamically binarized MNIST dataset (LeCun).
Revised: All experiments were done on the dynamically binarized MNIST dataset (LeCun).

Original: 5.2. S AMPLE Q UALITY AND D ATA LIKELIHOOD
Revised: 5.2. S AMPLE Q UALITY AND D ATA LIKELIHOOD

Original: We ﬁrst present samples from decoders trained with L2HMC, HMC and the ELBO (i.e. vanilla
Revised: We ﬁrst present samples from decoders trained with L2HMC, HMC and the ELBO (i.e. vanilla

Original: VAE).
Revised: VAE).

Original: Although higher log likelihood does not necessarily correspond to better samples (Theis et al., 2015), we can see in Figure 5, shown in the Appendix, that the decoder trained with L2HMC generates sharper samples than the compared methods.
Revised: Although higher log likelihood does not necessarily correspond to better samples (Theis et al., 2015), we can see in Figure 5, shown in the Appendix, that the decoder trained with L2HMC generates sharper samples than the compared methods.

Original: We now compare our method to HMC in terms of log-likelihood of the data.
Revised: We now compare our method to HMC in terms of log-likelihood of the data.

Original: As we previously stated, the marginal likelihood of a data point x ∈ X is not tractable as it requires integrating p ( x, z ) over a high-dimensional space.
Revised: As we previously stated, the marginal likelihood of a data point x ∈ X is not tractable as it requires integrating p ( x, z ) over a high-dimensional space.

Original: However, we can estimate it using annealed importance sampling (AIS; Neal (2001)).
Revised: However, we can estimate it using annealed importance sampling (AIS; Neal (2001)).

Original: Following Wu et al.
Revised: Following Wu et al.

Original: (2016), we evaluate our generative models on both training and held-out data.
Revised: (2016), we evaluate our generative models on both training and held-out data.

Original: In Figure 2, we plot the data’s log-likelihood against the number of gradient computation steps for both HMC-DGLM and L2HMC-DGLM.
Revised: In Figure 2, we plot the data’s log-likelihood against the number of gradient computation steps for both HMC-DGLM and L2HMC-DGLM.

Original: We can see that for a similar number of gradient computations, L2HMC-DGLM achieves higher likelihood for both training and held-out data.
Revised: We can see that for a similar number of gradient computations, L2HMC-DGLM achieves higher likelihood for both training and held-out data.

Original: This is a strong indication that L2HMC provides signiﬁcantly better posterior samples . 
Revised: This is a strong indication that L2HMC provides signiﬁcantly better posterior samples . [Figure]

Original: 5.2. I NCREASED E XPRESSIVITY OF THE POSTERIOR
Revised: 5.2. I NCREASED E XPRESSIVITY OF THE POSTERIOR

Original: In the standard VAE framework, approximate posteriors are often parametrized by a Gaussian, thus making a strong assumption of uni-modality.
Revised: In the standard VAE framework, approximate posteriors are often parametrized by a Gaussian, thus making a strong assumption of uni-modality.

Original: In this section, we show that using L2HMC to sample from the posterior enables learning of a richer posterior landscape.
Revised: In this section, we show that using L2HMC to sample from the posterior enables learning of a richer posterior landscape.

Original: Block Gibbs Sampling To highlight our ability to capture more expressive posteriors, we in-paint the top of an image using Block Gibbs Sampling using the approximate posterior or L2HMC.
Revised: Block Gibbs Sampling To highlight our ability to capture more expressive posteriors, we in-paint the top of an image using Block Gibbs Sampling using the approximate posterior or L2HMC.

Original: Formally, let x 0 be the starting image.
Revised: Formally, let x 0 be the starting image.

Original: We denote top or bottom-half pixels as x top0 and x bottom0 .
Revised: We denote top or bottom-half pixels as x top0 and x bottom0 .

Original: At each step t , we sample z ( t ) ∼ p ( z | x t ; θ ) , sample ˜ x ∼ p ( x | z t ; θ ) .
Revised: At each step t , we sample z ( t ) ∼ p ( z | x t ; θ ) , sample ˜ x ∼ p ( x | z t ; θ ) .

Original: We then set x top t +1 =
Revised: We then set x top t +1 =

Original: ˜ x top and x bottom t + = x bottom.
Revised: ˜ x top and x bottom t + = x bottom0 .

Original: We compare the results obtained by sampling from p ( z | x ; θ ) using q ψ (i.e. the[Figure] approximate posterior) vs. our trained sampler.
Revised: We compare the results obtained by sampling from p ( z | x ; θ ) using q ψ (i.e. the approximate posterior) vs. our trained sampler.

Original: The results are reported in Figure 3a.
Revised: The results are reported in Figure 3a.

Original: We can see that
Revised: We can see that

Original: L2HMC easily mixes between modes (3, 5, 8, and plausibly 9 in the ﬁgure) while the approximate posterior gets stuck on the same reconstructed digit (3 in the ﬁgure).
Revised: L2HMC easily mixes between modes (3, 5, 8, and plausibly 9 in the ﬁgure) while the approximate posterior gets stuck on the same reconstructed digit (3 in the ﬁgure).

Original: Visualization of the posterior After training a decoder with L2HMC, we randomly choose an element x 0 ∈ D and run 512 parallel L2HMC chains for 20 , 000 Metropolis-Hastings steps.
Revised: Visualization of the posterior After training a decoder with L2HMC, we randomly choose an element x 0 ∈ D and run 512 parallel L2HMC chains for 20 , 000 Metropolis-Hastings steps.

Original: We then ﬁnd the direction of highest variance, project the samples along that direction and show a histogram in Figure 3b.
Revised: We then ﬁnd the direction of highest variance, project the samples along that direction and show a histogram in Figure 3b.

Original: This plot shows non-Gaussianity in the latent space for the posterior.
Revised: This plot shows non-Gaussianity in the latent space for the posterior.

Original: Using our improved sampler enables the decoder to make use of a more expressive posterior, and enables the encoder to sample from this non-Gaussian posterior.
Revised: Using our improved sampler enables the decoder to make use of a more expressive posterior, and enables the encoder to sample from this non -Gaussian posterior.

Original:  F UTURE W ORK
Revised:  F UTURE W ORK

Original: The loss in Section 4.2 targets lag-one autocorrelation.
Revised: The loss in Section 4.2 targets lag-one autocorrelation.

Original: It should be possible to extend this to also target lag-two and higher autocorrelations.
Revised: It should be possible to extend this to also target lag-two and higher autocorrelations.

Original: It should also be possible to extend this loss to reward fast decay in the autocorrelation of other statistics of the samples, for instance the sample energy as well as the sample position.
Revised: It should also be possible to extend this loss to reward fast decay in the autocorrelation of other statistics of the samples, for instance the sample energy as well as the sample position.

Original: These additional statistics could also include learned statistics of the samples, combining beneﬁts of the adversarial approach of (Song et al., 2017) with the current work.
Revised: These additional statistics could also include learned statistics of the samples, combining beneﬁts of the adversarial approach of (Song et al., 2017) with the current work.

Original: Our learned generalization of HMC should prove complementary to several other research directions related to HMC.
Revised: Our learned generalization of HMC should prove complementary to several other research directions related to HMC.

Original: It would be interesting to explore combining our work with the use of HMC in a minibatch setting (Chen et al., 2014); with shadow Hamiltonians (Izaguirre & Hampton, 2004); with gradient pre-conditioning approaches similar to those used in Riemannian HMC (Girolami et al., 2009; Betancourt, 2013); with the use of alternative HMC accept-reject rules (Sohl-Dickstein et al., 2014; Berger et al., 2015); with the use of non-canonical Hamiltonian dynamics (Tripuraneni et al., 2016); with variants of AIS adapted to HMC proposals (Sohl-Dickstein & Culpepper, 2012); with the extension of HMC to discrete state spaces (Zhang et al., 2012); and with the use of alternative
Revised: It would be interesting to explore combining our work with the use of HMC in a minibatch setting (Chen et al., 2014); with shadow Hamiltonians (Izaguirre & Hampton, 2004); with gradient pre-conditioning approaches similar to those used in Riemannian HMC (Girolami et al., 2009; Betancourt, 2013); with the use of alternative HMC accept-reject rules (Sohl-Dickstein et al., 2014; Berger et al., 2015); with the use of non-canonical Hamiltonian dynamics (Tripuraneni et al., 2016); with variants of AIS adapted to HMC proposals (Sohl-Dickstein & Culpepper, 2012); with the extension of HMC to discrete state spaces (Zhang et al., 2012); and with the use of alternative

Original: Hamiltonian integrators (Creutz & Gocksch, 1989; Chao et al., 2015).
Revised: Hamiltonian integrators (Creutz & Gocksch, 1989; Chao et al., 2015).

Original: Finally, our work is also complementary to other methods not utilizing gradient information.
Revised: Finally, our work is also complementary to other methods not utilizing gradient information.

Original: For example, we could incorporate the intuition behind Multiple Try Metropolis schemes (Martino &
Revised: For example, we could incorporate the intuition behind Multiple Try Metropolis schemes (Martino &

Original: Read, 2013) by having several parametric operators and training each one when used.
Revised: Read, 2013) by having several parametric operators and training each one when used.

Original: In addition, one could draw inspiration from the adaptive literature (Haario et al., 2001; Andrieu & Thoms, 2008) or component-wise strategies (Gilks & Wild, 1992).
Revised: In addition, one could draw inspiration from the adaptive literature (Haario et al., 2001; Andrieu & Thoms, 2008) or component-wise strategies (Gilks & Wild, 1992).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we presented a general method to train expressive MCMC kernels parameterized with deep neural networks.
Revised: In this work we presented a general method to train expressive MCMC kernels parameterized with deep neural networks.

Original: Given a target distribution p , analytically known up to a constant, our method provides a fast-mixing sampler, able to efﬁciently explore the state space.
Revised: Given a target distribution p , analytically known up to a constant, our method provides a fast-mixing sampler, able to efﬁciently explore the state space.

Original: Our hope is that our method can be utilized in a “black-box” manner, in domains where sampling constitutes a huge bottleneck such as protein foldings (Sch¨utte et al., 1999) or physics simulations (Olsson, 1995).
Revised: Our hope is that our method can be utilized in a “black-box” manner, in domains where sampling constitutes a huge bottleneck such as protein foldings (Sch¨utte et al., 1999) or physics simulations (Olsson, 1995) A CKNOWLEDGMENTS

Revised: We would like to thank Ben Poole, Aditya Grover, David Belanger and Colin Raffel for insightful comments on the draft, Mohammad Norouzi for support and encouragement launching the project, and Jiaming Song for discussions and help running A-NICE-MC.

Original: Appendix A R EVERSE LEAPFROG OPERATOR
Revised: Appendix A R EVERSE LEAPFROG OPERATOR

Original: Let ξ = { x, v, d } in the extended state space with d = − 1 .
Revised: Let ξ = { x, v, d } in the extended state space with d = − 1 .

Original: Here, we describe the leapfrog updates for a single time step t , this consists of inverting the equations presented in the corresponding section.
Revised: Here, we describe the leapfrog updates for a single time step t , this consists of inverting the equations presented in the corresponding section.

Original: Let ζ 1 = { x, v, t } , we have:
Revised: Let ζ 1 = { x, v, t } , we have:

Original: [Equation]
Revised: [Equation]

Original: With the notation from Section
Revised: With the notation from Section

Original: [Equation]
Revised: [Equation]

Original: let ζ 2 (cid:44){ x m
Revised: let ζ 2 (cid:44)

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: ,v, t}
Revised: , t

Original: [Equation]
Revised: [Equation]

Original: Let us denote ζ 3 (cid:44) {x (cid:48) ¯ m
Revised: Let us denote ζ 3 (cid:44) {

Original: [Equation]
Revised: [Equation]

Original: , v, t } :
Revised: t } :

Original: [Equation]
Revised: [Equation]

Original: = [Equation] (cid:48)
Revised: Finally

Original: [Equation]
Revised: [Equation]

Original: Finally, the last update, with ζ 4 (cid:44) { x (cid:48)(cid:48) , ∂ x U ( x (cid:48)(cid:48)
Revised: the last update, with ζ 4 (cid:44) { x (cid:48)(cid:48) , ∂ x U ( x (cid:48)(cid:48)

Original: [Equation]
Revised: [Equation]

Original: , t } :
Revised: , t } :

Original: [Equation]
Revised: [Equation]

Original: It is important to note that to invert L θ , these steps should be ran for t from M to 1 .
Revised: It is important to note that to invert L θ , these steps should be ran for t from M to 1 .

Original: B D ETERMINANT OF THE J ACOBIAN
Revised: B D ETERMINANT OF THE J ACOBIAN

Original: Given the derivations (and notations) from Section 4, for the forward operator L θ , we can immediately compute the Jacobian:
Revised: Given the derivations (and notations) from Section 4, for the forward operator L θ , we can immediately compute the Jacobian:

Original: [Equation]
Revised: [Equation]

Original: Where ζ ti denotes the intermediary variable ζ i at time step t and d is the direction of ξ i.e. ξ = { x, v, d } .
Revised: Where ζ ti denotes the intermediary variable ζ i at time step t and d is the direction of ξ i.e. ξ = { x, v, d } .

Original: C E XPERIMENTAL DETAILS OF S ECTIONC. I MPLEMENTATION DETAILS
Revised: C E XPERIMENTAL DETAILS OF S ECTIONC. I MPLEMENTATION DETAILS

Original: First of all, we keep separate parameters for the network responsible for updating v and those updating x .
Revised: First of all, we keep separate parameters for the network responsible for updating v and those updating x .

Original: The architectures are the same.
Revised: The architectures are the same.

Original: Let us take the example of Q v , S v , T v .
Revised: Let us take the example of Q v , S v , T v .

Original: The time step t is given as input to the MLP, encoded as τ ( t ) = (cid:0) cos( 2 πtM ) , sin( 2 πtM ) (cid:1) .
Revised: The time step t is given as input to the MLP, encoded as τ ( t ) = (cid:0) cos( 2 πtM ) , sin( 2 πtM ) (cid:1) .

Original: σ ( · ) denotes the ReLU non-linearity.
Revised: σ ( · ) denotes the ReLU non-linearity.

Original: For n h hidden units per layer:
Revised: For n h hidden units per layer:

Original: • We ﬁrst compute h 1 =
Revised: • We ﬁrst compute h 1 =

Original: σ ( W 1 x + W 2 v + W 3 τ ( t ) + b ) ( h ∈ R n h ).
Revised: σ ( W 1 x + W 2 v + W 3 τ ( t ) + b ) ( h ∈ R n h ).

Original: • h 2 =
Revised: • h 2 =

Original: σ ( W 4 h + b 4 ) ∈ R n h • S v = λ s tanh
Revised: σ ( W 4 h + b 4 ) ∈ R n h • S v = λ s tanh

Original: ( W s 
Revised: ( W s h

Original: [Equation]
Revised: [Equation]

Original: In Section 5.1, the Q, S, T are neural networks with 2 hidden layers with 10 ( 100 for the 50 -d ICG) units and ReLU non-linearities.
Revised: In Section 5.1, the Q, S, T are neural networks with 2 hidden layers with 10 ( 100 for the 50 -d ICG) units and ReLU non-linearities.

Original: We train with Adam (Kingma & Ba, 2014) and a learning rate α = 10 − 3 .
Revised: We train with Adam (Kingma & Ba, 2014) and a learning rate α = 10 − 3 .

Original: We train for 5 , 000 iterations with a batch size of 200 .
Revised: We train for 5 , 000 iterations with a batch size of 200 .

Original: λ b was set to 0 for ICG and SCG and to 1 for MoG and Rough Well.
Revised: λ b was set to 0 for ICG and SCG and to 1 for MoG and Rough Well.

Original: For the MoG tasks, we train our sampler with a temperature parameter that we continuously anneal; we evaluate the trained sampler without using temperature. 
Revised: For the MoG tasks, we train our sampler with a temperature parameter that we continuously anneal; we evaluate the trained sampler without using temperature. [Figure]

Original: [Figure]
Revised: [Figure]

Original: C. A UTO - CORRELATION AND ESS
Revised: A UTO - CORRELATION AND ESS

Original: Let ( x τ ) τ ≤ T be a set of correlated samples converging to the distribution p with mean µ and covariance Σ .
Revised: Let ( x τ ) τ ≤ T be a set of correlated samples converging to the distribution p with mean µ and covariance Σ .

Original: We deﬁne auto-correlation at time t as:
Revised: We deﬁne auto-correlation at time t as:

Original: [Equation]
Revised: [Equation]

Original: We can now deﬁne effective sample size (ESS) as:
Revised: We can now deﬁne effective sample size (ESS) as:

Original: [Equation]
Revised: [Equation]

Original: Similar to Hoffman & Gelman (2014), we truncate the sum when the auto-correlation goes below 0 .
Revised: Similar to Hoffman & Gelman (2014), we truncate the sum when the auto-correlation goes below 0 .

Original: C. C OMPARISON WITH LAHMC
Revised: C. C OMPARISON WITH LAHMC

Original: We compare our trained sampler with LAHMC (Sohl-Dickstein et al., 2014).
Revised: We compare our trained sampler with LAHMC (Sohl-Dickstein et al., 2014).

Original: Results are reported in
Revised: Results are reported in

Original: Table 1.
Revised: Table 1.

Original: L2HMC largely outperforms LAHMC on all task.
Revised: L2HMC largely outperforms LAHMC on all task.

Original: LAHMC is also unable to mix between modes for the MoG task.
Revised: LAHMC is also unable to mix between modes for the MoG task.

Original: We also note that L2HMC could be easily combined with LAHMC, by replacing the leapfrog integrator of LAHMC with the learned one of L2HMC.
Revised: We also note that L2HMC could be easily combined with LAHMC, by replacing the leapfrog integrator of LAHMC with the learned one of L2HMC.

Original: Distribution
Revised: Distribution

Original: Gradient Evaluations
Revised: Gradient Evaluations

Original: [Equation] L2HMC
Revised: ESS-L2HMC

Original: ESS-LAHMC
Revised: ESS-LAHMC

Original: Ratio [Equation] ICG 156 .21 .7 .
Revised: Ratio 50 -d ICG 156 .21 .7 .

Original: Rough Well 12 .8 .1 .{+.
+}
Revised: Rough Well 12 .8 .1 {+.
+}

Original: [Equation]

Original: d SCG  16 .14 .
Revised: -d SCG  16 .14 .

Original: [Table]
Revised: [Table]

Original: D L2HMC-DGLM D. T RAINING ALGORITHM
Revised: D L2HMC-DGLM D. T RAINING ALGORITHM

Original: In this section, we present our training algorithm as well as a diagram explaining L2HMC-DGLM.
Revised: In this section, we present our training algorithm as well as a diagram explaining L2HMC-DGLM.

Original: For conciseness, given our operator L θ , we denote by K θ ( ·| x ) the distribution over next state given sampling of a momentum and direction and the Metropolis-Hastings step.
Revised: For conciseness, given our operator L θ , we denote by K θ ( ·| x ) the distribution over next state given sampling of a momentum and direction and the Metropolis-Hastings step.

Original: D. I MPLEMENTATION DETAILS OF L2HMC-DGLM
Revised: D. I MPLEMENTATION DETAILS OF L2HMC-DGLM

Original: Similar to our L2HMC training on unconditional sampling, we share weights across Q, S and T .
Revised: Similar to our L2HMC training on unconditional sampling, we share weights across Q, S and T .

Original: In addition, the auxiliary variable x (here the image from MNIST) is ﬁrst passed through a 2 -layer neural network, with softplus non-linearities and 512 hidden units.
Revised: In addition, the auxiliary variable x (here the image from MNIST) is ﬁrst passed through a 2 -layer neural network, with softplus non-linearities and 512 hidden units.

Original: This input is given to both
Revised: This input is given to both

Original: Input: 
Revised: Input: rate schedule

Original: Randomly initialize the decoder’s parameters φ and the approximate posterior ψ .
Revised: Randomly initialize the decoder’s parameters φ and the approximate posterior ψ .

Original: Initialize the parameters of the sampler θ with M leapfrog steps.
Revised: Initialize the parameters of the sampler θ with M leapfrog steps.

Original: for t = 0 to n iters − 1 do Randomly sample a minibatch
Revised: for t = 0 to n iters − 1 do Randomly sample a minibatch

Original: B from the dataset D .
Revised: B from the dataset D .

Original: L ELBO , L Sampler , L Decoder ← 0 for x ( b ) ∈ B do Sample ξ ( b )∼ r ( ·| x ( b ) ; ψ ) .
Revised: L ELBO , L Sampler , L Decoder ← 0 for x ( b ) ∈ B do Sample ξ ( b )∼ r ( ·| x ( b ) ; ψ ) .

Original: L ELBO ← p ( x ( b ) | z ( b ) 0 ; φ ) − KL( q ψ ( z | x ( b ) )
Revised: L ELBO ← p ( x ( b ) | z ( b ) 0 ; φ ) − KL( q ψ ( z | x ( b ) )

Original: || p ( z ))
Revised: || p ( z ))

Original: (cid:46) With ξ ( b ) 

Revised: (cid:46) With ξ ( b )

Original: = { z ( b ) 0 , v ( b ) 0 , d ( b ) 0 }
Revised: = { z ( b ) 0 , v ( b ) 0 , d ( b ) 0 }

Original: Deﬁne the energy function U x ( b ) ( z ) =
Revised: Deﬁne the energy function U x ( b ) ( z ) =

Original: − log p ( x ( b ) | z ; θ ) − log p ( z )
Revised: − log p ( x ( b ) | z ; θ ) − log p ( z )

Original: L Sampler ←λ ← (cid:113) Var( q ψ ( z ( b ) 

Revised: L Sampler ←λ ← (cid:113) Var( q ψ ( z ( b )

Original: | x ( b ) ) for j = 0 to J − 1 do ξ ( b ) j ← R ξ ( b ) j
Revised: | x ( b ) ) for j = 0 to J − 1 do ξ ( b ) j ← R ξ ( b ) j

Original: L Sampler ← L Sampler + (cid:96) λ ( ξ ( b ) j , FL
Revised: L Sampler ← L Sampler + (cid:96) λ ( ξ ( b ) j , FL

Original: θ ξ ( b ) j , A ( FL θ ξ ( b ) j | ξ ( b ) j ))
Revised: θ ξ ( b ) j , A ( FL θ ξ ( b ) j | ξ ( b ) j ))

Original: Set ξ ( b ) j +1 to FL θ ξ ( b ) j with probability A
Revised: Set ξ ( b ) j +1 to FL θ ξ ( b ) j with probability A

Original: ( FL θ ξ ( b ) j | ξ ( b ) j ) .
Revised: ( FL θ ξ ( b ) j | ξ ( b ) j ) .

Original: end for
Revised: end for

Original: L Decoder ← L Decoder + log p ( x ( b ) | z ( s ) J ; φ ) (cid:46) With ξ ( b ) J = { z ( b ) J , v ( b ) J , d ( b ) J } end for φ ← φ + α t ∇ φ L Decoder ψ ← ψ + α t ∇ ψ L ELBO θ ← θ + α t ∇ θ L Sampler end for
Revised: L Decoder ← L Decoder + log p ( x ( b ) | z ( s ) J ; φ ) (cid:46) With ξ ( b ) J = { z ( b ) J , v ( b ) J , d ( b ) J } end for φ ← φ + α t ∇ φ L Decoder ψ ← ψ + α t ∇ ψ L ELBO θ ← θ + α t ∇ θ L Sampler end for

Original: [Figure]
Revised: [Figure]

Original: networks {·} x and {·} v .
Revised: networks {·} x and {·} v .

Original: The architecture then consists of 2 hidden layers of 200 units and ReLU non-linearities.
Revised: The architecture then consists of 2 hidden layers of 200 units and ReLU non-linearities.

Original: For λ (scale parameter of the loss), we use the standard deviation of the approximate posterior.
Revised: For λ (scale parameter of the loss), we use the standard deviation of the approximate posterior.

Original: AIS Evaluation For each data point, we run 20 Markov Chains in parallel, 10 , 000 annealing steps with 10 leapfrogs per step and choose the step size for an acceptance rate of 0 . 65 .
Revised: AIS Evaluation For each data point, we run 20 Markov Chains in parallel, 10 , 000 annealing steps with 10 leapfrogs per step and choose the step size for an acceptance rate of 0 . 65 .

Original: D. MNIST S AMPLES
Revised: D. MNIST S AMPLES

Original: We show in Figure 5 samples from the three evaluated models: VAE (Kingma & Welling, 2013),
Revised: We show in Figure 5 samples from the three evaluated models: VAE (Kingma & Welling, 2013),

Original: HMC-DGLM (Hoffman, 2017) and L2HMC-DGLM.
Revised: HMC-DGLM (Hoffman, 2017) and L2HMC-DGLM.