Mean-Square Analysis with An Application to Optimal Dimension Dependence of Langevin Monte Carlo

Sampling algorithms based on discretizations of Stochastic Differential Equations1 ( SDEs ) compose a rich and popular subset of MCMC methods . This work pro-2 vides a general framework for the non-asymptotic analysis of sampling error in3 2-Wasserstein distance , which also leads to a bound of mixing time . The method4 applies to any consistent discretization of contractive SDEs . When applied to5 Langevin Monte Carlo algorithm , it establishes Õ ( √ d/ ) mixing time , without6 warm start , under the common log-smooth and log-strongly-convex conditions,7 plus a growth condition on the potential of target measures at infinity . This bound8 improves the best previously known Õ ( d/ ) result and is optimal in both dimension9 d and accuracy tolerance for log-smooth and log-strongly-convex target measures.10 Our theoretical analysis is further validated by numerical experiments.11 1 Introduction12 The problem of sampling statistical distributions has attracted considerable attention , not only in13 the fields of statistics and scientific computing , but also in machine learning ( Robert and Casella,14 2013 ; Andrieu et al. , 2003 ; Liu , 2008 ) ; for example , how various sampling algorithms scale with15 the dimension of the target distribution is a popular recent topic in statistical deep learning ( see16 discussions below for references ) . For samplers that can be viewed as discretizations of SDEs , the17 idea is to use an ergodic SDE whose equilibrium distribution agrees with the target distribution,18 and employ an appropriate numerical algorithm that discretizes ( the time of ) the SDE . The iterates19 of the numerical algorithm will approximately follow the target distribution when converged , and20 can be used for various downstream applications such as Bayesian inference and inverse problem21 ( Dashti and Stuart , 2017 ) . One notable example is the Langevin Monte Carlo algorithm ( LMC ) ,22 which corresponds to Euler-Maruyama discretization of overdamped Langevin equation . Its study23 dated back to at least the 90s ( Roberts et al. , 1996 ) but keeps on leading to important discoveries , for24 example , on non-asymptotics and dimension dependence , which are relevant to machine learning25 ( e.g. , Dalalyan ( 2017a , b ) ; Cheng et al . ( 2018a ) ; Durmus et al . ( 2019a , b ) ; Vempala and Wibisono26 ( 2019 ) ; Dalalyan and Riou-Durand ( 2020 ) ; Erdogdu and Hosseinzadeh ( 2020 ) ; Mou et al . ( 2019 ) ) .27 LMC is closely related to SGD too ( e.g. , Mandt et al . ( 2017 ) ) . Many other examples exist , based28 on alternative SDEs and different discretizations ( e.g. , Dalalyan and Riou-Durand ( 2020 ) ; Ma et al.29 ( 2021 ) ; Mou et al . ( 2021 ) ; Li et al . ( 2020 ) ; Roberts and Rosenthal ( 1998 ) ; Chewi et al . ( 2020 ) ; Shen30 and Lee ( 2019 ) ) .31 Quantitatively characterizing the non-asymptotic sampling error of numerical algorithms is usually32 critical for choosing the appropriate algorithm for a specific downstream application , for providing33 practical guidance on hyperparameter selection and experiment design , and for designing improved34 samplers . A powerful tool that dates back to ( Jordan et al. , 1998 ) is a paradigm of non-asymptotic35 error analysis , namely to view sampling as optimization in probability space , and it led to many36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . important recent results ( e.g. , Liu and Wang ( 2016 ) ; Dalalyan ( 2017a ) ; Wibisono ( 2018 ) ; Zhang37 et al . ( 2018 ) ; Frogner and Poggio ( 2020 ) ; Chizat and Bach ( 2018 ) ; Chen et al . ( 2018 ) ; Ma et al.38 ( 2021 ) ; Erdogdu and Hosseinzadeh ( 2020 ) ) . It works by choosing an objective functional , typically39 some statistical distances/diverges , and showing that the law of the iterates of sampling algorithms40 converges in that objective functional . However , the choice of the objective functional often needs to41 be customized for different sampling algorithms . For example , KL divergence works for LMC ( Cheng42 and Bartlett , 2018 ) , but a carefully hand-crafted cross term needs to be added to KL divergence for43 analyzing KLMC ( Ma et al. , 2021 ) . Even for the same underlying SDE , different discretization44 schemes exist and lead to different sampling algorithms , and the analyses of them had usually been45 case by case ( e.g. , Cheng et al . ( 2018b ) ; Dalalyan and Riou-Durand ( 2020 ) ; Shen and Lee ( 2019 ) ) .46 Therefore , it would be a desirable complement to have a unified , general framework to study the47 non-asymptotic error of SDE-based sampling algorithms.48 As an important member of the family of SDE-based sampling algorithms , Langevin Monte Carlo is49 widely used in practice . Its stochastic gradient version is implemented in common machine learning50 systems , such as Tensorflow ( Abadi et al. , 2016 ) , and is the off-the-shelf algorithm for large scale51 Bayesian inference . With the ever-growing size of parameter space , the non-asymptotic error of LMC52 is of central theoretical and practical interest , in particular , its dependence on the dimension of the53 sample space . The best current known upper bound of the mixing time in 2-Wasserstein distance for54 LMC is Õ ( d ) ( Durmus et al. , 2019b ) . Motivated by a recent result ( Chewi et al. , 2020 ) that shows55 better dimension dependence for a Metropolis-Adjusted improvement of LMC , we wonder if the56 current bound for ( unadjusted ) LMC is tight , and if not , what is the optimal dimension dependence ? 57 Our contribution We study a broad family of numerical algorithms that discretize SDEs that58 have a contraction property ( possibly after a coordinate transformation ) . For this type of problems,59 we revisit the classical mean-square analysis ( Milstein and Tretyakov , 2013 ) in numerical SDE60 literature and extend its the global error bound from finite time to infinite time . Same as in classical61 mean-square analysis , we show the global error is only half order lower than the order of local strong62 error ( p2 ) . We further obtain a Õ ( C 1 p2− 1 2 1 1 p2− 1 2 ) mixing time upper bound in 2-Wasserstein63 distance for the family of algorithms , where C is a constant containing various information of the64 underlying problem , e.g. , the dimension d.65 As an application of the general mixing time result , we study the widely used Langevin Monte66 Carlo algorithm ( LMC ) for sampling from a Gibbs distribution µ ∝ exp ( −f ( x ) ) , which is an67 Euler-Maruyama discretization of Langevin dynamics . Under the standard smoothness and strong-68 convexity assumptions , plus an additional linear growth condition on the third-order derivative of f ,69 we obtain a Õ ( √ d ) mixing time in 2-Wasserstein distance , which improves upon the previously best70 known Õ ( d ) result ( Durmus et al. , 2019b ) . For a comparison , note it was known that discretized71 kinetic Langevin dynamics can lead to √ d dependence on dimension ( Cheng and Bartlett , 2018 ; 72 Dalalyan and Riou-Durand , 2020 ) and some believe that it is the introduction of momentum that73 improves the dimension dependence , but our result shows that discretized overdamped Langevin ( no74 momentum ) can also have mixing time scaling like √ d. In fact , it is important to mention that it was75 recently shown that Metropolis-Adjusted Euler-Maruyama discretization of overdamped Langevin76 ( i.e. , MALA ) has an optimal dimension dependence of Õ ( √ d ) ( Chewi et al. , 2020 ) , while what we77 analyze here is the unadjusted version ( i.e. , LMC ) , and it has the same dimension dependence ( note78 however that our dependence is not as good as that for MALA ; more discussion in Section 4 ) . We79 also constructed an example that shows that the mixing time of LMC is at least Ω̃ ( √ d ) . Hence , our80 mixing time bound has the optimal dependence on both d and . Our theoretical analysis is further81 validated by empirical investigation of numerical examples.82 2 Preliminaries83 Notation Use the symbol x to denote a d-dimensional vector , and the plain symbol x to denote a84 scalar variable . Use ‖x‖ to denote the Euclidean norm of vector x . A numerical algorithm is denoted85 by A and its k-th iterate is denoted by x̄k . We slightly abuse notation by identifying measures with86 their density function w.r.t . Lebesgue measure . We use the convention Õ ( · ) = O ( · ) logO ( 1 ) ( · ) ,87 i.e. , the Õ ( · ) notation ignores the dependence on logarithmic factors . We use the notation Ω̃ ( · ) 88 similarly . Denote 2-Wasserstein distance by W2 ( µ1 , µ2 ) = ( inf ( X , Y ) ∼Π ( µ1 , µ2 ) E ‖X − Y ‖ 2 ) 1 2 ,89 where Π ( µ1 , µ2 ) is the set of couplings , i.e . all joint measures with X and Y marginals being90 µ1 and µ2 . Denote the target distribution by µ and the law of a random variable X by Law ( X ) .91 Finally , denote the mixing time of an sampling algorithm A converging to its target distribution µ in92 2-Wasserstein distance by τmix ( ; W2 ; A ) = inf { k ≥ 0|W2 ( Law ( x̄k ) , µ ) ≤ } .93 SDE for Sampling Consider a general SDE94 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-dimensional95 Wiener process . Under mild condition ( Pavliotis , 2014 , Theorem 3.1 ) , there exists a unique strong96 solution xt to Eq . ( 1 ) . Some SDEs admit geometric ergodicity , so that their solutions converge97 exponentially fast to a unique invariant distribution , and examples include the classical overdamped98 and kinetic Langevin dynamics , but are not limited to those ( e.g. , Mou et al . ( 2021 ) ; Li et al . ( 2020 ) ) .99 Such SDE are desired for sampling purposes , because one can set the target distribution to be the100 invariant distribution by choosing an SDE with an appropriate potential , and then solve the solution101 xt of the SDE and push the time t to infinity , so that ( approximate ) samples of the target distribution102 can be obtained . Except for a few known cases , however , explicit solutions of Eq . ( 1 ) are elusive and103 we have to resort to numerical schemes to simulate/integrate SDE . Such example schemes include,104 but are not limited to Euler-Maruyama method , Milstein methods and Runge-Kutta method ( e.g.,105 Kloeden and Platen ( 1992 ) ; Milstein and Tretyakov ( 2013 ) ) . With constant stepsize h and at k-th106 iteration , a typical numerical algorithm takes a previous iterate x̄k−1 and outputs a new iterate x̄k as107 an approximation of the solution xt of Eq . ( 1 ) at time t = kh.108 Langevin Monte Carlo Algorithm LMC algorithm is defined by the following update rule109 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 Euler-110 Maruyama discretization of the continuous overdamped Langevin dynamics dxt = −∇f ( xt ) dt+111 √ 2dBt , which converges to an equilibrium distribution µ ∼ exp ( −f ( x ) ) .112 Dalalyan ( 2017b ) provided a non-asymptotic analysis of LMC . An Õ ( d 2 ) mixing time bound in113 W2 for log-smooth and log-strongly-convex target measures ( Dalalyan , 2017a ; Cheng et al. , 2018a ; 114 Durmus et al. , 2019a ) has been established . It was further improved to Õ ( d ) under additional115 Lipschitz assumption on the Hessian of f ( Durmus et al. , 2019b ) . Mixing time bounds of LMC116 in other statistical distances/divergences have also been studied , including total variation distance117 ( Dalalyan , 2017b ; Durmus et al. , 2017 ) and KL divergence ( Cheng and Bartlett , 2018 ) .118 Classical Mean-Square Analysis A powerful framework for quantifying the global discretization119 error of a numerical algorithm for Eq . ( 1 ) , i.e. , ek = { E ‖xkh − x̄k‖ } 1 2 , is mean-square analysis120 ( e.g. , the monograph of Milstein and Tretyakov ( 2013 ) ) . Mean-square analysis studies how local121 integration error propagate and accumulate into global integration error ; in particular , if one-step122 ( local ) weak error and strong error ( both the exact solution xt and the numerical approximation start123 from the same initial value x ) satisfy124 ‖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 the125 global error can be bounded by ek ≤ C ( 1 + E ‖x0‖2 ) 1 2 hp2− 1 2 , k = 1 , 2 , · · · , K for some constant126 C > 0 dependent on Kh.127 Although classical mean-square analysis is only concerned with numerical integration error , sampling128 error can be also inferred . However , there is a limitation that prevents directly employing mean-square129 analysis in the non-asymptotic analysis of sampling algorithms . The bound of global error only holds130 in finite time because the constant C can grow exponentially as K increases , rendering the bound131 useless when K →∞.132 3 Mean-Square Analysis of Samplers Based on Contractive SDE133 In order to prepare for the analysis of sampling error , we first show that the finite time limitation of134 integration error analysis can be lifted if the SDE being discretized is contractive.135 More precisely , one bottleneck that prevents the results of classical mean-square analysis from136 extending to infinite time horizon , is the fact that the solution of a general SDE may not be bounded,137 and neither is its discretization . Note that local error ( Eq . ( 3 ) ) depends on the initial value . To go138 from local to global error , these ‘ initial ’ values correspond to iterates of numerical algorithms , which139 change from iteration to iteration and can be unbounded , hence when accumulated together , it is140 possible that the global error may blow up.141 Samplers considered here , on the other hand , are based on stochastic differential equations , each of142 which weakly converges to a limiting distributions . The solution of the underlying converging SDE,143 as it converges to the invariant measure , gradually inherits boundedness properties from the target144 measure . Thus , as long as the target measure has bounded 2nd-moment , a sampling algorithm based145 on a reasonable discretization of the SDE should also have bounded 2nd-moment . Motivated by this146 observation , we will assume the sampling algorithms we study are based on contractive SDEs , which147 is a sufficient condition to ensure the underlying SDE converges to a statistical distribution.148 Definition 3.1 . A stochastic differential equation is contractive if there exists a non-singular constant149 matrix A ∈ Rd×d , a constant β > 0 , such that any pair of solutions of the SDE satisfy150 ( E ∥∥A ( xt − yt ) ∥∥2 ) 12 ≤ ∥∥A ( x− y ) ∥∥ exp ( −βt ) , ( 4 ) where xt , yt are two solutions , driven by the same Brownian motion but evolved respectively from151 initial conditions x and y.152 Remark . As long as b and σ in ( 1 ) are not explicitly dependent on time , it suffices to find an153 arbitrarily small t0 > 0 and show ( 4 ) holds for all t < t0.154 Remark . Sometimes contraction is not easy to establish directly , but can be shown after an appro-155 priate coordinate transformation , see ( Dalalyan and Riou-Durand , 2020 , Proposition 1 ) for such a156 treatment for kinetic Langevin dynamics . The introduction of A permits such transformations.157 We now use contractivity to remove the finite time limitation . We will first need a lemma , which is a158 local ( short time ) result.159 Lemma 3.2 . ( Milstein and Tretyakov , 2013 , Lemma 1.3 ) Suppose b and σ in Eq . ( 1 ) are Lipschitz160 continuous . For two solutions xt , yt of Eq . ( 1 ) starting from x , y respectively , denote z : =161 ( xt − x ) − ( yt − y ) , then there exist C0 > 0 and h0 > 0 such that162 E ‖z‖2 ≤ C0 ‖x− y‖2 t , ∀x , y , 0 < t ≤ h0 . ( 5 ) Then we will have a sequence of results that connects sampling error ( a statistical property ) with163 local integration error ( a simulation property ) . This justifies our generic produce for non-asymptotic164 sampling error analysis , which only requires bounding the orders of local weak and strong integration165 errors ( in addition to establishing contractivity of the continuous dynamics ) .166 Theorem 3.3 . ( Global Integration Error , Infinite Time Version ) Suppose Eq . ( 1 ) is contractive with167 rate β and with respect to a non-singular matrix A ∈ Rd×d , with Lipschitz continuous b and σ,168 and there is a numerical algorithm A with step size h simulating the solution xt of the SDE , whose169 iterates are denoted by x̄k , k = 0 , 1 , · · · . Suppose there exists 0 < h0 ≤ 1 , C1 , C2 > 0 , D1 , D2 ≥170 0 , p1 ≥ 1 , 12 < p2 ≤ p1 − 1 2 such that for any 0 < h ≤ h0 , the algorithm A has , respectively , local171 weak and strong error of order p1 and p2 , defined as172 ∥∥E ( xh − x̄1 ) ∥∥ ≤ ( C1 +D1√E ‖x‖2 ) hp1 , ( E ‖xh − x̄1‖2 ) 1 2 ≤ ( C22 +D 2 2E ‖x‖ 2 ) 1 2 hp2 , ( 6 ) where xh solves Eq . ( 1 ) with any initial value x and x̄1 is the result of applying A to x for one step.173 If the solution of SDE xt and algorithm A both start from x0 , then for 0 < h ≤ h1 ,174 min { h0 , 1 4β , ( √ β 4 √ 2κAD2 ) 1 p2− 1 2 , ( β 8 √ 2κA ( D1+C0D2 ) ) 1 p2− 1 2 } , the global error ek is bounded as175 176 ek : = ( E‖xkh − x̄k‖2 ) 1 2 ≤ Chp2− 12 , k = 0 , 1 , 2 , · · · ( 7 ) where177 C = 2√ β κ2A ( C1 + C0C2 + √ 2U ( D1 + C0D2 ) √ β + C2 + √ 2D2U ) , ( 8 ) C0 is from Eq . ( 5 ) , κA is the condition number of matrix A and U2 , 4 ‖x0‖2 + 5Eµ ‖x‖2.178 Remark . We use the convention 1/0 =∞ when D1 = D2 = 0 . This is pertinent when a numerical179 algorithm A , e.g . LMC ( Lemma D.3 ) , produces bounded iterates . In such cases , the initial value in180 Eq . ( 6 ) are iterations of A and will be bounded , it then can be absorbed into C1 , C2 and we may set181 D1 = D2 = 0.182 Following Theorem 3.3 , we obtain the following non-asymptotic bound of the sampling error in W2:183 Theorem 3.4 . ( Non-Asymptotic Sampling Error Bound : General Case ) Under the same assump-184 tion and with the same notation of Theorem 3.3 , we have185 W2 ( Law ( x̄k ) , µ ) ≤ √ 2e−βkhW2 ( Law ( x0 ) , µ ) + √ 2Chp2− 1 2 , ∀0 < h ≤ h1 . A corollary of Theorem 3.4 is a bound on the mixing time of the sampling algorithm:186 Corollary 3.5 . ( Upper Bound of Mixing Time : General Case ) Under the same assumption and187 with the same notation of Theorem 3.3 , we have188 τmix ( ; W2 ; A ) ≤ max 1βh1 , 1β ( 2C ) 1 p2− 1 2 log 2 √ 2W2 ( Law ( x0 ) µ ) In particular , when high accuracy is needed , i.e. , < 2Chp2− 1 2 1 , we have189 τmix ( ; W2 ; A ) ≤ ( 2C ) 1 p2− 1 2 β 1 1 p2− 1 2 log 2 √ 2W2 ( Law ( x0 ) , µ ) = Õ C 1p2− 12 β 1 1 p2− 1 2 ( 9 ) Corollary 3.5 states how mixing time depends on the order of local ( strong ) error ( i.e. , p2 ) of a190 numerical algorithm . The larger p2 is , the shorter the mixing time of the algorithm is , in term of191 the dependence on accuracy tolerance parameter . It is important to note that for constant stepsize192 discretizations that are deterministic on the filtration of the driving Brownian motion and use only its193 increments , there is a strong order barrier , namely p2 ≤ 1.5 ( Rüemelin , 1982 ) ; however , methods194 involving multiple stochastic integrals ( e.g. , Kloeden and Platen ( 1992 ) ; Milstein and Tretyakov195 ( 2013 ) ) and randomization ( e.g. , Shen and Lee ( 2019 ) ) can yield a larger p2.196 The constant C defined in Eq . ( 7 ) typically contains rich information about the underlying SDE , e.g.197 dimension , Lipschitz constant of drift and noise diffusion , and the initial value x0 of the sampling198 algorithm . Through C , we can uncover the dependence of mixing time bound on various parameters,199 such as the dimension d. This will be exemplified with Langevin Monte Carlo in the next section.200 4 Non-Asymptotic Analysis of Langevin Monte Carlo Algorithm201 This section quantifies how LMC samples from Gibbs target distribution µ ∼ exp ( −f ( x ) ) that has202 a finite second moment , i.e. , ∫ Rd ‖x‖ 2 dµ < ∞ . Assume without loss of generality that the origin is203 a local minimizer of f , i.e . ∇f ( 0 ) = 0 ; this is for notational convenience in the analysis and can204 be realized via a simple coordinate shift , and it is not needed in the practical implementation . In205 addition , we assume the following two conditions hold:206 A 1 . ( Smoothness and Strong Convexity ) Assume f ∈ C2 and is L-smooth and m-strongly-convex,207 i.e . there exists 0 < m ≤ L such that mId 4 ∇2f ( x ) 4 LId , ∀x ∈ Rd.208 Denote the condition number of f by κ , Lm . The smoothness and strong-convexity assumption is209 the standard assumption in the literature of analyzing LMC algorithm ( Dalalyan , 2017a , b ; Cheng and210 Bartlett , 2018 ; Durmus et al. , 2019a , b ) .211 A 2 . ( Linear Growth of the 3rd-order derivative ) Assume f ∈ C3 and the operator ∇ ( ∆f ) grows212 at most linearly , i.e. , there exists a constant G > 0 such that ∥∥∇ ( ∆f ( x ) ) ∥∥ ≤ G ( 1 + ‖x‖ ) .213 Remark . The linear growth ( at infinity ) condition on∇∆f is actually not as restrictive as it appears,214 and in some sense even weaker than some classical condition for the existence of solutions to SDE.215 For example , a standard condition for ensuring the existence and uniqueness of a global solution to216 SDE is at most a linear growth ( at infinity ) of the drift ( Pavliotis , 2014 , Theorem 3.1 ) . If we consider217 monomial potentials , i.e. , f ( x ) = xp , p ∈ N+ , then the linear growth condition on∇∆f is met when218 p ≤ 4 , whereas the classical condition for the existence of solutions holds only when p ≤ 2.219 To apply mean-square analysis to study LMC algorithm , we will need to ensure the underlying220 Langevin dynamics is contractive , which we verify in Section C and D in the appendix . In addition,221 we work out all required constants to determine the C in Eq . 7 explicitly in the appendix . With all222 these necessary ingredients , we now invoke Theorem 3.4 and obtain the following result:223 Theorem 4.1 . ( Non-Asymptotic Error Bound : LMC ) Suppose Assumption 1 and 2 hold . LMC224 iteration x̄k+1 = x̄k − h∇f ( x̄k ) + √ 2hξk satisfies225 W2 ( Law ( x̄k ) , µ ) ≤ √ 2e−mkhW2 ( Law ( x0 ) , µ ) + √ 2CLMCh , 0 < h ≤ 1 4κL , k ∈ N ( 10 ) where CLMC = 10 ( L2+G ) m 3 2 √ 2d+m ( ‖x0‖2 + 1 ) = O ( √ d ) .226 Corollary 3.5 combined with the above result gives the following bound on the mixing time of LMC:227 Theorem 4.2 . ( Upper Bound of Mixing Time : LMC ) Suppose Assumption 1 and 2 hold . If running228 LMC from x0 , we then have229 τmix ( ; W2 ; LMC ) ≤ max { 4κ2 , 2CLMC m 1 } log 2 √ 2W2 ( Law ( x0 ) , µ ) where CLMC is the same in Theorem 4.1 . When high accuracy is needed , i.e. , ≤ CLMC2mκ2 , we have230 τmix ( ; W2 ; LMC ) ≤ 2CLMC m 1 log 2 √ 2W2 ( Law ( x0 ) , µ ) = Õ ( √ d ) . The Õ ( √ d ) mixing time bound in 2-Wasserstein distance improves upon the previous ones231 ( Dalalyan , 2017a ; Cheng and Bartlett , 2018 ; Durmus et al. , 2019b , a ) in the dependence of d and/or .232 If further assumingG = O ( L2 ) , we then haveCLMC = O ( κ2 √ m √ d ) and Thm.4.2 shows the mixing233 time is Õ ( κ2√ m √ d ) , which also improves the κ dependence in some previous results ( Dalalyan,234 2017a ; Cheng and Bartlett , 2018 ) in the regime m ≤ 1 . A brief comparison is summarized in Table 1.235 Optimality In fact , the Õ ( √ d ) mixing time of LMC has the optimal scaling one can expect . This236 is in terms of the dependence on d and , over the class of all log-smooth and log-strongly-convex237 target measures . To illustrate this , consider the following Gaussian target distribution whose potential238 is239 f ( x ) = m 2 d∑ i=1 x2i + L 2 2d∑ i=d+1 x2i , with m = 1 , L ≥ 4m . ( 11 ) We now establish a lower bound on the mixing time of LMC algorithm for this target measure.240 Theorem 4.3 . ( Lower Bound of Mixing Time ) Suppose we run LMC for the target measure defined241 in Eq . ( 11 ) from x0 = 12d , then for any choice of step size h > 0 within stability limit , we have242 τmix ( ; W2 ; LMC ) ≥ √ d 8 log √ d = Ω̃ ( √ d ) . Combining Theorem 4.2 and 4.3 , we see that mean-square analysis provides a tight bound for LMC.243 However , there is one limitation of our result – Assumption 2 , which is , although mild , still extra to the244 standard setup . Therefore , the gap between the upper bound and the lower bound of LMC algorithm245 over the entire family of log-smooth and log-strongly-convex target measures is not completely246 closed . We tend to believe that Assumption 2 may not be essential , but rather than an artifact of our247 proof technique . We hope to lift this restriction in future work.248 Comparison At least two sampling algorithms are closely related to LMC . One is Kinetic Langevin249 Monte Carlo algorithm ( KLMC ) , which is discretized kinetic/underdamped Langevin dynamics , and250 the other is Metropolis-Adjusted Langevin Algorithm ( MALA ) which uses the one-step update of251 LMC as a proposal and then accepts/rejects them with a Metropolis-Hastings algorithm.252 The Õ ( √ d ) mixing time in 2-Wasserstein distance of KLMC has been established for log-smooth253 and log-strongly-convex target measures in existing literature ( Cheng et al. , 2018b ; Dalalyan and254 Riou-Durand , 2020 ) . Due to its better dimension dependence over previously best known results of255 LMC , KLMC is understood to be the analog of Nesterov ’ s accelerated gradient method for sampling256 ( Ma et al. , 2021 ) . Our findings show that LMC is able to achieve the same mixing time , albeit under257 an additional growth-at-infinity condition . However , this does not say anything about whether/how258 KLMC accelerates LMC , as the optimality of KLMC bound is not yet clear . We also note KLMC has259 better condition number dependence , although the κ dependence in our bound may not be tight.260 In terms of MALA , a recent work ( Chewi et al. , 2020 ) establishes a Õ ( √ d ) mixing time in 2-261 Wasserstein distance with warm start , and the dimension dependence is shown to be optimal . We see262 that without the Metropolis adjustment , LMC can also achieve the optimal dimension dependence as263 MALA . But unlike LMC , MALA only has logarithmic dependence on 1 . Under warm-start condition,264 is it possible/how to improve the dependence of 1 for LMC , from polynomial to logarithmic ? This265 question is beyond the scope of this paper but worth further investigation.266 5 Numerical Examples267 This section numerically verifies our theoretical findings for LMC in Section 4 , with a particular268 focus on the dependence of the discretization error in Theorem 4.1 on dimension d and step size h.269 To this end , we consider two target measures specified by the following two potentials:270 f1 ( x ) = 1 2 ‖x‖2 + log d∑ i=1 exi and f2 ( x ) = 1 2 ‖x‖2 − 1 2d 1 2 d∑ i=1 cos ( d 1 4xi ) . ( 12 ) It is not hard to see that f1 is 2-smooth and 1-strongly convex , f2 is 32 -smooth and 1-strongly-271 convex . f2 is also used in ( Chewi et al. , 2020 ) to illustrate the optimal dimension dependence272 of MALA . Explicit expression of 2-Wasserstein distance between non-Gaussian distributions is273 1The dependence on κ is not readily available from Theorem 8 in Durmus et al . ( 2019b ) . 2The G = O ( L2 ) assumption is only for κ , m dependence . Removing it does not affect d , dependence . typically not available , instead , we use the Euclidean norm of the mean error as a surrogate because274 ∥∥Ex̄k − Eµx∥∥ ≤W2 ( Law ( x̄k ) , µ ) due to Jensen ’ s inequality . To obtain an accurate estimate of the275 ground truth , we run 108 independent LMC realizations using a tiny step size ( h = 0.001 ) , each till a276 fixed , long enough time , and use the empirical average to approximate Eµx.277 To study the dimension dependence of sampling error , we fix step size h = 0.1 , and for each278 d ∈ { 1 , 2 , 5 , 10 , 20 , 50 , 100 , 200 , 500 , 1000 } , we simulate 104 independent Markov chains using279 LMC algorithm for 100 iterations , which is long enough for the chain to be well-mixed . The mean280 and the standard deviation of the sampling error corresponding to the last 10 iterates are recorded.281 To study step size dependence of sampling error , dimension is fixed to be d = 10 . We experiment with282 step size h ∈ { 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 } × 10−1 . We fix a continuous time T = 20 , and run LMC283 algorithm for dTh e iterations for each h. The procedure is repeated 10 4 times with different random284 seeds to obtain independent samples . When the corresponding continuous time t = kh > 10 , we see285 from Eq . ( 10 ) that LMC is well converged and the sampling error is saturated by the discretization286 error . Therefore , for each h , we take the last d 10h e iterates and record the mean and standard deviation287 of their sampling error.288 The experiment results shown in Figure 1 are consistent with our theoretical analysis of the sampling289 error.Both linear dependence on √ d and h can be identified in and supported by the empirical290 evidence . Note results with smaller h are less accurate because one starts to see the error of empirical291 approximation due to finite samples . Experiments were conducted on a machine with a 2.20GHz292 Intel ( R ) Xeon ( R ) E5-2630 v4 CPU and an Nvidia GeForce GTX 1080 GPU.293 6 Conclusion294 This paper extends the mean-square analysis framework for analyzing the integration error of SDE295 to analyzing the sampling error in 2-Wasserstein distance . Corresponding mixing time bound296 unveils how a high-order numerical algorithm can help improve dependence on accuracy tolerance ,297 and potentially other parameters , such as the dimension . When applied to Langevin Monte Carlo298 algorithm , it obtains an improved and optimal Õ ( √ d/ ) bound , which was previously thought to be299 obtainable only with the addition of momentum.300 Here are some possible directions worth further investigations . ( i ) In data-intensive applications,301 stochastic gradients are typically used for better scalability . It seems natural to apply the mean-square302 analysis framework to study SDE-basd stochastic gradient MCMC methods ; ( ii ) Assumption 2 is303 likely to be an artifact of our analysis ; how to establish the optimal mixing time bound in the standard304 log-smooth and log-strongly-convex setup is still an open question ; ( iii ) Motivated by the recent305 result of MALA ( Chewi et al. , 2020 ) , it would be interesting to know whether the dependence on 1 306 can be improved to logarithmic , for example if LMC is initialized at a warm start.307 References308 Abadi , M. , Barham , P. , Chen , J. , Chen , Z. , Davis , A. , Dean , J. , Devin , M. , Ghemawat , S. , Irving , G.,309 Isard , M. , Kudlur , M. , Levenberg , J. , Monga , R. , Moore , S. , Murray , D. G. , Steiner , B. , Tucker,310 P. , Vasudevan , V. , Warden , P. , Wicke , M. , Yu , Y. , and Zheng , X . ( 2016 ) . Tensorflow : A system311 for large-scale machine learning . In 12th USENIX Symposium on Operating Systems Design and312 Implementation ( OSDI 16 ) , pages 265–283.313 Andrieu , C. , De Freitas , N. , Doucet , A. , and Jordan , M. I . ( 2003 ) . An introduction to mcmc for314 machine learning . Machine learning , 50 ( 1 ) :5–43.315 Chen , C. , Zhang , R. , Wang , W. , Li , B. , and Chen , L. ( 2018 ) . A unified particle-optimization frame-316 work for scalable bayesian sampling . In The Conference on Uncertainty in Artificial Intelligence.317 Cheng , X. and Bartlett , P. L. ( 2018 ) . Convergence of langevin mcmc in kl-divergence . PMLR 83,318 ( 83 ) :186–211.319 Cheng , X. , Chatterji , N. S. , Abbasi-Yadkori , Y. , Bartlett , P. L. , and Jordan , M. I . ( 2018a ) . Sharp con-320 vergence rates for langevin dynamics in the nonconvex setting . arXiv preprint arXiv:1805.01648.321 Cheng , X. , Chatterji , N. S. , Bartlett , P. L. , and Jordan , M. I . ( 2018b ) . Underdamped langevin mcmc:322 A non-asymptotic analysis . Proceedings of the 31st Conference On Learning Theory , PMLR.323 Chewi , S. , Lu , C. , Ahn , K. , Cheng , X. , Gouic , T. L. , and Rigollet , P. ( 2020 ) . Optimal dimension324 dependence of the metropolis-adjusted langevin algorithm . arXiv preprint arXiv:2012.12810.325 Chizat , L. and Bach , F. ( 2018 ) . On the global convergence of gradient descent for over-parameterized326 models using optimal transport . In Advances in neural information processing systems , pages327 3036–3046.328 Dalalyan , A. S. ( 2017a ) . Further and stronger analogy between sampling and optimization : Langevin329 monte carlo and gradient descent . Conference on Learning Theory , pages 678–689.330 Dalalyan , A. S. ( 2017b ) . Theoretical guarantees for approximate sampling from smooth and log-331 concave densities . Journal of the Royal Statistical Society : Series B ( Statistical Methodology ) ,332 79 ( 3 ) :651–676.333 Dalalyan , A. S. and Riou-Durand , L. ( 2020 ) . On sampling from a log-concave density using kinetic334 Langevin diffusions . Bernoulli , 26 ( 3 ) :1956–1988.335 Dashti , M. and Stuart , A. M. ( 2017 ) . The Bayesian Approach to Inverse Problems , pages 311–428.336 Springer International Publishing , Cham.337 Durmus , A. , Majewski , S. , and Miasojedow , B . ( 2019a ) . Analysis of langevin monte carlo via convex338 optimization . Journal of Machine Learning Research , 20:73–1.339 Durmus , A. , Moulines , E. , et al . ( 2017 ) . Nonasymptotic convergence analysis for the unadjusted340 langevin algorithm . Annals of Applied Probability , 27 ( 3 ) :1551–1587.341 Durmus , A. , Moulines , E. , et al . ( 2019b ) . High-dimensional bayesian inference via the unadjusted342 langevin algorithm . Bernoulli , 25 ( 4A ) :2854–2882.343 Erdogdu , M. A. and Hosseinzadeh , R. ( 2020 ) . On the convergence of langevin monte carlo : The344 interplay between tail growth and smoothness . arXiv preprint arXiv:2005.13097.345 Frogner , C. and Poggio , T. ( 2020 ) . Approximate inference with wasserstein gradient flows . In346 International Conference on Artificial Intelligence and Statistics.347 Jordan , R. , Kinderlehrer , D. , and Otto , F. ( 1998 ) . The variational formulation of the fokker–planck348 equation . SIAM journal on mathematical analysis , 29 ( 1 ) :1–17.349 Kloeden , P. E. and Platen , E. ( 1992 ) . Numerical solution of stochastic differential equations . Springer.350 Li , R. , Zha , H. , and Tao , M. ( 2020 ) . Hessian-free high-resolution nesterov acceleration for sampling.351 arXiv preprint arXiv:2006.09230.352 Liu , J. S. ( 2008 ) . Monte Carlo strategies in scientific computing . Springer Science & Business Media.353 Liu , Q. and Wang , D. ( 2016 ) . Stein variational gradient descent : A general purpose bayesian inference354 algorithm . In Advances in neural information processing systems , pages 2378–2386.355 Ma , Y.-A. , Chatterji , N. S. , Cheng , X. , Flammarion , N. , Bartlett , P. L. , and Jordan , M. I . ( 2021 ) . Is356 there an analog of Nesterov acceleration for gradient-based MCMC ? Bernoulli , 27 ( 3 ) :1942 – 1992.357 Mandt , S. , Hoffman , M. D. , and Blei , D. M. ( 2017 ) . Stochastic gradient descent as approximate358 bayesian inference . Journal of Machine Learning Research , 18 ( 134 ) :1–35.359 Milstein , G. N. and Tretyakov , M. V. ( 2013 ) . Stochastic numerics for mathematical physics . Springer360 Science & Business Media.361 Mou , W. , Flammarion , N. , Wainwright , M. J. , and Bartlett , P. L. ( 2019 ) . Improved bounds for362 discretization of langevin diffusions : Near-optimal rates without convexity . arXiv preprint363 arXiv:1907.11331.364 Mou , W. , Ma , Y.-A. , Wainwright , M. J. , Bartlett , P. L. , and Jordan , M. I . ( 2021 ) . High-order365 langevin diffusion yields an accelerated mcmc algorithm . Journal of Machine Learning Research,366 22 ( 42 ) :1–41.367 Pavliotis , G. A . ( 2014 ) . Stochastic processes and applications : diffusion processes , the Fokker-Planck368 and Langevin equations , volume 60 . Springer.369 Robert , C. and Casella , G. ( 2013 ) . Monte Carlo statistical methods . Springer Science & Business370 Media.371 Roberts , G. O. and Rosenthal , J. S. ( 1998 ) . Optimal scaling of discrete approximations to langevin372 diffusions . Journal of the Royal Statistical Society : Series B ( Statistical Methodology ) , 60 ( 1 ) :255–373 268.374 Roberts , G. O. , Tweedie , R. L. , et al . ( 1996 ) . Exponential convergence of langevin distributions and375 their discrete approximations . Bernoulli , 2 ( 4 ) :341–363.376 Rüemelin , W. ( 1982 ) . Numerical treatment of stochastic differential equations . SIAM Journal on377 Numerical Analysis , 19 ( 3 ) :604–613.378 Shen , R. and Lee , Y. T. ( 2019 ) . The randomized midpoint method for log-concave sampling . In379 Advances in Neural Information Processing Systems , pages 2098–2109.380 Vempala , S. and Wibisono , A . ( 2019 ) . Rapid convergence of the unadjusted langevin algorithm:381 Isoperimetry suffices . In Advances in Neural Information Processing Systems , pages 8092–8104.382 Wibisono , A . ( 2018 ) . Sampling as optimization in the space of measures : The langevin dynamics as383 a composite optimization problem . In Conference On Learning Theory , pages 2093–3027.384 Zhang , R. , Chen , C. , Li , C. , and Carin , L. ( 2018 ) . Policy optimization as wasserstein gradient flows.385 In International Conference on Machine Learning , pages 5737–5746.386 Checklist387 1 . For all authors ... 388 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s389 contributions and scope ? [ Yes ] 390 ( b ) Did you describe the limitations of your work ? [ Yes ] Limitation is discussed in the391 paragraph below Table 1.392 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] To the393 best of our knowledge , this work does not have any potential negative societal impact.394 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to395 them ? [ Yes ] 396 2 . If you are including theoretical results ... 397 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] Assumptions398 are fully , clearly stated in theorems.399 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] Complete proofs of400 all theoretical results are provided in supplementary materials.401 3 . If you ran experiments ... 402 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-403 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] All are provided404 in supplementary materials.405 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they406 were chosen ) ? [ Yes ] Our experiments do not need training . We specify all the hyperpa-407 rameters needed to run the experiments.408 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-409 ments multiple times ) ? [ Yes ] See Figure 1.410 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type411 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] We report the hardware used to run412 our experiments.413 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 414 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 415 ( b ) Did you mention the license of the assets ? [ N/A ] 416 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ N/A ] 417 418 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re419 using/curating ? [ N/A ] 420 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable421 information or offensive content ? [ N/A ] 422 5 . If you used crowdsourcing or conducted research with human subjects ... 423 ( a ) Did you include the full text of instructions given to participants and screenshots , if424 applicable ? [ N/A ] 425 ( b ) Did you describe any potential participant risks , with links to Institutional Review426 Board ( IRB ) approvals , if applicable ? [ N/A ] 427 ( c ) Did you include the estimated hourly wage paid to participants and the total amount428 spent on participant compensation ? [ N/A ] 429

Summary: This paper first gives an analysis framework for controlling 2-Wasserstein distance of sampling algorithms coming from discretizing a continuous Langevin dynamics. This is achieved by first considering the bias (local weak error) and mse (local strong error) of single step update.  The multi-step error accumulate exponentially, but together with the contraction guarantee, the maximum sampling error after arbitrary steps could be controlled. Next, this paper introduces an assumption on the linear growth of the 3rd-order derivative which is shown to improve the convergence rate in terms of dependency of dimension. The convergence analysis is carried out based on the general analysis framework.