Adaptive Thermostats for Noisy Gradient Systems

Introduction

Stochastic thermostats are powerful tools for sampling probability measures on high-dimensional spaces. These methods combine an extended dynamics with degenerate stochastic perturbation to ensure ergodicity. The traditional use of thermostats in molecular dynamics is to sample a well-specified equilibrium system involving a known force field which is the gradient of a potential energy function. Recently, however, these techniques have become increasingly popular for problems of more general form, including the following:

multiscale models in which the forces are obtained by approximate sampling in another scale regime ;

nonequilibrium physical models in which the potential energy function either is evolving or does not completely specify the system ;

Bayesian machine learning applications in which a dataset defines an objective function which leads to an effective force law .

In this article, we consider thermostats and numerical methods for sampling an underlying probability measure in the presence of error, under the assumption that the errors are random with a simple distributional form and unknown, but constant or slowly varying, parameters. In the cases considered, these methods are simple to implement, robust, and efficient.

The main tool that we employ in this article is the general concept of a thermostat as a (stochastic) distributional control for a dynamical system. These methods originate in molecular dynamics, and it is simplest to explain them in that context. Classical molecular dynamics tracks the motion of individual atoms determined by Newton’s law in the microcanonical ($NVE$) ensemble, where energy (i.e., the Hamiltonian of the system) is always conserved . However, constant energy is not the appropriate setting of a real-world laboratory environment. In most cases, one wishes instead to sample the canonical ($NVT$) ensemble, where temperature, as an intensive variable, is conserved, by using thermostat techniques .

The idea of a thermostat is to modify dynamics so that a prescribed invariant measure is sampled. There are competing aims in this type of work. For example, one may wish to perturb the underlying Newtonian dynamics minimally, so that temporal correlations are preserved, or one may be interested in sampling rare events in a system with metastable states; thus a variety of methods have been developed. The most obvious proposals, and also the oldest, are Brownian and Langevin dynamics. In Brownian (sometimes called “overdamped Langevin”) dynamics, the system is

where ${{\bm{q}}}$ represents a $3N$-dimensional vector of time-dependent random variables, ${\rm d{\bf W}}$ represents a vector of infinitesimal Wiener increments, $\beta$ is a positive parameter (proportional to the reciprocal temperature), $U$ is the potential energy function, and $\lambda$ is a free parameter which represents a time-rescaling. It can be shown that this system (1) ergodically samples the Gibbs–Boltzmann probability distribution $\bar{\rho}_{\beta}\propto{\rm exp}(-\beta U)$. For simplicity, we assume that the configurations ${{\bm{q}}}$ are restricted to a compact and simply connected domain $\Omega_{{{\bm{q}}}}$. In molecular dynamics applications, the starting point is the potential energy function, which is usually assumed to be a semiempirical formula constructed from primitive functions via an a priori parameter fitting procedure. Alternatively, one may assume that it is the probability distribution that is specified and that the potential energy is constructed from it via

which, of course, requires that $\rho>0$. In many applications it is found that the use of a first order dynamics such as (1) is inefficient or introduces unphysical dynamical properties, and one employs, instead, the Langevin dynamics method:

Again, $\gamma$ in these equations is a free parameter, termed the “friction constant”. It is related to the timescale on which the variables of the system interact with particles of a fictitious extended “bath”, but it cannot be associated with a simple time-rescaling of the equations of motion and is thus different from $\lambda$ in (1). It is a little more involved to show that (2)–(3) ergodically samples the distribution with density $\rho_{\beta}\propto{\rm exp}(-\beta H({{\bm{q}}},{{\bm{p}}}))$, where $H({{\bm{q}}},{{\bm{p}}})={{{\bm{p}}}}^{T}{{\bm{M}}}^{-1}{{{\bm{p}}}}/2+U({{\bm{q}}})$. In molecular dynamics, the $3N\times 3N$ matrix ${{\bm{M}}}$ is typically diagonal and contains the masses of atoms, ${{{\bm{p}}}}$ represents the momentum vector, and $H$ is the Hamiltonian or energy function. In more general settings, the masses and friction coefficient may be treated as free parameters, and by computing long trajectories of (2)–(3), one may obtain averages with respect to $\bar{\rho}_{\beta}({{\bm{q}}})$; i.e., if $\{\left({{\bm{q}}}(\tau),{{{\bm{p}}}}(\tau)\right):\tau\geq 0\}$ is a path generated by solving the SDE system (2)–(3), one has, for suitable test functions $\phi({{\bm{q}}})$, and under certain conditions on the potential energy function $U$ ,

where ${\rm d}\omega_{{\bm{q}}}={\rm d}{{\bm{q}}}_{1}{\rm d}{{\bm{q}}}_{2}\dots\rm d{{\bm{q}}}_{N}$. In other words, the projected path defines a sampler for the density $\bar{\rho}_{\beta}$.

Langevin dynamics can thus be seen as an extended system which allows sampling to be performed in a reduced cross section of phase space by marginalization over long trajectories; this is the essential property of a thermostat. Other types of thermostats include Nosé–Hoover–Langevin (NHL) dynamics and various generalized schemes (see, e.g., ). In these methods, one adds additional auxiliary variables which are meant to control the dynamics (via a negative feedback loop), and the auxiliary variables are then further coupled to stochastic processes of Ornstein–Uhlenbeck type which can provide ergodicity . (Note that the use of purely deterministic approaches, such as Nosé–Hoover, results in ergodicity issues .) The use of auxiliary variables can provide a degree of flexibility in the design of the thermostat, for example, allowing the treatment of systems arising in fluid dynamics or imposing an isokinetic constraint . Very recently, we have further generalized the NHL method to obtain pairwise Nosé–Hoover–Langevin (PNHL), which is a momentum-conserving thermostat and thus applicable to the simulation of hydrodynamic behavior in complex fluids and polymers in mesoscales .

2 Noisy Gradients

The gradient (or Hamiltonian) structure is essential to the nature of all the methods described above since it is only by use of this feature that the underlying Fokker–Planck equation can be shown to have the desired steady state solution. However, in many applications, in particular multiscale modelling, the force is corrupted by significant approximation error and cannot be viewed as the gradient of a single global potential function. One imagines a large extended system involving configurational variables ${{\bm{q}}}$ and ${{\bm{y}}}$, with $({{\bm{q}}},{{\bm{y}}})\in\Omega_{{{\bm{q}}}}\times\Omega_{{{\bm{y}}}}$ (compact), and an overall distribution described by a Gibbs–Boltzmann density

In practice most systems constructed in this way, for example, those arising in mixed quantum and classical molecular models , will admit very substantial errors in the forces; that is,

Depending on the method of computation, it may be reasonable to assume that the errors $\Delta^{k}$ are normally distributed with zero mean, which is justified by the central limit theorem , but the variance of the errors is generally not known and will be dependent on the location ${{\bm{q}}}$ where they are computed; thus we would expect

where $\bm{\Sigma}^{k}({{\bm{q}}})$ is an unknown covariance matrix. It should be noted that the assumption of the errors being Gaussian distributed is also often adopted in Bayesian inverse problems and elsewhere.

The most straightforward approach to the problem is to first treat the estimation problem for $\bm{\Sigma}^{k}$ separately, by some means, and then to use this within a standard Brownian or Langevin dynamics algorithm. The difficulty is that this requires a high level of local accuracy in the calculations, which is likely to be burdensome and involve redundant computation. What we would prefer to do is to resolve the correct target distribution by a global calculation.

This problem has recently been encountered in the data science community, where it has attracted considerable attention . To illustrate, we consider the problem of Bayesian sampling , where one is interested in correctly drawing states from a posterior probability density defined as

where ${{\bm{\theta}}}$ is the parameter vector of interest, $\mathbf{X}$ represents the entire dataset, and, $\pi(\mathbf{X}|{{\bm{\theta}}})$ and $\pi({{\bm{\theta}}})$ represent the likelihood and prior distributions, respectively. In these applications, the distribution parameters are interpreted as the configuration variables (${{\bm{\theta}}}\equiv{{\bm{q}}}$). We introduce a potential energy $U({{\bm{\theta}}})$ by defining $\pi({{\bm{\theta}}}|\mathbf{X})\propto\exp(-\beta U({{\bm{\theta}}}))$; thus taking the logarithm of (5) gives

Assuming the data are independent and identically distributed (i.i.d.), the logarithm of the likelihood distribution can then be calculated as

where $N$ is the size of the entire dataset.

However, in machine learning applications, one often finds that directly sampling with the entire large-scale dataset is computationally infeasible. For instance, standard Markov chain Monte Carlo (MCMC) methods require the calculation of the acceptance probability and the creation of informed proposals based on the whole dataset, while the gradient is evaluated through the whole dataset in the hybrid Monte Carlo (HMC) method , again resulting in severe computational complexity.

3 Sampling Methods for Noisy Gradients

In the original SGLD method, samples are generated by Brownian dynamics,

where $\mathbf{R}_{n}$ is a vector of i.i.d. standard normal random variables. It should be emphasized that $\Delta t_{n}$ is a sequence of stepsizes decreasing to zero . Although a central limit theorem associated with the decreasing stepsize sequence was established by Teh et al. , a fixed stepsize is often preferred in practice, which is the choice in this article as in Vollmer et al. , where a modified SGLD (mSGLD) is introduced:

is the covariance matrix of the noisy force.

A stochastic gradient Hamiltonian Monte Carlo (SGHMC) method was also proposed very recently by Chen et al. , which incorporates a parameter-dependent diffusion matrix $\bm{\Sigma}({{\bm{q}}})$ (i.e., the covariance matrix of the noisy force). $\bm{\Sigma}({{\bm{q}}})$ is intended to effectively offset the stochastic perturbation of the gradient. However, it is very difficult to accommodate $\bm{\Sigma}({{\bm{q}}})$ in practice; moreover, as pointed out in , poor estimation of it may have a significant adverse influence in correctly sampling the target distribution unless the stepsize is small enough.

These problems challenge the conventional mechanism of thermostats. An article of Jones and Leimkuhler provides an alternative means of tackling this problem by showing that Nosé–Hoover dynamics is able to adaptively dissipate excess heat pumped into the system while maintaining the Gibbs (canonical) distribution. In the setting of systems involving a driving stochastic perturbation, the adaptive Nosé–Hoover method is referred to as Ad-NH, with similar generalizations of Nosé–Hoover–Langevin (Ad-NHL) and Langevin dynamics (Ad-Langevin) available. An idea equivalent to Ad-Langevin was very recently applied in the setting of Bayesian sampling for use in data science calculations by Ding et al. , which they referred to as the stochastic gradient Nosé–Hoover thermostat (SGNHT). It showed significant advantages over alternative techniques such as SGHMC . However, the numerical method used by Ding et al. is not optimal, neither in terms of its accuracy (measured per unit work) nor its stability (measured by the largest usable stepsize).

Although extended systems have been increasingly popular in molecular simulations, the mathematical analysis of the order of convergence, specifically in terms of the bias in averaged quantities computed using numerical trajectories, is not fully understood. Using a splitting approach, we propose in this article an alternative numerical method for Ad-Langevin simulation that substantially improves on the existing schemes in the literature in terms of accuracy, robustness, and overall numerical efficiency.

The rest of the article is organized as follows. In Section 2, we describe the construction of adaptive formulations for noisy gradients including the Ad-Langevin/SGNHT method. Section 3 considers the construction of numerical methods for solving the SDEs. Numerical experiments are performed in Section 4. Our experiments are of a more limited nature in comparison with those of Ding et al. , but we believe them to be representative of performance on a significant class of problems. Finally, we summarize our findings in Section 5.

Adaptive Thermostats for Noisy Gradients

In this section, we discuss the construction of thermostats to approximate samples with respect to the target measure (i.e., the correct marginalized Gibbs density) if the covariance matrix of the noisy force is constant, i.e., $\bm{\Sigma}({{\bm{q}}})=\sigma^{2}{{\bm{I}}}$ ($\sigma$ is a constant positive quantity). The procedure was outlined in the paper of Jones and Leimkuhler and relies on the fact that a fixed amplitude noise perturbation engenders a shift of the auxiliary variable in the extended stationary distribution associated with the Nosé–Hoover thermostat.

If the system is not coming from a Newtonian dynamics model, then it is unclear that we need to rely on second order dynamics for this purpose. To see why this is the case, we explain what goes wrong if we try to use first order dynamics. In what follows, we assume that the covariance matrix of the noisy force is constant, although we ultimately intend to apply the method more generally (see recent work on a novel covariance-controlled adaptive Langevin thermostat that can handle parameter-dependent noise in ). Even in the constant $\sigma$ case it is a nontrivial problem to extract statistics related to a particular target temperature, since we do not assume that $\sigma$ is known.

For $\sigma$ constant, let us first consider the SDE

and seek $\chi(\cdot)$ so that an extended Gibbs distribution with density of the form $\psi({{\bm{q}}},\xi)=\bar{\rho}_{\beta}({{\bm{q}}})\varphi(\xi)$ is (ergodically) preserved. The variable $\xi$ is an auxiliary variable. We do not generally care what its distribution is, but it is crucial that

the overall density is in product form, and

$\varphi(\xi)\geq 0$ is normalizable and of a simple, easily sampled form.

These conditions ensure that we can easily average out over the auxiliary variable to compute the averages of functions of ${{\bm{q}}}$ which are of greatest interest.

Proposition 1. Let $\chi({{\bm{q}}})=-\beta^{-1}\Delta U({{\bm{q}}})+\|\nabla U({{\bm{q}}})\|^{2}$; then (13)–(14) preserves the modified Gibbs distribution

Proof. The Fokker–Planck equation corresponding to (13)–(14) is

Proposition 1 tells us that if we can solve system (13)–(14), under an assumption of ergodicity, we can compute averages with respect to the target Gibbs distribution without actually knowing the value of $\sigma$. $\sigma$ could be observed retrospectively by simply averaging $\xi$ during simulation, since $\langle\xi\rangle=\beta\sigma^{2}/2$.

The problem is that the dynamics (13)–(14) is not quite what we want. A typical numerical method for this system might be constructed based on modification of the Euler–Maruyama method:

however, observe that this method requires separate knowledge of $\nabla U({{\bm{q}}})$ and $\sigma$, which is generally impossible a priori, as we assume that the force is polluted by unknown noise. The form of the equations means that we evaluate the product of $\xi$ and the deterministic force, on the one hand, and the random perturbation, on the other hand, separately, and these contributions are independently scaled by $\Delta t$ and $\sqrt{\Delta t}$, respectively.

To adaptively control the invariant distribution, we consider the following second order formulation, which was first introduced in the paper of Jones and Leimkuhler :

A similar system (SGNHT) was used by Ding et al. , who also explored its application to three examples from machine learning. These experiments demonstrated that Ad-Langevin has superior performance compared to SGHMC in various applications, confirming the importance of adaptively dissipating additional noise in sampling. However, there remain two important issues that we wish to address in this article: (1) the underlying dynamics of the Ad-Langevin method is not clear due to the presence of the stochastically perturbed gradient; (2) little attention has been paid to the design of optimal numerical methods for implementing Ad-Langevin with attention to stability and numerical efficiency.

One may wonder why the artificial noise is needed (i.e., $\sigma_{\rm A}\neq 0$), since we are assuming the presence of noise in the gradient itself. The reason is as follows: in defining a numerical method for the noisy gradient system, the force (including the random perturbation) will in general be multiplied by $\Delta t$, where $\Delta t$ is the timestep. On the other hand, the Itō rule implies that the scaling of random perturbations in an SDE should be by a factor proportional to $\sqrt{\Delta t}$; thus, effectively, if we are to relate the thermostatted method to a standard SDE, the standard deviation of the noise is reduced by multiplication by the factor $\sqrt{\Delta t}$. The noise perturbation introduced at each timestep (and the effective diffusion) is thus reduced for small stepsizes and it is therefore important to inject additional artificial noise in order to stabilize the invariant distribution. A rewriting of the Ad-Langevin system as a standard Itō SDE system makes clear the relation between the different terms

Let us note the main features of the dynamics (19):

The invariant distribution for the given system may be directly obtained by study of its Fokker–Planck equation. Following , it is straightforward to show that (19) has the following invariant distribution:

where $Z$ is the normalizing constant and

where $\sigma_{\rm F}=\sigma\sqrt{\Delta t}$. Observe that this means that if $\sigma_{\rm A}=0$, then, as ${\rm lim}_{\Delta t\rightarrow 0}\sigma_{\rm F}=0$, we find that $\xi$ tends to a variable which is normally distributed with mean zero. Alternatively, if $\sigma_{\rm A}\neq 0$, one would obtain

where $\beta^{-1}\mu^{-1}$ is the variance and the symbol $\overset{\mathscr{L}}{\to}$ indicates that $\xi$ converges in probability law to a normally distributed random variable with the indicated parameters. The order of the limits here is important: $t\rightarrow\infty$ first (to reach the invariant distribution), then $\Delta t\rightarrow 0$.

The ergodicity of (19) with respect to the distribution indicated above can easily be demonstrated by reference to Hörmander’s condition for hypoellipticity following the method in , as for Langevin dynamics. The only additional step is to verify that the noise propagates into the $\xi$ variable, which follows due to its strong coupling to the momenta.

Numerical Methods for Adaptive Thermostats

Since stochastic systems in most of the cases cannot be solved “exactly”, splitting methods are often adopted in practice. For instance here, the vector field of the Ad-Langevin/SGNHT (17) can be split into four pieces which are denoted as “A”, “B”, “O”, and “D”, in such a way that each piece can be solved “exactly”,

Clearly parts “A” and “D” can be solved “exactly”. As mentioned previously, the underlying dynamics for “B” is

The “O” or “Ornstein–Uhlenbeck” part is usually stated with $\xi$ a positive constant, in which case the solution is found to be

where $\mathbf{p}(0)$ is the initial value of the variable and $\mathbf{R}$ is a vector of i.i.d. standard normal random variables. However, the same formula (23) is easily seen to be valid for $\xi<0$, since the quantity $(1-e^{-2\xi\Delta t})/(2\xi)$ is strictly greater than zero unless $\xi=0$. (The proof is obtained by following the standard procedure .) When $\xi=0$, one can simply replace $(1-e^{-2\xi\Delta t})/(2\xi)$ by its well-defined asymptotic limit,

The generators associated with each piece are defined, respectively, as

Overall, the generator of the Ad-Langevin/SGNHT (17) system can be written as

The flow map (or phase space propagator) of the system can be written in the shorthand notation

where the exponential map here denotes the solution operator. Approximations of ${\cal F}_{t}$ can be obtained as products (taken in different arrangements) of exponentials of the splitting terms. For example, the phase space propagation of the method proposed by Ding et al. for the Ad-Langevin/SGNHT (17) system (denoted as “SGNHT-N”) can be written as

and $\exp\left(\Delta t\mathcal{L}_{f}\right)$ represents the phase space propagator associated with the corresponding vector field $f$. Because of its nonsymmetric structure, one anticipates first order convergence to the invariant measure (for any choice of $\sigma$). Due to the naming of the component parts, the SGNHT-N method may be denoted by “PAD”.

Overall, the SGNHT-N/PAD integration method is as follows:

We propose symmetric alternative methods, such as the following symmetric Ad-Langevin/ SGNHT (SGNHT-S) splitting method:

where exact solvers for parts “B” and “O” derived above are applied. The SGNHT-S method may be referred to as “BADODAB”, where it should be noted that the various operations are symmetrically applied and the steplengths are uniform and span the interval $\Delta t$. Other symmetric splittings are considered below.

The SGNHT-S numerical integration method may be written as

The force computed at the end of each timestep can be reused at the start of the next step; thus only one force calculation is needed in SGNHT-S at each timestep, the same as for SGNHT-N. In practice, one could replace the exponential and square root operations in the exact solver of the “O” part by their respective well-defined asymptotic expansions to reduce the computational cost.

The analysis of the accuracy of ergodic averages (averages with respect to the invariant measure) in stochastic numerical methods can be performed using the framework of long-time Talay–Tubaro expansion, as developed in . In what follows we compare the order of convergence of the two Ad-Langevin/SGNHT methods with a clean gradient.

For a splitting method described by $\mathcal{L}=\mathcal{L}_{\alpha}+\mathcal{L}_{\beta}+\dots+\mathcal{L}_{\zeta}$, we define the effective operator $\hat{\mathcal{L}}^{{\dagger}}$ associated with the perturbed system obtained using the numerical method with stepsize $\Delta t$ by the relation

This operator can be computed using the Baker–Campbell–Hausdorff (BCH) expansion and can thus be viewed as a perturbation of the exact Fokker–Planck operator $\mathcal{L}^{{\dagger}}$:

for some perturbation operators $\mathcal{L}^{{\dagger}}_{i}$.

for some correction functions $f_{i}$ satisfying $\langle f_{i}\rangle=0$.

Substituting $\hat{\mathcal{L}}^{{\dagger}}$ and $\hat{\rho}$ into the stationary Fokker–Planck equation

by equating first order terms in $\Delta t$.

According to the BCH expansion, for (noncommutative) linear operators $X$ and $Y$, we have

The notation $[X,Y]=XY-YX$ denotes the commutator of operators $X$ and $Y$.

These equations demonstrate that for nonsymmetric splitting methods, there typically exists a nonzero term $\mathcal{L}^{{\dagger}}_{1}\propto[X,Y]\neq 0$, while the condition $\mathcal{L}^{{\dagger}}_{1}=0$, implying $f_{1}=0$, is automatically satisfied for symmetric splitting methods; thus, for observables $\phi({{\bm{q}}},{{\bm{p}}},\xi)$, assuming the asymptotic expansion holds, the computed average would be of order two

where $\langle\cdot\rangle$ denotes the average with respect to the target invariant distribution. Therefore, the SGNHT-S method (28) would have second order convergence for all the observables.

We can work out the leading operator $\mathcal{L}^{{\dagger}}_{1}$ associated with the nonsymmetric SGNHT-N/PAD method (26) of Ding et al. ,

2 Superconvergence Property

Recently, it has been demonstrated in the setting of Langevin dynamics that a particular symmetric splitting method (“BAOAB”), which requires only one force calculation per step, is fourth order for configurational quantities in the ergodic limit and in the limit of large friction .

Following the standard procedure described in Section 3.1, we obtain the following PDE associated with the BADODAB method:

where $\mathcal{L}^{{\dagger}}$ is the exact Fokker–Planck operator

whose action on the extended invariant measure reads as

The equation is very complicated, and we have no direct means of solving it. However, the additional variable $\xi$ has mean $\hat{\gamma}$. If we suppose that $\mu$ is large, then the variance of $\xi$ will be small. In this case we can consider the approximation obtained by replacing functions of $\xi$ in the PDE (35) by their corresponding averages

We use this as part of an averaging of the stationary Fokker–Planck equation with respect to the auxiliary variable. That is, we project the Fokker–Planck equation and its solution by integrating with respect to the Gaussian distribution of $\xi$ in the ergodic limit. We can think of this is as defining a sort of “subspace projection”; it is related to the Galerkin method that is widely used in solving high-dimensional linear systems and PDEs, including Fokker–Planck equations . In this case, we apply the projection operator

where $\nu$ is an arbitrary function, to the PDE (35). Effectively, this results in the reduced equation

where the operator $\check{\mathcal{L}}^{{\dagger}}$ is just the operator $\mathcal{L}^{{\dagger}}$ reduced by the action of the projection, and which acts on functions of $q$ and $p$; this is nothing other than the corresponding adjoint generator of Langevin dynamics. Likewise, $\hat{f}_{2}$ is now a function of $q$ and $p$ only. The right-hand side simplifies to

where $\rho_{\beta}$ is the Gibbs (canonical) density (${\rm exp}(-\beta H(q,p))$).

We consider the high friction limit ($\hat{\gamma}\rightarrow\infty$) and expand $\hat{f}_{2}$ in a series involving the reciprocal friction $\varepsilon=1/\hat{\gamma}$,

with each function $\hat{f}_{2,i}$ satisfying $\langle\hat{f}_{2,i}\rangle=0$. Dividing (35) by the friction coefficient $\hat{\gamma}$, we obtain

We take the high thermal mass limit ($\mu\rightarrow\infty$) in such a way that $\varepsilon=1/\mu=1/\hat{\gamma}$. The use of this limit yields the following terms of the expansion of the right-hand side in powers of $\varepsilon$. Defining

Furthermore, by equating powers of the reciprocal friction $\varepsilon$, we can solve a sequence of equations

to obtain the leading term $\hat{f}_{2,0}$, i.e.,

which leads to the average of configurational observables $\phi(q)$ with respect to the invariant measure as

Thus, for configurational observables the BADODAB method has fourth order convergence to the invariant measure in the large friction and thermal mass limits (i.e., $\varepsilon\rightarrow 0$),

Numerical Experiments

In this section, we conduct a variety of numerical experiments to compare the performance of the different schemes presented in this article.

Before we compare various methods in machine learning applications (i.e., with a noisy gradient), we first demonstrate the order of convergence of various splitting methods with a clean gradient.

A popular model of an $N$-body system with pair interactions based on a spring with rest length (i.e., pendulum) was used, a standard if simplified model of molecular dynamics. The total potential energy of the system is defined as

where $r_{ij}=\|{{\bm{q}}}_{i}-{{\bm{q}}}_{j}\|$ denotes the distance between two particles $i$ and $j$, and $\varphi(r_{ij})$ represents the pair potential energy

We first compare the two SGNHT methods on controlling two configurational quantities: configurational temperature and average potential energy. The configurational temperature , which, as the kinetic temperature, should in principle be equal to the target temperature, can be defined as

where the angle brackets denote the averages, and $\nabla_{i}U$ and $\nabla^{2}_{i}U$ represent the gradient and Laplacian of the potential energy $U$ with respect to the position of particle $i$, respectively (see more discussions in ).

As shown in Figure 1, with the help of the dashed order lines, we can see that SGNHT-N and SGNHT-S show first and second order convergence, respectively, as expected. It is clear that SGNHT-S has not only at least one order of magnitude improvement in accuracy in both observables, but also much greater robustness over the SGNHT-N method, which becomes completely unstable at around $\Delta t=0.08$. The results on the configurational temperature and average potential energy are rather similar; therefore in what follows we present only average potential energy results.

2 Bayesian Inference

In this subsection we compare methods in a classical Bayesian inference model in one dimension, i.e., to estimate the mean of a normal distribution with known variance . More precisely, given $N$ i.i.d. samples from a normal distribution, $x_{i}\sim\mathcal{N}(\check{\mu},\hat{\sigma}^{2})$, where it should be noted that $\check{\mu}$ is the true mean, when we draw samples with known $\hat{\sigma}^{2}$ and a uniform prior distribution ranging from $-N/2$ to $N/2$, we are able to calculate the posterior distribution of the mean in a closed form

where $\hat{x}=\sum^{N}_{i=1}x_{i}/N$. In the context of stochastic gradient approximation, we have

Note that stepsizes for SGNHT (second order dynamics) and SGLD (first order dynamics) based methods are not directly comparable—as mentioned in the stepsize of a first order dynamics method like Euler–Maruyama when viewed as the limiting discretization of a Langevin integrator corresponds to $\Delta t^{2}/2$, where $\Delta t$ is the stepsize of the Langevin method. However, in our experiments we are uninterested in the time-dynamics of the system and care only about the invariant measure. Therefore the important relationship is the error in thermodynamic averages in comparison with the number of timesteps (work), which quantifies the efficiency of a given method. The stepsize is just an arbitrary parameter which allows for refinement of the statistical calculation.

Between the two SGNHT methods, SGNHT-S (the new scheme being proposed here) is obviously superior to SGNHT-N: the latter starts to show significant deviation from the true distribution at $\Delta t=0.02$, while the distribution of the former still looks well matched to the true one at $\Delta t=0.03$. Our observations are confirmed by Figure 5, where the mean absolute error (MAE) of the distribution of the two SGNHT methods is plotted. The MAE, which can be thought of as a relative error in distribution, is defined as

where $\bar{N}$ denotes the number of intervals, which was chosen as 100. $\omega_{i}$ and $\hat{\omega}_{i}$ represent the observed frequency in bin $i$ and the exact expected frequency, respectively . As can be seen, the stability threshold of SGNHT-N was around $\Delta t=0.03$, beyond which the system became unstable, as highlighted in the figure (in which case the system blew up, resulting in a 100% MAE). Once again, SGNHT-S not only shows an order of magnitude better accuracy but also has a much greater robustness than SGNHT-N. In particular, for defined accuracy, the SGNHT-S method is able to use double the stepsize compared to SGNHT-N, which means a remarkable 50% improvement in overall numerical efficiency as defined in .

3 Bayesian Logistic Regression

Following , we also investigate the performance of different methods for a more complicated Bayesian logistic regression model. The data $y_{i}\in\{-1,1\}$ were modelled by

Following the same procedure in the Bayesian inference example (Section 4.2), we can calculate the noisy force and then plug it into different thermostats for sampling.

In our numerical experiments, we considered the $d=3$ case with $N=1000$ data points. We chose the dataset to be

Of the two SGNHT methods, the SGNHT-S method again shows not only at least an order of magnitude improvement on accuracy but also much better robustness than the other: SGNHT-N became unstable just above $\Delta t=0.02$. Remarkably, the SGNHT-S method at $\Delta t=0.1$ still achieves better accuracy than the SGLD method at $\Delta t=0.01$. In other words, the method we propose here gives more than a 90% improvement in overall numerical efficiency compared to one of the most popular methods in the literature. For fixed accuracy, the SGNHT-S method can use almost four times the stepsize of the SGNHT-N method (i.e., an improvement of about 75% in overall numerical efficiency).

Conclusions

We have reviewed a variety of methods in stochastic gradient systems with applications in machine learning. We have provided a theoretical discussion on the foundation (underlying dynamics) of those stochastic gradient systems, which has been lacking in the literature. We have also proposed a new symmetric splitting (at least second order) method in SGNHT (SGNHT-S/BADODAB), which substantially improves the accuracy and robustness compared to a nonsymmetric splitting (first order) method (SGNHT-N) proposed recently in the literature. Furthermore, we have demonstrated that under certain conditions the SGNHT-S/BADODAB method can inherit the superconvergence property recently discovered in integrators for Langevin dynamics, i.e., fourth order convergence to the invariant measure for configurational averages.

By conducting various numerical experiments, we have demonstrated that the two SGNHT methods outperform the popular SGLD method and its variant mSGLD. In particular, the SGNHT-S method can use up to ten times the stepsize of SGLD, which implies a remarkable more than 90% improvement in overall numerical efficiency. Between the two SGNHT methods, the SGNHT-S method can use almost four times the stepsize of SGNHT-N for defined accuracy (i.e., about a 75% improvement in overall numerical efficiency).

It should be noted that in certain cases, it may be desirable to employ a Metropolis–Hastings procedure in order to remove the discretization bias . However, we emphasize that the correction is not without computational cost, particularly as the dimension is increased , and the results of and of the current article demonstrate that high accuracy with respect to the invariant distribution is often achievable using traditional numerical integration techniques, thus in many cases entirely eliminating the necessity of Metropolis–Hastings corrections (see more discussions in ). Moreover, we mention that the methods of this article can in principle be combined with Metropolis–Hastings algorithms if it is necessary to completely eliminate the discretization bias.

Acknowledgements

The authors thank Ben Goddard, Charles Matthews, Tony Shardlow, Zhanxing Zhu, and Konstantinos Zygalakis for stimulating discussions and valuable suggestions. The authors further thank the anonymous referees for their comments, which substantially contributed to the presentation of our results. XS gratefully acknowledges the financial support from the University of Edinburgh and China Scholarship Council.

References