Exact Langevin Dynamics with Stochastic Gradients
Abstract
Stochastic gradient Markov Chain Monte Carlo algorithms are popular samplers for approximate inference, but they are generally biased.
We show that many recent versions of these methods (e.g. Chen et al.
(2014)) cannot be corrected using Metropolis-Hastings rejection sampling, because their acceptance probability is always zero.
We can ﬁx this by employing a sampler with realizable backwards trajectories, such as Gradient-Guided
Monte Carlo (Horowitz, 1991), which generalizes stochastic gradient Langevin dynamics (Welling and Teh, 2011) and Hamiltonian Monte Carlo.
We show that this sampler can be used with stochastic gradients, yielding nonzero acceptance probabilities which can be computed even across multiple steps
1. Introduction
Exact posterior inference in many Bayesian models is intractable, which has led to the development of diﬀerent approximate inference schemes.
One popular class of such algorithms are Markov Chain Monte Carlo (MCMC) methods.
These simulate a Markov chain, constructed so that its stationary distribution is the true posterior (Geyer, 1992).
Among the
MCMC methods, Hamiltonian Monte Carlo (HMC) has achieved recognition as being the “gold standard”, mostly for its favorable convergence properties (Neal, 2012).
While HMC yields great performance, its adoption in many real applications is hindered by the fact that it requires passes over the whole data set to compute one gradient step, which is prohibitively expensive in many settings (Chen et al., 2014).
Especially for largedata Bayesian models such as Bayesian neural networks (BNNs), stochastic gradient MCMC (SG-MCMC) methods are more popular, due to their favorable computational costs and their similarity to stochastic gradient descent (SGD).
Arguably, the most popular SG-MCMC method for BNNs in recent years has been stochastic gradient HMC (Chen et al., 2014,
SGHMC), although its predecessor stochastic gradient Langevin dynamics (Welling and
Teh, 2011, SGLD) also enjoys considerable following.
Since all of these approaches are approximate, in order to ensure that they sample from the true posterior, rejection sampling is often necessary, for instance, following the
Metropolis-Hastings (MH) scheme (Metropolis et al., 1953; Hastings, 1970).
However, in this work, we draw attention to the unfortunate fact that SGHMC does not allow for MH rejection sampling, since its acceptance probability is always zero.
We then revisit gradientguided Monte Carlo (GGMC) (Horowitz, 1991), which can overcome this problem.
We show that GGMC provides a uniﬁed framework for SGLD and SGHMC and that it yields nonzero acceptance probabilities, which can even be computed across multiple steps.
Our contributions are as follows:
• We show that the SG-MCMC schemes used in recent work (Chen et al., 2014; Wenzel et al., 2020) always have acceptance probability zero, because the backward trajectory is not realizable in the symplectic Euler-Maruyama integrator they employ (Sec. 3).
• We revisit GGMC, a variant of HMC.
It generalises HMC and SGLD, always has positive acceptance probability and, with correction, accepts stochastic gradients (Sec. 4).
The scheme was independently discovered in the statistics (Neal, 2012; Horowitz, 1991) and molecular dynamics literatures (Bussi and Parrinello, 2007; Leimkuhler and Matthews, 2015).
The latter have shown that it simulates Langevin dynamics.
• We examine an augmentation scheme to calculate the acceptance probabilities across multiple stochastic steps (Sec. 5, by Zhang et al.
This allows us to use stochastic gradients, but sample from the exact posterior, by doing one MH correction using the exact likelihood after many stochastic gradient steps.
Based on our ﬁndings, we believe that work in Bayesian deep learning should discard the Euler-Maruyama scheme used by SGHMC and its variants (Ding et al., 2014; Ma et al., 2015).
Instead, to have conﬁdence in our sampling, we should use reversible integrators (Matthews and Weare, 2018; Leimkuhler et al., 2019; Zhang et al., 2020, this work) and keep track of their acceptance probability.
2. Background: Hamiltonian Monte Carlo
We wish to sample from a target distribution π ( θ )
∝ ˜
π
( θ ) , which we can only evaluate up to a normalizing constant.
HMC ﬁrst augments the parameter space θ ∈ R d with a momentum m ∈ R d of the same dimension.
It constructs a Markov chain that alternates: re-sampling the momentum, and then following an ordinary diﬀerential equation (ODE), speciﬁed by the Hamiltonian H ( θ , m ) , for a given amount of time.
The stationary distribution of this
Markov chain (with exact simulation) is the Boltzmann distribution
[Equation]
Here U ( θ ) is the potential energy, which we relate to the unnormalised target distribution by U ( θ )
= − log ˜ π
( θ ) .
The kinetic energy K ( m ) =
12 m T M − 1 m makes the momenta m Gaussian, and thus easy to sample.
M is a positive deﬁnite (usually diagonal) preconditioner, the mass matrix.
T ∈ R
+ is a positive temperature.
For HMC, T = 1 , but this changes in the next section.
By construction, the tempered target ˜ π ( θ ) 1 /T is a marginal of the Boltzmann distribution.
We may write the ODE’s full state as s =
( θ , m ) .
Unfortunately, we cannot exactly simulate the ODE for most posteriors, so we must discretise it with an integrator.
This unavoidably introduces errors into the simulation, which have to be controlled if we want to guarantee that the sampling is accurate.
The standard way to do this is via the Metropolis-Hastings acceptance probability.
For this, we deﬁne the transition kernel: g ( s ∗ | s n )
= 1 if the approximate simulation starting at s n = ( θ n , m n ) ends at s ∗ =
( θ ∗ , m ∗ ) , otherwise 0.
Then, we accept the new sample (i.e. set s n +
= s ∗ ) with probability min (1 , MH ( s ∗ | s n )) , where
[Equation]
Otherwise we reject the sample, that is, set s n +
= s n .
We negate the momentum so the integrator applied backwards would retrace the same steps.
Afterwards, we are free to negate the momentum again without a rejection step, since K ( − m ) = K ( m ) , so it leaves the target distribution (1) invariant (Neal, 2012, eq. 5.20).
For regular HMC, this can be ignored because the momentum gets resampled immediately.
3. SG-MCMC acceptance probabilities are zero
All the Markov chains we consider in this paper are discretisations of Langevin’s stochastic diﬀerential equations (SDEs) which, borrowing notation from Wenzel et al.
(2020), are
[Equation]
[Equation]
[Equation]
Here, γ ∈ R + is a friction coeﬃcient and W is a standard Wiener process with covariance
I .
If γ = 0 , these become Hamilton’s equations, used for HMC.
Setting d m = 0 , we obtain the overdamped Langevin equation, which discretises to the Metropolis-adjusted Langevin algorithm (MALA; e.g. Roberts and Rosenthal (1998)).
For time step h , the symplectic Euler-Maruyama scheme in Chen et al.
(2014, SGHMC) and Wenzel et al.
(2020) simulates these SDEs as follows.
Using noise (cid:15) ∼ N ( 0 , I ) ,
[Equation]
Like in HMC, it is the discretisation, not the SDE, that determines our simulation’s
Markov transition kernel g .
Unfortunately, we will see that the backward transition is unrealizable, so its density g ( θ n , − m n | θ n +1 , − m n +1 ) is almost always zero.
Because this density appears in the numerator of eq.
(2), the resulting acceptance probability for the Euler-Maruyama scheme is zero as well.
Theorem 1 For any starting point ( θ n , m n ) and step size h > 0 , we sample ( θ n +1 , m n +1 ) using equations (5, 6).
Then, the backward transition density g ( θ n , − m n | θ n +1 , − m n +1 ) is zero, with probability 1 over the randomness of the forward transition.
Proof Consider the transition distribution, g ( θ ∗ , m ∗ | θ , m ) .
Because it has been generated from equations (5, 6), any realization from g must satisfy them.
In particular, it has to satisfy eq.
(6), i.e. θ ∗ =
θ + h M − 1 m ∗ .
Thus, for the purposes of the acceptance probability (eq. 2), we can deﬁne the density of g as zero outside of its support,
[Equation]
otherwise
[Equation]
Now, let us check whether this condition can simultaneously hold for the forward and backward transitions in eq.
That is, is it possible that
[Equation]
Solving the linear system we can see that, since M is positive deﬁnite and h > 0 , eq.
(8) is true if and only if m n = m n +1 .
Since the step size is h > 0 , this only happens when the Gaussian random draw (cid:15) in eq.
(5) is equal to v =
( h M 2 γT ) − 1 / 2 ( hγ m n + h ∇
θ
U ( θ n )) .
Because the covariance of (cid:15) is the identity, this is a probability zero event.
Thus, with probability 1, (cid:15) (cid:54) = v , which implies m n (cid:54) = m n +1 .
We know the left-hand-side expression of eq.
(8) is true, so that implies θ n (cid:54)
=
θ n +1 + h M − 1 ( − m n ) , which implies by eq.
(7) that g ( θ n , − m n | θ n +1 , − m n +1 ) = 0 .
The easiest way to ensure m n = m n +1 , such that the acceptance probability remains nonzero, is to set the time step h = 0 .
This, however, makes the sampler useless.
If we take the limit h → 0 instead, the Euler-Maruyama scheme gets arbitrarily close to the true SDE trajectory.
Crucially however, per theorem 1, the acceptance probability remains 0 for any h > 0 .
Thus, it is impossible to use the acceptance probability to monitor the discretisation error.
The Euler-Maruyama scheme cannot satisfy detailed balance.
So far, we have not considered stochastic gradients, but this result includes them: it still holds if we substitute an arbitrary g n for ∇ θ
U ( θ n ) .
4. GGMC solves the problem by using a diﬀerent integrator
To get a nonzero acceptance probability, we can just use a time-reversible integrator.
The one we present here has been discovered several times in statistics (Horowitz, 1991; Neal, 1993) and molecular dynamics (Bussi and Parrinello, 2007; Leimkuhler and Matthews, 2015).
We call it Gradient-guided Monte Carlo (GGMC), following Horowitz (1991), because it operates similarly to random walk Metropolis-Hastings, but with a gradient inﬂuencing the random walk’s direction.
Neal (1993, sec. 5.1) calls it the “stochastic dynamics” method, or “Langevin
Monte Carlo (LMC) with partial momentum refreshment” (Neal, 2012).
Leimkuhler and
Matthews (2015) call this integrator “OBABO” because of the order of its steps.
The integrator splits Langevin dynamics (eqs. 4–3) into three parts, which can be solved exactly, and applies them in sequence (Leimkuhler and Matthews, 2015).
Let a = e − γh , and
[Equation] (cid:15) [Equation] ∼ N ( 0 , I ) be Gaussian random variables with identity covariance.
Then
[Equation]
Steps O.1 and O.2 are partial momentum refreshment, and steps B.1, A and B.2 form one leapfrog step.
Thus, the discretisation error is locally O ( h 3 ) (Neal, 2012).
Establishing the trajectory discretisation error is complicated, because it requires solving a partial diﬀerential equation (PDE).
However, all found PDE solutions indicate that the error is O ( h 2 ) (Leimkuhler and Matthews, 2015;
Sivak et al., 2014), the same as HMC.
This scheme generalizes several popular samplers.
To obtain HMC, we may repeat BAB many times, and use O as momentum refreshment.
Alternatively, setting γ = 0 implies a = 1 , so the O steps disappear, and the integrator becomes one leapfrog step.
Adding a separate momentum refreshment then also yields HMC.
To obtain MALA 3 or SGLD (Welling and Teh, 2011), we take γ → ∞ .
This implies a → 0 , which means that step O. entirely ignores the incoming momentum m n and just samples a fresh one.
Because of this, we do not need to calculate the next momentum m n +1 at the end of a step, or represent the momentum at all.
Thus, steps B.2 and O.2 fall out.
We may also fold steps O.1 and
B.1 into A to obtain the more common single-step expression for SGLD.
Note that there are other modern methods that simulate Langevin dynamics while preserving reversibility.
For instance, Matthews and Weare (2018) propose a variant of
ABOBA that corrects for the (known or estimated) covariance of the gradient noise, similar to SGHMC, to get the correct O ( h 2 ) statistics.
Also, Zhang et al.
(2020) develop an integrator that uses time-asymmetric gradients but, nevertheless, is realizable in the reverse direction and thus has a nonzero acceptance probability.
Acceptance probability.
The key ingredient is the ratio of backward and forward transition probabilities, which we derive in appendix A.
The resulting acceptance probability is
[Equation]
This is the same expression as in HMC or Langevin MC, were we to regard only the inner leapfrog step as the transition to accept or reject.
Interestingly, the acceptance probability does not at all depend on the permanence a of the momentum.
5. Acceptance probabilities for multiple and stochastic steps
The inventors of HMC, Duane et al.
(1987), already observed that the kinetic and potential energies used to specify the transition probability g need not be the same ones used to calculate the acceptance probability.
Only the latter will determine what distribution we actually sample from.
That is, the sampler (be it HMC, or GGMC) may use stochastic or incorrect gradients: so long as the MH acceptance probability is calculated with respect to the target ˜
π ( θ ) 1 /T
e − K ( m ) /T , the resulting algorithm will estimate averages with respect to the target distribution.
This would seem to be irrelevant in practice because, naïvely, we need to calculate the exact acceptance probability after every (stochastic) GGMC step.
This would bring the overall algorithm’s per-step complexity to linear in the size of the data set, which is precisely what we want to avoid when using stochastic gradients.
Instead, perhaps we can calculate the MH correction only after many GGMC steps.
Naïvely, this is impractical, since even the two-step transition probability is an intractable integral:
g 2 ( s n +2 | s n ) = (cid:82) g ( s n +2 | s n +1 )
g ( s n +1 | s n )d s n +1 .
However, it turns out we can just multiply the per-step probabilities.
Starting at s n , after N steps of MCMC, we can accept the resulting sample s n + N with probability min (1 , MH ( s n + N | s n ))
(Zhang et al., 2020), where
[Equation]
Due to telescopic cancellation, we need only evaluate the exact likelihood at the end of many steps.
For example, we could follow 10 epochs of stochastic gradients, and then spend 1 epoch evaluating the likelihood exactly for the MH correction step.
This can be combined with other algorithms (Neal, 2012; Sohl-Dickstein et al., 2014, sec. 5.4) to reduce the rejection probability.
Zhang et al.
(2020, Lemma 2) propose this deferred MH scheme and show its correctness for any sampler with a homogeneous Markov chain.
We can extend it to inhomogeneous
Markov chains, such as the cosine schedule (Zhang et al., 2019) by considering the following scheme.
We augment the state of the Markov chain with N random variables, { ζ i } Ni =1 , that will determine the randomness of the intermediate steps.
Their target distributions have to be time-symmetric, i.e. p ζ i ( · ) =
p ζ N − i +1 ( · ) for all i .
Before an MCMC multi-step, we sample all the ζ i from their target distribution.
Then, we use them to provide all the randomness for the intermediate steps, so each step is just the result of a deterministic transformation:
[Equation]
The functions f i have to be time-symmetric as well ( f i =
f N − i +1 ).
Otherwise, the backward transition is not realizable using the same ζ i as the forward one, so the acceptance probability is zero.
If they are time-symmetric, the resulting acceptance probability is eq.
In this view, the cosine schedule (Zhang et al., 2019) has acceptance probability zero, even if a backward-realizable integrator is used.
The cosine schedule is not a time-symmetric transformation: the schedule of learning rates is high at the beginning and low at the end of a cycle.
Wu et al.
(2020) employ a similar construction to use MCMC samplers as normalizing ﬂows.
Neal (1996, Section 3.5.1) proposed the “partial gradient” construction: a variant of this for HMC and dataset minibatching in particular, motivated by the same operatorsplitting view of integrators as the leapfrog.
6. Conclusion
We have shown that two popular SG-MCMC samplers cannot be corrected using rejection sampling, since their acceptance probabilities are zero.
We have then revisited GGMC, which solves this problem and generalizes SGLD and HMC.
Finally, we have also presented a way to compute the acceptance probabilities of GGMC across several steps, leading to a practical method for drawing exact posterior samples using only stochastic gradients.
It is time for the Bayesian deep learning community to let go of the Euler-Maruyama scheme.
Instead, we should explore the breadth of integrators for Langevin dynamics which we have inherited from molecular dynamics and statistics.
We have shown that correctness is achievable.
Thus, we should demand correctness from our samplers.
Appendix A. Acceptance probability calculation for GGMC
The procedure used by Bussi and Parrinello (2007) is to start with the joint distribution of the noises (cid:15) , (cid:15) (cid:48) and transform it to the distribution over θ n +1 , m n +1 | θ n , m n using the density transformation formula with the determinant of the Jacobian.
However, in this case, there is somewhat more illustrative way.
Consider
[Equation]
, and thus
[Equation]
, (12) and similarly for r (cid:48) and (cid:15) (cid:48) .
Note that r
= m n +1 / 4 − √ a m n , and r (cid:48)
=
m n +1 − √ a m n +3 / 4 .
By substituting r , r (cid:48) into eqs.
(B.1) and (B.2), we can write θ n +1 , m n +1 as
[Equation]
[Equation]
The reciprocal of the determinant of the Jacobian is thus
[Equation]
which is a constant, so it will be eliminated when dividing by the backward transition density
[Equation]
Putting this all together
[Equation]
the forward transition density is
[Equation]
By substituting θ n +1 ↔ θ n , m n → − m n +1 , m n +1 / 4 → − m n +3 / 4 , m n +3 / 4 → − m n +1 / 4 , m n +1 → − m n , we can obtain the backward transition probability
[Equation]
Dividing them, the constant terms from the Gaussians and h − 1 | M | disappear
[Equation]
and we obtain
[Equation]
where K ( v ) = 12 v T M − 1 v is the kinetic energy deﬁned in section 3.
This simpliﬁes greatly, to
[Equation]
The log ratio of the Boltzmann densities (eq. (1)) is
[Equation]
Now, to obtain the acceptance probability in eq.
(2) we add eqs.
(19) and (20), obtaining
[Equation]
which is exactly the same as eq.