Original: M ETA -L EARNING FOR S TOCHASTIC G RADIENT MCMC
Revised: M ETA -L EARNING FOR S TOCHASTIC G RADIENT MCMC

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Stochastic gradient Markov chain Monte Carlo (SG-MCMC) has become increasingly popular for simulating posterior samples in large-scale Bayesian modeling.
Revised: Stochastic gradient Markov chain Monte Carlo (SG-MCMC) has become increasingly popular for simulating posterior samples in large-scale Bayesian modeling.

Original: However, existing SG-MCMC schemes are not tailored to any speciﬁc probabilistic model, even a simple modiﬁcation of the underlying dynamical system requires signiﬁcant physical intuition.
Revised: However, existing SG-MCMC schemes are not tailored to any speciﬁc probabilistic model even a simple modiﬁcation of the underlying dynamical system requires signiﬁcant physical intuition

Original: This paper presents the ﬁrst meta-learning algorithm that allows automated design for the underlying continuous dynamics of an SG-MCMC sampler.
Revised: This paper presents the ﬁrst meta learning algo - rithm that allows automated design for the underlying continuous dynamics of an SG-MCMC sampler.

Original: The learned sampler generalizes Hamiltonian dynamics with state-dependent drift and diffusion, enabling fast traversal and efﬁcient exploration of energy landscapes.
Revised: The learned sampler generalizes Hamiltonian dynamics with state-dependent drift and diffusion, enabling fast traversal and efﬁcient exploration of energy landscapes.

Original: Experiments validate the proposed approach on learning tasks with Bayesian fully connected neural networks, Bayesian convolutional neural networks and Bayesian recurrent neural networks, showing that the learned sampler out-performs generic, hand-designed SG-MCMC algorithms, and generalizes to different datasets and larger architectures.
Revised: Experiments validate the proposed approach on learning tasks with Bayesian fully connected neural networks, Bayesian convolutional neural networks and Bayesian recurrent neural networks, showing that the learned sampler out-performs generic, hand-designed SG-MCMC algorithms and generalizes to different datasets and larger architectures

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: There is a resurgence of research interests in Bayesian deep learning (Graves, 2011; Blundell et al., 2015; Hernández-Lobato & Adams, 2015; Hernandez-Lobato et al., 2016; Gal & Ghahramani, 2016; Ritter et al., 2018), which applies Bayesian inference to neural networks for better uncertainty estimation.
Revised: There is a resurgence of research interests in Bayesian deep learning (Graves, 2011; Blundell et al., 2015; Hernández-Lobato & Adams, 2015; Hernandez-Lobato et al., 2016; Gal & Ghahramani, 2016; Ritter et al., 2018), which applies Bayesian inference to neural networks for better uncertainty estimation.

Original: It is crucial for e.g. better exploration in reinforcement learning (Deisenroth & Rasmussen, 2011; Depeweg et al., 2017), resisting adversarial attacks (Feinman et al., 2017; Li & Gal, 2017;
Revised: It is crucial for e.g. better exploration in reinforcement learning (Deisenroth & Rasmussen, 2011; Depeweg et al., 2017), resisting adversarial attacks (Feinman et al., 2017; Li & Gal, 2017;

Original: Louizos & Welling, 2017) and continual learning (Nguyen et al., 2018).
Revised: Louizos & Welling, 2017) and continual learning (Nguyen et al., 2018).

Original: A popular approach to performing Bayesian inference on neural networks is stochastic gradient Markov chain Monte Carlo (SG-MCMC), which adds properly scaled Gaussian noise to a stochastic gradient ascent procedure (Welling & Teh, 2011).
Revised: A popular approach to performing Bayesian inference on neural networks is stochastic gradient Markov chain Monte Carlo (SG-MCMC), which adds properly scaled Gaussian noise to a stochastic gradient ascent procedure (Welling & Teh, 2011).

Original: Recent advances in this area further introduced optimization techniques such as pre-conditioning (Ahn et al., 2012; Patterson & Teh, 2013), annealing (Ding et al., 2014) and adaptive learning rates (Li et al., 2016a;
Revised: Recent advances in this area further introduced optimization techniques such as pre-conditioning (Ahn et al., 2012; Patterson & Teh, 2013), annealing (Ding et al., 2014) and adaptive learning rates (Li et al., 2016a;

Original: Chen et al., 2016).
Revised: Chen et al., 2016).

Original: All these efforts have made SG-MCMC highly scalable to many deep learning tasks, including shape and texture modeling in computer vision (Li et al., 2016b) and language modeling with recurrent neural networks (Gan et al., 2017).
Revised: All these efforts have made SG-MCMC highly scalable to many deep learning tasks, including shape and texture modeling in computer vision (Li et al., 2016b) and language modeling with recurrent neural networks (Gan et al., 2017).

Original: However, inventing novel dynamics for SG-MCMC requires signiﬁcant mathematical work to ensure the sampler stationary distribution is the target distribution, which is less friendly to practitioners.
Revised: However, inventing novel dynamics for SG-MCMC requires signiﬁcant mathematical work to ensure the sampler’s stationary distribution is the target distribution, which is less friendly to practitioners.

Original: Furthermore, many of these algorithms are designed as a generic sampling procedure, and the associated physical mechanism might not be best suited for sampling neural network weights.
Revised: Furthermore, many of these algorithms are designed as a generic sampling procedure, and the associated physical mechanism might not be best suited for sampling neural network weights.

Original: This paper aims to automate the SG-MCMC proposal design by introducing meta-learning techniques (Schmidhuber, 1987; Bengio et al., 1992; Naik & Mammone, 1992; Thrun & Pratt, 1998).
Revised: This paper aims to automate the SG-MCMC proposal design by introducing meta-learning techniques (Schmidhuber, 1987; Bengio et al., 1992; Naik & Mammone, 1992; Thrun & Pratt, 1998).

Original: The general idea is to train a learner on one or multiple tasks in order to acquire common knowledge that generalizes to future tasks.
Revised: The general idea is to train a learner on one or multiple tasks in order to acquire common knowledge that generalizes to future tasks.

Original: Recent applications of meta-learning include learning to transfer knowledge to unseen few-shot learning tasks (Santoro et al., 2016; Ravi & Larochelle, 2017;
Revised: Recent applications of meta-learning include learning to transfer knowledge to unseen few-shot learning tasks (Santoro et al., 2016; Ravi & Larochelle, 2017;

Original: Finn et al., 2017), and learning algorithms such as gradient descent (Andrychowicz et al., 2016; Li &Malik, 2017;
Revised: Finn et al., 2017), and learning algorithms such as gradient descent (Andrychowicz et al., 2016; Li &

Original:  Wichrowska et al., 2017), Bayesian optimization (Chen et al., 2017) and reinforcement learning (Duan et al., 2016; Wang et al., 2016).
Revised: Malik, 2017; Wichrowska et al., 2017), Bayesian optimization (Chen et al., 2017) and reinforcement learning (Duan et al., 2016; Wang et al., 2016).

Original: Unfortunately these advances cannot be directly transfered to the world of MCMC samplers, as a naive neural network parameterization of the transition kernel does not guarantee the posterior distribution to be the stationary distribution of the sampler.
Revised: Unfortunately, these advances cannot be directly transferred to the world of MCMC samplers, as a naive neural network parameterization of the transition kernel does not guarantee the posterior distribution to be the stationary distribution of the sampler.

Original: We present to the best of our knowledge the ﬁrst attempt towards meta-learning an SG-MCMC algorithm.
Revised: We present to the best of our knowledge the ﬁrst attempt towards meta-learning an SG-MCMC algorithm.

Original: Concretely, our contribution include:
Revised: Concretely, our contribution includes:

Original: • An SG-MCMC sampler that extends Hamiltonian dynamics with learnable diffusion and curl matrices.
Revised: • An SG-MCMC sampler that extends Hamiltonian dynamics with learnable diffusion and curl matrices.

Original: Once trained, the sampler can generalize to different datasets and architectures.
Revised: Once trained, the sampler can generalize to different datasets and architectures.

Original: • Extensive evaluation of the proposed sampler on Bayesian fully connected neural networks, Bayesian convolutional neural networks and Bayesian recurrent neural networks, with comparisons to popular SG-MCMC schemes based on e.g. Hamiltonian Monte Carlo (Chen et al., 2014) and pre-conditioned Langevin dynamics (Li et al., 2016a).
Revised: • Extensive evaluation of the proposed sampler on Bayesian fully connected neural networks, Bayesian convolutional neural networks and Bayesian recurrent neural networks, with comparisons to popular SG-MCMC schemes based on e g. Hamiltonian Monte Carlo (Chen et al., 2014) and pre-conditioned Langevin dynamics (Li et al., 2016a).

Original:  B ACKGROUND : A COMPLETE FRAMEWORK FOR SG-MCMC
Revised:  B ACKGROUND : A COMPLETE FRAMEWORK FOR SG-MCMC

Original: Consider sampling from a target density π ( θθθ ) that is deﬁned by an energy function : U ( θθθ ) , θθθ ∈ R D , π ( θθθ ) ∝ exp( − U ( θθθ )) .
Revised: Consider sampling from a target density π ( θθθ ) that is deﬁned by an energy function : U ( θθθ ) , θθθ ∈ R D , π ( θθθ ) ∝ exp( − U ( θθθ )) .

Original: In this paper we focus on this sampling task with a Bayesian modeling set-up, i.e. given observed data D = { ooo n } Nn =1 , we deﬁne a probabilistic model p ( D ,θθθ ) =
Revised: In this paper, we focus on this sampling task with a Bayesian modeling set-up, i.e. given observed data D = { ooo n } Nn =1 , we deﬁne a probabilistic model p ( D ,θθθ ) =

Original: (cid:81)
Revised: (cid:81)

Original: Nn =1 p ( ooo n | θθθ ) p ( θθθ ) , and we want samples fromtarget the  density deﬁned as posterior distribution π ( θθθ ) = p ( θθθ |D ) .
Revised: Nn =1 p ( ooo n | θθθ ) p ( θθθ ) , and we want samples from the target density deﬁned as posterior distribution π ( θθθ ) = p ( θθθ |D ) .

Original: We use Bayesian neural networks as an illustrating example, in this case, ooo n = ( xxx n ,yyy n ) , the prior p ( θθθ ) is a Gaussian N ( θθθ ;000 , λ − 1 III{+
 I I I+} ) , and the energy function is deﬁned as
Revised: We use Bayesian neural networks as an illustrating example, in this case, ooo n = ( xxx n ,yyy n ) , the prior p ( θθθ ) is a Gaussian N ( θ θ θ ;0 0 0 , λ − {+
 I I I+} ) , and the energy function is deﬁned as

Original: [Equation]
Revised: [Equation]

Original:  with (cid:96) ( yyy, ˆ yyy ) usually deﬁned as the (cid:96) 2 loss for regression or the cross-entropy loss for classiﬁcation.
Revised: || 22 , (1) with (cid:96) ( yyy, ˆ yyy ) usually deﬁned as the (cid:96) 2 loss for regression or the cross-entropy loss for classiﬁcation.

Original: A typical MCMC sampler constructs a Markov chain with a transition kernel , and corrects the proposed samples with Metropolis-Hastings (MH) rejection steps.
Revised: A typical MCMC sampler constructs a Markov chain with a transition kernel , and corrects the proposed samples with Metropolis-Hastings (MH) rejection steps.

Original: Some of these methods, e.g. Hamiltonian
Revised: Some of these methods, e.g. Hamiltonian

Original: Monte Carlo (HMC) (Duane et al., 1987; Neal et al., 2011), further augment the state space as zzz =
Revised: Monte Carlo (HMC) (Duane et al., 1987; Neal et al., 2011), further augment the state space as zzz =

Original: ( θθθ,rrr ) with auxiliary variables rrr , and sample from the augmented distribution π ( zzz )
Revised: ( θθθ,rrr ) with auxiliary variables rrr , and sample from the augmented distribution π ( zzz )

Original: ∝ exp ( − H ( zzz )) , with the Hamiltonian H ( zzz ) =
Revised: ∝ exp ( − H ( zzz )) , with the Hamiltonian H ( zzz ) =

Original: U ( θθθ ) + g ( rrr ) such that (cid:82) exp( − g ( rrr ))
Revised: U ( θθθ ) + g ( rrr ) such that (cid:82) exp( − g ( rrr ))

Original: drrr = C .
Revised: drrr = C .

Original: Thus, marginalizing out the auxiliary variable rrr will not affect the stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) .
Revised: Thus, marginalizing out the auxiliary variable rrr will not affect the stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) .

Original: For deep learning tasks, the observed dataset D often contains thousands, if not millions, of instances, making MH rejection steps computationally prohibitive.
Revised: For deep learning tasks, the observed dataset D often contains thousands, if not millions, of instances, making MH rejection steps computationally prohibitive.

Original: Fortunately this is mitigated by SG-MCMC, whose transition kernel is implicitly deﬁned by a stochastic differential equation (SDE) that leaves the target density invariant (Welling & Teh, 2011;
Revised: Fortunately this is mitigated by SG-MCMC, whose transition kernel is implicitly deﬁned by a stochastic differential equation (SDE) that leaves the target density invariant (Welling & Teh, 2011;

Original: Ahn et al., 2012; Patterson & Teh, 2013; Chen et al., 2014;
Revised: Ahn et al., 2012; Patterson & Teh, 2013; Chen et al., 2014;

Original: Ding et al., 2014).
Revised: Ding et al., 2014).

Original: Such  Markov process is called Itô diffusion governed by the continuous-time SDEs:
Revised: Such a Markov process is called Itô diffusion governed by the continuous-time SDEs:

Original: [Equation]
Revised: [Equation]

Original: with fff ( zzz ) the deterministic drift , WWW ( t ) the Wiener process, and DDD ( zzz ) the diffusion matrix.
Revised: with fff ( zzz ) the deterministic drift , WWW ( t ) the Wiener process, and DDD ( zzz ) the diffusion matrix.

Original: As a simple example, Langevin dynamics considers zzz = θθθ , fff ( θθθ ) =
Revised: As a simple example, Langevin dynamics considers zzz = θθθ , fff ( θθθ ) =

Original: −∇ θθθ U ( θθθ ) and DDD ( θθθ )
Revised: −∇ θθθ U ( θθθ ) and DDD ( θθθ )

Original: =
Revised: =

Original: I .
Revised: I .

Original: Then using forward Euler discretization with step-size η the update rule of the parameters is
Revised: Then using forward Euler discretization with step-size η the update rule of the parameters is

Original: [Equation]
Revised: [Equation]

Original: Stochastic gradient Langevin dynamics (SGLD, Welling & Teh, 2011) proposed an approximation to (3), by replacing the exact gradient ∇ θθθ U ( θθθ ) with an estimate using a mini-batch of datapoints:
Revised: Stochastic gradient Langevin dynamics (SGLD, Welling & Teh, 2011) proposed an approximation to (3), by replacing the exact gradient ∇ θθθ U ( θθθ ) with an estimate using a mini-batch of datapoints:

Original: [Equation]
Revised: [Equation]

Original: Therefore SGLD can be viewed as a stochastic gradient descent (SGD) that adds in a properly scaled Gaussian noise term.
Revised: Therefore SGLD can be viewed as a stochastic gradient descent (SGD) that adds in a properly scaled Gaussian noise term.

Original: Similarly SGHMC (Chen et al., 2014) is closely related to momentum
Revised: Similarly, SGHMC (Chen et al., 2014) is closely related to momentum

Original: SGD (see appendix).
Revised: SGD (see appendix).

Original: Furthermore, MH rejection steps are usually dropped in SG-MCMC when a carefully selected discretization step-size is in use.
Revised: Furthermore, MH rejection steps are usually dropped in SG-MCMC when a carefully selected discretization step-size is in use.

Original: Therefore SG-MCMC has the same computational complexity as many stochastic optimization algorithms, making it highly scalable for sampling posterior distributions of neural network weights conditioned on big datasets.
Revised: Therefore SG-MCMC has the same computational complexity as many stochastic optimization algorithms, making it highly scalable for sampling posterior distributions of neural network weights conditioned on big datasets.

Original: Ma et al.
Revised: Ma et al.

Original: (2015) derived a framework of SG-MCMC samplers using advanced statistical mechanics (Yin & Ao, 2006; Shi et al., 2012), which explicitly parameterizes the drift f f f ( z z z ) :
Revised: (2015) derived a framework of SG-MCMC samplers using advanced statistical mechanics (Yin & Ao, 2006; Shi et al., 2012), which explicitly parameterizes the drift fff ( zzz ) :

Original: [Equation]
Revised: [Equation]

Original: with QQQ ( zzz ) the curl matrix, DDD ( zzz ) the diffusion matrix and ΓΓΓ( zzz ) a correction term.
Revised: with QQQ ( zzz ) the curl matrix, DDD ( zzz ) the diffusion matrix and ΓΓΓ( zzz ) a correction term.

Original: Remarkably Ma et al.
Revised: Remarkably Ma et al.

Original: (2015) showed the completeness of their framework:
Revised: (2015) showed the completeness of their framework:

Original: π ( zzz ) ∝ exp( − H ( zzz )) is a stationary distribution of the SDE (2)+(5) for any pair of positive semi-deﬁnite matrix DDD ( zzz ) and skew-symmetric matrix QQQ ( zzz ) ; 2.
Revised: π ( zzz ) ∝ exp( − H ( zzz )) is a stationary distribution of the SDE (2)+(5) for any pair of positive semi-deﬁnite matrix DDD ( zzz ) and skew-symmetric matrix QQQ ( zzz ) ; 2.

Original: for any Itô diffusion process that has the unique stationary distribution π ( zzz ) , under mild conditions there exist DDD ( zzz ) and QQQ ( zzz ) matrices such that the process is governed by (2)+(5).
Revised: for any Itô diffusion process that has the unique stationary distribution π ( zzz ) , under mild conditions there exist DDD ( zzz ) and QQQ ( zzz ) matrices such that the process is governed by (2)+(5).

Original: As a consequence, the construction of an SG-MCMC algorithm reduces to deﬁning the state-space zzz and the DDD  , QQQ  matrices.
Revised: As a consequence, the construction of an SG-MCMC algorithm reduces to deﬁning the state-space zzz and the DDD ( zzz ) , QQQ ( zzz ) matrices.

Original: Indeed Ma et al.
Revised: Indeed Ma et al.

Original: (2015) also casted existing SG-MCMC samplers within the framework, and proposed an improved version of SG-Riemannian-HMC.
Revised: (2015) also cast existing SG-MCMC samplers within the framework, and proposed an improved version of SG-Riemannian-HMC.

Original: In general, an appropriate design of these two matrices leads to signiﬁcant improvement on mixing as well as reduction of sample bias (Li et al., 2016a; Ma et al., 2015).
Revised: In general, an appropriate design of these two matrices leads to signiﬁcant improvements in mixing as well as reduction of sample bias (Li et al., 2016a; Ma et al., 2015).

Original: However, this design has been historically based on strong physical intuitions from statistical mechanics (Duane et al., 1987; 
Revised: However, this design has been historically based on strong physical intuitions from statistical mechanics (Duane et al., 1987; Neal et al., 2011;

Original: Neal et al., 2011; Ding et al., 2014).
Revised: Ding et al., 2014).

Original: Therefore, it can still be difﬁcult for practitioners to understand and engineer the sampling method that is best suited to their machine learning tasks.
Revised: Therefore, it can still be difﬁcult for practitioners to understand and engineer the sampling method that is best suited to their machine learning tasks.

Original: In the next section we will describe our recipe on meta-learning an SG-MCMC sampler of the form (2)+(5).
Revised: In the next section, we will describe our recipe on meta-learning an SG-MCMC sampler of the form (2)+(5).

Original: Before the presentation, we emphasize that the completeness result of the framework is beneﬁcial for our meta-learning task.
Revised: Before the presentation, we emphasize that the completeness result of the framework is beneﬁcial for our meta-learning task.

Original: On  one hand, as meta-learning searches the best algorithm for a given set of tasks, it is crucial that the search space is large enough to contain many useful candidates.
Revised: On the one hand, as meta-learning searches the best algorithm for a given set of tasks, it is crucial that the search space is large enough to contain many useful candidates.

Original: On the other hand, some form of “correctness” guarantee is often required to achieve better generalization to test tasks that might not be very similar to the training tasks.
Revised: On the other hand, some form of “correctness” guarantee is often required to achieve better generalization to test tasks that might not be very similar to the training tasks.

Original: Ma et al.
Revised: Ma et al.

Original: (2015)’s completeness result indicates that our proposed method searches SG-MCMC samplers in the biggest subset of all Itô diffusion processes such that each instance is a valid posterior sampler.
Revised: (2015)’s completeness result indicates that our proposed method searches SG-MCMC samplers in the biggest subset of all Itô diffusion processes such that each instance is a valid posterior sampler.

Original: Therefore, the proposed meta-learning algorithm has the best from both worlds, indeed our experiments show that the learned sampler is superior to a number of other baseline SG-MCMC methods.
Revised: Therefore, the proposed meta-learning algorithm has the best from both worlds, indeed our experiments show that the learned sampler is superior to a number of other baseline SG-MCMC methods.

Original:  M ETA - LEARNING FOR SG-MCMC
Revised:  M ETA - LEARNING FOR SG-MCMC

Original: This section presents a meta-learning approach to learn an SG-MCMC proposal from data.
Revised: This section presents a meta-learning approach to learn an SG-MCMC proposal from data.

Original: Our aim is to design an appropriate parameterization of DDD ( zzz ) and QQQ ( zzz ) , so that the sampler can be trained on simple tasks with a meta-learning procedure, and generalize to more complicated densities.
Revised: Our aim is to design an appropriate parameterization of DDD ( zzz ) and QQQ ( zzz ) , so that the sampler can be trained on simple tasks with a meta-learning procedure, and generalize to more complicated densities.

Original: For simplicity, we only augment the state-space by one extra variable ppp called momentum (Duane et al., 1987;
Revised: For simplicity, we only augment the state-space by one extra variable ppp called momentum (Duane et al., 1987;

Original: Neal et al., 2011), although generalization to e.g. thermostat variable (Ding et al., 2014) is straight-forward.
Revised: Neal et al., 2011), although generalization to e.g. thermostat variable (Ding et al., 2014) is straightforward.

Original: Thus, the augmented state-space is zzz =
Revised: Thus, the augmented state-space is zzz =

Original: ( θθθ,ppp ) (i.e. rrr = ppp ), and the Hamiltonian is deﬁned as H ( zzz ) =
Revised: ( θθθ,ppp ) (i.e. rrr = ppp ), and the Hamiltonian is deﬁned as H ( zzz ) =

Original: U ( θθθ ) + 12 ppp T ppp with identity mass matrix.
Revised: U ( θθθ ) + 12 ppp T ppp with identity mass matrix.

Original: 3. E FFICIENT PARAMETERIZATION OF DIFFUSION AND CURL MATRICES
Revised: 3. E FFICIENT PARAMETERIZATION OF DIFFUSION AND CURL MATRICES

Original: For neural networks, the dimensionality of θθθ can be at least tens of thousands.
Revised: For neural networks, the dimensionality of θθθ can be at least tens of thousands.

Original: Thus, training and applying full DDD ( zzz ) and QQQ ( zzz ) matrices can cause  huge computational burden, let alone gradient computations required by ΓΓΓ( zzz ) .
Revised: Thus, training and applying full DDD ( zzz ) and QQQ ( zzz ) matrices can cause a huge computational burden, let alone gradient computations required by ΓΓΓ( zzz ) .

Original: To address this, we deﬁne the preconditioning matrices as follows:
Revised: To address this, we deﬁne the preconditioning matrices as follows:

Original: [Equation]
Revised: [Equation]

Original: Here fff θ D and fff θ Q are neural network parameterized functions that will be detailed in section 3.2, and c is a small positive constant.
Revised: Here fff φ D and fff φ Q are neural network parameterized functions that will be detailed in section 3.2, and c is a small positive constant.

Original: We choose DDD f and QQQ f to be diagonal for fast computation, although future work can explore low-rank matrix solutions.
Revised: We choose DDD f and QQQ f to be diagonal for fast computation, although future work can explore low-rank matrix solutions.

Original: From Ma et al.
Revised: From Ma et al.

Original: , our design has the unique stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) if fff φ D is non-negative for all zzz .
Revised: , our design has the unique stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) if fff φ D is non-negative for all zzz .

Original: We discuss the role of each precondition matrix for better intuition.
Revised: We discuss the role of each precondition matrix for better intuition.

Original: The curl matrix QQQ ( zzz ) in (2) mainly controls the deterministic drift forces introduced by the energy gradient ∇ θθθ U ( θθθ ) (as seen in many HMC-like procedures and Eq.5).
Revised: The curl matrix QQQ ( zzz ) in (2) mainly controls the deterministic drift forces introduced by the energy gradient ∇ θθθ U ( θθθ ) (as seen in many HMC-like procedures and in eq.

Original: Usually we only have access to the stochastic gradient ∇ θθθ ˜ U
Revised: Usually, we only have access to the stochastic gradient ∇ θθθ ˜ U

Original: ( θθθ ) through data sub-sampling, and here DDD  acts as the friction to counter for the associated noise that mainly affects the momentum ppp .
Revised: ( θθθ ) through data sub-sampling, and here DDD ( zzz ) acts as the friction to counter for the associated noise that mainly affects the momentum ppp .

Original: This explains the design of the diffusion matrix DDD ( zzz ) that uses DDD f ( zzz ) to control the amount of friction and injected noise to the momentum.
Revised: This explains the design of the diffusion matrix DDD ( zzz ) that uses DDD f ( zzz ) to control the amount of friction and injected noise to the momentum.

Original: Furthermore,  DDD f ( zzz ) should also account for the pre-conditioning effect introduced by QQQ f ( zzz ) , e.g, when the magnitude of QQQ f  is large, we need higher friction correspondingly.
Revised: Furthermore, [Figure] DDD f ( zzz ) should also account for the pre-conditioning effect introduced by QQQ f ( zzz ) , e.g, when the magnitude of QQQ f ( zzz ) is large, we need higher friction correspondingly.

Original: This explains the squared term fff φ Q ( zzz ) (cid:12) fff
Revised: This explains the squared term fff φ Q ( zzz ) (cid:12) fff

Original: φ Q ( zzz ) in DDD f  design.
Revised: φ Q ( zzz ) in DDD f ( zzz ) design.

Original: Thescaling positive  constant α is heuristically selected following (Chen et al., 2014;[Figure] Ma et al., 2015) (see appendix).
Revised: The positive scaling constant α is heuristically selected following (Chen et al., 2014; Ma et al., 2015) (see appendix).

Original: Finally the extra term ΓΓΓ( zzz ) =
Revised: Finally, the extra term ΓΓΓ( zzz ) =

Original: [ΓΓΓ θθθ ( zzz ) , ΓΓΓ ppp ( zzz )] is responsible for compensating the changes introduced by preconditioning matrices QQQ ( zzz ) and DDD ( zzz ) .
Revised: [ΓΓΓ θθθ ( zzz ) , ΓΓΓ ppp ( zzz )] is responsible for compensating the changes introduced by preconditioning matrices QQQ ( zzz ) and DDD ( zzz ) .

Original: The discretized dynamics of the state zzz =
Revised: The discretized dynamics of the state zzz =

Original: ( θθθ,ppp ) with step-size η and stochastic gradient ∇ θθθ ˜ U 
Revised: ( θθθ,ppp ) with step-size η and stochastic gradient ∇ θ θθ ˜ U ( θθθ ) are ppp t +1 = (1 − ηDDD [Equation] ( zzz t ))

Revised: [Equation]

Revised: − ηQQQ

Original:  ( θθθ ) are
Revised: [Equation] ( zzz [Equation] ) ∇

Revised: [Equation]

Revised: t

Revised: [Equation]

Revised: [Equation]

Revised: + η ΓΓΓ

Revised: [Equation]

Revised: ( zzz t )

Revised: + (cid:15)(cid:15)(cid:15), (cid:15)(cid:15)(cid:15)

Revised: [Equation]

Revised: N (

Revised: [Equation]

Revised: 2 ηDDD f ( zzz t )) , θθθ t +1 = θθθ t +

Original: [Equation]
Revised: [Equation]

Original: Note that we use a modiﬁed forward Euler discretization (Neal et al., 2011) here, and the computation graph of eq.
Revised: f (ˆ zzz t ) ppp t +1 + η ΓΓΓ [Equation] (ˆ zzz t ) , zzz t = [Equation]

Revised: θθθ t

Revised: [Equation]

Revised: t ] , ˆ

Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: t +1 ] .

Revised: We use a modiﬁed forward Euler discretization (Neal et al., 2011) here, and the computation graph of eq.

Original: (7) is visualized in the right part of Figure 1 (see appendix for SGHMC discretized updates).
Revised: (7) is visualized in the right part of Figure 1 (see appendix for SGHMC discretized updates).

Original: Again we see that QQQ f  is responsible for the acceleration of θθθ , and from the (1 − ηDDD f ) term in the update equation of ppp , we see that DDD f  controls the friction introduced to the momentum.
Revised: Again we see that QQQ f ( zzz ) is responsible for the acceleration of θθθ , and from the (1 − ηDDD f ( zzz )) term in the update equation of ppp , we see that DDD f ( zzz ) controls the friction introduced to the momentum.

Original: Observing that  the noisy gradient is approximately Gaussian distributed in the big data setting, Ma et al.
Revised: Note that in the big-data setting, the noisy gradient is approximately Gaussian distributed with mean 000 and variance VVV ( θθθ ) . Observing this, Ma et al.

Original: (2015) further suggested a correction scheme to counter for stochastic gradient noise, which samples the Gaussian noise (cid:15)(cid:15)(cid:15) ∼ N (000 , 2 ηDDD f ( zzz ) − η 2 ˜ B ( θθθ )) with an empirical estimate of thegradient variance ˜ B ( θθθ )  instead.
Revised: (2015) further suggested a correction scheme to counter for stochastic gradient noise, which samples the Gaussian noise (cid:15)(cid:15)(cid:15) ∼ N (000 , 2 ηDDD f ( zzz ) − η 2 ˜ B ( θθθ )) with an empirical estimate of the variance ˜ B ( θθθ ) ≈ QQQ f ( zzz ) VVV ( θθθ ) QQQ Tf ( zzz ) instead.

Original: These corrections can be dropped when the discretization step-size η is small, therefore we do not consider them in our experiments.
Revised: These corrections can be dropped when the discretization step-size η is small, therefore, we do not consider them in our experiments.

Original: 3. C HOICES OF INPUTS TO THE NEURAL NETWORKS
Revised: 3. C HOICES OF INPUTS TO THE NEURAL NETWORKS

Original: We now present detailed functional forms for fff φ Q and fff φ D .
Revised: We now present detailed functional forms for fff φ Q and fff φ D .

Original: When designing these, our goal was to achieve a good balance between generalization power and computational efﬁciency.
Revised: When designing these, our goal was to achieve a good balance between generalization power and computational efﬁciency.

Original: Recall that the curl matrix QQQ  mainly controls the drift of the dynamics, and the desired behavior is the fast traverse through low density regions.
Revised: Recall that the curl matrix QQQ ( zzz ) mainly controls the drift of the dynamics, and the desired behavior is the fast traverse through low-density regions.

Original: One useful source of information to identify this is the energy function U ( θθθθθθθθθ ) .
Revised: One useful source of information to identify this is the energy function U ( θθθθθθθθθ ) .

Original: 1 We also include the momentum p i to the inputs of fff φ Q , allowing the QQQ  matrix to observe the velocity information of the θθθ i .
Revised: 1 We also include the momentum p i to the inputs of fff φ Q , allowing the QQQ ( zzz ) matrix to observe the velocity information of the θθθ i .

Original: We further add an offset β to QQQ  to prevent the vanishing of this matrix.
Revised: We further add an offset β to QQQ ( zzz ) to prevent the vanishing of this matrix.

Original: Putting all of them together, we deﬁne the i th element of fff φ Q as
Revised: Putting all of them together, we deﬁne the i th element of fff φ Q as

Original: [Equation]
Revised: [Equation]

Original: The corresponding ΓΓΓ term requires both ∂ θ i f φ Q ( U ( θθθ ) , p i ) and ∂ p i f φ Q ( U ( θθθ ) , p i ) .
Revised: The corresponding ΓΓΓ( zzz ) term requires both ∂ θ i f φ Q ( U ( θθθ ) , p i ) and ∂ p i f φ Q ( U ( θθθ ) , p i ) .

Original: The energy gradient ∂ θ i U ( θθθ ) also appears in (7), so it remains to compute ∂ U fff φ Q , which, along with ∂ p i f φ Q ( U ( θθθ ) , p i ) , can be obtained by automatic differentiation (Abadi et al., 2015).
Revised: The energy gradient ∂ θ i U ( θθθ ) also appears in (7), so it remains to compute ∂ U fff φ Q , which, along with ∂ p i f φ Q ( U ( θθθ ) , p i ) , can be obtained by automatic differentiation (Abadi et al., 2015).

Original: Matrix DDD  is responsible for the friction and the stochastic gradient noise, which are crucial for better exploration around high density regions.
Revised: Matrix DDD ( zzz ) is responsible for the friction and the stochastic gradient noise, which are crucial for better exploration around high-density regions.

Original: Therefore, we also add the energy gradient ∂ θ i U ( θθθ ) to the inputs, meaning that the i th element of fff φ D is
Revised: Therefore, we also add the energy gradient ∂ θ i U ( θθθ ) to the inputs, meaning that the i th element of fff φ D is

Original: [Equation]
Revised: [Equation]

Original: By  construction of the DDD  matrix, the ΓΓΓ vector only requires ∇ ppp DDD f without computing any higher order information.
Revised: By the construction of the DDD ( zzz ) matrix, the ΓΓΓ vector only requires ∇ ppp DDD f without computing any higher order information.

Original: In practice both U ( θθθ ) and ∂ θ i U ( θθθ ) are replaced by their stochastic estimates ˜
Revised: In practice, both U ( θθθ ) and ∂ θ i U ( θθθ ) are replaced by their stochastic estimates ˜

Original: U ( θθθ ) and ∂ θ i ˜ U
Revised: U ( θθθ ) and ∂ θ i ˜ U

Original: ( θθθ ) , respectively.
Revised: ( θθθ ) , respectively.

Original: To keep the scale of the inputs roughly the same across tasks, we rescale all the inputs using statistics computed by simulating the sampler with random initialized fff φ D and fff φ Q .
Revised: To keep the scale of the inputs roughly the same across tasks, we rescale all the inputs using statistics computed by simulating the sampler with randomly initialized fff φ D and fff φ Q .

Original: When the computational budget is limited, we replace the exact gradient computation required by ΓΓΓ( zzz ) with ﬁnite difference approximations.
Revised: When the computational budget is limited, we replace the exact gradient computation required by ΓΓΓ( zzz ) with ﬁnite difference approximations.

Original: We refer the reader to the appendix for details.
Revised: We refer the reader to the appendix for details.

Original: 3. L OSS FUNCTION DESIGN FOR META - LEARNING
Revised: 3. L OSS FUNCTION DESIGN FOR META - LEARNING

Original: Another challenge is to design a meta-learning procedure for the sampler to encourage faster convergence and low bias on test tasks.
Revised: Another challenge is to design a meta-learning procedure for the sampler to encourage faster convergence and low bias on test tasks.

Original: To achieve these goals we propose two loss functions that we named as  cross-chain loss and  in-chain loss .
Revised: To achieve these goals we propose two loss functions that we named as the cross-chain loss and the in-chain loss .

Original: From now on we consider the discretized dynamics and deﬁne q t ( θθθ |D ) as the marginal distribution of the random variable θθθ at time t .
Revised: From now on we consider the discretized dynamics and deﬁne q t ( θθθ |D ) as the marginal distribution of the random variable θθθ at time t .

Original: Cross-chain loss We introduce cross-chain loss that encourages the sampler to converge faster.
Revised: Cross-chain loss We introduce cross-chain loss that encourages the sampler to converge faster.

Original: Since the sampler is guaranteed to have the unique stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) , fast convergence means that KL[ q t || π ] is close to zero when t is small.
Revised: Since the sampler is guaranteed to have the unique stationary distribution π ( θθθ ) ∝ exp( − U ( θθθ )) , fast convergence means that KL[ q t || π ] is close to zero when t is small.

Original: Therefore this KL-divergence becomes a sensible objective to minimize, which is equivalent to maximizing the variational lowerbound (or ELBO): L t VI ( q t )
Revised: Therefore this KL-divergence becomes a sensible objective to minimize, which is equivalent to maximizing the variational lowerbound (or ELBO): L t VI ( q t )

Original: =
Revised: =

Original: − E q t [ U ( θθθ )]
Revised: − E q t [ U ( θ θ θ )]

Original: + H
Revised: + H

Original: [ q t ] (Jordan et al., 1999; Beal, 2003).
Revised: [ q t ] (Jordan et al., 1999; Beal, 2003).

Original: We further make the objective doubly stochastic: (1) the energy term is further approximated by its stochastic estimates ˜
Revised: We further make the objective doubly stochastic: (1) the energy term is further approximated by its stochastic estimates ˜

Original: U ( θθθ ) ; (2) we use Monte Carlo variational inference (MCVI) (Ranganath et al., 2014;
Revised: U ( θθθ ) ; (2) we use Monte Carlo variational inference (MCVI, Ranganath et al., 2014;

Original: Blundell et al., 2015) which estimates the lower-bound with samples θθθ kt ∼ q t ( θθθ t |D ) , k
Revised: Blundell et al., 2015) which estimates the lower-bound with samples θθθ kt ∼ q t ( θθθ t |D ) , k

Original: = 1 , ..., K .
Revised: = 1 , ..., K .

Original: These samples { θθθ kt } K,Tk =1 ,t =1 are obtained by simulating K parallel Markov chains with the sampler, and the cross-chain loss is deﬁned by accumulating the lower-bounds through time:
Revised: These samples { θθθ kt } K,Tk =1 ,t =1 are obtained by simulating K parallel Markov chains with the sampler, and the cross-chain loss is deﬁned by accumulating the lower-bounds through time:

Original: [Equation]
Revised: [Equation]

Original: By minimizing this objective, we can improve the convergence of the sampler, especially at the early times of the Markov chain.
Revised: By minimizing this objective, we can improve the convergence of the sampler, especially at the early times of the Markov chain.

Original: The objective also takes the sampler bias into account because the two distributions will match when the KL-divergence is minimized.
Revised: The objective also takes the sampler bias into account because the two distributions will match when the KL-divergence is minimized.

Original: In-chain loss For very big neural networks, simulating multiple Markov chains is prohibitively expensive.
Revised: In-chain loss For very big neural networks, simulating multiple Markov chains is prohibitively expensive.

Original: The issue is mitigated by thinning that collects samples for every τ step (after burn-in), which effectively draws samples from the averaged distribution ¯ q (θθθ |D ) =
Revised: The issue is mitigated by thinning that collects samples for every τ step (after burn-in), which effectively draws samples from the averaged distribution ¯ q (

Original: [Equation]
Revised: [Equation]

Original: [Equation] (cid:80)
Revised: D

Original: [Equation]
Revised: [Equation]

Original: T/τ (cid:99) s
Revised: =

Original: [Equation]
Revised: [Equation]

Original: q sτ ( θθθ ) . The in-chain loss is, therefore, deﬁned as the ELBO evaluated at the averaged distribution ¯ q , which is then approximated by Monte Carlo with samples ΘΘΘ kT,τ = {θθθ ksτ }
Revised: The in-chain loss is, therefore, deﬁned as the ELBO evaluated at the averaged distribution ¯ q , which is then approximated by Monte Carlo with samples ΘΘΘ kT,τ = {

Original: [Equation]
Revised: [Equation]

Original: T/
Revised: ksτ

Original: [Equation]
Revised: [Equation]

Original: (cid:99)

Original: s = obtained by thinning:
Revised: obtained by thinning:

Original: [Equation]
Revised: [Equation]

Original: Gradient approximation We leverage the recently proposed Stein gradient estimator (Li & Turner, 2018) to estimate the intractable gradients ∇ φ log q t ( θθθ ) for cross-chain loss and ∇ φ log ¯ q ( θθθ ) for in-chain loss.
Revised: Gradient approximation We leverage the recently proposed Stein gradient estimator (Li & Turner, 2018) to estimate the intractable gradients ∇ φ log q t ( θθθ ) for cross-chain loss and ∇ φ log ¯ q ( θθθ ) for in-chain loss.

Original: Precisely, by the chain rule we have ∇ φ log q t ( θθθ ) = ∇ φ θθθ ∇ θθθ log q t ( θθθ ) , so it remains to estimate the gradients GGG = ( ∇ θθθ 1 t log q t ( θθθ 1 t ) , . .
Revised: Precisely, by the chain rule, we have ∇ φ log q t ( θθθ ) = ∇ φ θθθ ∇ θθθ log q t ( θθθ ) , so it remains to estimate the gradients GGG = ( ∇ θθθ 1 t log q t ( θθθ 1 t ) , . .

Original: , ∇ θθθ Kt log q t ( θθθ Kt )) T at the sampled locations { θθθ kt } Kk =1 ∼ q t .
Revised: , ∇ θθθ Kt log q t ( θθθ Kt )) T at the sampled locations { θθθ kt } Kk =1 ∼ q t .

Original: The recipe ﬁrst constructs a kernel matrix KKK with KKK ij =
Revised: The recipe ﬁrst constructs a kernel matrix KKK with KKK ij =

Original: K ( θθθ it ,θθθ j t ) , and then estimates the gradients by GGG ≈ − ( KKK + λIII ) − 1 (cid:104)∇ ,KKK (cid:105) , where (cid:104)∇ ,KKK (cid:105) ij =
Revised: K ( θθθ it ,θθθ jt ) , and then estimates the gradients by GGG ≈ − ( KKK + λIII ) − 1 (cid:104)∇ ,KKK (cid:105) , where (cid:104)∇ ,KKK (cid:105) ij =

Original: (cid:80)
Revised: (cid:80)

Original: Kk =1 ∂ θθθ kt ( j ) K ( θθθ kt ,θθθ it ) .
Revised: Kk =1 ∂ θθθ kt ( j ) K ( θθθ kt ,θθθ it ) .

Original: In our experiments, we use the RBF kernel, and the corresponding gradient estimator has a simple analytic form that can be computed efﬁciently in O
Revised: In our experiments, we use the RBF kernel, and the corresponding gradient estimator has a simple analytic form that can be computed efﬁciently in O

Original: ( K 2 D + K 3 ) time (usually K (cid:28) D ).
Revised: ( K 2 D + K 3 ) time (usually K (cid:28) D ).

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Since the development of SGLD (Welling & Teh, 2011), SG-MCMC has been increasingly popular for posterior sampling on big data.
Revised: Since the development of SGLD (Welling & Teh, 2011), SG-MCMC has been increasingly popular for posterior sampling on big data.

Original: In detail, Chen et al.
Revised: In detail, Chen et al.

Original: (2014) scaled up HMC with stochastic gradients, Ding et al.
Revised: (2014) scaled up HMC with stochastic gradients, Ding et al.

Original: (2014) further augmented the state space with a temperature auxiliary  variable, and Springenberg et al.
Revised: (2014) further augmented the state space with an auxiliary temperature variable, and Springenberg et al.

Original: (2016) improved robustness through scale adaptation.
Revised: (2016) improved robustness through scale adaptation.

Revised: The SG-MCMC

Revised: [Figure]

Revised: i

Revised: [Figure]

Revised: r g

Original: The SG-MCMC      extensions to Riemannian Langevin dynamics and HMC (Girolami & Calderhead, 2011) have also been proposed (Patterson & Teh, 2013; Ma et al., 2015).
Revised: [Figure] c [Figure] StartEnd      6      extensions to Riemannian Langevin dynamics and HMC (Girolami & Calderhead, 2011) have also been proposed (Patterson & Teh, 2013; Ma et al., 2015).

Original: Our proposed sampler architecture further generalizes SG-Riemannian-HMC as it decouples the design of DDD ( zzz ) and QQQ ( zzz ) matrices, and the detailed functional form of these two matrices are also learned from data.
Revised: Our proposed sampler architecture further generalizes SG-Riemannian-HMC as it decouples the design of DDD ( zzz ) and QQQ ( zzz ) matrices, and the detailed functional form of these two matrices are also learned from data.

Original: Our approach is closely related to the recent line of work on learning optimization algorithms.
Revised: Our approach is closely related to the recent line of work on learning optimization algorithms.

Original: Speciﬁcally, Andrychowicz et al.
Revised: Speciﬁcally, Andrychowicz et al.

Original: (2016) trained a recurrent neural network (RNN) based optimizer [Figure] i v e r g e n c [Figure]
Revised: (2016) trained a recurrent neural network (RNN) based optimizer that transfers to similar tasks with supervised learning.

Original: StartEnd      6      that transfers to similar tasks with supervised learning.

Original: Later Chen et al.
Revised: Later Chen et al.

Original: (2017) generalized this approach to Bayesian optimization (Brochu et al., 2010; Snoek et al., 2012) which is gradient-free.
Revised: (2017) generalized this approach to Bayesian optimization (Brochu et al., 2010; Snoek et al., 2012) which is gradient-free.

Original: We do not use RNN in our approach as it cannot be represented within the framework of Ma et al.
Revised: We do not use RNNs in our approach as it cannot be represented within the framework of Ma et al.

Original: We leave the combination of learnable RNN proposals to future work.
Revised: We leave the combination of learnable RNN proposals to future work.

Original: Also Li & Turner (2018) presented an initial attempt to meta-learn an approximate inference algorithm, which simply combined the stochastic gradient and the Gaussian noise with a neural network.
Revised: Also Li & Turner (2018) presented an initial attempt to meta-learn an approximate inference algorithm, which simply combined the stochastic gradient and the Gaussian noise with a neural network.

Original: Thus the stationary distribution of that sampler (if it exists) is only an approximation to the exact posterior.
Revised: Thus the stationary distribution of that sampler (if it exists) is only an approximation to the exact posterior.

Original: On the other hand, the proposed sampler (with η → 0 ) is guaranteed to be correct by the complete framework (Ma et al., 2015).
Revised: On the other hand, the proposed sampler (with η → 0 ) is guaranteed to be correct by the complete framework (Ma et al., 2015).

Original: Very recently Wu et al.
Revised: Very recently Wu et al.

Original: (2018) discussed that short-horizon meta-objectives for learning optimizers can cause a serious issue for long-time generalization.
Revised: (2018) discussed that short-horizon meta-objectives for learning optimizers can cause a serious issue for long-time generalization.

Original: We found this bias is less severe in our approach, again due to the fact that the learned sampler is provably correct.
Revised: We found this bias is less severe in our approach, again due to the fact that the learned sampler is provably correct.

Original: Recent research also considered improving HMC with a trainable transition kernel.
Revised: Recent research also considered improving HMC with a trainable transition kernel.

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

Original: (2015) improved upon vanilla HMC by introducing a trainable re-sampling distribution for the momentum.
Revised: (2015) improved upon vanilla HMC by introducing a trainable re-sampling distribution for the momentum.

Original: Song et al.
Revised: Song et al.

Original: (2017) parameterized the HMC transition kernel with a trainable invertible transformation called non-linear independent components estimation (NICE) (Dinh et al., 2014), and train it with Wasserstein adversarial training (Arjovsky et al., 2017).
Revised: (2017) parameterized the HMC transition kernel with a trainable invertible transformation called non-linear independent components estimation (NICE, Dinh et al., 2014), and train it with Wasserstein adversarial training (Arjovsky et al., 2017).

Original: Levy et al.
Revised: Levy et al.

Original: (2018) generalized
Revised: (2018) generalized

Original: HMC by augmenting the state space with a binary direction variable, and they parameterized the transition kernel with a non-volume preserving invertible transformation that is inspired by  real-valued non-volume preserving (RealNVP) ﬂows (Dinh et al., 2017).
Revised: HMC by augmenting the state space with a binary direction variable, and they parameterized the transition kernel with a non-volume preserving invertible transformation that is inspired by the real-valued non-volume preserving (RealNVP) ﬂows (Dinh et al., 2017).

Original: The sampler is trained with  expected squared jump distance (Pasarica & Gelman, 2010).
Revised: The sampler is trained with the expected squared jump distance (Pasarica & Gelman, 2010).

Original: We note that adversarial training is less reliable for high dimensional data, thus it is not considered in this paper.
Revised: We note that adversarial training is less reliable for high dimensional data, thus it is not considered in this paper.

Original: Also the jump distance does not explicitly take the sampling bias and convergence speed into account.
Revised: Also, the jump distance does not explicitly take the sampling bias and convergence speed into account.

Original: More importantly, the purpose of these approaches is to directly improve the HMC-like sampler on the target distribution , and with NICE/RealNVP parametrization it is difﬁcult to generalize the sampler to densities of different dimensions.
Revised: More importantly, the purpose of these approaches is to directly improve the HMC-like sampler on the target distribution , and with NICE/RealNVP parametrization it is difﬁcult to generalize the sampler to densities of different dimensions.

Original: In contrast, our goal is to learn an SG-MCMC sampler that can later be transferred to sample from different Bayesian neural network posterior distributions, which will typically have different dimensionality and include tens of thousands of random variables.
Revised: In contrast, our goal is to learn an SG-MCMC sampler that can later be transferred to sample from different Bayesian neural network posterior distributions, which will typically have different dimensionality and include tens of thousands of random variables.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate the meta-learned SG-MCMC sampler, which is referred to as NNSGHMC or the meta sampler in the following.
Revised: We evaluate the meta-learned SG-MCMC sampler, which is referred to as NNSGHMC or the meta sampler in the following.

Original: Detailed test set-ups are reported in the appendix.
Revised: Detailed test set-ups are reported in the appendix.

Original: Code is available at https://github.com/FirstAuthor/MetaSGMCMC .
Revised: The code is available at https://github.com/WenboGong/MetaSGMCMC .

Original:  5. S YNTHETIC EXAMPLE: SAMPLING G AUSSIAN VARIABLES WITH NOISY GRADIENTS
Revised: [Figure] [Table] 5. S YNTHETIC EXAMPLE

Original: We ﬁrst consider sampling Gaussian variables to demonstrate fast convergence and low bias of the meta sampler.
Revised: We ﬁrst consider sampling Gaussian variables to demonstrate fast convergence and low bias of the meta sampler.

Original: To mimic stochastic gradient settings, we manually inject Gaussian noise with unit variance to the gradient as suggested by (Chen et al., 2014).
Revised: To mimic stochastic gradient settings, we manually inject Gaussian noise with unit variance to the gradient as suggested by (Chen et al., 2014).

Original: The training density is a 10D Gaussian with randomly generated diagonal covariance matrix, and the test density is a 20D Gaussian.
Revised: The training density is a 10D Gaussian with randomly generated diagonal covariance matrix, and the test density is a 20D Gaussian.

Original: For evaluation, we simulate K = 50 parallel chains for T = 12 , 000 steps.
Revised: For evaluation, we simulate K = 50 parallel chains for T = 12 , 000 steps.

Original: Then we follow Ma et al.
Revised: Then we follow Ma et al.

Original: (2015) to evaluate the sampler’s bias measured by the KL divergence from the empirical estimate to the ground truth.
Revised: (2015) to evaluate the sampler’s bias measured by the KL divergence from the empirical estimate to the ground truth.

Original: Results are visualized on the left panel of Figure 2, showing that the meta sampler[Figure] [Table] both converges much faster and achieves lower bias compared to SGHMC.
Revised: Results are visualized on the left panel of Figure 2, showing that the meta sampler both converges much faster and achieves lower bias compared to SGHMC.

Original: The effective sample size 2 for SGHMC and NNSGHMC are 22 and 59 , again indicating better efﬁciency of the meta sampler.
Revised: The effective sample sizefor SGHMC and NNSGHMC are 22 and 59 , again indicating better efﬁciency of the meta sampler.

Original: For illustration purposes, we also plot in the other two panels the trajectory ofa samples by simulating NNSGHMC (middle) and SGHMC (right) on a 2D Gaussian for  ﬁxed amount of time ηT .
Revised: For illustration purposes, we also plot in the other two panels the trajectory of samples by simulating NNSGHMC (middle) and SGHMC (right) on a 2D Gaussian for a ﬁxed amount of time ηT .

Original: This conﬁrms that the meta sampler explores more efﬁciently and is less affected by the injected noise.
Revised: This conﬁrms that the meta sampler explores more efﬁciently and is less affected by the injected noise.

Original: 5. B AYESIAN FEEDFORWARD NEURAL NETWORKS
Revised: 5. B AYESIAN FEEDFORWARD NEURAL NETWORKS

Original: Next we consider Bayesian neural network classiﬁcation on MNIST data with three generalization tests: network architecture generalization (NT), activation function generalization (AF) and dataset generalization (Data).
Revised: Next, we consider Bayesian neural network classiﬁcation on MNIST data with three generalization tests: network architecture generalization (NT), activation function generalization (AF) and dataset generalization (Data).

Original: In all tests the sampler is trained with a 1-hidden layer multi-layer perceptron (MLP) (20 units, ReLU activation) as the underlying model for the target distribution π ( θθθ ) .
Revised: In all tests, the sampler is trained with a 1-hidden layer multi-layer perceptron (MLP) (20 units, ReLU activation) as the underlying model for the target distribution π ( θθθ ) .

Original: We also report long-time horizon generalization results, meaning that the simulation time steps in test time is much longer than that of training (cf.
Revised: We also report long-time horizon generalization results, meaning that the simulation time steps in test time are much longer than that of training (cf.

Original: Andrychowicz et al., 2016).
Revised: Andrychowicz et al., 2016).

Original: Algorithms in comparison include
Revised: Algorithms in comparison include

Original: SGLD (Welling & Teh, 2011), SGHMC (Chen et al., 2014) and preconditioned SGLD (PSGLD,
Revised: SGLD (Welling & Teh, 2011), SGHMC (Chen et al., 2014) and preconditioned SGLD (PSGLD,

Original: Li et al., 2016a).
Revised: Li et al., 2016a).

Original: Note that PSGLD uses RMSprop-like preconditioning techniques (Tieleman &
Revised: Note that PSGLD uses RMSprop-like preconditioning techniques (Tieleman &

Original: Hinton, 2012) that require moving average estimates of the gradient’s second moments.
Revised: Hinton, 2012) that require moving average estimates of the gradient’s second moments.

Original: Therefore the underlying dynamics of PSGLD cannot be represented within our framework (6).
Revised: Therefore the underlying dynamics of PSGLD cannot be represented within our framework (6).

Original: Thus we mainly focus on comparisons with SGLD and SGHMC, and leave the PSGLD results as reference.
Revised: Thus we mainly focus on comparisons with SGLD and SGHMC, and leave the PSGLD results as reference.

Original: The discretization step-sizes for the samplers are tuned on the validation dataset for each task.
Revised: The discretization step-sizes for the samplers are tuned on the validation dataset for each task.

Original: Architecture generalization (NT)
Revised: Architecture generalization (NT)

Original: In this test we use the trained sampler to draw samples from the posterior distribution of a 2-hidden layer MLP with 40 units and ReLU activations.
Revised: In this test we use the trained sampler to draw samples from the posterior distribution of a 2-hidden layer MLP with 40 units and ReLU activations.

Original: Figure 3 shows  learning curves of test error and negative test log-likelihood (NLL) for 100 epochs, where the ﬁnal performance is reported in Table 1.
Revised: Figure 3 shows the learning curves of test error and negative test log-likelihood (NLL) for 100 epochs, where the ﬁnal performance is reported in Table 1.

Original: Overall NNSGHMC achieves the fastest convergence even when compared with PSGLD.
Revised: Overall NNSGHMC achieves the fastest convergence even when compared with PSGLD.

Original: It has the lowest test error compared to SGLD and SGHMC.
Revised: It has the lowest test error compared to SGLD and SGHMC.

Original: NNSGHMC’s ﬁnal test LL is on par with SGLD and slightly worse than PSGLD, but it is still better than SGHMC. 
Revised: NNSGHMC’s ﬁnal test LL is on par with SGLD and slightly worse than PSGLD, but it is still better than SGHMC. [Figure]

Original: Architecture + Activation function generalization (NT+AF)
Revised: Architecture + Activation function generalization (NT+AF)

Original: Next we replace NT’s test network’s activation function with sigmoid and re-run the same test as before.
Revised: Next we replace NT’s test network’s activation function with sigmoid and re-run the same test as before.

Original: Again results in Figureand Table 1 show that NNSGHMC converges faster than others for both test error and NLL.
Revised: Again results in Figureand Table 1 show that NNSGHMC converges faster than others for both test error and NLL.

Original: It also achieves the best NLL results among all samplers, and the same test error as SGHMC.
Revised: It also achieves the best NLL results among all samplers, and the same test error as SGHMC.

Original: Architecture + Dataset generalization (NT+Data)
Revised: Architecture + Dataset generalization (NT+Data)

Original: In this test we split MNIST into training task (classifying digits 0-4) and test task (digits 5-9).
Revised: In this test we split MNIST into training task (classifying digits 0-4) and test task (digits 5-9).

Original: The meta sampler is trained with the smaller MLP,[Figure] and it is evaluated on the task with the larger MLP with NT’s architecture.
Revised: The meta sampler is trained with the smaller MLP, and it is evaluated on the task with the larger MLP with NT’s architecture.

Original: Thus, the meta sampler are trained without any knowledge of the test task’s training and test data.
Revised: Thus, the meta sampler is trained without any knowledge of the test task’s training and test data.

Original: From Figure 3, we see that NNSGHMC, although a bit slower at  start, catches up quickly and proceeds to lower error.
Revised: From Figure 3, we see that NNSGHMC, although a bit slower at the start, catches up quickly and proceeds to lower error.

Original: The difference between these samplers NLL results is marginal, and NNSGHMC is on par with PSGLD.
Revised: The difference between these samplers NLL results is marginal, and NNSGHMC is on par with PSGLD.

Original: Learned strategies For better intuition, we visualize in Figure 4 the contours of f f f φ D and f f f φ Q .
Revised: Learned strategies For better intuition, we visualize in Figure 4 the contours of fff φ D and fff φ Q .

Original: From the left panel, fff φ Q has learned a nearly linear strategy w.r.t.
Revised: From the left panel, fff φ Q has learned a nearly linear strategy w.r.t.

Original: the energy and small variations w.r.t.
Revised: the energy and small variations w.r.t.

Original: the momentum.
Revised: the momentum.

Original: This enables the sampler for fast traversal through low density (high energy) regions and better exploration at high density (low energy) area.
Revised: This enables the sampler for fast traversal through low density (high energy) regions and better exploration at high density (low energy) area.

Original: The strategy learned for the diffusion matrix DDD  is rather interesting.
Revised: The strategy learned for the diffusion matrix DDD ( zzz ) is rather interesting.

Original: Recall that DDD  is parametrized by both fff φ D and fff φ Q (cid:12) fff φ Q (Eq.6).
Revised: Recall that DDD ( zzz ) is parametrized by both fff φ D and fff φ Q (cid:12) fff φ Q (eq.

Original: Since Figure 4 (left) indicates that fff φ Q is large in high energy regions, the amount of friction is adequate, so fff φ D tends to reduce its output to maintain momentum (see the middle plot).
Revised: Since Figure 4 (left) indicates that fff φ Q is large in high energy regions, the amount of friction is adequate, so fff φ D tends to reduce its output to maintain momentum (see the middle plot).

Original: By contrast, at low energy regions fff φ D increases to add friction in order to prevent divergence.
Revised: By contrast, at low energy regions, fff φ D increases to add friction in order to prevent divergence.

Original: The right panel visualizes the interactions between the momentum and the mean gradient − 1 N ∇ θθθ U ( θθθ ) at a ﬁxed energy level.
Revised: The right panel visualizes the interactions between the momentum and the mean gradient − 1 N ∇ θθθ U ( θθθ ) at a ﬁxed energy level.

Original: This indicates that the meta sampler has learned a strategy to prevent overshoot by producing large friction, indeed fff φ D returns large values when the signs of the momentum and the gradient differ.
Revised: This indicates that the meta sampler has learned a strategy to prevent overshoot by producing large friction, indeed fff φ D returns large values when the signs of the momentum and the gradient differ.

Original: 5. B AYESIAN CONVOLUTIONAL NEURAL NETWORKS
Revised: 5. B AYESIAN CONVOLUTIONAL NEURAL NETWORKS

Original: Following the setup of BNN MNIST experiments, we also test our algorithm on convolutional neural networks (CNNs) for CIFAR-10 (Krizhevsky, 2009) classiﬁcation, again with three generalization tasks (NT, AF and Data).
Revised: Following the setup of BNN MNIST experiments, we also test our algorithm on convolutional neural networks (CNNs) for CIFAR-10 (Krizhevsky, 2009) classiﬁcation, again with three generalization tasks (NT, AF and Data).

Original: The meta sampler is trained using a smaller CNN with two convolutional layers ( 3 × 3 × 3 × 8 and 3 × 3 × 8 × 8 ) and one fully connected (fc) layer (50 hidden units).
Revised: The meta sampler is trained using a smaller CNN with two convolutional layers ( 3 × 3 × 3 × 8 and 3 × 3 × 8 × 8 ) and one fully connected (fc) layer (50 hidden units).

Original: ReLU activations, and max-pooling operators of size 2 are applied after each convolutional layer.
Revised: ReLU activations, and max-pooling operators of size 2 are applied after each convolutional layer.

Original: The meta sampler is trained using 100 “meta-epochs”, where each “meta-epoch” has 5 data epochs.
Revised: The meta sampler is trained using 100 “meta-epochs”, where each “meta-epoch” has 5 data epochs.

Original: At the beginning of each “meta-epoch”, a “replay” technique inspired by experience replay (Lin, 1993) is utilised (see appendix).
Revised: At the beginning of each “meta-epoch”, a “replay” technique inspired by experience replay (Lin, 1993) is utilized (see appendix).

Original: The discretization step-sizes are tuned on a validation dataset for each task.
Revised: The discretization step-sizes are tuned on a validation dataset for each task.

Original: Architecture generalization (NT)
Revised: Architecture generalization (NT)

Original: The test CNN has two convolutional layers ( 3 × 3 × 3 × 16 and 3 × 3 × 16 × 16 ) and one fc layer (100 hidden units), resulting in roughly 4 × dimenality of θθθ .
Revised: The test CNN has two convolutional layers ( 3 × 3 × 3 × 16 and 3 × 3 × 16 × 16 ) and one fc layer (100 hidden units), resulting in roughly 4 × dimenality of θθθ .

Original: Figure 5 shows that the meta sampler achieves the fastest learning at the ﬁrst 10 epochs, and continues to have better performance in both test accuracy and NLL.
Revised: Figure 5 shows that the meta sampler achieves the fastest learning at the ﬁrst 10 epochs, and continues to have better performance in both test accuracy and NLL.

Original: Interestingly, PSGLD slows down quickly after 3 epochs, and it converges to a worse answer.
Revised: Interestingly, PSGLD slows down quickly after 3 epochs, and it converges to a worse answer.

Original: The best performance over 200 epochs is shown in Table 2, where the meta sampler is a clear winner in both accuracy and NLL.
Revised: The best performance over 200 epochs is shown in Table 2, where the meta sampler is a clear winner in both accuracy and NLL.

Original: This demonstrates that our sampler indeed converges faster and has found a better posterior mode. 
Revised: This demonstrates that our sampler indeed converges faster and has found a better posterior mode. [Figure] [Table]

Original: Architecture + Activation function generalization (NT+AF)
Revised: Architecture + Activation function generalization (NT+AF)

Original: We use the same CNN architecture as in NT but replace the ReLU activations with sigmoid.
Revised: We use the same CNN architecture as in NT but replace the ReLU activations with sigmoid.

Original: Figure 5 and Table 2 show that the meta sampler again has better convergence speed and the best ﬁnal performance.
Revised: Figure 5 and Table 2 show that the meta sampler again has better convergence speed and the best ﬁnal performance.

Original: Architecture + Dataset generalization(NT+Data) We split CIFAR-10 according to labels 0-4 as the training task and 5-9 as the test task.
Revised: Architecture + Dataset generalization (NT+Data) We split CIFAR-10 according to labels 0-4 as the training task and 5-9 as the test task.

Original: We also used the same CNN architecture as in NT.
Revised: We also used the same CNN architecture as in NT.

Original: From
Revised: From

Original: Figure 5 and Table 2, the meta sampler consistently achieves the fastest convergence speed.
Revised: Figure 5 and Table 2, the meta sampler consistently achieves the fastest convergence speed.

Original: It also achieves similar accuracy as SGHMC, but it has slightly worse test NLL compared to SGHMC.[Figure] [Table] [Table]
Revised: It also achieves similar accuracy as SGHMC, but it has slightly worse test NLL compared to SGHMC.

Original:  B AYESIAN RECURRENT NEURAL NETWORKS
Revised: 5. B AYESIAN RECURRENT NEURAL NETWORKS

Original: Lastly we consider a more challenging setup: sequence modeling with Bayesian RNNs.
Revised: Lastly, we consider a more challenging setup: sequence modeling with Bayesian RNNs.

Original: Here a single datum is a sequence ooo n = { xxx 1 n , ...,xxx Tn } and the log-likelihood is deﬁned as log p ( ooo n | θθθ ) =
Revised: Here a single datum is a sequence ooo n = { xxx 1 n , ...,xxx Tn } and the log-likelihood is deﬁned as log p ( ooo n | θθθ ) =

Original: (cid:80)
Revised: (cid:80)

Original: Tt =1 log p ( xxx nt | xxx n 1 , . . .
Revised: Tt =1 log p ( xxx nt | xxx n 1 , . . .

Original: ,xxx nt − 1 ,θθθ ) , with each of the conditional densities produced by a gated recurrent unit (GRU) network (Cho et al., 2014).
Revised: ,xxx nt − 1 ,θθθ ) , with each of the conditional densities produced by a gated recurrent unit (GRU) network (Cho et al., 2014).

Original: We consider four polyphonic music datasets for this task:
Revised: We consider four polyphonic music datasets for this task:

Original: Piano-midi (Piano) as training data, and Nottingham (Nott), MuseData (Muse) and JSB chorales (JSB) for evaluation.
Revised: Piano-midi (Piano) as training data, and Nottingham (Nott), MuseData (Muse) and JSB chorales (JSB) for evaluation.

Original: The meta sampler is trained on a small GRU with 100 hidden states.
Revised: The meta sampler is trained on a small GRU with 100 hidden states.

Original: At test time we follow Chen et al.
Revised: At test time, we follow Chen et al.

Original: and set the step-size to η = 0 .
Revised: and set the step-size to η = 0 .

Original: We found SGLD signiﬁcantly under-performs, so instead, we report the performances of two optimizers. Adam (Kingma & Ba, 2014) and Santa (Chen et al., 2016), taken from Chen et al.
Revised: We found SGLD signiﬁcantly under-performs, so instead, we report the performances of two optimizers, Adam (Kingma & Ba, 2014) and Santa (Chen et al., 2016), taken from Chen et al.

Original: Again, these two optimizers use moving average schemes which are out of the scope of our framework, so we mainly compare the meta sampler with SGHMC and leave the others as references.
Revised: Again, these two optimizers use moving average schemes which are out of the scope of our framework, so we mainly compare the meta sampler with SGHMC and leave the others as references.

Original: The meta sampler is tested on the four datasets using 200 unit GRU.
Revised: The meta sampler is tested on the four datasets using 200 unit GRU.

Original: So for Piano this corresponds to architecture generalization only.
Revised: So for Piano this corresponds to architecture generalization only.

Original: From Figure 6 we see that the meta sampler achieves faster convergence compared to SGHMC, at the same time it achieves similar speed as Santa at early stages.
Revised: From Figure 6 we see that the meta sampler achieves faster convergence compared to SGHMC, at the same time it achieves similar speed as Santa at early stages.

Original: All the samplers achieve best results close to Santa on Piano.
Revised: All the samplers achieve best results close to Santa on Piano.

Original: The meta sampler successfully generalizes to the other three datasets, demonstrating faster convergence than SGHMC consistently, and better ﬁnal performance on Muse.
Revised: The meta sampler successfully generalizes to the other three datasets, demonstrating faster convergence than SGHMC consistently, and better ﬁnal performance on Muse.

Original: Interestingly, the meta sampler’s ﬁnal results on Nott and
Revised: Interestingly, the meta sampler’s ﬁnal results on Nott and

Original: JSB are slightly worse than other samplers.
Revised: JSB are slightly worse than other samplers.

Original: Presumably these two datasets are very different from
Revised: Presumably, these two datasets are very different from

Original: Muse and Piano, therefore, the energy landscape is less similar to the training density (see appendix).
Revised: Muse and Piano, therefore, the energy landscape is less similar to the training density (see appendix).

Original: Speciﬁcally JSB is a dataset with much shorter sequences.
Revised: Speciﬁcally, JSB is a dataset with much shorter sequences.

Original: And in this case, SGHMC also exhibits over-ﬁtting but to a smaller degree.
Revised: And in this case, SGHMC also exhibits over-ﬁtting but to a smaller degree.

Original: Therefore, we further test the meta sampler on JSB without the offset β in fff φ Q to reduce the acceleration (denoted as NNSGHMC-s).
Revised: Therefore, we further test the meta sampler on JSB without the offset β in fff φ Q to reduce the acceleration (denoted as NNSGHMC-s).

Original: Surprisingly, NNSGHMC-s  convergences in similar speeds as the original one, but with less amount of over-ﬁtting and better ﬁnal test NLL 8.40 .
Revised: Surprisingly, NNSGHMC-s [Figure] convergences in similar speeds as the original one, but with less amount of over-ﬁtting and better ﬁnal test NLL 8.40 .

Original:  C ONCLUSIONS AND FUTURE WORK
Revised:  C ONCLUSIONS AND FUTURE WORK

Original: We have presented a meta-learning algorithm that can learn an SG-MCMC sampler on simpler tasks and generalizes to more complicated densities in high dimensions.
Revised: We have presented a meta-learning algorithm that can learn an SG-MCMC sampler on simpler tasks and generalizes to more complicated densities in high dimensions.

Original: Experiments on Bayesian MLPs,
Revised: Experiments on Bayesian MLPs,

Original: Bayesian CNNs and Bayesian RNNs conﬁrmed the strong generalization of the trained sampler to  long-time horizon as well as across datasets and network architectures.
Revised: Bayesian CNNs and Bayesian RNNs conﬁrmed the strong generalization of the trained sampler to the long-time horizon as well as across datasets and network architectures.

Original: Future work will focus on[Figure] better designs for both the sampler and the meta-learning procedure.
Revised: Future work will focus on better designs for both the sampler and the meta-learning procedure.

Original: For the former, temperature variable augmentation as well as moving average estimation will be explored.
Revised: For the former, temperature variable augmentation as well as moving average estimation will be explored.

Original: For the latter, better loss functions will be proposed for faster training, e g. by reducing the unrolling steps of the sampler during training.
Revised: For the latter, better loss functions will be proposed for faster training, e.g. by reducing the unrolling steps of the sampler during training.

Original: Finally, the automated design of generic MCMC algorithms that might not be derived from continuous Markov processes remains an open challenge
Revised: Finally, the automated design of generic MCMC algorithms that might not be derived from continuous Markov processes remains an open challenge. A CKNOWLEDGMENTS

Revised: We thank Shixiang Gu, Mark Rowland and Cheng Zhang for comments on the manuscript.

Revised: We also appreciate Changyou Chen for providing the experiment results of Bayesian RNN (Chen et al., 2016).

Revised: Wenbo Gong is supported by the CSC-Cambridge Trust Scholarship.

Original: A C OMPARING MOMENTUM SGD AND SGHMC
Revised: A C OMPARING MOMENTUM SGD AND SGHMC

Original: Similar to the relationship between SGLD and SGD, SGHMC is closely related SGD with momentum (SGD-M).
Revised: Similar to the relationship between SGLD and SGD, SGHMC is closely related SGD with momentum (SGD-M).

Original: First in HMC, the state space is augmented with an additional momentum variable denoted as ppp ∈ R D .
Revised: First in HMC, the state space is augmented with an additional momentum variable denoted as ppp ∈ R D .

Original: We assume an identity mass matrix associated with that momentum term.
Revised: We assume an identity mass matrix associated with that momentum term.

Original: Then the corresponding drift fff ( θθθ,ppp ) and diffusion matrix DDD are:
Revised: Then the corresponding drift fff ( θθθ,ppp ) and diffusion matrix DDD are:

Original: [Equation]
Revised: [Equation]

Original: where CCC is a positive deﬁnite matrix called friction coefﬁcient.
Revised: where CCC is a positive deﬁnite matrix called friction coefﬁcient.

Original: Thus, HMC’s continuous-time dynamics is governed by the following SDE:
Revised: Thus, HMC’s continuous-time dynamics is governed by the following SDE:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The discretized update rule (with simple Euler discretization) of HMC with step-size η is
Revised: The discretized update rule (with simple Euler discretization) of HMC with step-size η is

Original: [Equation]
Revised: [Equation]

Original:  If stochastic gradient ∇ ˜ U ( θθθ ) is used, we need to replace the covariance matrix of (cid:15)(cid:15)(cid:15) with 2 η ( CCC − ˆ BBB ) where ˆ BBB is the variance estimation of the gradients.
Revised: [Equation] If stochastic gradient ∇ ˜ U ( θθθ ) is used, we need to replace the covariance matrix of (cid:15)(cid:15)(cid:15) with 2 η ( CCC − ˆ BBB ) where ˆ BBB is the variance estimation of the gradients.

Original: On the other hand, the update equations of SGD with momentum (SGD-M) are the following:
Revised: On the other hand, the update equations of SGD with momentum (SGD-M) are the following:

Original: [Equation]
Revised: [Equation]

Original: where k and l are called momentum discount factor and learning rate , respectively.
Revised: where k and l are called momentum discount factor and learning rate , respectively.

Original: Also we can rewrite the SGHMC update equations by setting ηppp t = vvv t , l = η 2 , k = (1 − ηCCC ) ,
Revised: Also we can rewrite the SGHMC update equations by setting ηppp t = vvv t , l = η 2 , k = (1 − ηCCC ) ,

Original: [Equation]
Revised: [Equation]

Original: Thus, the discretized SGHMC updates can be viewed as the SGD-M update injected with carefully controlled Gaussian noise.
Revised: Thus, the discretized SGHMC updates can be viewed as the SGD-M update injected with carefully controlled Gaussian noise.

Original: Therefore, the hyperparameter of SGHMC can be heuristically chosen based on the experience of SGD-M and vice versa.
Revised: Therefore, the hyperparameter of SGHMC can be heuristically chosen based on the experience of SGD-M and vice versa.

Original: Neal et al.
Revised: Neal et al.

Original: (2011) showed that in practice, simple Euler discretization for HMC simulation might cause divergence, therefore advanced discretization schemes such as Leapfrog and modiﬁed Euler are recommended.
Revised: (2011) showed that in practice, simple Euler discretization for HMC simulation might cause divergence, therefore advanced discretization schemes such as Leapfrog and modiﬁed Euler are recommended.

Original: We use modiﬁed Euler discretization in our implementation of SGHMC and the meta sampler, resulting in the following update:
Revised: We use modiﬁed Euler discretization in our implementation of SGHMC and the meta sampler, resulting in the following update:

Original: [Equation]
Revised: [Equation]

Original: B F INITE DIFFERENCE APPROXIMATION FOR THE G AMMA VECTOR
Revised: B F INITE DIFFERENCE APPROXIMATION FOR THE G AMMA VECTOR

Original: The main computational burden is the gradient computation required by ΓΓΓ( zzz ) vector.
Revised: The main computational burden is the gradient computation required by ΓΓΓ( zzz ) vector.

Revised: To simplify notations we will write DDD = DDD ( zzz ) and QQQ =

Revised: QQQ ( zzz ) .

Original: From the parametrization of  QQQ and DDD matrix in  (6), for θθθ,ppp ∈
Revised: From the parametrization of the QQQ and DDD matrices in eq. (6), for θ θ θ,p p p ∈

Original: R D we have ΓΓΓ( zzz ) =
Revised: R D we have Γ Γ Γ( z z z ) =

Original: [ΓΓΓ θθθ , ΓΓΓ ppp ] .
Revised: [Γ Γ Γ θθθ , Γ Γ Γ ppp ] .

Original: For the ﬁrst term ΓΓΓ θθθ , we have
Revised: For the ﬁrst term Γ Γ Γ θθθ , we have

Original: [Equation]
Revised: [Equation]

Original: Due to the two-stage update of Euler integrator, at time t, we have f t − 1 φ Q ,i =
Revised: Due to the two-stage update of Euler integrator, at time t, we have f t − 1 φ Q ,i =

Original: f φ Q ,i ( U ( θθθ t − 1 ) , p t − 1 i ) , ˆ f t − 1 φ Q ,i = f φ Q ,i ( U ( θθθ t − 1 ) , p t i ) and f t − 1 φ D ,i = f φ D ,i ( U ( θθθ t − 1 ) , p t − 1 i , ∇ θ t − 1 i
Revised: f φ Q ,i ( U ( θθθ t − 1 ) , p t − 1 i ) , ˆ f t − 1 φ Q ,i = f φ Q ,i ( U ( θθθ t − 1 ) , p t i ) and f t − 1 φ D ,i = f φ D ,i ( U ( θθθ t − 1 ) , p t − 1 i , ∇ θ t − 1 i

Original: U ( θθθ t − 1 )) .
Revised: U ( θθθ t − 1 )) .

Original: Thus a proper ﬁnite difference method requires f φ Q ,i ( U ( θθθ t ) , p t − 1 i ) , which is not exactly the history from the previous time.
Revised: Thus a proper ﬁnite difference method requires f φ Q ,i ( U ( θθθ t ) , p t − 1 i ) , which is not exactly the history from the previous time.

Original: Therefore we further approximate it using delayed estimate:
Revised: Therefore we further approximate it using delayed estimate:

Original: [Equation]
Revised: [Equation]

Original: Similarly [Equation] the ΓΓΓ ppp term expands as
Revised: Similarly, the ΓΓΓ ppp term expands as

Original: [Equation]
Revised: [Equation]

Original: We further approximate ∂f φQ,i ∂U ( θθθ ) by the following
Revised: We further approximate ∂f φQ,i ∂U ( θθθ ) by the following

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: This only requires the storage of previous QQQ matrix.
Revised: This only requires the storage of previous QQQ matrix.

Original: However, ∂f φD,i ∂p i requires one further forward pass to obtain ˆ f t − 1 φ D ,i = f
Revised: However, ∂f φD,i ∂p i requires one further forward pass to obtain ˆ f t − 1 φ D ,i = f

Original: φ D ,i ( U ( θθθ t ) , p t − 1 i , ∇ θ ti U ( θθθ t )) , thus [Equation] we have
Revised: φ D ,i ( U ( θθθ t ) , p t − 1 i , ∇ θ ti U ( θθθ t )) , thus, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore the proposed ﬁnite difference method only requires one more forward passes to compute ˆ fff t − 1 φ D and instead, save 3 back-propagations.
Revised: Therefore the proposed ﬁnite difference method only requires one more forward passes to compute ˆ fff t − 1 φ D and instead, save 3 back-propagations.

Original: As back-propagation is typically more expensive than forward pass, our approach reduces running time drastically, especially when the sampler are applied to large neural network.
Revised: As back-propagation is typically more expensive than forward pass, our approach reduces running time drastically, especially when the sampler are applied to large neural network.

Original: Time complexity ﬁgures Every SG-MCMC method (including the meta sampler) requires ∇ θθθ  U ( θθθ ) .
Revised: Time complexity ﬁgures Every SG-MCMC method (including the meta sampler) requires ∇ θθθ ˜ U ( θθθ ) .

Original: The main burden is the forward pass and back-propagation through the DDD  and QQQ matrix, where the latter one has been replaced by the proposed ﬁnite difference scheme.
Revised: The main burden is the forward pass and back-propagation through the DDD ( zzz ) and QQQ ( zzz ) matrices, where the latter one has been replaced by the proposed ﬁnite difference scheme.

Original: The time complexity is O ( HD ) for both forward pass and ﬁnite difference with H the number of hidden units in the neural network of the meta sampler.
Revised: The time complexity is O ( HD ) for both forward pass and ﬁnite difference with H the number of hidden units in the neural network of the meta sampler.

Original: Parallel computation with GPUs improves real-time speed, indeed in our MNIST experiment the meta sampler spends roughly 1.5x time when compared with SGHMC.
Revised: Parallel computation with GPUs improves real-time speed, indeed in our MNIST experiment the meta sampler spends roughly 1.5x time when compared with SGHMC.

Original: During meta sampler training, the Stein gradient estimator requires the kernel matrix inversion which is O ( K 3 ) for cross-chain training.

Original: In practice, we only run a few parallel Markov chains

Original: K = 20 ∼ 50 , thus, this will not incur huge computation cost.

Original: For in-chain loss the computation can also be reduced with proper thinning schemes.

Original: C D ETAILS OF THE S TEIN GRADIENT ESTIMATOR
Revised: C D ETAILS OF THE S TEIN GRADIENT ESTIMATOR

Original: For a distribution q ( θθθ ) that is implicitly deﬁned by a generative procedure, the density q ( θθθ ) is often intractable.
Revised: For a distribution q ( θθθ ) that is implicitly deﬁned by a generative procedure, the density q ( θθθ ) is often intractable.

Original: Li & Turner (2018) derived the Stein gradient estimator that estimates G = ( ∇ θθθ 1 log q ( θθθ 1 ) , · · · ∇ θθθ K log q ( θθθ K ))
Revised: Li & Turner (2018) derived the Stein gradient estimator that estimates G = ( ∇ θθθ 1 log q ( θθθ 1 ) , · · · ∇ θθθ K log q ( θθθ K ))

Original: T on samples θθθ 1 , ...,θθθ K ∼ q ( θθθ ) .
Revised: T on samples θθθ 1 , ...,θθθ K ∼ q ( θθθ ) .

Original: There are two different ways to derive this gradient estimator, here we brieﬂy introduce one of them, and refer the readers to Li &
Revised: There are two different ways to derive this gradient estimator, here we brieﬂy introduce one of them, and refer the readers to Li &

Original: Turner (2018) for details.
Revised: Turner (2018) for details.

Original: We start by introducing Stein’s identity (Stein, 1972; 1981; Gorham & Mackey, 2015; Liu et al., 2016).
Revised: We start by introducing Stein’s identity (Stein, 1972; 1981; Gorham & Mackey, 2015; Liu et al., 2016).

Original: Let h : R d × 1 → R d (cid:48) × 1 be a differentiable multivariate test function which maps θθθ to a column vector h ( θθθ ) =
Revised: Let h : R d × 1 → R d (cid:48) × 1 be a differentiable multivariate test function which maps θθθ to a column vector h ( θ θ θ ) =

Original: [ h 1 ( θθθ ) , h 2 ( θθθ ) , ..., h d (cid:48)
Revised: [ h 1 ( θ θ θ ) , h 2 ( θ θ θ ) , ..., h d (cid:48)

Original: ( θθθ )] T .
Revised: ( θ θ θ )] T .

Original: One can use integration by parts to show the following Stein’s identity when a boundary condition lim || θθθ ||→ ∞ q ( θθθ ) h ( θθθ ) = 0 is assumed for the test function
Revised: One can use integration by parts to show the following Stein’s identity when a boundary condition lim || θθθ ||→ ∞ q ( θθθ ) h ( θθθ ) = 0 is assumed for the test function:

Original: [Equation]
Revised: [Equation]

Original: This boundary condition holds for almost any test function if q has sufﬁciently fast-decaying tails (e.g. Gaussian tails).
Revised: This boundary condition holds for almost any test function if q has sufﬁciently fast-decaying tails (e.g. Gaussian tails).

Original: Li & Turner (2018) proposed the Stein gradient estimator for ∇ θθθ log q ( θθθ ) by inverting a Monte Carlo (MC) version of Stein’s identity (23):
Revised: Li & Turner (2018) proposed the Stein gradient estimator for ∇ θ θ θ log q ( θθθ ) by inverting a Monte Carlo (MC) version of Stein’s identity (23):

Original: [Equation]
Revised: [Equation]

Original: Then G is obtained by ridge regression (with || · || F the Frobenius norm of a matrix)
Revised: Then G is obtained by ridge regression (with || · || F the Frobenius norm of a matrix)

Original: [Equation]
Revised: [Equation]

Original: which has an analytical solution
Revised: which has an analytical solution

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Here θθθ k ( j ) denotes the j th element of vector θθθ k .
Revised: Here θθθ k ( j ) denotes the j th element of vector θθθ k .

Original: One can show that the RBF kernel satisﬁes Stein’s identity (Liu et al., 2016).
Revised: One can show that the RBF kernel satisﬁes Stein’s identity (Liu et al., 2016).

Original: In this case h ( θθθ ) =
Revised: In this case h ( θθθ ) =

Original: K ( θθθ, · ) , d (cid:48) =
Revised: K ( θθθ, · ) , d (cid:48) =

Original: + ∞ and by the reproducing kernel property, h (θθθ ) T h ( θθθ (cid:48) )
Revised: + ∞ and by the reproducing kernel property, h (

Original: = (cid:104)K ( θθθ, · ) , K ( θθθ (cid:48) , · ) (cid:105) H =
Revised: [Equation] H =

Original: K ( θθθ,θθθ (cid:48) ) .
Revised: K ( θθθ,θθθ (cid:48) ) .

Original: Li & Turner (2018) also show that the Stein gradient estimator can be obtained by minimizing a Monte Carlo estimate of the kernelized Stein discrepancy (Chwialkowski et al., 2016;
Revised: Li & Turner (2018) also show that the Stein gradient estimator can be obtained by minimizing a Monte Carlo estimate of the kernelized Stein discrepancy (Chwialkowski et al., 2016;

Original: Liu et al., 2016).
Revised: Liu et al., 2016).

Original:  It is well-known for kernel methods that a better choice of the kernel can greatly improve the performance.
Revised: The kernel choice It is well-known for kernel methods that a better choice of the kernel can greatly improve the performance.

Original: However, optimal kernels are often problem speciﬁc, and they are generally difﬁcult to obtain.
Revised: However, optimal kernels are often problem speciﬁc, and they are generally difﬁcult to obtain.

Original: Recently, a popular approach for kernel design is to compose a simple kernel (e.g. RBF kernel) on features extracted from a deep neural network.
Revised: Recently, a popular approach for kernel design is to compose a simple kernel (e.g. RBF kernel) on features extracted from a deep neural network.

Original: Representative work include deep kernel learning for Gaussian processes (Wilson et al., 2016), and adversarial approaches to learn kernel parameters (Li et al., 2017; Bi ´ nkowski et al., 2018).
Revised: Representative work include deep kernel learning for Gaussian processes (Wilson et al., 2016), and adversarial approaches to learn kernel parameters (Li et al., 2017; Bi ´ nkowski et al., 2018).

Original: Unfortunately, both approaches do not scale very well to our application as θθθ has at least tens of thousands of dimensions.
Revised: Unfortunately, both approaches do not scale very well to our application as θθθ has at least tens of thousands of dimensions.

Original: Furthermore, they both considered kernel learning for observed data, while in our case θθθ is a latent variable to be inferred.
Revised: Furthermore, they both considered kernel learning for observed data, while in our case θθθ is a latent variable to be inferred.

Original: Therefore it remains a research question on how to learn kernels on latent variables efﬁciently, and addressing this question is out of the scope of the paper.
Revised: Therefore it remains a research question on how to learn kernels on latent variables efﬁciently, and addressing this question is out of the scope of the paper.

Original: Instead, we follow Liu &Wang (2016); Li &
Revised: Instead, we follow Liu &

Original:  Turner (2018) to use RBF kernel for the gradient estimator.
Revised: Wang (2016); Li & Turner (2018) to use RBF kernel for the gradient estimator.

Original: Other kernels can be trivially adapted to our method.
Revised: Other kernels can be trivially adapted to our method.

Original: We expect even better performance if an optimal kernel is in use, but we leave the investigation to future work.
Revised: We expect even better performance if an optimal kernel is in use, but we leave the investigation to future work.

Revised: Time complexity ﬁgures During meta sampler training, the Stein gradient estimator requires the kernel matrix inversion which is O ( K 3 ) for cross-chain training.

Revised: In practice, we only run a few parallel Markov chains K = 20 ∼ 50 , thus, this will not incur huge computation cost.

Revised: For in-chain loss the computation can also be reduced with proper thinning schemes.

Original: D I MPLEMENTATION DETAILS OF THE TRAINING LOSS
Revised: D I MPLEMENTATION DETAILS OF THE TRAINING LOSS

Original: We visualize on the left panel of Figure 7 the unrolled computation scheme.
Revised: We visualize on the left panel of Figure 7 the unrolled computation scheme.

Original: We apply truncated back-propagate through time (BPTT) to train the sampler.
Revised: We apply truncated back-propagate through time (BPTT) to train the sampler.

Original: Speciﬁcally, we manually stop the gradient ﬂow through the input of DDD and QQQ matrix to avoid computing higher order gradients.
Revised: Speciﬁcally, we manually stop the gradient ﬂow through the input of DDD and QQQ matrices to avoid computing higher order gradients.

Original: We also illustrate cross-chain in-chain training on the right panel of Figure 7.
Revised: We also illustrate cross-chain in-chain training on the right panel of Figure 7.

Original: Cross-chain training encourages both fast convergence and low bias, provided that the samples are taken from parallel chains.
Revised: Cross-chain training encourages both fast convergence and low bias, provided that the samples are taken from parallel chains.

Original: On the other hand, in-chain training encourages sample diversity inside a chain.
Revised: On the other hand, in-chain training encourages sample diversity inside a chain.

Original: In practice, we might consider thinning the chains when performing in-chain training.
Revised: In practice, we might consider thinning the chains when performing in-chain training.

Original: Empirically this improves the Stein gradient estimator’s accuracy as the samples are spread out.
Revised: Empirically this improves the Stein gradient estimator’s accuracy as the samples are spread out.

Original: Computationally, this also prevents inverting big matrices for the Stein gradient estimator, and reduces the number of backpropagation operations.
Revised: Computationally, this also prevents inverting big matrices for the Stein gradient estimator, and reduces the number of backpropagation operations.

Original: Another trick we applied is parallel chain sub-sampling: if all the chains are used, then there is less encouragement of singe chain mixing, since the parallel chain samples can be diverse enough already to give reasonable gradient estimate.
Revised: Another trick we applied is parallel chain sub-sampling: if all the chains are used, then there is less encouragement of singe chain mixing, since the parallel chain samples can be diverse enough already to give reasonable gradient estimate.

Original: E I NPUT PRE - PROCESSING
Revised: E I NPUT PRE - PROCESSING

Original: One potential challenge is that for different tasks and problem dimensions, the energy function, momentum and energy gradient can have very different scales and magnitudes.
Revised: One potential challenge is that for different tasks and problem dimensions, the energy function, momentum and energy gradient can have very different scales and magnitudes.

Original: This affects the meta sampler’s generalization, for example, if training and test densities have completely different energy scales, then the meta sampler is likely to produce wrong strategies.
Revised: This affects the meta sampler’s generalization, for example, if training and test densities have completely different energy scales, then the meta sampler is likely to produce wrong strategies.

Original: This is especially the case when the meta sampler is generalized to much bigger networks or to very different datasets.
Revised: This is especially the case when the meta sampler is generalized to much bigger networks or to very different datasets.

Original: To mediate this issue, we propose to pre-process the inputs to both fff φ D and fff φ Q networks to make it at similar scale as those in training task.
Revised: To mediate this issue, we propose to pre-process the inputs to both fff φ D and fff φ Q networks to make it at similar scale as those in training task.

Original: Recall that the energy function is U ( θθθ ) =
Revised: Recall that the energy function is U ( θθθ ) =

Original: − (cid:80) Nn =1 log p ( yyy n | xxx n ,θθθ ) − log p ( θθθ ) where the prior log p ( θθθ ) is often  isotropic Gaussian distribution.
Revised: − (cid:80) Nn =1 log p ( yyy n | xxx n ,θθθ ) − log p ( θθθ ) where the prior log p ( θθθ ) is often an isotropic Gaussian distribution.

Original: Thus the energy function scale linearly w.r.t both the dimensionality of θθθ and the total number of observations N .
Revised: Thus the energy function scale linearly w.r.t both the dimensionality of θθθ and the total number of observations N .

Original: Often the energy function is further approximated using mini-batches of M datapoints.
Revised: Often the energy function is further approximated using mini-batches of M datapoints.

Original: Putting them together, we propose pre-processing the energy as
Revised: Putting them together, we propose pre-processing the energy as

Original: [Equation]
Revised: [Equation]

Original: (26) where D train and D test are the dimensionality of θθθ in the training task and the test task, respectively.
Revised: (26) where D train and D test are the dimensionality of θθθ in the training task and the test task, respectively.

Original: Importantly, for RNNs N represents the total sequence length, namely N = (cid:80) N data n =1 T n , where
Revised: Importantly, for RNNs N represents the total sequence length, namely N = (cid:80) N data n =1 T n , where

Original: N data is the total number of sequences and T n is the sequence length for a datum xxx n .
Revised: N data is the total number of sequences and T n is the sequence length for a datum xxx n .

Original: We also deﬁne M accordingly.
Revised: We also deﬁne M accordingly.

Original: The momentum and energy gradient magnitudes are estimated by simulating a randomly initialized meta sampler for short iterations.
Revised: The momentum and energy gradient magnitudes are estimated by simulating a randomly initialized meta sampler for short iterations.

Original: With these statistics we normalize both the momentum and the energy gradient to have roughly zero mean and unit variance.
Revised: With these statistics we normalize both the momentum and the energy gradient to have roughly zero mean and unit variance.

Original: F E XPERIMENT S ETUP F. T OY E XAMPLE
Revised: F E XPERIMENT S ETUP F. T OY E XAMPLE

Original: We train our meta sampler on a 10D uncorrelated Gaussian with mean (3 , ..., 3) and randomly generated covariance matrix.
Revised: We train our meta sampler on a 10D uncorrelated Gaussian with mean (3 , ..., 3) and randomly generated covariance matrix.

Original: We do not set any offset and additional frictions, i.e. α = 0 and β = 0 .
Revised: We do not set any offset and additional frictions, i.e. α = 0 and β = 0 .

Original: The noise estimation matrix ˜ BBB are set to be 0 for both meta sampler and SGHMC.
Revised: The noise estimation matrix ˜ BBB are set to be 0 for both meta sampler and SGHMC.

Original: To mimic stochastic gradient, we manually inject Gaussian noise with zero mean and unit variance into ∇ θθθ ˜ U ( θθθ )
Revised: To mimic stochastic gradient, we manually inject Gaussian noise with zero mean and unit variance into ∇ θθθ ˜ U ( θθθ )

Original: = ∇ θθθ U ( θθθ )+ (cid:15)(cid:15)(cid:15),(cid:15)(cid:15)(cid:15) ∼ N (000 ,III ) .
Revised: = ∇ θθθ U ( θθθ )+ (cid:15)(cid:15)(cid:15),(cid:15)(cid:15)(cid:15) ∼ N (000 ,III ) .

Original: The functions fff φ D and fff φ Q are represented by 1-hidden-layer
Revised: The functions fff φ D and fff φ Q are represented by 1-hidden-layer

Original: MLPs with 40 hidden units.
Revised: MLPs with 40 hidden units.

Original: For training task, the meta sampler step size is 0.01.
Revised: For training task, the meta sampler step size is 0.01.

Original: The initial positions are drawn from Uniform ([0 , 6] D ) .
Revised: The initial positions are drawn from Uniform ([0 , 6] D ) .

Original: We train our sampler for 100 epochs and each epochs consists 4 x 100 steps.
Revised: We train our sampler for 100 epochs and each epochs consists 4 x 100 steps.

Original: For every 100 steps, we updates our QQQ and DDD  using Adam optimizer with learning rate 0.0005.
Revised: For every 100 steps, we updates the QQQ and DDD matrices using Adam optimizer with learning rate 0.0005.

Original: Then we continue the updated sampler with last position and momentum until4 sub-epochs are ﬁnished.
Revised: Then we continue the updated sampler with last position and momentum until sub-epochs are ﬁnished.

Original: We re-initialize the momentum and position.
Revised: We re-initialize the momentum and position.

Original: We use both cross-chain and in-chain losses.
Revised: We use both cross-chain and in-chain losses.

Original: The Stein Gradient estimator uses RBF kernel with bandwidth chosen to be 0.5 times the median-heuristic estimated value.
Revised: The Stein Gradient estimator uses RBF kernel with bandwidth chosen to be 0.5 times the median-heuristic estimated value.

Original: We unroll the Markov Chain for 20 steps before we manually stop the gradient.
Revised: We unroll the Markov Chain for 20 steps before we manually stop the gradient.

Original: For cross-chain training, we take sampler across chain for each 2 time steps.
Revised: For cross-chain training, we take sampler across chain for each 2 time steps.

Original: For in-Chain, we discard initial 50 points for burn-in and sub-sample the chain with batch size 5.
Revised: For in-Chain, we discard initial 50 points for burn-in and sub-sample the chain with batch size 5.

Original: We thin the samples for every 3 steps.
Revised: We thin the samples for every 3 steps.

Original: For both training and evaluation, we run 50 parallel Markov Chains.
Revised: For both training and evaluation, we run 50 parallel Markov Chains.

Original: The test task is to draw samples from a 20D correlated Gaussian with with mean (3 , ..., 3) and randomly generated covariance matrix.
Revised: The test task is to draw samples from a 20D correlated Gaussian with with mean (3 , ..., 3) and randomly generated covariance matrix.

Original: The step size is 0.025 for both meta sampler and SGHMC.
Revised: The step size is 0.025 for both meta sampler and SGHMC.

Original: To stabilize the meta sampler we also clamp the output values of fff φ Q within [ − 5 , 5] .
Revised: To stabilize the meta sampler we also clamp the output values of fff φ Q within [ − 5 , 5] .

Original: The friction matrix for SGHMC is selected as III .
Revised: The friction matrix for SGHMC is selected as III .

Original: F. B AYESIAN MLP MNIST
Revised: F. B AYESIAN MLP MNIST

Original: In MNIST experiment, we apply input pre-processing on energy function as in (26) and scale energy gradient by 70.
Revised: In MNIST experiment, we apply input pre-processing on energy function as in (26) and scale energy gradient by 70.

Original: Also, we scale up fff φ D by 50 to account for sum of stochastic noise.
Revised: Also, we scale up fff φ D by 50 to account for sum of stochastic noise.

Original: The offset α is selected as 0 .
Revised: The offset α is selected as 0 .

Original: 01 η as suggested by Chen et al.
Revised: 01 η as suggested by Chen et al.

Original: (2014), where η = (cid:113) lrN with lr the per-batch learning rate.
Revised: (2014), where η = (cid:113) lrN with lr the per-batch learning rate.

Original: We also turn off the off-set and noise estimation, i.e. β = 0 and ˜ B B B = 0 .
Revised: We also turn off the off-set and noise estimation, i.e. β = 0 and ˜ B B B = 0 .

Original: We runparallel chains for both training and evaluation.
Revised: We runparallel chains for both training and evaluation.

Original: We only adopt the cross chain training with thinning samplers of 5 times step.
Revised: We only adopt the cross chain training with thinning samplers of 5 times step.

Original: We also use the ﬁnite difference technique during evaluation to speed-up computations.
Revised: We also use the ﬁnite difference technique during evaluation to speed-up computations.

Original: F.2. A RCHITECTURE G ENERALIZATION (NT)
Revised: F.2. A RCHITECTURE G ENERALIZATION (NT)

Original: We train the meta sampler on a smaller BNN with architecture 784-20-10 and ReLU activation function, then test it on a larger one with architecture 784-40-40-10.
Revised: We train the meta sampler on a smaller BNN with architecture 784-20-10 and ReLU activation function, then test it on a larger one with architecture 784-40-40-10.

Original: In both cases the batch size is 500 following Chen et al.
Revised: In both cases the batch size is 500 following Chen et al.

Original: Both fff φ D and fff φ Q are parameterized by 1-hidden-layer MLPs with 10 units.
Revised: Both fff φ D and fff φ Q are parameterized by 1-hidden-layer MLPs with 10 units.

Original: The per-batch learning rate is 0.007.
Revised: The per-batch learning rate is 0.007.

Original: We train the sampler for 100 epochs and each one consists of 7 sub-epochs.
Revised: We train the sampler for 100 epochs and each one consists of 7 sub-epochs.

Original: For each sub-epoch, we run the sampler for 100 steps.
Revised: For each sub-epoch, we run the sampler for 100 steps.

Original: We re-initialize θθθ and momentum after each epoch.
Revised: We re-initialize θθθ and momentum after each epoch.

Original: To stabilize the meta sampler in evaluation, we ﬁrst run the meta sampler with small per-batch learning rate 0 .
Revised: To stabilize the meta sampler in evaluation, we ﬁrst run the meta sampler with small per-batch learning rate 0 .

Original: 0085 for 3 data epochs and clamp the QQQ values.
Revised: 0085 for 3 data epochs and clamp the QQQ values.

Original: After, we increase the per-batch learning rate to 0 .
Revised: After, we increase the per-batch learning rate to 0 .

Original: 018 with clipped fff φ Q .
Revised: 018 with clipped fff φ Q .

Original: The learning rate for SGHMC is 0 .
Revised: The learning rate for SGHMC is 0 .

Original: 01 for all times.
Revised: 01 for all times.

Original: For SGLD and PSGLD, they are 0 .
Revised: For SGLD and PSGLD, they are 0 .

Original: 2 and 1 . 4 × 10 − 3 respectively.
Revised: 2 and 1 . 4 × 10 − 3 respectively.

Original: These step-sizes are tuned on MNIST validation data.
Revised: These step-sizes are tuned on MNIST validation data.

Original: F.2. NT + A CTIVATION FUNCTION GENERALIZATION
Revised: F.2. NT + A CTIVATION FUNCTION GENERALIZATION

Original: We modify the test network’s activation function to sigmoid .
Revised: We modify the test network’s activation function to sigmoid .

Original: We use almost the same settings as in network generalization tests, except that the per-batch learning rates are tuned again on validation data.
Revised: We use almost the same settings as in network generalization tests, except that the per-batch learning rates are tuned again on validation data.

Original: For the meta sampler and SGHMC, they are 0 .
Revised: For the meta sampler and SGHMC, they are 0 .

Original: 18 and 0 . 15 .
Revised: 18 and 0 . 15 .

Original: For SGLD and PSGLD, they areand 1 .
Revised: For SGLD and PSGLD, they areand 1 .

Original: 3 × 10 − 2 .
Revised: 3 × 10 − 2 .

Original: F.2. NT + D ATASET G ENERALIZATION
Revised: F.2. NT + D ATASET G ENERALIZATION

Original: We train the meta sampler on ReLU network with architecture 784-20-5 to classify images 0-4, and test the sampler on ReLU network 784-40-40-5 to classify images 5-9.
Revised: We train the meta sampler on ReLU network with architecture 784-20-5 to classify images 0-4, and test the sampler on ReLU network 784-40-40-5 to classify images 5-9.

Original: The settings are mostly the same as in network architecture generalization for both training and evaluation.
Revised: The settings are mostly the same as in network architecture generalization for both training and evaluation.

Original: One exception is again the per-batch learning rate for PSGLD, which is tuned as 1 . 3 × 10 − 3 .
Revised: One exception is again the per-batch learning rate for PSGLD, which is tuned as 1 . 3 × 10 − 3 .

Original: Note that even though we use the same per-batch learning rate as before, the discretization step-size is now different due to smaller training dataset, thus, α will be automatically adjusted accordingly.
Revised: Note that even though we use the same per-batch learning rate as before, the discretization step-size is now different due to smaller training dataset, thus, α will be automatically adjusted accordingly.

Original: F. B AYESIAN C ONVOLUTIONAL N EURAL N ETWORK ON CIFAR- 
Revised: F. B AYESIAN C ONVOLUTIONAL N EURAL N ETWORK ON CIFAR-

Original: CIFAR-10 dataset contains 50,000 training images with 10 labels and 10,000 test images.
Revised: CIFAR-10 dataset contains 50,000 training images with 10 labels and 10,000 test images.

Original: We train our meta sampler using smaller CNN classiﬁer with two convolutional layer ( 3 × 3 × 3 × 8 and 3 × 3 × 8 × 8 , no padding) and one fc layer of 50 hidden units.
Revised: We train our meta sampler using smaller CNN classiﬁer with two convolutional layer ( 3 × 3 × 3 × 8 and 3 × 3 × 8 × 8 , no padding) and one fc layer of 50 hidden units.

Original: Therefore the dimensionality of θθθ is 15 , 768 .
Revised: Therefore the dimensionality of θθθ is 15 , 768 .

Original: The training sampler discretization step-size η is (cid:113) 0 .
Revised: The training sampler discretization step-size η is (cid:113) 0 .

Original: and scaling term is α = 0 .
Revised: and scaling term is α = 0 .

Original: 005 η .
Revised: 005 η .

Original: To make it analogous to optimization methods, we call 0 .
Revised: To make it analogous to optimization methods, we call 0 .

Original: 0007 as per-batch learning rate and 0 .
Revised: 0007 as per-batch learning rate and 0 .

Original: as friction coefﬁcient .
Revised: as friction coefﬁcient .

Original: The fff φ Q and fff φ D are deﬁned by 2-layer MLPs with 10 hidden units.
Revised: The fff φ Q and fff φ D are deﬁned by 2-layer MLPs with 10 hidden units.

Original: We set the offset values to 0 for both QQQ and DDD .
Revised: We set the offset values to 0 for both QQQ and DDD .

Original: Further, we scale up the output of DDD f ( zzz ) by 10 and its gradient input ∇ U ( θθθ ) by 100.
Revised: Further, we scale up the output of DDD f ( zzz ) by 10 and its gradient input ∇ U ( θθθ ) by 100.

Original: We scale up the energy input U ( θθθ ) to both fff φ Q and fff φ D by 5.
Revised: We scale up the energy input U ( θθθ ) to both fff φ Q and fff φ D by 5.

Original: We train our meta sampler using 100 “meta epoch” with 5 data epoch and 500 batch size.
Revised: We train our meta sampler using 100 “meta epoch” with 5 data epoch and 500 batch size.

Original: Within each ”meta epoch”, we repeat the following computation for 10 times: we run 50 parallel chains using the meta sampler for 50 iterations (0.5 dataset epoch), compute the loss function, and update the meta sampler’s parameters using Adam.
Revised: Within each ”meta epoch”, we repeat the following computation for 10 times: we run 50 parallel chains using the meta sampler for 50 iterations (0.5 dataset epoch), compute the loss function, and update the meta sampler’s parameters using Adam.

Original: We manually stop the gradient after 20 iterations.
Revised: We manually stop the gradient after 20 iterations.

Original: Then we start the next sub-epoch using the last θθθ and ppp .
Revised: Then we start the next sub-epoch using the last θθθ and ppp .

Original: After we ﬁnish all sub-epoch, we re-initialize the θθθ and ppp using replay techniques with probability 0.15.
Revised: After we ﬁnish all sub-epoch, we re-initialize the θθθ and ppp using replay techniques with probability 0.15.

Original: The sub-sample chain number for in-chain loss is set to 5.
Revised: The sub-sample chain number for in-chain loss is set to 5.

Original: F.3. T HE REPLAY TECHNIQUE
Revised: F.3. T HE REPLAY TECHNIQUE

Original: Experience replay (Lin, 1993) is a technique broadly used in reinforcement learning literature.
Revised: Experience replay (Lin, 1993) is a technique broadly used in reinforcement learning literature.

Original: Inspired by this, in Bayesian CNN experiments we train the meta sampler in a similar way, and we found this replay technique particularly useful for more complicated dataset like CIFAR-10.
Revised: Inspired by this, in Bayesian CNN experiments we train the meta sampler in a similar way, and we found this replay technique particularly useful for more complicated dataset like CIFAR-10.

Original: At the beginning of each “meta epoch”, each chain is initialized  with a speciﬁc state randomly chosen from a replay pool with a pre-deﬁned replay probability, or with a random state sampled from a 
Revised: At the beginning of each “meta epoch”, each chain is initialized either with a speciﬁc state randomly chosen from a replay pool, or with a random state sampled from a Gaussian distribution.

Original: Gaussian distribution.
Revised: We use a pre-deﬁned replay probability to control the replay strategy.

Original: The replay pool is updated after each sub-epoch, and it has a queue-like data structure of constant size, so that the old states are replaced by the new ones.
Revised: The replay pool is updated after each sub-epoch, and it has a queue-like data structure of constant size, so that the old states are replaced by the new ones.

Original: Therefore, this replay technique is useful for both short-time and long-time horizon generalization.
Revised: Therefore, this replay technique is useful for both short-time and long-time horizon generalization.

Original: On one hand, the meta sampler can continue with previous states, allowing it to accommodate long-time horizon behavior.
Revised: On one hand, the meta sampler can continue with previous states, allowing it to accommodate long-time horizon behavior.

Original: On the other hand, due to non-zero probability of random restart, the meta sampler can learn a better strategy for fast convergence.
Revised: On the other hand, due to non-zero probability of random restart, the meta sampler can learn a better strategy for fast convergence.

Original: Therefore with this replay technique, the sampler can observe both burn-in and roughly-converged behavior, and this balance is controlled by the replay probability.
Revised: Therefore with this replay technique, the sampler can observe both burn-in and roughly-converged behavior, and this balance is controlled by the replay probability.

Original: F.3. A RCHITECTURE GENERALIZATION (NT)
Revised: F.3. A RCHITECTURE GENERALIZATION (NT)

Original: For architecture generalization, the test CNN has two convolutional layer ( 3 × 3 × 3 × 16 and 3 × 3 × 16 × 16 , no padding) and one fully connected layer with 100 hidden units.
Revised: For architecture generalization, the test CNN has two convolutional layer ( 3 × 3 × 3 × 16 and 3 × 3 × 16 × 16 , no padding) and one fully connected layer with 100 hidden units.

Original: Thus, the dimensionality of θθθ is 61,478, roughly 4 times of the training dimension.
Revised: Thus, the dimensionality of θθθ is 61,478, roughly 4 times of the training dimension.

Original: We run 20 parallel chains in test time.
Revised: We run 20 parallel chains in test time.

Original: We split the 50,000 training images into 45,000 training and 5,000 validation images, and tune the discretization step-size of each sampling and optimization methods on the validation set for 80 epochs.
Revised: We split the 50,000 training images into 45,000 training and 5,000 validation images, and tune the discretization step-size of each sampling and optimization methods on the validation set for 80 epochs.

Original: For test, we run the tuned samplers/optimizers for 200 data epoch (roughlytimes longer than training) to ensure convergence.
Revised: For test, we run the tuned samplers/optimizers for 200 data epoch (roughlytimes longer than training) to ensure convergence.

Original: For the meta sampler, the per-batch rate is 0 . 003 .
Revised: For the meta sampler, the per-batch rate is 0 . 003 .

Original: For SGHMC, the per-batch is also 0.003 with friction coefﬁcient 0 .
Revised: For SGHMC, the per-batch is also 0.003 with friction coefﬁcient 0 .

Original: For SGLD, the per-batch learning rate is 0 . 15 .
Revised: For SGLD, the per-batch learning rate is 0 . 15 .

Original: PSGLD uses 1 .
Revised: PSGLD uses 1 .

Original: 3 × 10 − 3 as learning rate and 0 .
Revised: 3 × 10 − 3 as learning rate and 0 .

Original: 99 as moving average term.
Revised: 99 as moving average term.

Original: For optimization methods, we use learning rate 0.002 for Adam and 0.003 for SGD-M.
Revised: For optimization methods, we use learning rate 0.002 for Adam and 0.003 for SGD-M.

Original: The momentum term is 0.9.
Revised: The momentum term is 0.9.

Original: To prevent overﬁtting, we use weight penalty with coefﬁcient 0 . 001 .
Revised: To prevent overﬁtting, we use weight penalty with coefﬁcient 0 . 001 .

Original: F.3. NT + S IGMOID GENERALIZATION
Revised: F.3. NT + S IGMOID GENERALIZATION

Original: The test CNN has same architecture as in NT, except that it replaces all ReLU activation functions with sigmoid activations.
Revised: The test CNN has same architecture as in NT, except that it replaces all ReLU activation functions with sigmoid activations.

Original: We ﬁx all other parameters for sampling method and only re-tune the step sizes using same setup as in NT.
Revised: We ﬁx all other parameters for sampling method and only re-tune the step sizes using same setup as in NT.

Original: The per-batch rate for meta sampler, SGHMC, SGLD and PSGLD are 0 . 1 , 0 .
Revised: The per-batch rate for meta sampler, SGHMC, SGLD and PSGLD are 0 . 1 , 0 .

Original: 5 and 0 .
Revised: 5 and 0 .

Original: 005 respectively.
Revised: 005 respectively.

Original: For optimization methods, the step size for Adam and
Revised: For optimization methods, the step size for Adam and

Original: SGD-M are 0 .
Revised: SGD-M are 0 .

Original: 002 and 0 .
Revised: 002 and 0 .

Original: 03 respectively.
Revised: 03 respectively.

Original: F.3. NT + D ATASET GENERALIZATION
Revised: F.3. NT + D ATASET GENERALIZATION

Original: We split the CIFAR-10 training and test dataset according to the labels.
Revised: We split the CIFAR-10 training and test dataset according to the labels.

Original: We use training data with labels 0-4 for meta sampler training, training data with labels 5-9 for test CNN training, and test data with labels 5-9 for test CNN evaluation.
Revised: We use training data with labels 0-4 for meta sampler training, training data with labels 5-9 for test CNN training, and test data with labels 5-9 for test CNN evaluation.

Original: Thus, the meta sampler has no access to the test task’s training and test data during sampler training.
Revised: Thus, the meta sampler has no access to the test task’s training and test data during sampler training.

Original: We train our sampler using the same scaling terms as in NT but reduce the discretization step-size to 0 . 0005 .
Revised: We train our sampler using the same scaling terms as in NT but reduce the discretization step-size to 0 . 0005 .

Original: The rest setup is the same as in NT.
Revised: The rest setup is the same as in NT.

Original: We use the same test CNN architecture and ReLU activation as in NT, and tune the learning rate using validation data.
Revised: We use the same test CNN architecture and ReLU activation as in NT, and tune the learning rate using validation data.

Original: The step size for the meta sampler, SGHMC, SGLD and PSGLD are 0 . 0015 , 0 . 005 ,
Revised: The step size for the meta sampler, SGHMC, SGLD and PSGLD are 0 . 0015 , 0 . 005 ,

Original: [Table]
Revised: [Table]

Original: 0 . 2 and
Revised: 0 . 2 and

Original: F. B AYESIAN RNN
Revised: F. B AYESIAN RNN

Original: The Piano data is selected as the training task, which is further split into training, validation and test subsets.
Revised: The Piano data is selected as the training task, which is further split into training, validation and test subsets.

Original: We use batch-size 1, meaning that the energy and the gradient are estimated on a single sequence.
Revised: We use batch-size 1, meaning that the energy and the gradient are estimated on a single sequence.

Original: The meta sampler uses similar neural network architectures as in MNIST tests.
Revised: The meta sampler uses similar neural network architectures as in MNIST tests.

Original: The training and evaluation per-batch learning rate for all the samplers is set to be 0.001 following Chen et al.
Revised: The training and evaluation per-batch learning rate for all the samplers is set to be 0.001 following Chen et al.

Original: We train the meta sampler for 40 epochs with 7 sub-epochs with only cross chain loss.
Revised: We train the meta sampler for 40 epochs with 7 sub-epochs with only cross chain loss.

Original: Each sub-epochs consists 70 iterations.
Revised: Each sub-epochs consists 70 iterations.

Original: We scale the DDD output by 20 and set α = 0 .
Revised: We scale the DDD output by 20 and set α = 0 .

Original: 002 η , where η is deﬁned in the same way as before.
Revised: 002 η , where η is deﬁned in the same way as before.

Original: We use zero offset during training, i.e. β = 0 .
Revised: We use zero offset during training, i.e. β = 0 .

Original: We apply input pre-processing for both fff φ D and fff φ Q .
Revised: We apply input pre-processing for both fff φ D and fff φ Q .

Original: To prevent divergence of the meta sampler at early training stage.
Revised: To prevent divergence of the meta sampler at early training stage.

Original: We also set the constant of c = 100 to the f φ D .
Revised: We also set the constant of c = 100 to the f φ D .

Original: For dataset generalization, we tune the off-set value based on Piano validation set and transfer the tuned setting β = − 1 .
Revised: For dataset generalization, we tune the off-set value based on Piano validation set and transfer the tuned setting β = − 1 .

Original: 5 to the other three datasets.
Revised: 5 to the other three datasets.

Original: For Piano architecture generalization, we do not tune any hyper-parameters including β and use exactly same settings as training.
Revised: For Piano architecture generalization, we do not tune any hyper-parameters including β and use exactly same settings as training.

Original: Exact gradient is used in RNN experiments instead of computing ﬁnite differences.
Revised: Exact gradient is used in RNN experiments instead of computing ﬁnite differences.

Original: G RNN DATASET DESCRIPTION
Revised: G RNN DATASET DESCRIPTION

Original: We list some data statistics in Table 3 which roughly indicates the similarity between datasets.
Revised: We list some data statistics in Table 3 which roughly indicates the similarity between datasets.

Original: Piano dataset is the smallest in terms of data number, however, the averaged sequence length is the largest.
Revised: Piano dataset is the smallest in terms of data number, however, the averaged sequence length is the largest.

Original: Muse dataset is similar to Piano in sequence length and energy scale but much larger in terms of data number.
Revised: Muse dataset is similar to Piano in sequence length and energy scale but much larger in terms of data number.

Original: On the other hand, Nott dataset has very different energy scale compared to the other three.
Revised: On the other hand, Nott dataset has very different energy scale compared to the other three.

Original: This potentially makes the generalization much harder due to inconsistent energy scale fed into fff φ Q and fff φ D .
Revised: This potentially makes the generalization much harder due to inconsistent energy scale fed into fff φ Q and fff φ D .

Original: For JSB, we notice a very short sequence length on average, therefore the GRU model is more likely to over-ﬁt.
Revised: For JSB, we notice a very short sequence length on average, therefore the GRU model is more likely to over-ﬁt.

Original: Indeed, some algorithms exhibits signiﬁcant over-ﬁtting behavior on JSB dataset compared to other data (Santa is particularly severe).
Revised: Indeed, some algorithms exhibits signiﬁcant over-ﬁtting behavior on JSB dataset compared to other data (Santa is particularly severe).

Original: H A DDITIONAL P LOTS H. S HORT RUN COMPARISON
Revised: H A DDITIONAL P LOTS H. S HORT HORIZON PERFORMANCE COMPARISONS

Original: We also run the samplers using the same settings as in MNIST experiments for a short period of time (500 iterations).
Revised: We also run the samplers using the same settings as in MNIST experiments for a short period of time (500 iterations).

Original: We also compare to other optimization methods including Momentum SGD (SGD-M) and Adam.
Revised: We also compare to other optimization methods including momentum SGD (SGD-M) and Adam.

Original: We use the same per-batch learning rate for SGD-M and SGHMC as in MNIST experiment.
Revised: We use the same per-batch learning rate for SGD-M and SGHMC as in MNIST experiment.

Original: For Adam, we use 0.002 for ReLU and 0.01 for Sigmoid network.
Revised: For Adam, we use 0.002 for ReLU and 0.01 for Sigmoid network.

Original: The results are shown in Figure 8.
Revised: The results are shown in Figure 8.

Original: Meta sampler and Adam achieves the fastest convergence speed.
Revised: Meta sampler and Adam achieves the fastest convergence speed.

Original: This again conﬁrms the faster convergence of the meta sampler especially at initial stages.
Revised: This again conﬁrms the faster convergence of the meta sampler especially at initial stages.

Original: We also provide additional contour plots (Figure 9) for MNIST experiments to demonstrate the strategy learned by fff φ D for reference.
Revised: We also provide additional contour plots (Figure 9) for MNIST experiments to demonstrate the strategy learned by fff φ D for reference.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]