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 manuscript considers the unadjusted Langevin Monte Carlo (LMC) algorithm, and performs non-asymptotic analysis of its convergence with respect to the 2 Wasserstein distance. The main contribution is a mixing time bound of O(d^0.5/\epilson), which improves upon the existing O(d/\epsilon) bound of Durmus and Moulines. The O(d^0.5/\epilson) is also shown to be optimal for the classes of probability distributions considered.