Sqrt(d) Dimension Dependence of Langevin Monte Carlo

( √ d/ ) mixing time bound for LMC , without warm start , under the common log-smooth and log-strongly-convex conditions , plus a growth condition on the 3rd-order derivative of the potential of target measures . This bound improves the best previously known Õ ( d/ ) result and is optimal ( in terms of order ) in both dimension d and accuracy tolerance for target measures satisfying the aforementioned assumptions . Our theoretical analysis is further validated by numerical experiments . 1 INTRODUCTION . The problem of sampling statistical distributions has attracted considerable attention , not only in the fields of statistics and scientific computing , but also in machine learning ( Robert & Casella , 2013 ; Andrieu et al. , 2003 ; Liu , 2008 ) ; for example , how various sampling algorithms scale with the dimension of the target distribution is a popular recent topic in statistical deep learning ( see discussions below for references ) . For samplers that can be viewed as discretizations of SDEs , the idea is to use an ergodic SDE whose equilibrium distribution agrees with the target distribution , and employ an appropriate numerical algorithm that discretizes ( the time of ) the SDE . The iterates of the numerical algorithm will approximately follow the target distribution when converged , and can be used for various downstream applications such as Bayesian inference and inverse problem ( Dashti & Stuart , 2017 ) . One notable example is the Langevin Monte Carlo algorithm ( LMC ) , which corresponds to an Euler-Maruyama discretization of the overdamped Langevin equation . Its study dated back to at least the 90s ( Roberts et al. , 1996 ) but keeps on leading to important discoveries , for example , on non-asymptotics and dimension dependence , which are relevant to machine learning ( e.g. , Dalalyan ( 2017a ; b ) ; Cheng et al . ( 2018a ) ; Durmus et al . ( 2019 ) ; Durmus & Moulines ( 2019 ) ; Vempala & Wibisono ( 2019 ) ; Dalalyan & Riou-Durand ( 2020 ) ; Li et al . ( 2019 ) ; Erdogdu & Hosseinzadeh ( 2021 ) ; Mou et al . ( 2019 ) ; Lehec ( 2021 ) ) . LMC is closely related to SGD too ( e.g. , Mandt et al . ( 2017 ) ) . Many other examples exist , based on alternative SDEs and/or different discretizations ( e.g. , Dalalyan & Riou-Durand ( 2020 ) ; Ma et al . ( 2021 ) ; Mou et al . ( 2021 ) ; Li et al . ( 2020 ) ; Roberts & Rosenthal ( 1998 ) ; Chewi et al . ( 2021 ) ; Shen & Lee ( 2019 ) ) . Quantitatively characterizing the non-asymptotic sampling error of numerical algorithms is usually critical for choosing the appropriate algorithm for a specific downstream application , for providing practical guidance on hyperparameter selection and experiment design , and for designing improved samplers . A powerful tool that dates back to ( Jordan et al. , 1998 ) is a paradigm of non-asymptotic error analysis , namely to view sampling as optimization in probability space , and it led to many important recent results ( e.g. , Liu & Wang ( 2016 ) ; Dalalyan ( 2017a ) ; Wibisono ( 2018 ) ; Zhang et al . ( 2018 ) ; Frogner & Poggio ( 2020 ) ; Chizat & Bach ( 2018 ) ; Chen et al . ( 2018 ) ; Ma et al . ( 2021 ) ; Erdogdu & Hosseinzadeh ( 2021 ) ) . It works by choosing an objective functional , typically some statistical distances/divergences , and showing that the law of the iterates of sampling algorithms converges in that objective functional . However , the choice of the objective functional often needs to be customized for different sampling algorithms . For example , KL divergence works for LMC ( Cheng & Bartlett , 2018 ) , but a carefully hand-crafted cross term needs to be added to KL divergence for analyzing KLMC ( Ma et al. , 2021 ) . Even for the same underlying SDE , different discretization schemes exist and lead to different sampling algorithms , and the analyses of them had usually been case by case ( e.g. , Cheng et al . ( 2018b ) ; Dalalyan & Riou-Durand ( 2020 ) ; Shen & Lee ( 2019 ) ) . Therefore , it would be a desirable complement to have a unified , general framework to study the non-asymptotic error of SDE-based sampling algorithms . Toward this goal , an alternative approach to analysis has recently started attracting attention , namely to resort to the numerical analysis of SDE integrators ( e.g. , Milstein & Tretyakov ( 2013 ) ; Kloeden & Platen ( 1992 ) ) and quantitatively connect the integration error to the sampling error . One remarkable work in this direction is Li et al . ( 2019 ) , which will be discussed in greater details later on . The main tool of analysis in this paper will be a strengthened version ( in specific aspects that will be clarified soon ) of the result in Li et al . ( 2019 ) . Although this analysis framework is rather general and applicable to a broad family of numerical methods that discretize contractive1 SDEs , the main innovation focuses on a specific sampling algorithm , namely LMC , which is widely used in practice . Its stochastic gradient version is implemented in common machine learning systems , such as Tensorflow ( Abadi et al. , 2016 ) , and is the off-the-shelf algorithm for large scale Bayesian inference . With the ever-growing size of parameter space , the non-asymptotic error of LMC is of central theoretical and practical interest , in particular , its dependence on the dimension of the sample space . The best current known upper bound of the mixing time in 2-Wasserstein distance for LMC is Õ ( d ) ( Durmus & Moulines , 2019 ) . Motivated by a recent result ( Chewi et al. , 2021 ) that shows better dimension dependence for a Metropolis-Adjusted improvement of LMC , we will investigate if the current bound for ( unadjusted ) LMC is tight , and if not , what is the optimal dimension dependence . Our contributions . The main contribution of this work is an improved Õ ( √ d ) mixing time upper bound for LMC in 2-Wasserstein distance , under reasonable regularity assumptions . More specifically , we study LMC for sampling from a Gibbs distribution dµ ∝ exp ( −f ( x ) ) dx . Under the standard smoothness and strong-convexity assumptions , plus an additional linear growth condition on the third-order derivative of the potential ( which also shares connections to popular assumptions in the frontier literature ) , our bound improves upon the previously best known Õ ( d ) result ( Durmus & Moulines , 2019 ) in terms of dimension dependence . For a comparison , note it was known that discretized kinetic Langevin dynamics can lead to √ d dependence on dimension ( Cheng & Bartlett , 2018 ; Dalalyan & Riou-Durand , 2020 ) and some believe that it is the introduction of momentum that improves the dimension dependence , but our result shows that discretized overdamped Langevin ( no momentum ) can also have mixing time scaling like √ d. In fact , it is important to mention that recently shown was that Metropolis-Adjusted Euler-Maruyama discretization of overdamped Langevin ( i.e. , MALA ) has an optimal dimension dependence of Õ ( √ d ) ( Chewi et al. , 2021 ) , while what we analyze here is the unadjusted version ( i.e. , LMC ) , and it has the same dimension dependence ( note however that our dependence is not as good as that for MALA ; more discussion in Section 4 ) . We also constructed an example which shows that the mixing time of LMC is at least Ω̃ ( √ d ) . Hence , our mixing time bound has the optimal dependence on both d and , in terms of order , for the family of target measures satisfying those regularity assumptions . Our theoretical analysis is further validated by empirical investigation of numerical examples . A minor contribution of this work is the error analysis framework that we use . It is based on the classical mean-square analysis ( Milstein & Tretyakov , 2013 ) in numerical SDE literature , however extended from finite time to infinite time . It is a minor contribution because this extension was 1possibly after a coordinate transformation already pioneered in the milestone work of Li et al . ( 2019 ) , although we will develop a refined version . Same as in classical mean-square analysis and in Li et al . ( 2019 ) , the final ( sampling in this case ) error is only half order lower than the order of local strong integration error ( p2 ) . This will lead to a Õ ( C 1 p2− 1 2 / 1 p2− 1 2 ) mixing time upper bound in 2-Wasserstein distance for the family of algorithms , where C is a constant containing various information of the underlying problem , e.g. , the dimension d. Nevertheless , the following two are new to this paper : ( i ) We weakened the requirement on local strong and weak errors . More precisely , Li et al . ( 2019 ) requires uniform bounds on local errors , but this could be a nontrivial requirement for SDE integrators ; the improvement here only requires non-uniform bounds ( although establishing the same result consequently needs notably more efforts , these are included in this paper too ) . ( ii ) The detailed expressions of our bounds are not the same as those in Li et al . ( 2019 ) ( even if local errors could be uniformly bounded ) , and as we are interested in dimension-dependence of LMC , we work out constants and carefully track their dimension-dependences . Bounds and constants in Li et al . ( 2019 ) might not be specifically designed for tightly tracking dimension dependences , as the focus of their seminal paper was more on dependence ; consequently , its general error bound only led to a Õ ( d ) -dependence in mixing time when applied to LMC ( see Example 1 in Li et al . ( 2019 ) ) , whereas our result leads to Õ ( √ d ) . 2 PRELIMINARIES . Notation Use symbol x to denote a d-dim . vector , and plain symbol x to denote a scalar variable . ‖x‖ denotes the Euclidean vector norm . A numerical algorithm is denoted by A and its k-th iterate by x̄k . Slightly abuse notation by identifying measures with their density function w.r.t . Lebesgue measure . Use the convention Õ ( · ) = O ( · ) logO ( 1 ) ( · ) and Ω̃ ( · ) = Ω ( · ) logO ( 1 ) ( · ) , i.e. , the Õ ( · ) /Ω̃ ( · ) notation ignores the dependence on logarithmic factors . Use the notation Ω̃ ( · ) similarly . Denote 2-Wasserstein distance by W2 ( µ1 , µ2 ) = ( inf ( X , Y ) ∼Π ( µ1 , µ2 ) E ‖X − Y ‖ 2 ) 1 2 , where Π ( µ1 , µ2 ) is the set of couplings , i.e . all joint measures with X and Y marginals being µ1 and µ2 . Denote the target distribution by µ and the law of a random variableX by Law ( X ) . Finally , denote the mixing time of an sampling algorithm A converging to its target distribution µ in 2-Wasserstein distance by τmix ( ; W2 ; A ) = inf { k ≥ 0|W2 ( Law ( x̄k ) , µ ) ≤ } . SDE for Sampling Consider a general SDE dxt = b ( t , xt ) dt+ σ ( t , xt ) dBt ( 1 ) where b ∈ Rd is a drift term , σ ∈ Rd×l is a diffusion coefficient matrix and Bt is a l-dimensional Wiener process . Under mild condition ( Pavliotis , 2014 , Theorem 3.1 ) , there exists a unique strong solution xt to Eq . ( 1 ) . Some SDEs admit geometric ergodicity , so that their solutions converge exponentially fast to a unique invariant distribution , and examples include the classical overdamped and kinetic Langevin dynamics , but are not limited to those ( e.g. , Mou et al . ( 2021 ) ; Li et al . ( 2020 ) ) . Such SDE are desired for sampling purposes , because one can set the target distribution to be the invariant distribution by choosing an SDE with an appropriate potential , and then solve the solution xt of the SDE and push the time t to infinity , so that ( approximate ) samples of the target distribution can be obtained . Except for a few known cases , however , explicit solutions of Eq . ( 1 ) are elusive and we have to resort to numerical schemes to simulate/integrate SDE . Such example schemes include , but are not limited to Euler-Maruyama method , Milstein methods and Runge-Kutta method ( e.g. , Kloeden & Platen ( 1992 ) ; Milstein & Tretyakov ( 2013 ) ) . With constant stepsize h and at k-th iteration , a typical numerical algorithm takes a previous iterate x̄k−1 and outputs a new iterate x̄k as an approximation of the solution xt of Eq . ( 1 ) at time t = kh . Langevin Monte Carlo Algorithm LMC algorithm is defined by the following update rule x̄k = x̄k−1 − h∇f ( x̄k−1 ) + √ 2hξk , k = 1 , 2 , · · · ( 2 ) where { ξk } k∈Z > 0 are i.i.d . standard d-dimensional Gaussian vectors . LMC corresponds to an EulerMaruyama discretization of the continuous overdamped Langevin dynamics dxt = −∇f ( xt ) dt+√ 2dBt , which converges to an equilibrium distribution µ ∼ exp ( −f ( x ) ) . Dalalyan ( 2017b ) provided a non-asymptotic analysis of LMC . An Õ ( d 2 ) mixing time bound in W2 for log-smooth and log-strongly-convex target measures ( Dalalyan , 2017a ; Cheng et al. , 2018a ; Durmus et al. , 2019 ) has been established . It was further improved to Õ ( d ) under an additional Hessian Lipschitz condition ( Durmus & Moulines , 2019 ) . Mixing time bounds of LMC in other statistical distances/divergences have also been studied , including total variation distance ( Dalalyan , 2017b ; Durmus & Moulines , 2017 ) and KL divergence ( Cheng & Bartlett , 2018 ) . Classical Mean-Square Analysis A powerful framework for quantifying the global discretization error of a numerical algorithm for Eq . ( 1 ) , i.e. , ek = { E ‖xkh − x̄k‖ } 1 2 , is mean-square analysis ( e.g. , Milstein & Tretyakov ( 2013 ) ) . Mean-square analysis studies how local integration error propagate and accumulate into global integration error ; in particular , if one-step ( local ) weak error and strong error ( both the exact and numerical solutions start from the same initial value x ) satisfy ‖Exh − Ex̄1‖ ≤C1 ( 1 + E ‖x‖2 ) 1 2 hp1 , ( local weak error ) ( E ‖xh − x̄1‖2 ) 1 2 ≤C2 ( 1 + E ‖x‖2 ) 1 2 hp2 , ( local strong error ) ( 3 ) over a time interval [ 0 , Kh ] for some constants C1 , C2 > 0 , p2 ≥ 12 and p1 ≥ p2 + 1 2 , then the global error is bounded by ek ≤ C ( 1 + E ‖x0‖2 ) 1 2 hp2− 1 2 , k = 1 , · · · , K for some constant C > 0 dependent on Kh . Although classical mean-square analysis is only concerned with numerical integration error , sampling error can be also inferred . However , there is a limitation that prevents its direct employment in analyzing sampling algorithms : the global error bound only holds in finite time because the constant C can grow exponentially as K increases , rendering the bound useless when K →∞ .

Summary: The paper is concerned with the non-asymptotic analysis of SDE-based sampling algorithms and has two main contributions. Firstly, it improves upon the general framework of [Li et al. (2019)](https://proceedings.neurips.cc/paper/2019/hash/7d265aa7147bd3913fb84c7963a209d1-Abstract.html) and, in particular, does not require uniform boundedness assumptions for the local strong/weak errors. This means their first result, Theorem 3.3, is easier to apply to numerical schemes for contractive SDEs and can already be seen to unify previous works (such as the analysis of KLMC in [Dalalyan et al. (2020)](https://projecteuclid.org/journals/bernoulli/volume-26/issue-3/On-sampling-from-a-log-concave-density-using-kinetic-Langevin/10.3150/19-BEJ1178.short)). Secondly, the authors consider the Unadjusted Langevin Algorithm (ULA) and derive a 2-Wasserstein bound of $\mathcal{O}(\sqrt{d} h)$, which has optimal dependence on both the dimension $d$ and step size $h$. To achieve this, they impose an additional smoothness condition on $f$ that is less restrictive than having the Hessian $\nabla^2 f$ Lipschitz continuous. The authors support this theory with numerical examples, where they demonstrate that a lower bound on the 2-Wasserstein error scales as $\mathcal{O}(\sqrt{d})$ and $\mathcal{O}(h)$.