Global Non-convex Optimization with Discretized Diffusions
Murat A. Erdogdu, Lester Mackey, Ohad Shamir
Introduction
Consider the unconstrained and possibly non-convex optimization problem
Recent studies have shown that the Langevin algorithm – in which an appropriately scaled isotropic Gaussian vector is added to a gradient descent update – globally optimizes whenever the objective is dissipative ( for ) with a Lipschitz gradient . Remarkably, these globally optimized objectives need not be convex and can even be multimodal. The intuition behind the success of the Langevin algorithm is that the stochastic optimization method approximately tracks the continuous-time Langevin diffusion which admits the Gibbs measure – a distribution defined by – as its invariant distribution. Here, is an inverse temperature parameter, and when is large, the Gibbs measure concentrates around its modes. As a result, for large values of , a rapidly mixing Langevin algorithm will be close to a global minimum of . In this case, rapid mixing is ensured by the Lipschitz gradient and dissipativity. Due to its simplicity, efficiency, and well-understood theoretical properties, the Langevin algorithm and its derivatives have found numerous applications in machine learning [see, e.g., 28, 6].
In this paper, we prove an analogous global optimization property for the Euler discretization of any smooth and dissipative diffusion and show that different diffusions are suitable for solving different classes of convex and non-convex problems. Our non-asymptotic analysis, based on a multidimensional version of Stein’s method, establishes explicit bounds on both integration and optimization error. Our contributions can be summarized as follows:
For any function , we provide explicit \mathcal{O}\mathopen{}\mathclose{{}\left(\frac{1}{\epsilon^{2}}}\right) bounds on the numerical integration error of discretized dissipative diffusions. Our bounds depend only on simple properties of the diffusion’s coefficients and Stein factors, i.e., bounds on the derivatives of the associated Poisson equation solution.
For pseudo-Lipschitz , we derive explicit first through fourth-order Stein factor bounds for every fast-coupling diffusion with smooth coefficients. Since our bounds depend on Wasserstein coupling rates, we provide user-friendly, broadly applicable tools for computing these rates. The resulting computable integration error bounds recover the known Markov chain Monte Carlo convergence rates of the Langevin algorithm in both convex and non-convex settings but apply more broadly.
We introduce new explicit bounds on the expected suboptimality of sampling from a diffusion. Together with our integration error bounds, these yield computable and convergent bounds on global optimization error. We demonstrate that improved optimization guarantees can be obtained by targeting distributions other than the standard Gibbs measure.
We show that different diffusions are appropriate for different objectives and detail concrete examples of global non-convex optimization enabled by our framework but not covered by the existing Langevin theory. For example, while the Langevin diffusion is particularly appropriate for dissipative and hence quadratic growth , we show alternative diffusions are appropriate for “heavy-tailed” with subquadratic or sublinear growth.
We emphasize that, while past work has assumed the existence of finite Stein factors , focused on deriving convergence rates with inexplicit constants , or concentrated singularly on the Langevin diffusion , the goals of this work are to provide the reader with tools to (a) check the appropriateness of a given diffusion for optimizing a given objective and (b) compute explicit optimization and integration error bounds based on easily accessed properties of the objective and chosen diffusion. The rest of the paper is organized as follows. Section 1.1 surveys related work. Section 2 provides an introduction to diffusions and their use in optimization and reviews our notation. Section 3 provides explicit bounds on integration error in terms of Stein factors and on Stein factors in terms of simple properties of and the diffusion. In Section 4, we provide explicit bounds on optimization error by targeting Gibbs and non-Gibbs invariant measures and discuss how to obtain better optimization error using non-Gibbs invariant measures. We give concrete examples of applying these tools to non-convex optimization problems in Section 5 and conclude in Section 6.
The Euler discretization of the Langevin diffusion is commonly termed the Langevin algorithm and has been studied extensively in the context of sampling from a log concave distribution. Non-asymptotic integration error bounds for the Langevin algorithm are studied in . A representative bound follows from combining the ergodicity of the diffusion with a discretization error analysis and yields error in steps for the strongly log concave case and steps for the general log concave case .
Our work is motivated by a line of research that uses the Langevin algorithm to globally optimize non-convex functions. Gelfand and Mitter established the global convergence of an appropriate variant of the algorithm, and Raginsky et al. subsequently used optimal transport theory to prove optimization and integration error bounds. For example, provides an integration error bound of after \mathcal{O}\big{(}\frac{1}{\epsilon^{4}}\text{poly}(\log(\frac{1}{\epsilon}))\frac{1}{\lambda_{*}}\big{)} steps under the quadratic-growth assumptions of dissipativity and a Lipschitz gradient; the estimate involves the inverse spectral gap parameter , a quantity that is often unknown and sometimes exponential in both inverse temperature and dimension. Gao et al. obtained similar guarantees for stochastic Hamiltonian Monte Carlo algorithms for empirical and population risk minimization under a dissipativity assumption with rate estimates. In this work, we accommodate “heavy-tailed” objectives that grow subquadratically and trade the often unknown and hence inexplicit spectral gap parameter of for the more user-friendly distant dissipativity condition (Prop. 3.4) which provides a straightforward and explicit certification of fast coupling and hence the fast mixing of a diffusion. For distantly dissipative diffusions, the size of our error bounds is driven primarily by a computable distance parameter; in the Langevin setting, an analogous quantity is studied in place of the spectral gap in the contemporaneous work of .
Cheng et al. provide integration error bounds for sampling with the overdamped Langevin algorithm under a distant strong convexity assumption (a special case of distant dissipativity). The authors build on the results of and establish error in steps. We consider general distantly dissipative diffusions and establish an integration error of in steps under mild assumptions on the objective function and smoothness of the diffusion.
Vollmer et al. used the solution of the Poisson equation in their analysis of stochastic Langevin gradient descent, invoking the bounds of Pardoux and Veretennikov [24, Thms. 1 and 2] to obtain Stein factors. However, Thms. 1 and 2 of yield only inexplicit constants and require bounded diffusion coefficients, a strong assumption violated by the examples treated in Section 5. Chen et al. considered a broader range of diffusions but assumed, without verification, that Stein factor and Markov chain moment were universally bounded by constants independent of all problem parameters. One of our principal contributions is a careful enumeration of the dependencies of these Stein factors and Markov chain moments on the objective and the candidate diffusion. Our convergence analysis builds on the arguments of , and our Stein factor bounds rely on distant and uniform dissipativity conditions for -Wasserstein rate decay and the smoothing effect of the Markov semigroup . Our Stein factor results significantly generalize the existing bounds of by accommodating pseudo-Lipschitz objectives and quadratic growth in the covariance coefficient and deriving the first four Stein factors explicitly.
Optimization with Discretized Diffusions: Preliminaries
associated with our objective . Inserting and into the formula (2.2) we obtain
which reduces to . We emphasize that the choice of the Gibbs measure is arbitrary, and we will consider other measures that yield superior guarantees for certain minimization problems.
In practice, the diffusion 2.1 cannot be simulated in continuous time and is instead approximated by a discrete-time numerical integrator. We will show that a particular discretization, the Euler method, can be used as a global optimization algorithm for various families of convex and non-convex . The Euler method is the most commonly used discretization technique due to its explicit form and simplicity; however, our analysis can be generalized to other numerical integrators as well. For the Euler discretization of the SDE (2.1) corresponds to the Markov chain updates
where is the step size, and is an isotropic Gaussian vector that is independent from . This update rule defines a Markov chain which typically has an invariant measure that is different from the invariant measure of the continuous time diffusion. However, when the step size is sufficiently small, the difference between two invariant measures becomes small and can be quantitatively characterized [see, e.g., 22]. Our optimization algorithm is simply to evaluate the function at each Markov chain iterate and report the point with the smallest function value.
and bound each term on the right-hand side separately. The integration error—which captures both the short-term non-stationarity of the chain and the long-term bias due to discretization—is the subject of Section 3; we develop explicit bounds using techniques that build upon . The expected suboptimality quantifies how well exact samples from minimize on average. In Section 4, we extend the Gibbs measure Langevin diffusion bound of Raginsky et al. to more general invariant measures and associated diffusions and demonstrate the benefits of targeting non-Gibbs measures.
We say a function is pseudo-Lipschitz continuous of order if it satisfies
Explicit Bounds on Integration Error
We develop our explicit bounds on integration error in three steps. In Theorem 3.1, we bound integration error in terms of the polynomial growth and dissipativity of diffusion coefficients (Conditions 1 and 2) and Stein factors bounds on the derivatives of solutions to the diffusion’s Poisson equation (Condition 3). Condition 3 is a common assumption in the literature but is typically not verified. To address this shortcoming, Theorem 3.2 shows that any smooth, fast-coupling diffusion admits finite Stein factors expressed in terms of diffusion coupling rates (Condition 4). Finally, in Section 3.1, we provide user-friendly tools for explicitly bounding those diffusion coupling rates. We begin with our conditions.
For , the diffusion (2.1) satisfies the dissipativity condition
Our final condition concerns the solution of the Poisson equation (also known as the Stein equation in the Stein’s method literature) associated with our candidate diffusion.
The function solves the Poisson equation with generator (3.2)
is pseudo-Lipschitz of order with constant , and has -th order derivative with degree- polynomial growth for , i.e.,
The coefficients govern the regularity of the Poisson equation solution and are termed Stein factors in the Stein’s method literature. Although variants of Condition 3 have been assumed in previous work , we emphasize that this assumption is not easily verified, and frequently only empirical evidence is provided as justification for the assumption . We will ultimately derive explicit expressions for the Stein factors for a wide variety of diffusions and functions , but first we will use the Stein factors to bound the integration error of our discretized diffusion.
Let Conditions 1, 2 and 3 hold for some . For any even integerIn a typical example where is bounded by a quadratic polynomial, we have and . We also remind the reader that the double factorial is of order . and a step size satisfying ,
This integration error bound, proved in Appendix A, is \mathcal{O}\big{(}\frac{1}{\eta M}+\eta\big{)} since the higher order term can be combined with the dominant term yielding as . We observe that one needs \mathcal{O}\mathopen{}\mathclose{{}\left(\epsilon^{-2}}\right) steps to reach a tolerance of . Theorem 3.1 seemingly makes no assumptions on the objective function , but in fact the dependence on is present in the growth parameters, the Stein factors, and the polynomial degree of the Poisson equation solution. For example, we will show in Theorem 3.2 that the polynomial degree is upper bounded by that of the objective function . To characterize the function classes covered by Theorem 3.1, we next turn to dissecting the Stein factors.
While verifying Conditions 2 and 1 for a given diffusion is often straightforward, it is not immediately clear how one might verify Condition 3. As our second principal contribution, we derive explicit values for the Stein factors for any smooth and dissipative diffusion exhibiting fast -Wasserstein decay:
where infimum is taken over all couplings between and . We further define the relative rates
Assume that Conditions 1, 2 and 4 hold and that is pseudo-Lipschitz continuous of order with, for , at most degree- polynomial growth of its -th order derivatives. Then, Condition 3 is satisfied with Stein factors
The proof of Theorem 3.2 is given in Section B and relies on the explicit transition semigroup derivative bounds of . We emphasize that, to provide finite Stein factors, Theorem 3.2 only requires -Wasserstein decay and allows the -Wasserstein rate to grow. An integrable Wasserstein rate is an indication that a diffusion mixes quickly to its stationary distribution. Hence, Theorem 3.2 suggests that, for a given , one should select a diffusion that mixes quickly to a stationary measure that, like the Gibbs measure Section 2, has modes at the minimizers of . We explore user-friendly conditions implying fast Wasserstein decay in Section 3.1 and detailed examples deploying these tools in Section 5. Crucially for the “heavy-tailed” examples given in Section 5, Theorem 3.2 allows for an unbounded diffusion coefficient , unlike the classic results of .
A simple condition that leads to exponential and -Wasserstein decay is uniform dissipativity 3.11. The next result from (see also [2, Sec. 1], [15, Thm. 10]) makes the relationship precise.
In the Gibbs measure Langevin case, where and , uniform dissipativity is equivalent to the strong convexity of . As we will see in Section 5, the extra degree of freedom in the diffusion coefficient will allow us to treat non-convex and non-strongly convex functions .
A more general condition leading to exponential -Wasserstein decay is the distant dissipativity condition 3.12. The following result of builds upon the pioneering analyses of Eberle [11, Cor. 2] and Wang [27, Thm. 2.6] to provide explicit Wasserstein decay.
for , , and has Wasserstein rate for
Conveniently, both uniform and distant dissipativity imply our dissipativity condition, Condition 2. The Prop. 3.4 rates feature the distance-dependent parameter . In the pre-conditioned Langevin Gibbs setting ( and constant) when is the negative log likelihood of a multimodal Gaussian mixture, in 3.12 represents the maximum distance between modes . When is relatively small, the convergence of the diffusion towards its stationary distribution is rapid, and the non-uniformity parameter is small; when is relatively large, the parameter grows exponentially in , as would be expected due to infrequent diffusion transitions between modes.
Our next result, proved in Appendix D, provides a user-friendly set of sufficient conditions for verifying distant dissipativity and hence exponential Wasserstein decay in practice.
holds for , , then, for any inverse temperature , the diffusion with drift and diffusion coefficients and has stationary density and satisfies 3.12 with , , , and .
Explicit Bounds on Optimization Error
To convert our integration error bounds into bounds on optimization error, we now turn our attention to bounding the expected suboptimality term of 2.5. To characterize the expected suboptimality of sampling from a measure with modes matching the minima of , we generalize a result due to Raginsky et al. . The original result [25, Prop. 3.4] was designed to analyze the Gibbs measure Section 2 and demanded that be smooth, in the sense that . Our next proposition, proved in Appendix C, is designed for more general measures and importantly relaxes the smoothness requirements on .
Suppose is the stationary density of an -dissipative diffusion (Condition 2) with global maximizer . If takes the generalized Gibbs form for and , we have
More generally, if for some and and all , then
When , is the Gibbs measure, and the bound 4.1 exactly recovers [25, Prop. 3.4]. The generalized Gibbs measures with allow for improved dependence on the inverse temperature when . Note however that, for , the distributions also require knowledge of the optimal value . In certain practical settings, such as neural network optimization, it is common to have . When is unknown, a similar analysis can be carried out by replacing with an estimate, and the bound (4.1) still holds up to a controllable error factor.
By combining Prop. 4.1 with Theorem 3.1, we obtain a complete bound controlling the global optimization error of the best Markov chain iterate.
Instantiate the assumptions and notation of Theorem 3.1 and Prop. 4.1. If the diffusion has the generalized Gibbs stationary density , then
Finally, we demonstrate that, for quadratic functions, the generalized Gibbs expected suboptimality bound 4.1 can be further refined to remove the dependence.
The bound (4.4) applies to any with level set (i.e., ) volume proportional to .
Applications to Non-convex Optimization
We next provide detailed examples of verifying that a given diffusion is appropriate for optimizing a given objective, using either uniform dissipativity (Prop. 3.3) or our user-friendly distant dissipativity conditions (Prop. 3.5). When the Gibbs measure Langevin diffusion is used, our results yield global optimization when is strongly convex (condition 3.11 with and ) or has strongly convex tails (condition 3.14 with ). To highlight the value of non-constant diffusion coefficients, we will focus on “heavy-tailed” examples that are not covered by the Langevin theory.
We begin with a pedagogical example of selecting an appropriate diffusion and verifying our global optimization conditions. Fix and consider , a simple non-convex objective which exhibits sublinear growth in and hence does not satisfy dissipativity (Condition 2) when paired with the Gibbs measure Langevin diffusion (). To target the Gibbs measure Section 2 with inverse temperature , we choose the diffusion with coefficients and for and . This choice satisfies Condition 1 with \lambda_{b}=\mathcal{O}\mathopen{}\mathclose{{}\left(1}\right), \lambda_{\sigma}=\mathcal{O}\mathopen{}\mathclose{{}\left(\gamma^{-1/2}}\right), and \lambda_{a}=\mathcal{O}\mathopen{}\mathclose{{}\left(\gamma^{-1}}\right) with respect to and Condition 2 with and . In fact, this diffusion satisfies uniform dissipativity,
Figure 1 illustrates the value of this designed diffusion over gradient descent and the standard Langevin algorithm. Here, , , the inverse temperature , the step size , and each algorithm is run from the initial point . We observe that the Langevin algorithm diverges, and gradient descent requires thousands of iterations to converge while the designed diffusion converges to the region of interest after 15 iterations.
2 Non-convex learning with linear growth
Next consider the canonical learning problem of regularized loss minimization with
We then show that from Prop. 3.5 is bounded and that, for suitable loss choices, is bounded and Lipschitz so that 3.14 holds with and sufficiently large.
Conclusion
In this paper, we showed that the Euler discretization of any smooth and dissipative diffusion can be used for global non-convex optimization. We established non-asymptotic bounds on global optimization error and integration error with convergence governed by Stein factors obtained from the solution of the Poisson equation. We further provided explicit bounds on Stein factors for large classes of convex and non-convex objective functions, based on computable properties of the objective and the diffusion. Using this flexibility, we designed suitable diffusions for optimizing non-convex functions not covered by the existing Langevin theory. We also demonstrated that targeting distributions other than the Gibbs measure can give rise to improved optimization guarantees.
References
Appendix A Proof of Theorem 3.1: Integration error of discretized diffusions
Denoting by and using the integral form Taylor’s theorem on around the previous iterate , and taking expectations, we obtain
The first term on the right hand side can be written as
where in the last step, we used the fact that is independent from and that odd moments of are 0. Similarly for the second and the third terms, we obtain respectively
By combining these with (3.3), we find that (A.1) can be written as
Finally, dividing each term by , averaging over , and using the triangle inequality, we reach the bound
For the first term on the right hand side, using Condition 3 and Lemma A.2, we can write
where we used Young’s inequality in the second step and Lemma A.2 in the last step.
The second term in the above inequality can be bounded by
where in the last step we used Lemma A.2.
Similarly, the third and the fourth terms in the inequality (A.9) can be bounded as
We first bound the expectation in the above integral
Using Condition 1, Lemma E.1 and , we obtain
Therefore, the last term in (A.9) can be bounded by
the right hand side of (A.19) can be bounded by
Combining the above bounds in (A.10), (A.11), (A.13), (A.14) and (A.21) and applying them on (A.9), we reach the final bound
It is well known that the dissipativity condition on the second moment carries directly to the higher order moments . The following lemma will be useful when we bound the higher order moments of the discretized diffusion.
For , we have the following relation
Further, assume that Conditions 1 and 2 hold, and . Then,
The proof for the first statement easily follows from the following expression,
For second statement, we use the first statement with and Conditions 1 and 2. First, we consider the case and write
Using the inequality given in Lemma E.3 twice, we obtain
Same calculation yields a similar expression for the case . Generalizing, we obtain the following formula,
A.2 Proof of Lemma A.2: Markov Chain Moment Bounds
Let the Conditions 1 and 2 hold. For , denote by an even integer satisfying . If the step size satisfies
First, we handle the even moments. For , we write
A compact and more interpretable estimate for can be obtained as follows,
where we use a looser bound to ensure that the right hand side is larger than 1.
The above analysis only covers the even moments so far. For any integer , denote by an even integer that is not smaller than . Then, by the Hölder’s inequality we write
Appendix B Proof of Theorem 3.2: Stein Factor Bounds
Our proof of Theorem 3.2 will use a representation of the Poisson equation solution in terms of the transition semigroup, i.e.,
coupled with the following bounds on the derivatives of the semigroup. See for a proof of the above representation.
Furthermore, satisfies the degree- polynomial growth bound
To establish the first Stein factor bound , we combine the representation (B.3), the triangle inequality, and the definition of pseudo-Lipschitzness to find that
Invoking the pseudo-Lipschitz constant for (B.4) now yields the first Stein factor bound.
For each additional Stein factor, the dominated convergence theorem will enable us to differentiate under the integral sign. For the second Stein factor , using the second derivative of the representation (B.3) and the bound (B.6), we obtain
The final bound is obtained by taking the supremum over and , i.e.,
For the third Stein factor , using the third derivative of the representation (B.3), and the bounds (B.7) and (B.8), we obtain
Lastly, for the fourth Stein factor using the fourth derivative of the representation (B.3) together with the bounds (B.9) and (B.11),
The final result follows from taking a supremum over :
Appendix C Proofs of Expected Suboptimality Bounds
We begin by proving the more general claim 4.2. Our dissipativity assumption together with the diffusion moment bounds in implies that . Moreover, as noted in the proof of [25, Prop. 3.4], the differential entropy is bounded by that of a multivariate Gaussian with the same second moments:
Meanwhile, . Our smoothness assumption, a polar coordinate transform, and the integral identity of [16, 3.326 2] imply that
The result 4.2 now follows by summing the estimates C.2 and C.4.
Now consider the case in which . By design, is also a global minimizer of with . Therefore, by Taylor’s theorem, we have for each
The generalized Gibbs result 4.1 now follows from the general claim 4.2 and Jensen’s inequality as for . ∎
Using the variable change ,, the above equals
where is the Gamma function. Substituting back , and noting that for all , we get that the above equals
Appendix D Proof of Prop. 3.5: User-friendly Wasserstein decay for Gibbs measures
Appendix E Auxiliary Lemmas
For which is independent from , we have
The exact expressions for the quadratic form moments can be found in . We simply use the properties of Frobenius norm to obtain a compact upper bound. ∎
where in the last step, we used and the Bernoulli inequality
For and , we have .
The derivative of the polynomial has roots at , and a root at . Therefore, for , attains its minimum value at . We choose so that
where for the last step, we use for and . ∎