Improved Bounds for Discretization of Langevin Diffusions: Near-Optimal Rates without Convexity

07/25/2019 ∙ by Wenlong Mou, et al. ∙ 1

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 O(η) to O(η^2) 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.

READ FULL TEXT VIEW PDF
POST COMMENT

Comments

There are no comments yet.

Authors

page 1

page 2

page 3

page 4

This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.

1 Introduction

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)

(1)

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

(2)

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?

Our contributions:

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 Requirecontraction 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 Requiremixing Additionalassumptions
Dalalyan (2017) No No None
Alfonsi et al. (2015) No No Second-ordersmooth drift
Dalalyan and Karagulyan (2017),Durmus and Moulines (2017) Yes - Yes Second-ordersmooth drift
Cheng and Bartlett (2018),Ma et al. (2018) No - Yes strong convexityoutside 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 convexityoutside a ball
This paper No No Second-ordersmooth drift
Table 1: Comparison between discretization of Langevin diffusion and sampling algorithms.
Related work:

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.

Notation:

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

(3)
Assumption 2 (Smooth drift term).

There is a finite constant such that

(4)
Assumption 3 (Distant dissipativity).

There exist strictly positive constants such that

(5)
Assumption 4 (Smooth Initialization).

The initializations and , for the processes (1) and (2) respectively, are drawn from a density such that

(6)

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:

Theorem 1.

Consider the original Langevin diffusion (1) under Assumptions 1,  2,  3 and 4. Then there are universal constants such that for any and all times , the KL error of the Euler-Maruyama discretization (2) is bounded as

(7)

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:

Theorem 2.

Consider the original Langevin diffusion (1) under Assumptions 1,  2, and 4, and suppose that for some non-negative function . Then for any stepsize and time , the KL error of the Euler-Maruyama discretization (2) is bounded as

(8)

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

(9)

With this definition, we have the following:

Corollary 1.

Consider a density of the form such that:

  1. The gradient satisfies Assumptions 12, and 3.

  2. The distribution defined by satisfies a log-Sobolev inequality with constant .

Then under the initialization Assumption 4, for any , the unadjusted Langevin algorithm (2) with drift and step size has mixing times bounded as:

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

In this section, we provide a high-level overview of the three main steps that comprise the proof of Theorem 1; the subsequent Sections 34, and 5 provide the details of these steps.

Step 1:

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

(10)

See Lemma 2 for details.

Step 2:

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:

Proposition 1.

Under Assumptions 12 and 4, for any , we have

(11)
Step 3:

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.

Proposition 2.

Under Assumptions 12 and 4, for and , we have

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.

  • Without Assumption 3, if the drift term is negative gradient of a function with , the second and fourth moment can still be bounded (see Lemma 12). The proof uses the moment bounds of along the path. Collecting these results yields Eq (8), which completes our sketch of the proof of Theorem 2.

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

(12)

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

(13)

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

.

Lemma 1.

The density of the process defined in (12) satisfies the PDE

(14)

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:

(15)

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.

(16a)
(16b)
(16c)
Proof of equation (16a):

We show:

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):

We have:

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.

Lemma 2.

Suppose that the densties and satisfy the Fokker-Planck equations (13) and (14), respectively. Then

(17)

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

(18)

The reminder term is relatively easy to control, since it contains a factor, which is already of order . More formally, we have:

Lemma 3.

Let us define . We have under Assumptions 1 and 2:

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.

(a) (b)
Figure 1: (a) Non-smooth initial distribution leads to scaling. (b) Smooth initial distribution: Oscillation cancels and leads to error.

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

(19)

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.

Lemma 4.

For all , we have

and consequently,

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:

Lemma 5.

Under Assumption 1, the following bounds hold for all :

(20a)
(20b)

See Section 4.4 for the proof of this lemma.

Combining the Taylor expansion (18) with the bounds from Lemma 4 and 5 yields the bound claimed in Proposition 1.

4.3 Proof of Lemma 4

We prove here Lemma 4 which controls the dominant term of the decomposition of in (19). Recall is expressed in term of the gradient of the Gaussian density:

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: