We present an improved analysis of the Euler-Maruyama discretization of the Langevin diffusion. Our analysis does not require global contractivity, and yields polynomial dependence on the time horizon. Compared to existing approaches, we make an additional smoothness assumption, and improve the existing rate from to in terms of the KL divergence. This result matches the correct order for numerical SDEs, without suffering from exponential time dependence. When applied to algorithms for sampling and learning, this result simultaneously improves all those methods based on Dalayan’s approach.
Improved Bounds for Discretization of Langevin Diffusions: Near-Optimal Rates without Convexity
|Wenlong Mou||Nicolas Flammarion||Martin J. Wainwright||Peter L. Bartlett|
|Department of Electrical Engineering and Computer Sciences|
|Department of Statistics|
|The Voleon Group|
July 29, 2019
In recent years, the machine learning and statistics communities have witnessed a surge of interest in the Langevin diffusion process, and its connections to stochastic algorithms for sampling and optimization. The Langevin diffusion in is defined via the Itô stochastic differential equation (SDE)
where is a standard -dimensional Brownian motion, and the function is known as the drift term. For a drift term of the form for some differentiable function , the Langevin process (1) has stationary distribution with density ; moreover, under mild growth conditions on , the diffusion converges to this stationary distribution as . See Pavliotis (2014) for more background on these facts, which underlie the development of sampling algorithms based on discretizations of the Langevin diffusion. Diffusive processes of this nature also play an important role in understanding stochastic optimization; in this context, the Gaussian noise helps escaping shallow local minima and saddle points in finite time, making it especially useful for non-convex optimization. From a theoretical point of view, the continuous-time process is attractive to analyze, amenable to a range of tools coming from stochastic calculus and Brownian motion theory (Revuz and Yor, 1999). However, in practice, an algorithm can only run in discrete time, so that the understanding of discretized versions of the Langevin diffusion is very important.
The discretization of SDEs is a central topic in the field of scientific computation, with a wide variety of schemes proposed and studied (Kloeden, 1992; Higham, 2001). The most commonly used discretization is the Euler-Maruyama discretization: parameterized by a step size , it is defined by the recursion
Here the sequence is formed of -dimensional standard Gaussian random vectors.
From past work, the Euler-Murayama scheme is known to have first-order accuracy under appropriate smoothness conditions. In particular, the Wasserstein distances for between the original Langevin diffusion and the discretized version decays as as decays to zero, with the dimension and time horizon captured in the order notation (see, e.g., Alfonsi et al., 2015). When the underlying dependence on the time horizon is explicitly calculated, it can grow exponentially, due to the underlying Grönwall inequality. If the potential is both suitably smooth and strongly convex, then the scaling with remains first-order, and the bound becomes independent of time (Durmus and Moulines, 2017; Dalalyan and Karagulyan, 2017). These bounds, in conjunction with the coupling method, have been used to bound the mixing time of the unadjusted Langevin algorithm (ULA) for sampling from strongly-log-concave densities. Moreover, this bound aligns well with the classical theory of discretization for ordinary differential equations (ODEs), where finite-time discretization error may suffer from bad dependence on , and either contraction assumptions or symplectic structures are needed in order to control long-time behavior (Iserles, 2009).
On the surface, it might seem that SDEs pose greater numerical challenges than ODEs; however, the presence of randomness actually has been shown to help in the long-term behavior of discretization. Dalalyan (2017) showed that the pathwise Kullback-Leibler (KL) divergence between the original Langevin diffusion (1) and the Euler-Maruyma discretization (2) is bounded as with only smoothness conditions. This result enables comparison of the discretization with the original diffusion over long time intervals, even without contraction. The discretization techniques of Dalalyan (2017) serve as a foundation for a number of recent papers on sampling and non-convex learning, including the papers (Raginsky et al., 2017; Tzen et al., 2018; Liang and Su, 2017).
On the other hand, this bound on the KL error is likely to be loose in general. Under suitable smoothness conditions, standard transportation inequalities (Bolley and Villani, 2005) guarantee that such a KL bound can be translated into a -bound in Wasserstein distance. Yet, as mentioned in the previous paragraph, the Wasserstein rate should be under enough smoothness assumption. This latter result either requires assuming contraction or leads to exponential time dependence, leading naturally to the question: can we achieve best of both worlds? That is, is it possible to prove a Wasserstein bound without convexity or other contractivity conditions?
In this paper, we answer the preceding question in the affirmative: more precisely, we close the gap between the correct rate for the Euler-Maruyama method and the linear dependence on time horizon, without any contractivity assumptions. As long as the drift term satisfies certain first and second-order smoothness, as well as growth conditions at far distance, we show the KL divergence between marginal distributions of equation (1) and equation (2), at any time , is bounded as . Note that this bound is non-asymptotic, with polynomial dependence on all the smoothness parameters, and linear dependence on . As a corollary of this improved discretization bound, we give improved bounds for using the unadjusted Langevin algorithm (ULA) for sampling from a distribution satisfying a log-Sobolev inequality. In addition, our improved discretization bound improves a number of previous results on non-convex optimization and inference, all of which are based on the discretized Langevin diffusion.
In the proof of our main theorem, we introduce a number of new techniques. A central challenge is how to study the evolution of time marginals of the interpolation of discrete-time Euler algorithm, and in order to do so, we derive a Fokker-Planck equation for the interpolated process, where the drift term is the backward conditional expectation of at the previous step, conditioned on the current value of . The difference between this new drift term for the interpolated process and itself can be much smaller than the difference between at two time points. Indeed, taking the conditional expectation cancels out the bulk of the noise terms, assuming the density from the previous step is smooth enough. We capture the smoothness of density at the previous step by its Fisher information, and develop De Bruijn-type inequalities to control the Fisher information along the path. Combining this regularity estimate with suitable tail bounds leads to our main result. We suspect that our analysis of this interpolated process and associated techniques for regularity estimates may be of independent interest.
|Paper||Require contraction||111We only listed time horizon dependence for methods that guarantee discretization error between continuous-time and discrete-time for any time. If the proof requires mixing and does not give the difference between the one-time distributions, we mark it as “-”.Time||222The distances are measured in . If the original bound is shown for KL, it is transformed into using transportation inequalities, resulting in the same rate. We mark with * if the original bound was shown in KL Step size||Require mixing||Additional assumptions|
|Alfonsi et al. (2015)||No||No||Second-order smooth drift|
|Dalalyan and Karagulyan (2017), Durmus and Moulines (2017)||Yes||-||Yes||Second-order smooth drift|
|Cheng and Bartlett (2018), Ma et al. (2018)||No||-||Yes||strong convexity outside a ball|
|Cheng et al. (2018a), Bou-Rabee et al. (2018) 333For Hamiltonian Monte-Carlo, which is based on discretization of ODE, instead of SDE.||No||-||Yes||strong convexity outside a ball|
|This paper||No||No||Second-order smooth drift|
Recent years have witnessed a flurry of activity in statistics and machine learning on the Langevin diffusion and related stochastic processes. A standard application is sampling from a density of the form based on an oracle that returns the pair for any query point . In the log-concave case, algorithms for sampling under this model are relatively well-understood, with various methods for discretization and variants of Langevin diffusion proposed in order to refine the dependence on dimension, accuracy level and condition number (Dalalyan, 2017; Durmus and Moulines, 2017; Cheng et al., 2018b; Lee et al., 2018b; Mangoubi and Vishnoi, 2018; Dwivedi et al., 2018).
When the potential function is non-convex, the analysis of continuous-time convergence and the discretization error analysis both become much more involved. When the potential satisfies a logarithmic Sobolev inequalities, continuous-time convergence rates can be established (see e.g. Markowich and Villani, 2000), and these guarantees have been leveraged for sampling algorithms (Bernton, 2018; Wibisono, 2018; Ma et al., 2018). Coupling-based results for the Wasserstein distance have also been shown for variants of Langevin diffusion (Cheng et al., 2018a; Bou-Rabee et al., 2018). Beyond sampling, the global convergence nature of Langevin diffusion has been used in non-convex optimization, since the stationary distribution is concentrated around global minima. Langevin-based optimization algorithms have been studied under log-Sobolev inequalities (Raginsky et al., 2017), bounds on the Stein factor (Erdogdu et al., 2018); in addition, accelerated methods have been studied (Chen et al., 2019). The dynamics of Langevin algorithms have also been studied without convergence to stationarity, including exiting times (Tzen et al., 2018), hitting times (Zhang et al., 2017), exploration of a basin-of-attraction (Lee et al., 2018a), and statistical inference using the path (Liang and Su, 2017). Most of the works in non-convex setting are based on the discretization methods introduced by Dalalyan (2017).
Finally, in a concurrent and independent line of work, Fang and Giles (2019) also studied a multi-level sampling algorithm without imposing a contraction condition, and obtained bounds for the mean-squared error; however, their results do not give explicit dependence on problem parameters. Since the proofs involve bounding the moments of Radon-Nikodym derivative, their results may be exponential in dimension, as opposed to the polynomial-dependence given here.
We let denote the Euclidean norm of a vector . For a matrix we let denote its spectral norm. For a function , we let denote its Jacobian evaluated at . We use to denote the law of random variable . We define the constant . When the variable of the integrand is not explicitly written, integrals are taking with respect to the Lebesgue measure: in particular, for an integrable function , we use as a shorthand for .
2 Main Results
We now turn to our main results, beginning with our assumptions and a statement of our main theorem. We then develop and discuss a number of corollaries of these main results.
2.1 Statement of main results
Our main results involve three conditions on the drift term , and one on the initialization:
Assumption 1 (Lipschitz drift term).
There is a finite constant such that
Assumption 2 (Smooth drift term).
There is a finite constant such that
Assumption 3 (Distant dissipativity).
There exist strictly positive constants such that
Assumption 4 (Smooth Initialization).
Note that no contractivity assumption on the drift term
is imposed. Rather, we use the notion of distant dissipativity, which
is substantially weaker; even this assumption is relaxed in
Theorem 2. The initialization
condition (6) is clearly satisfied by the
standard Gaussian density, but
Assumption 4 allows for other densities
with quadratic tail behavior.
With these definitions, the main result of this paper is the following:
If we track only the dependence on , the result (7) can be summarized as a bound of the form . This result should be compared to the bound obtained by Dalalyan (2017) using only Assumption 2. It is also worth noticing that the term only comes with the third order derivative bound, which coincides with the Wasserstein distance result, based on a coupling proof, as obtained by Durmus and Moulines (2017) and Dalalyan and Karagulyan (2017). However, these works do not study separately the discretization error of the discrete process and assume contractivity.
Note that Assumption 3 can be substantially relaxed when the drift is negative gradient of a function. Essentially, we only require this function to be non-negative, along with the smoothness assumptions. In such case, we have the following discretization error bound:
Once again tracking only the dependence on , the bound (8) can be summarized as . This bound has weaker dependency on , but it holds for any non-negative potential function without any growth conditions.
When the problem of sampling from a target distribution is considered, the above bounds applied to the drift term yield bounds in TV distance, more precisely via the convergence of the Fokker-Planck equation and the Pinsker inequality (Dalalyan, 2017). Instead, in this paper, so as to obtain a sharper result, we directly combine the result of Theorem 1 with the analysis of Cheng and Bartlett (2018). A notable feature of this strategy is that it completely decouples analyses of the discretization error and of the convergence of the continuous-time diffusion process. The convergence of the continuous-time process is guaranteed when the target distribution satisfies a log-Sobolev inequality (Toscani, 1999; Markowich and Villani, 2000).
Given an error tolerance and a distance function , we define the associated mixing time of the discretized process
With this definition, we have the following:
Consider a density of the form such that:
The distribution defined by satisfies a log-Sobolev inequality with constant .
The set of distributions satisfying a log-Sobolev inequality (Gross, 1975) includes strongly log-concave distributions (Bakry and Émery, 1985) as well as perturbations thereof (Holley and Stroock, 1987). For example, it includes distributions that are strongly log-concave outside of a bounded region, but non-log-concave inside of it, as analyzed in some recent work (Ma et al., 2018). Under the additional smoothness Assumption 2, we obtain an improved mixing-time compared to of Ma et al. (2018). On the other hand, we obtain the same mixing time in as the papers (Durmus and Moulines, 2017; Dalalyan and Karagulyan, 2017) but under weaker assumptions on the target distribution—namely, those that satisfy a log-Sobolev inequality as opposed to being strongly log-concave.
2.2 Overview of proof
First, we construct a continuous-time interpolation of the discrete-time process , and prove that its density satisfies an analogue of the Fokker-Plank equation (see Lemma 1). The elliptic operator of this equation is time-dependent, with a drift term given by the backward conditional expectation of the original drift term . By direct calculation, the time derivative of the KL divergence between the interpolated and the original Langevin diffusion is controlled by the mean squared difference between the drift terms of the Fokker-Planck equations for the original and the interpolated processes, namely the quantity
See Lemma 2 for details.
Our next step is to control the mean-squared error term (10). Compared to the MSE bound obtained from the Girsanov theorem by Dalalyan (2017), our bound has an additional backward conditional expectation inside the norm. Directly pulling this latter outside the norm by convexity inevitably entails a KL bound due to fluctuations of the Brownian motion. However, taking the backward expectation cancels out most of the noises, as long as the distribution of the initial iterate at each step is smooth enough. This geometric intuition is explained precisely in Section 4.1, and concretely implemented in Section 4.2. The following proposition summarizes the main conclusion from Steps 1 and 2:
The third step is to bound the moments of and , so as to control the right-hand side of equation (11). In order to bound the Fisher information term , we prove an extended version of the De Brujin formula for the Fokker-Planck equation of (see Lemma 6). It bounds the time integral of by moments of . Since Proposition 1 requires control of the Fisher information at the grid points , we bound the integral at time by the one at time ; see Lemma 7 for the precise statement. Combining these results, we obtain the following bound of the averaged Fisher information.
It remains to bound the moments of along the path. By Proposition 1 and Proposition 2, the second and fourth moment of are used. With different assumptions on the drift term, different moments bounds can be established, leading to Theorem 1 and Theorem 2, respectively.
Under distant dissipativity (Assumption 3), the -th moment of this process can be bounded from above, for arbitrary value of . (see Lemma 11). The proof is based on the Burkholder-Davis-Gundy inequality for continuous martingales. Collecting these results yields Eq (7), which completes our sketch of the proof of Theorem 1.
3 Interpolation, KL Bounds and Fokker-Planck Equation
Following Dalalyan (2017), the first step of the proof is to construct a continuous-time interpolation for the discrete-time algorithm (2). In particular, we define a stochastic process over the interval via
Let be the natural filtration associated with the Brownian motion . Conditionally on , the process is a Brownian motion with constant drift and starting at . This interpolation has been used in past work (Dalalyan, 2017; Cheng and Bartlett, 2018). In their work, the KL divergence between the law of processes and is controlled, via a use of the Girsanov theorem, by bounding Radon-Nikodym derivatives. This approach requires controlling the quantity for . It is should be noted that it scales as , due to the scale of oscillation of Brownian motions.
In our approach, we overcome this difficulty by only considering the KL divergence of the one-time marginal laws . Let us denote the densities of and with respect to Lebesgue measure in by and , respectively. It is well-known that when is Lipschitz, then the density satisfies the Fokker-Planck equation
where denotes the Laplacian operator. On the other hand, the interpolated process is not Markovian, and so does not have a semigroup generator. For this reason, it is difficult to directly control the KL divergence between it and the original Langevin diffusion. In the following lemma, we construct a different partial differential equation that is satisfied by .
The density of the process defined in (12) satisfies the PDE
where is a time-varying drift term.
See Section 3.1 for the proof of this lemma. The key observation is that, conditioned on the -field , the process is a Brownian motion with constant drift, whose conditional density satisfies a Fokker-Planck equation. Taking the expectation on both sides, and interchanging the integral with the derivatives, we obtain the Fokker-Planck equation for the density unconditionally.
In Lemma 1, we have a Fokker-Planck equation with time-varying coefficients; it is satisfied by the one-time marginal densities of the continuous-time interpolation for (2). This representation provides convenient tool for bounding the time derivative of KL divergence, a task to which we turn in the next section.
3.1 Proof of Lemma 1
We first consider the conditional distribution of , conditioned on . At time , it starts with an atomic mass (viewed as Dirac -function at point , which is a member of the tempered distribution space (see, e.g., Rudin, 1991). Its derivatives and Hessian are well-defined as well.) For , this conditional density follows the Fokker-Planck equation for a Brownian motion with constant drift:
where the partial derivatives are in terms of the dummy variable . Take expectations of both sides of (15). By interchanging derivative and integration, we obtain the following identities. Rigorous justification are provided below.
Proof of equation (16a):
Applying Lemma 14 in Appendix D, we can show that the density has a tail decaying as . We then note that is equal to the semigroup generator of the conditional Brownian motion with constant drift, which also decays exponentially with , in a small neighborhood of , for fixed . So the quantity has a dominating function of the form of in a small neighborhood of . Combining with the dominated convergence theorem justifies step (i).
Proof of equation (16b):
In order to justify step (i), we first note that, according to Assumption 1, both of the functions and grow at most linearly in , for fixed . By the rapid decay of the tail of shown in Lemma 14, and the decay of the tail of obtained by elementary results on the Gaussian density, we have a dominating function of the form of . This justifies by the dominated convergence theorem. Then simply follows from the Bayes rule.
Proof of equation(16c):
We similarly have:
Note that for any density function . Since is a quadratic function in the variable , its gradient is linear (it also grows at most linearly with ), and its Laplacian is constant. Therefore, we have a dominating function of form for the integrand, which guarantees the interchange between the integral and the Laplacian operator. This leads to .
Combining these identities yields
where for .
4 Controlling the KL divergence: Proof of Proposition 1
We now turn to the proof of Proposition 1, which involves bounding the derivative . We first compute the derivative using the Fokker-Planck equation established in Lemma 1, and then upper bound it by a regularity estimate of the density and moment bounds on . The key geometric intuition underlying our argument is the following: if the drift is second-order smooth and the initial distribution at each step is also smooth, most of the Gaussian noise is cancelled out, and only higher-order terms remain. This intuition is fleshed out in Section 4.1.
In the following lemma, we give an explicit upper bound on the KL divergence between the one-time marginal distributions of the interpolated process and the original diffusion, based on Fokker-Planck equations derived above.
See Appendix A for the proof of this claim.
It is worth noting the key difference between our approach and the method of Dalalyan (2017), which is based on the Girsanov theorem. His analysis controls the KL divergence via the quantity , a term which scales as even for the simple case of the Ornstein-Uhlenbeck process. Indeed, the Brownian motion contributes to an oscillation in , dominating other lower-order terms. By contrast, we control the KL divergence using the quantity . Observe that is exactly the backward conditional expectation of conditioned on the value of . Having the conditional expectation inside (rather than outside) the norm enables the lower-order oscillations to cancel out.
In the remainder of this section, we focus on bounding the integral on the right-hand side of (17). Since the difference between and comes mostly from an isotropic noise, we may expect it to mostly cancel out. In order to exploit this intuition, we use the third-order smoothness condition (see Assumption 2) so as to perform the Taylor expansion
The reminder term is relatively easy to control, since it contains a factor, which is already of order . More formally, we have:
See Appendix A for the proof of this claim.
It remains to control the first order term. From Assumption 1, the Jacobian norm is at most ; accordingly, we only need to control the norm of the vector . It corresponds to the difference between the best prediction about the past of the path and the current point, given the current information. Herein lies the main technical challenge in the proof of Proposition 1, apart from the construction of the Fokker-Planck equation for the interpolated process. Before entering the technical details, let us first provide some geometric intuition for the argument.
4.1 Geometric Intuition
Suppose that we were dealing with the conditional expectation of , conditioned on ; in this case, the Gaussian noise would completely cancel out (see (12)). However, we are indeed reasoning backward, and itself is dependent with the Gaussian noise added to this process. It is unclear whether the cancellation occurs when computing . In fact, it occurs only under particular situations, which turn out to typical for the discretized process.
Due to the dependence between and Gaussian noise, we cannot expect cancellation to occur in general. Figure 1(a) illustrates an extremal case, where the initial distribution at time is an atomic mass. When we condition on the value at as well, the process behaves like a Brownian bridge. Consequently, it makes no difference whether the conditional expectation is inside or outside the norm: in either case, there is a term of the form , which scales as .
On the other hand, as illustrated in Figure 1(b), if the initial distribution is uniform over some region, the initial point is almost equally likely to be from anywhere around , up to the drift term, and most of the noise gets cancelled out. In general, if the initial distribution is smooth, locally it looks almost uniform, and similar phenomena should also hold true. Thus we expect to be decomposed into terms coming from the drift and terms coming from the smoothness of the initial distribution.
4.2 Upper Bound via Integration by Parts
With this intuition in hand, we now turn to the proof itself. In order to leverage the smoothness of the initial distribution, we use integration by parts to move the derivatives onto the density of . From Bayes’ formula, we have
Since the density is a Gaussian centered at with fixed covariance, the gradient with respect to is the density itself times a linear factor , with an additional factor depending on the Jacobian of . This elementary fact motivates a decomposition whose goal is to express as the sum of the conditional expectation of and some other terms which are easy to control. More precisely, in order to expose a gradient of the Gaussian density, we decompose the difference into three parts, namely , where
We define the conditional expectations for and control the three terms separately.
Let us denote by the -dimensional standard Gaussian density. The first term can directly be expressed in term of the gradient of :
where we used the chain rule and . Thus, applying integration by parts, we write in a revised form.
For all , we have
See Section 4.3 for the proof of this lemma.
It is clear from Lemma 4 that a regularity estimates on the moments of gives an estimates on the squared integral. Such a bound with reasonable dimension dependence is nontrivial to obtain. This is postponed to section 5.
The remaining two terms are relatively easy to control, as summarized in the following:
Under Assumption 1, the following bounds hold for all :
See Section 4.4 for the proof of this lemma.
4.3 Proof of Lemma 4
where is the -dimensional standard Gaussian density. We first note the tail of the Gaussian density is trivial, and the tail of is justified by Appendix D. Therefore we obtain applying integration by parts:
Then, applying the Cauchy-Schwartz inequality yields
This last inequality concludes the proof of Lemma 4.
4.4 Proof of Lemma 5
Recall that this lemma provides bounds on the remaining two terms and