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 derives an improved 2-Wasserstein convergence bound for the Langevin algorithm (ULA). As noted in the paper, the folk wisdom is that the previously known error bound for the ULA is O(d/epsilon), compared to the underdamped MCMC O(sqrt{d}/epsilon). The paper claims to improve the ULA bound in W-2 distance, matching the case of the underdamped bound. See my detailed comments for questions.