The Poisson Midpoint Method for Langevin Dynamics:  Provably Efficient Discretization for Diffusion Models

Authors: saravanan  kandasamy
dheeraj  nagaraj
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-07 12:41:33

Abstract: Langevin Dynamics is a Stochastic Differential Equation (SDE) central to sampling and generative modeling and is implemented via time discretization. Langevin Monte Carlo (LMC), based on the Euler-Maruyama discretization, is the simplest and most studied algorithm. LMC can suffer from slow convergence - requiring a large number of steps of small step-size to obtain good quality samples. This becomes stark in the case of diffusion models where a large number of steps gives the best samples, but the quality degrades rapidly with smaller number of steps. Randomized Midpoint Method has been recently proposed as a better discretization of Langevin dynamics for sampling from strongly log-concave distributions. However, important applications such as diffusion models involve non-log concave densities and contain time varying drift. We propose its variant, the Poisson Midpoint Method, which approximates a small step-size LMC with large step-sizes. We prove that this can obtain a quadratic speed up of LMC under very weak assumptions. We apply our method to diffusion models for image generation and show that it maintains the quality of DDPM with 1000 neural network calls with just 50-80 neural network calls and outperforms ODE based methods with similar compute.

Main paper: 
 introduction: 
 The task of sampling from a target distribution is central to Bayesian inference, generative modeling, differential privacy and theoretical computer science [48, 20, 17, 25] . Sampling algorithms, based on the discretization of a stochastic differential equation (SDE) called the Langevin Dynamics, are widely used. The straightforward time discretization (i.e., Euler Maruyama discretization) of Langevin dynamics, called Langevin Monte Carlo (LMC), is popular due to its simplicity. The convergence properties of LMC have been studied extensively in the literature under various conditions on the target distribution [8, 12, 11, 45, 15, 33, 7, 1, 6, 9, 16, 30, 53, 5] . LMC can suffer from slow convergence to the target distribution, and often requires a large number of steps with a very fine time discretization (i.e., small step-size), making it prohibitively expensive.The Poisson Midpoint Method introduced in this paper approximates multiple steps of small step-size Euler-Maruyama discretization with one step of larger step-size via stochastic approximation. In the case of LMC, we show that our method (called PLMC) converges to the target as fast as LMC with a much smaller step-size without any additional assumptions such as isoperimetry or strong log concavity (up to a small additional error term). This is a variant of the Randomized Midpoint Method (RLMC) studied in the literature [40, 18, 50] (see Section 1.1 for comparison).Diffusion models are state-of-the-art in generating new samples of images and videos given samples [20, 43, 34] . These start with a Gaussian noise vector and evolve it through the time-reversal of the SDE called the Ornstein-Uhlenbeck process. The time reversed process can be written as an SDE (Langevin Dynamics with a time dependent drift) or as an ODE (see Section 2) . The DDPM scheduler [20] , which discretizes the SDE, obtains the best quality images with a small step-size and a large number of steps (usually 1000 steps). However, its quality degrades with larger step-sizes and a small number of steps (say 100 steps). Schedulers such as DDIM ( [42] ), DPM-Solver ( [28, 29] ), PNDM ( [27] ) which solve the ODE via numerical methods perform much better than DDPM with a small number of steps. However, it is noted that they do not match the performance of DDPM with 1000 steps over many datasets ([42, 28, 41] ). Poisson Midpoint Method gives a scheduler for the time-reversed SDE which maintains the quality of DDPM with 1000 steps, with just 50-80 steps.Euler Maruyama discretization of SDEs is known to be inefficient and many powerful numerical integration techniques have been studied extensively ( [23, 32, 3, 26] ). However, higher order methods such as the Runge-Kutta method require the existence and boundedness of higher order derivatives of the drift. The Randomized Midpoint Method for LMC (RLMC) was introduced for strongly logconcave sampling [40] and was further explored in [50, 18] . It was shown that RLMC, under certain conditions, can sample with a larger step size for fewer steps compared to Euler Maruyama and yet obtain the same accuracy. RLMC is popular due to its simplicity, and ease of implementation and does not require higher order bounded derivatives of the drift function. However, the current theoretical results are restricted to the case of strongly log-concave sampling, whereas non-log-concave sampling is of immense practical interest.(1) We design the Poisson Midpoint Method which discretizes SDEs by approximating K-steps of Euler-Maruyama discretization with step-size α K by just one step with step-size α. We show a strong error bound between these two processes under general conditions in Theorem 1 (no assumption on mixing, smoothness etc). This is based on a Central Limit Theorem (CLT) based method in [10] to analyze stochastic approximations of LMC.(2) We apply our method to LMC to obtain PLMC. We show that it achieves a speed-up in sampling for both Overdamped LMC (OLMC) and Underdamped LMC (ULMC) whenever LMC mixes, without additional assumptions such as isoperimetry or strong-log concavity.(3) When the target obeys the Logarithmic Sobolev Inequalities (LSI), we show that PLMC achieves a quadratic speed up for both OLMC and ULMC. Prior works on midpoint methods [40, 18, 50] only considered strongly log-concave distributions. We also show an improvement in computational complexity for ULMC from 1 ϵ 2/3 to 1 √ ϵ to achieve ϵ error2 . (4) Empirically, we show that our technique can match the quality of DDPM Scheduler with 1000 steps with fewer steps, achieving up to 4x gains in compute. Over multiple datasets, our method outperforms ODE based schedulers such as DPM-Solver and DDIM in terms of quality.X 0:T denotes (X t ) 0≤t≤T and X K(0:T ) denotes (X tK ) 0≤t≤T . I denotes identity matrix over R k×k whenever k is clear from context. For any vector x ∈ R k , ∥x∥ denotes its Euclidean norm. For any a, b ∈ Z and a > b, we take the b t=a to be 0, and the product b t=a to be 1. Underdamped Langevin Dynamics happens in R 2d . Here, we take vectors named X (along with subscripts and superscripts) as X = [U V ]⊺ where U, V ∈ R d also carry the same subscripts and superscripts (e.g: X4 corresponds to Ũ4 , Ṽ4 ). In this case I represents identity matrix in R 2d×2d and I d denotes the identity matrix in R d×d . For any random variable X, we let Law(X) denote its probability measure. By TV(µ, ν) and KL µ ν we denote the total variation distance and KL divergence (respectively) between two probability measure µ, ν. O, Ω, Θ are standard Kolmogorov complexity notations whereas Õ, Ω, Θ are same as O, Ω, Θ up to poly-logarithmic factors in the problem parameters such as 1 ϵ , 1 α , K, T, d. 
problem setup: 
 Given a random vector X 0 ∈ R d , consider an iterative, discrete time process (X t ) t∈N∪{0} over R d , with step-size α > 0 given by:EQUATIONWhere A α , G α , Γ α are d × d matrix valued functions of the step-size α and (Z t ) t≥0 i.i.d.∼ N (0, I d ). b : R d → R d is the drift. Call this process S(A, G, Γ, b, α). We consider Overdamped Langevin Monte Carlo (OLMC), Underdamped Langevin Monte Carlo (ULMC) and DDPMs as key examples.Overdamped Langevin Monte Carlo: Consider Overdamped Langevin Dynamics for some F : R d → R:EQUATIONHere B τ is the standard Brownian motion in R d . Under mild conditions on F and X0 , Law( Xτ )τ →∞ → π * where π ⋆ (X) ∝ exp(-F (X)) is the stationary distribution.Picking A α = I, G α = αI, Γ α = √ 2αI and b(x, tα) = -∇F (x) in Equation (1) gives us Euler-Maruyama discretization of Overdamped Langevin Dynamics: X t in (1) approximates Xαt [38, 35] . OLMC is the canonical algorithm for sampling and has been studied under assumptions such as log-concavity of π ⋆ [8, 12, 11] or that π ⋆ satisfies isoperimetric inequalities [45, 15, 33, 7, 1] .Underdamped Langevin Monte Carlo occurs in 2d dimensions. We writeX t = [U t V t ] ⊺ ∈ R 2dwhere U t ∈ R d is the position and V t ∈ R d is the velocity. Fix a damping factor γ > 0. We take: A h := I d 1 γ (1 -e -γh )I d 0 e -γh I d , G h := 1 γ (h -1 γ (1 -e -γh ))I d 0 1 γ (1 -e -γh )I d 0 b(X t , tα) := -∇F (U t ) 0 Γ 2 h := 2 γ h -2 γ (1 -e -γhThe stationary distribution of the SDE is given by π ⋆ (U, V ) ∝ exp(-F (U )-∥V ∥ 2 2 ) [13, 9] . ULMC is popular in the literature and has been analyzed in the strongly log-concave setting [6, 9, 16] and under isoperimetry conditions [30, 53] . We refer to [53] for a complete literature review.Denoising Diffusion Models: In this case the stochastic differential equation is given by:EQUATIONThis also admits an equivalent characteristic ODE given below:EQUATIONSee [43, 5, 4] for further details. Here p τ is the probability density of e -t X * + √ 1 -e -2t Z where X * is drawn from the target and Z is drawn from N (0, I) independently. The drift ∇ log p τ is learned via neural networks for discrete time instants τ 0 , . . . , τ n-1 (usually n = 1000). In practice, the iterations are written in the form 3 : X t+1 = a t X t + b t ∇ log p τt (X t ) + σ t Z t where a t , b t , σ t are chosen for best performance. Aside from the original choice in [20] , many others have been proposed ( [2, 42] ). Since A α , G α , Γ α in Equation (1) are time independent, we provide a variant of the Poisson Midpoint Method to suit DDPMs in Section A.1, along with a few other optimizations.Scaling Relations: We impose the following scaling relations on the matrices A h , G h , Γ h for every h ∈ R + , n ∈ N, which are satisfied by both OLMC and ULMC. Whenever b(•) is a constant function, these ensure that K steps of Equation (1) with step-size α/K is the same as 1 step with step-size α in distribution.(A h ) n = A hn ; ( n-1 i=0 (A h ) i )G h = G hn ; n-1 i=0 (A h ) i Γ 2 h (A ⊺ h ) i = Γ 2 hnRandomized Midpoint Method and Stochastic Approximation We first illustrate the Randomized Midpoint Method [40, 18, 50] by applying it to OLMC (to obtain RLMC) to motivate our method (the Poisson Midpoint Method) and explain why we expect a quadratic speed up shown in Section 3. Overdamped Langevin Dynamics (2) satisfies:EQUATIONTaking X t as the approximation to Xtα , LMC approximates the integral (t+1)α tα ∇F ( Xs )ds with α∇F (X t ), giving a 'biased' estimator to the integral (conditioned on X t = Xtα ). This gives the LMC updates X t+1 = X t -α∇F (X t ) + √ 2αZ t ; Z t ∼ N (0, I). RLMC chooses a uniformly random point in the interval [tα, (t + 1)α] instead of initial point tα as described below:Let u ∼ Unif([0, 1]), Z t,1 , Z t,2 ∼ N (0, I) be independent and define the midpoint X t+u := X t - uα∇F (X t ) + √ 2uαZ t,1(notice X t+u is an approximation for X(t+u)α ). The RLMC update is:X t+1 = X t -α∇F (X t+u ) + √ 2uαZ t,1 + 2(1 -u)αZ t,2 .Notice that √ 2uαZ t,1 + 2(1 -u)αZ t,2 |u, X t ∼ N (0, I), and E[α∇F (X t+u )|X t , Z t,1 , Z t,2 ] = α 1 0 F (X t+s )ds which is a better approximation of the integral than α∇F (X t ). Therefore RLMC provides a nearly unbiased approximation to the updates in Equation (6) .Intuitively, we expect that reducing the bias leads to a quadratic speed-up. Let Z, Z ′ ∼ N (0, 1) and independent. For ϵ small enough it is easy to show that, KL Law(Z + ϵ) Law(Z) = Θ(ϵ 2 ) whereas KL Law(Z + ϵZ ′ ) Law(Z) = Θ(ϵ 4 ). We hypothesize that Z t + error in integral is closer to N (0, I) when the error term has a small mean and a large variance (as in RLMC) than when it has a large mean but 0 variance (as in LMC). However, rigorous analysis of RLMC has only been done under assumptions like strong log-concavity of the target distribution. This is due to the fact that Z t is dependent on the error in the integral, disallowing the argument above.Our method, PLMC, circumvents these issues by considering a discrete set of midpoints{0, 1 K , . . . , K-1 K } instead of [0, 1].It picks each midpoint with probability 1 K independently, allowing us to prove results under more general conditions using the intuitive ideas described above. Thus, our method is a variant of RLMC which is amenable to more general mathematical analysis. The OPTION 2 of our method (see below) makes this connection clearer. PLMC is naturally suited to DDPMs since the drift function is trained only for a discrete number of timesteps (see Section 2).We introduce the Poisson Midpoint Method (PLMC) which approximates K steps of S(A, G, Γ, b, α K ) with step-size α K with one step of which has a step-size α. We denote this by PS(A, G, Γ, b, α, K) and let its iterates be denoted by ( XtK ) t≥0 or (X P t ) t≥0 . Suppose H t,i ∈ {0, 1} be any binary sequence and Z tK+i be a sequence of i.i.d. N (0, I) for t, i ∈ N ∪ {0}, 0 ≤ i ≤ K -1. Given XtK , we define the interpolation:EQUATIONNote that this interpolation is cheap since every one of XtK+i can be computed with just one evaluation of the function b(). We then define the refined iterates for 0 ≤ i ≤ K -1 for a given t as:EQUATIONWe pick H t,i based on the following two options, independent of Z tK+i , X0 :OPTION 1: H t,i are i.i.d. Ber( 1 K ).OPTION 2: Let u t ∼ Unif({0, . . . , K -1}) i.i.d. and H t,i := 1(u t = i).Remark 1. We call our method Poisson Midpoint method since in OPTION 1 the set of midpoints {αt + αi K : H t,i = 1} converges to a Poisson process over [αt, α(t + 1)] as K → ∞.The Algorithm and Computational Complexity The algorithm PS(A, G, Γ, b, α, K) computes XK(t+1) given XKt in one step by unrolling the recursion given in Equation (8) . For the sake of clarity, we will relabel XtK to be X P t to stress the fact that it is the t-th iteration of PS(A, G, Γ, b, α, K).Step 1:Generate I t = {i 1 , . . . , i N } ⊆ {0, . . . , K -1} such that H t,i = 1 iff i ∈ I t . Let i 1 < i 2 • • • < i N hold.When N = 0, we take this to be the empty set.Step 2: Let M 0 := 0 and let W t,k be a sequence of i.i.d. N (0, I) random vectors. For k = 1, . . . , N, N + 1, we take:M k = A α(i k -i k-1 ) K M k-1 + Γ α(i k -i k-1 ) K W t,k We use the convention that i 0 = 0, i N +1 = K -1, A 0 = I and Γ 0 = 0. Step 3: For k = 1, . . . , N , compute XtK+i k := A αi k K X P t + G αi k K b(X P t , αt) + M kStep 4:Corr := K N k=1 G (K-1-i k )α K (b( XtK+i , (tK+i)α K ) -b(X P t , αt))Step 5:X P t+1 = A α X P t +G α b(X P t , αt)+M N +1 +Corr (9)That is, the algorithm first generates the random mid-points H t,i , computes the interpolation XtK+i only when H t,i = 1 and then computes b( XtK+i , (tK+i)α) for these points. These computations are then combined to compute X(t+1)K . In most applications, it is computationally easy to generate Gaussian random vectors and perform vector operations such as summation. However, the evaluation of the drift function b() is expensive. Therefore, in this work, we consider the number of evaluations of the drift function as the measure of computational complexity. The following proposition establishes that each iteration of PLMC requires 2 evaluations of b() in expectation.Proposition 1. When the scaling relations hold (Section 2.1), the trajectory (X P t ) t≥0 in Equation (9) has the same joint distribution as the trajectory ( XtK ) t≥0 given in Equation (8) . In expectation, one step of PS(A, G, Γ, b, α, K) requires two evaluations of the function b(•).We call PS(A, G, Γ, b, α, K) as PLMC whenever S(A, G, Γ, b, α K ) is either OLMC or ULMC.1.ES 2p (∇F ) ≤ C p (LαN ) 2p ES 2p (V ) + C p T L p (d + log K) p + 2 p N p-1 E(S 2 (∇F )) p (50) 2. ES 2p (V ) ≤ C p (αN ) 2p ES 2p (∇F ) + T γ 2 L (d + log K) p + 2 p N p-1 E(S 2 (V )) p (51)Proof. We consider a t := ∇F ( Ũ(t-1)K ), b t := Ṽ(t-1)K . Let N satisfy the condition in the Lemma. Since T is a power of 2, clearly, N divides T . Let T k = {(k -1)N + 1, . . . , kN } be as defined in Lemma 26. Consider the following upper bound derived using the results of Lemma 23EQUATION)Applying Lemma 26, we conclude:EQUATIONSimilarly, we have:EQUATION)Applying Lemma 26, we conclude:EQUATIONCombining Lemmas 24 and 22 gives the following theorem.Theorem 4. Fix p ≥ 1. There exist constants C p , c p , cp > 0 such that whenever:γ ≥ C p √ L, αγ < c p , α 3p-1 T p-1 L 2p γ p+1< cp , the following results hold:S 2p (∇F ) ≤ C p γ 2p-1 α E∥ Ṽ0 ∥ 2p + E|(F (Ψ 0 ) -F (Ψ T )) + | p + 1 + C p T γ 4p L p + (γαT ) p-1 γ 2p (d + log K) p S 2p (V ) ≤ C p 1 γα E∥ Ṽ0 ∥ 2p + E|(F (Ψ 0 ) -F (Ψ T )) + | p + 1 + C p T γ 2p L p + (γαT ) p-1 (d + log K) p 
main results: 
 Theorem 1 gives an upper bound for the KL divergence of the trajectory generated by PS(A, G, Γ, b, α, K) to the one generated by S(A, G, Γ, b, α K ). We refer to Section C for its proof. We note that Theorem 1 does not make any mixing or smoothness assumptions on b(•) and that it can handle time dependent drifts. We refer to Section 4 for a proof sketch and discussion. Theorem 1. Let X t be the iterates of S(A, G, Γ, b, α K ) and X P t be the iterates of PS(A, G, Γ, b, α, K) with OPTION 1. Suppose that X P 0 = X 0 . Let XtK+i be the iterates in Equation (8) . Define random variables:BtK+i := Γ -1 α K G α K [b( XtK+i, tα + iα K ) -b( XtK+i, tα + iα K )] βtK+i := ∥KΓ -1 α K G α K [b( XtK+i, tα + αi K ) -b( XtK , tα)]∥Then, for some universal constant C and any r > 1:EQUATIONWe now apply Theorem 1 to the case of OLMC and ULMC under additional assumptions, with the proofs in Sections E and F respectively. Assumption 1.F : R d → R d is L-smooth (i.e., ∇F is L-Lipschitz).x * is its global minimizer. Assumption 2. The initialization X 0 is such that E∥X 0 -x * ∥ 14 < C 14 init d 7 . The assumptions above are very mild and standard in the literature. Specifically, Assumption 2 shows that the initialization is close to global optimum by O( √ d) up to 14th moments. For instance, this is satisfied when the initialization is a standard Gaussian variable with mean µ satisfying ∥µ -x * ∥ = O( √ d).Specifically this is true when µ = 0 and ∥x* ∥ = O( √ d). This is a weak assumption which is implied from common initialization assumptions in the literature as listed below. It can be replaced with the assumptions in [53, Appendix D and Lemma 27] which considers Gaussian initializations with the right variance and mean. The original randomized midpoint method work [40] considers initializing at x * whereas [45] considers a Gaussian initialization with the right variance at a local minimum of F .We do not make any assumptions regarding isoperimetry of the target distribution π ⋆ (x) ∝ exp(-F (x)). Theorem 2 (OLMC). Consider the setting of Theorem 1 with OLMC under Assumptions 1 and 2. There exists constants c 1 , c 2 > 0 such that whenever αL < c 1 andEQUATION) Remark 2. There are lower order terms hidden in the O() notation. These are explicated in Equation (40) in the appendix. The next theorem gives a similar guarantee for ULMC and the lower order terms are explicated in Equation 60 in the appendix. Theorem 3 (ULMC). Consider the setting of Theorem 1 with ULMC under Assumptions 1 and 2. Suppose that x * is the global minimizer of F . There exist constants C 1 , c 1 , c 2 such that wheneverEQUATIONOLMC and ULMC are sampling algorithms which output approximate samples (X T and U T respectively) from the distribution with density π ⋆ ∝ e -F . Given ϵ > 0, prior works give upper bounds on T and the corresponding step-size α as a function of ϵ to achieve guarantees such as KL Law(X T ) π ⋆ ≤ ϵ 2 or TV(Law(X T ), π ⋆ ) ≤ ϵ. By Pinsker's inequality TV 2 ≤ 2KL therefore we guarantees for KL ≤ ϵ 2 to those for TV ≤ ϵ as is common in the literature.Quadratic Speedup Let T = Θ(1/α) as is standard. Choosing K = Θ(1/α), our method applied to OLMC achieves a KL divergence of O(α 2 ) to OLMC with step-size α 2 . Similarly, our method applied to ULMC achieves a KL divergence of O(α 4 ) to ULMC with step-size α 2 . Whenever the KL divergence of OLMC (resp. ULMC) output to π ⋆ , with step-size η is Ω(η) (resp Ω(η 2 )) Theorem 2 (resp. Theorem 3) demonstrates a quadratic speed up.To show the generality of our results, we combine Theorems 2 and 3 with convergence results for OLMC /ULMC in the literature ( [45, 53] ) when π ⋆ satisfies the Logarithmic Sobolev Inequality with constant λ (λ-LSI). We obtain convergence bounds for the last iterate of PLMC to π ⋆ under the same conditions. λ-LSI is more general than strong log-concavity (λ-strongly log-concave π ⋆ satisfies λ-LSI). It is stable under bounded multiplicative perturbations of the density [21] and Lipschitz mappings. λ-LSI condition has been widely used to study sampling algorithms beyond log-concavity. We present our results in Table 1 and refer to Section G for the exact conditions and results. for TV ≤ ϵ 4 Proof Sketch Sketch for Theorem 1 For the proof of Theorem 1, we follow the recipe given in [10] in order to analyze the stochastic approximations of LMC -where only an unbiased estimator for the drift function is known. The bias variance decomposition in Lemma 1, shows that the iterations of PS(A, G, Γ, α, K) can be written in the same form of as the iterations of S(A, G, Γ, α K ):XtK+i+1 = A α K XtK+i + G α K b( XtK+i) + Γ α K ZtK+iWhere ZtK+i := Z tK+i + B tK+i + S tK+i , B tK+i is the 'bias' with a non-zero conditional mean, and S tK+i is the variance with 0 conditional mean (conditioned on XtK+i ). They are independent of Z tK+i conditioned on XtK+i . Note that the sequence ( ZtK+i ) t,i is neither i.i.d. nor Gaussian. If it was a sequence of i.i.d. N (0, I), then this is exactly same as S(A, G, Γ, α K ). The main idea behind the proof of Theorem 1 is that due to data-processing inequality, it is sufficient to show that ( ZtK+i ) t,i is close to a sequence of i.i.d. Gaussian random vectors in KL-divergence. The bias term can be shown to lead to an error bounded by t,i E∥B tK+i ∥ 2 , which roughly corresponds to the KL divergence between N (B tK+i , I) and N (0, I). We then show that Z tK+i +S tK+i | Z0:tK+i-1 is close in distribution to N (0, I). In order to achieve this, we first modify the Wasserstein CLT established in [51] to show that Z tK+i + S tK+i is close in distribution to N (0, I + Σ t,i ) when conditioned on Z0:tK+i-1 where Σ t,i is the conditional covariance of S tK+i . This CLT step gives us the error of the formT -1 s=0 K-1 i=0 CE β 10 sK+i K 3 + β 6 sK+i K 2 + β 2r sK+i K in Theorem 1.We then use the standard formula for KL divergence between Gaussians to bound the distance between N (0, I + Σ t,i ) to N (0, I). This accounts for the fact that the Gaussian noise considered has a slightly higher variance than I. This leads to the leading term CEβ 4 sK+i K 2 .Sketch for Theorem 2 Applying Theorem 1 to OLMC , note that the term β tKi depends on how far the coarse estimate XtK+i is from the true value XtK+i . Indeed, under the smoothness assumption on F we show that:β 2p tK+i ≤ ( L 2 αK 2 ) p sup 0≤j≤K-1 ∥ XtK+j -XtK ∥ 2p. Thus:E∥βtK+i∥ 2p ≲ L 2p α 3p K p E∥∇F ( XtK )∥ 2p + L 2p α 2p K p d p .Therefore, the proof reduces to boundingT -1 t=0 E∥∇F ( XtK )∥ 2p .We observe that X(t+1)K = XtK -α(∇F ( XtK ) + ∆ t ) + √ 2αZ t where ∆ t is a small error term appearing due to Poisson Midpoint Method. Notice that this is approximately stochastic gradient descent on F with a large noise √ 2α. Therefore, using the taylor approximation of F , we can show that:EF ( X(t+1)K ) -EF ( XtK ) ≲ -α∥∇F ( XtK )∥ 2 + αd + o(αd) =⇒ T -1 t=0 E∥∇F ( XtK )∥ 2 ≲ F (X0) -infx F (x) α + LT d + o(LT d)The following sophisticated bound derived in this work is novel to the best of our knowledge: T -1 t=0 E∥∇F ( XtK )∥ 2p ≲ L p-1 E (F (X0) -infx F (x)) p α + T L p d p (1 + (αLT )T -1 t=0 E∥∇F ( ŨtK )∥ 2p ≲ γ 2p γα E∥ Ṽ0∥ 2p + E|F (Ψ0) -F (x * )| p + T γ 4p L p + (γαT ) p-1 γ 2p d p T -1 t=0 E∥ ṼtK ∥ 2p ≲ 1 γα E∥ Ṽ0∥ 2p + E|(F (Ψ0) -F (x * )| p + T γ 2p L p + (γαT ) p-1 d p 
perform the update:: 
 Xt-K = A K Xt + K i=1 B K,i b( Xt , t) + K i=1 C K,i Z i + K i=1 pH i B K,i b( X t-i+1 , t -i + 1) -b( Xt , t)whereX t-τ = A τ Xt + τ i=1 B τ,i b( Xt , t) + τ j=1 C τ,K-j+1 Z[j], ∀τ ∈ [K -1]. 5. Return Xt-K . InterpolationConstants(t, k): 1. A k ← a t • a t-1 • a t-k+1 . 2. B k,i = a t-k+1 • • • a t-i • b t-i+1 , if i < k b t-k+1 , if i = k for each i ∈ [k]. 3. C k,i = a t-k+1 • • • a t-i • σ t-i+1 , if i < k σ t-k+1 , if i = k for each i ∈ [k]. 4. Return (A k , B k , C k ). Remark 3.In the theoretically analyzed algorithm given in Section 2.2, we take the midpoints to be {0, 1 K , . . . , K-1 K }. But here, we take the midpoints to be { 1 K , . . . , K-1 K } since the Xt = Xt , making the correction termpH i B K,i b( X t-i+1 , t -i + 1) -b( Xt , t) = 0 when i = 0. 
experiments: 
 We now present experiments to evaluate Poisson Midpoint Method as a training-free scheduler for diffusion models. We consider the Latent Diffusion Model (LDM) [39] for CelebAHQ 256, LSUN Churches, LSUN Bedrooms and FFHQ datasets using the official (PyTorch) codebase and checkpoints. We compare the sample quality of the Poisson Midpoint Method against established methods such as DDPM, DDIM and DPM-Solver, varying the number of neural network calls (corresponding to the drift b(x, t)) used to generate a single image.To evaluate the quality, we generate 50k images for each method and number of neural network calls and compare it with the training dataset. We use Fréchet Inception Distance (FID) [19] metric for LSUN Churches and LSUN Bedrooms. For CelebAHQ 256 and FFHQ, we use Clip-FID, a more suitable metric as it is known that FID may exhibit inconsistencies with human evaluations datasets outside of Imagenet [24] . We refer to Section A.3 in the appendix for further details.We refer to the ODE based sampler with η = 0 (see [42] ) setting as DDIM and use the implementation in [39] . We generate images for number of neural network calls ranging from 20 to 500. For DPM-Solver, we port the official codebase of [28] to generate images for different numbers of neural network calls ranging from 10 to 100 using MultistepDPMSolver and tune the hyperparameter 'order' over {2, 3} and 'skip_type' over {'logSNR', 'time_uniform', 'time_quadratic'} for each instance to obtain the best possible FID score. This ensures that the baseline is competitive.For the sake of clarity, we will call all SDE based methods, including DDIM with η > 0 (see [42] ) as DDPM. The DDPM scheduler has many different proposals for coefficients a t , b t , c t (see Section 2, [42, 2] ), apart from the original proposal in the work of [20] . Based on these proposals, we consider three different variants of DDPM in our experiments (see Section A.5 for exact details). This choice can have a significant impact on the performance (See Figure 1 in the Appendix) for a given number of denoising diffusion steps. For the Poisson Midpoint Method, we implement the algorithm shown in Section A.1 for number of diffusion steps ranging from 20 to 500, corresponding to 40 to 750 neural network calls (see Section A.2). This approximates K steps of the 1000 step DDPM with a single step. For both Poisson Midpoint Method and DDPM, we plot the results from the best variant in Table 2 for a given number of neural network calls and refer to Section A.5 for the numbers of all variants. Poisson Midpoint Method incurs additional noise in each iteration due to the randomness introduced by H i . This can lead to a large error when K is large. When K is large, we reduce the variance of the Gaussian noise to compensate as suggested in the literature (see Covariance correction in [10] and [31, Equation 9] ). We refer to Section A.5 for full details.We refer to the outcome of our empirical evaluations in Table 2 . The first column compares the performance against DDPM. We see that for all the datasets considered, Poisson Midpoint Method can match the quality of the DDPM sampler with 1000 neural network calls with just 40-80 neural network calls. Observe that for CelebA, LSUN-Church and FFHQ datasets, the performance of DDPM degrades rapidly with lower number of steps, showing the advantage our method in this regime. However, a limitation of our work is that the quality of our method degrades rapidly at around 40-50 neural network calls. We believe this is because our stochastic approximation breaks down with larger step-sizes and further research is needed to mitigate this.The second column compares the performance of our method against ODE based methods. It is known in the literature that DDPM with 1000 steps outperforms DDIM and DPM-Solver in terms of the quality for a large number of models and datasets [41, 42] . Thus, in terms of quality, Poisson midpoint method with just 50-80 neural network calls outperforms ODE based methods with a similar amount of compute. Note that we optimize the performance of DPM-Solver over 6 different variants as mentioned above to maintain a fair comparison. However, in the very low compute regime (~10 steps), DPM-Solver remains the best choice. 
conclusion: 
 We introduce the Poisson Midpoint Method, which efficiently discretizes Langevin Dynamics and theoretically demonstrates quadratic speed up over Euler-Maruyama discretization under general conditions. We apply our method to diffusion models for image generation, and show that our method maintains the quality of 1000 step DDPM with just 50-80 neural network calls. This outperforms ODE based methods such as DPM-Solver in terms of quality, with a similar amount of compute. Future work can explore variants of Poisson midpoint method with better performance when fewer than 50 neural network calls are used. An interesting theoretical direction would be to derive convergence bounds for algorithms such as DDPM which have a time dependent drift function. Future work can also consider convergence rates of PLMC under conditions such as the Poincare Inequality and whenever ∇F is Hölder continuous instead of Lipschitz continuous. 
societal impact: 
 Our work considers an efficient numerical discretization schemes for making diffusion model inference more efficient. Publicly available, pre-trained diffusion models are very impactful and have significant risk of abuse. In addition to theoretical guarantees, our work considers empirical experiments to evaluate the inference efficiency on publicly available, widely used diffusion models over curated datasets. We do not foresee any significant positive or negative social impact of our work.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Details regarding the GPU, CPU and the RAM are provided in the paper.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: The paper is purely technical and we do not foresee any societal impact due to this.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification:We use already open sourced models and datasets, widely used in the research community.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification:Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: A.1 PseudocodeThe Algorithm Let a t , b t and σ t be the coefficients of the DDPM denoising step at time t as per Section 2. Let N denote the number of train steps. The main paper used notations and conventions based on the LMC literature. Here we follow different conventions and notations to connect with the diffusion model literature.1. We take the dynamics to go backward in time (i.e., compute X t-1 from X t ) 2. We use b(X t , t) instead of b(X t , αt) to denote the neural network estimate of∇ log p τ N -t-1 (X t ).DDPM scheduler samples X N -1 ∼ N (0, I) and iteratively computes X 0 as:X t-1 = a t X t + b t b(X t , t) + σ t Z t ; Z t ∼ N (0, I)We take N = 1000 as is standard in the literature. The Poisson Midpoint Method approximates K steps of the iteration above, with a single step. That is, it obtains an approximation for X t-K directly from X t . The number of iterations deployed by Poisson Midpoint Method to sample X 0 is N/K. We compute the number of neural network calls required in Section A.2. The exact description of Poisson Midpoint Method is given below.PoissonMidpointMethod( Xt , t, K, b(, ), OPTION):1. Let (A k , B k , C k ) ← InterpolationConstants(t, k), for every k ∈ [K]. 2. If OPTION = 1: Define H i ∼ Ber(1/K) i.i.d for every i ∈ [K -1]. Define p = K. If OPTION = 2: Define H i ← 1(u = i) where u ∼ Unif({1, . . . , K -1}) i.i.d. Define p = K -1. 3. For every i ∈ [K], Z i ∼ N (0, I).While each step of Poisson Midpoint Method approximates K steps of DDPM, it can use more than one neural network calls. Since neural networks calls are the most computationally expensive parts of diffusion models (4 orders of magnitude larger FLOPs compared to other computations), we compute the number of neural network calls for the Poisson Midpoint Method given K, N for both OPTION 1 and OPTION 2.For OPTION 1, b( Xt , t) is always computed. b( Xt-i+1 , t -i + 1) needs to be computed iff H i = 1. It is easy to show that E|{i ∈ [K -1] : H i = 1}| = K-1K . Thus, the expected number of neural network calls per step is 2 -1 K . The total number number of neural network calls is N K (2 -1 K ). OPTION 2 always computes evaluates b() twice per step. Therefore, the number of neural network calls is 2N K .FID scores between two image datasets can vary greatly depending on the image processing details which do not have much bearing on visual characteristics (such as png vs jpeg) [37] . Thus, we utilize the standardized Clean-FID codebase [37] to compute the scores for all the datasets. As is standard, we generate 50k images for all our evaluations. For CelebAHQ 256, we evaluate the CLIP-FID scores against the combined 30k training and validation samples of the CelebA 1024x1024 dataset [22] , where the CLIP-FID is invoked by using the flags model_name = clip_vit_b_32 and mode = clean [37] . For FFHQ 256, we utilize the precomputed statistics of [37] on the 1024 resolution trainval70k split with the flags dataset_split = trainfull and mode=clean. As for LSUN Churches [49] , we utilize the precomputed statistics of [37] with the flags dataset_split = trainfull and mode=clean to calculate the FID. And for LSUN Bedrooms [49] , we downloaded the training split of [39] which consists of 3,033,042 images and used the codebase of [37] to compute the FID between two folders.For our experiments, we utilized a high-performance computational setup consisting of 1a , Variant 2 outperforms the rest of the variants for steps smaller than 125, while Variant 1 achieves the best performance for larger steps. For FFHQ 256, similar phase transitions can be observed in Figure 1c at around 50-100 steps. For LSUN Bedrooms, we observed that Variant 2 outperforms 1 and 3 on all the regimes. Our experiments were carried out on all the variants and on all the steps, and the best scores are considered for comparison with Poisson Midpoint Method in Table 2 . Let α t , β t , ᾱt be as defined in [20] .Variant 1: The first variant uses the following closed form expression to perform the update.EQUATION)Variant 2: The second variant of DDPM uses the default coefficients of DDIM as provided in the codebase of [39] , with η = 1 and σ 2 t = (1-αt)•βt+1 (1-αt+1) . Variant 3: The third variant of DDPM also uses the default coefficients of DDIM with η = 1, but uses a modified lower variance [31, 10] . Specifically, for Variant 1, we consider four sub-variants 1a, 1b, 1c, 1d that respectively correspond to variances σ 2 t = βt 1+i•βt for i = {-1, 0, 1, 2}, with the rest of the coefficients a t , b t defined as in (13) . Note that the variance is inversely proportional to i. We observe that lower variance leads to faster convergence when the number of diffusion steps is small. However, higher variances yield slightly better scores at higher steps, better quality at the expense of increased neural network evaluations. This phenomenon can be observed in Tables 3, 4 , 5 and 6 where the optimal score for each step-size/neural network calls is highlighted in bold. Therefore, we evaluate samples on all of the above mentioned variants and choose the best variant for every number of steps/neural network calls. We see that whenever K is small (2-4), then OPTION 1 uses fewer neural network calls than OPTION 2. However, the without replacement sampling technique in OPTION 2 incurs lower variance error in the updates, allowing better convergence with lower number of steps (See Section A.2). Thus, we use OPTION 2 for smaller number of steps and OPTION 1 for larger number of steps.σ 2 t = (1 -α t ) • β t+1 .For unconditional image generation of high-resolution images, third order (order=3) Multistep DPM-Solvers with uniform time steps are recommended. To ensure a competitive benchmark, we evaluate the DPM-Solver on six different settings by tuning the parameters 'order' and 'skip_type'. Our evaluations are shown in Table 7 , where the best score for every setting, highlighted in bold, are chosen in our plots for comparison. We use the official code of [37] with parameters 'algorithm_type = dpmsolver' 'method = multistep' and vary the parameters 'order' and 'skip_type' accordingly. The proof relies on the Wasserstein CLT based approach introduced in [10] to compare two discrete time stochastic processes. For the sake of clarity we consider the drift function b(x, t) to be b(x) (i.e., time invariant). However, the proof goes through for time varying drifts as well.C.1 The Bias-Variance Decomposition:In the Lemma below, we will rewrite the update equations for XtK+i given in Equation ( 8) in the same form as the update equations for S(A, G, Γ, α K ) given in Equation ( 1). The lemma follows by re-arranging the terms in Equation ( 8). Lemma 1. We can writeXtK+i+1 = A α K XtK+i + G α K b( XtK+i ) + Γ α K ZtK+iWhere ZtK+i := Z tK+i + B tK+i + S tK+i such that 'bias' B tK+i is defined asB tK+i := Γ -1 α K G α K [b( XtK+i ) -b( XtK+i )]and the variance S tK+i is defined as:S tK+i := K(H t,i -1 K )Γ -1 α K G α K [b( XtK+i ) -b( XtK )]By Equation (1), the iterates of S(A, G, Γ, α K ) satisfy:EQUATIONBy Lemma 1, we can write down the interpolating process as:EQUATIONNote that while Z tK+i are i.i.d. N (0, I) vectors, ZtK+i can be non-Gaussian and non-i.i.d. Below we will show that the KL divergence between the joint laws of (X tK+i ) t,i and ( XtK+i ) t,i can be bounded by bounding the KL divergence between (Z tK+i ) t,i and ( ZtK+i ) t,i . We now state the following standard results from information theory. Lemma 2 (Chain Rule for KL divergence). Let p, q be two probability distributions over X × Y where X and Y are polish spaces. Let p x , q x denote their respective marginals over X and let p y|x , q y|x denote the conditional distribution over Y conditioned on x ∈ X . Then, We refer to [44, Lemma 4.18] for a proof. Lemma 3 (Data Processing Inequality). Let F : R k1 → R k2 be any measurable function. Let P, Q be any probability distributions over R k1 . Then, the following inequality holds:KL p q = KL p x q x + E x∼px KL p y|x q y|xKL F # P F # Q ≤ KL P QThe lemma below follows from the fact that Equation ( 14) and Equation ( 15) have the same functional form.Lemma 4. Suppose S(A, G, Γ, α K ) and PS(A, G, Γ, α, K) are initialized at the same point X 0 . That is, X 0 = X0 . Then, there exists a measurable function F T such that the following hold almost surely.(X τ ) 0≤τ ≤T = F T (X 0 , Z 0 , . . . , Z T -1 )( Xτ ) 0≤τ ≤T = F T (X 0 , Z0 , . . . , ZT -1 )By Lemma 3 and Lemma 4, we have:EQUATIONIn the last step, we have used the chain rule (Lemma 2).We now seek to control the individual KL-divergences in the RHS of Equation ( 16) via Wasserstein distances. We re-state [10, Lemma 26]: Lemma 5. Suppose Z ∼ N (0, σ 2 I). Let A and B be random variables independent of Z. Then,KL Law(Z + A) Law(Z + B) ≤ 1 2σ 2 W 2 2 (A, B) .The Lemma below follows from the triangle inequality for Wasserstein distance. Lemma 6. Let A, B, C be random vectors over R k with finite second moments. Then,W 2 (Law(A), Law(B + C)) ≤ W 2 (Law(A), Law(B)) + E∥C∥ 2From Lemma 1, we show that Zτ = Z τ + B τ + S τ . Conditioned on Z0:τ-1 , X 0 , Z τ ∼ N (0, I) and is independent of B τ , S τ . We now use Gaussian splitting: if A 1 , A 2 ∼ N (0, I) i.i.d, then A1+A2 √ 2 ∼N (0, I). Therefore, letting Z τ,1 , Z τ,2 be i.i.d from N (0, I) independent of Z0:τ-1 , X 0 , S τ , B τ , we can write:EQUATION2 ), Law(EQUATIONIn the second step, we have used the inequality in Lemma 5. In the third step we have used Lemma 6.In Equation ( 17), we note that Z τ,2 is Gaussian but S τ need not be Gaussian (but has zero mean). In the lemma stated below, we show that Z τ,2 + √ 2S τ | Z0:τ-1 , X 0 is close to a Gaussian of the same variance by adapting the arguments given in the proof of [51, Lemma 1.6 ]. The proof is defered to Section D Lemma 7. Suppose N is random vector such that the following conditions are satisfied:1. ∥N∥ ≤ β almost surely, EN = 0 and ENN ⊺ = Σ.2. N takes its value in a one dimensional sub-space almost surely. Suppose Z ∼ N (0, I) is independent of N. Let P = Law( √ I + ΣZ) and Q = Law(Z + N). Denoting ν := Tr(Σ). Then,EQUATIONWe also have the crude bound:EQUATIONThe next lemma demonstrates the convexity of the Wasserstein distance, which is a straightforward consequence of the Kantorovich duality [47] . Lemma 8 (Convexity of Wasserstein Distance). Suppose µ is a measure over R d and Q(•, •) be a kernel over R d with respect to some arbitrary measurable space Ω and let M be a probability measure over Ω. That is Q(•, ω) is a probability distribution over R d for every ω ∈ Ω.W 2 µ, Q(•, ω)dM (ω) ≤ W 2 (µ, Q(•, ω)) dM (ω)Lemma 9 (Wasserstein Distance Between Gaussians, [36] ).EQUATION2 )Lemma 10. Let P := Law(Z Kt+i,2 ) = N (0, I) ,Q := Law(Z Kt+i,2 + √ 2(S Kt+i )|X 0 , Z0:Kt+i-1 )Define the random variableEQUATIONProof. Note that E[√ 2S tK+i | Z0:Kt+i-1 , Y 0:Kt+i-1 , X 0 ] = 0Now, define the random variable βEQUATIONClearly, β is measurable with respect to the sigma algebra of Z0:Kt+i-1 , Z 0:Kt+i-1 , X 0 . By the definition of S tK+i , it is clear that ∥S tK+i ∥ ≤ β almost surely. Define the conditional covariance (only in this proof) to be:EQUATIONThus, almost surely:Tr(Σ) ≤ 2β 2 K √2S tK+i takes its values in a one-dimensional sub-space almost surely when conditioned on Z0:Kt+i-1 , Z 0:Kt+i-1 , X 0 . It is independent of Z tK+i conditioned on Z0:Kt+i-1 , Z 0:Kt+i-1 , X 0 .Let us now define the following random probability distributions measurable with respect to the sigma algebra of Z0:Kt+i-1 , X 01. Q := Law(Z Kt+i,2 + √ 2S Kt+i | Z0:Kt+i-1 , Z 0,Kt+i-1 , X 0 ) 2. P := Law( √ I + ΣZ Kt+i,2 | Z0:Kt+i-1 , Z 0,Kt+i-1 , X 0 )First, note that by Lemma 8 and Jensen's inequality, we have:EQUATIONFirst consider W 2 2 P , P . Conditioned on Z0:Kt+i-1 , Z 0,Kt+i-1 , X 0 , Σ has at-most one non-zero eigenvalue. Without loss of generality, we can take Σ = νe 1 e ⊺ 1 for the calculations below. From Lemma 9, we conclude that:EQUATION)In the last step, we have used the following inequality which follows from the mean-value theorem:√ 1 + ν ≥ 1 + ν 2 -ν 2 8 .We check that the conditions for Lemma 7 hold for P, Q (almost surely conditioned on Z0:Kt+i-1 , X 0 ) and ν ≤ β 2 K to conclude:EQUATIONCombining this with Equation ( 24), we conclude that for any r ≥ 1:EQUATIONWe combine Equations ( 16) and (17) with Lemma 10 to conclude the result of Theorem 1.The crude bound in Equation ( 19) follows via a naive coupling argument. We will now sketch a proof Equation ( 18) in Lemma 7 by showcasing how to modify the proof of [52, Lemma 1.6]. Our proof deals with a specialized case compared to [51, Lemma 1.6] . This allows us to derive a stronger result. In this section, by 'original proof', we refer to the proof in [52] .Since N is supported on a single dimensional sub-space almost surely, we take this direction to be e 1 almost surely without loss of generality. We thus take the covariance matrix of N to be Σ = νe 1 e ⊺ 1 . We generate jointly distributed random vectors Z, Z ′ ∈ R d as follows: Let ⟨Z, e j ⟩ are i.i.d. standard normal random variables for j = 2, . . . , d and ⟨Z ′ , e j ⟩ = ⟨Z, e j ⟩ almost surely. We generate ⟨Z ′ , e 1 ⟩ to be standard normal independent of (⟨Z, e j ⟩) j≥2 . Let ⟨N, e 1 ⟩ be independent of Z ′ . We draw ⟨Z, e 1 ⟩ to be standard normally distributed and Wasserstein-2 optimally coupled to ⟨Z ′ ,e1⟩+⟨N,e1⟩ √ 1+ν and independent of all other random variables mentioned above. We can easily check that ( √ I + ΣZ, Z ′ + N) as defined above is a coupling between P, Q.Via this coupling, we conclude that:EQUATIONLet Z 1 := ⟨Z, e 1 ⟩ and Z ′ 1 := ⟨Z ′ , e 1 ⟩, N 1 := ⟨N, e 1 ⟩. We define m = 1+ 1 ν (and note throughout the proof that a value of infinity can be easily handled). We see thatZ ′ 1 +N1 √ 1+ν = 1 -1 m Z ′ 1 + 1 -1 m N 1Note that 1 -1 m N is denoted by Y in the original proof and the bound is∥Y ∥ ≤ β √ n instead of ∥N∥ ≤ β in this proof.Consider the function f (x), which denotes the ratio of density function ofEQUATIONWhere N ′ 1 is an i.i.d copy of N 1 and r(n) := 1 2(n 2 -1) -1 2 log(1 + 1 n 2 -1 ). Define Q 1 = -(N 2 1 +(N ′ 1 ) 2 )+2mN1N ′ 1 2(m+1) + 1 2(m 2 -1) -r(m).We modify the estimates in Lemma 4.4 and Lemma 4.5 in the original proof in the following:Lemma 11. 1. |Q 1 | ≤ m|N1N ′ 1 | m+1 + β 2 m+1 + 1 2(m 2 -1) almost surely 2. E[Q 1 ] = -1 2(m 2 -1) -r(m) 3. E[Q 2 1 ] ≤ mβ 2 +2m 2 +12(m 2 -1) 2Proof.1. Note that r(m) ≤ 1 2(m 2 -1) . By triangle inequality, we have:EQUATION2. A direct calculation shows the identity for EQ 1 .3. Now, consider EQ 2 1 . We follow the proof of [52, Lemma 4.5], to conclude the following inequalities. Since EQ 1 ≤ 0 and 1 2(m 2 -1) -r(m) ≥ 0, we have:EQUATIONNote that the parameter σ i found in the original proof satisfies σ i = 1 in our proof. We now proceed with the proof of Lemma 7. DefineR(Q) = exp(Q) -1 -Q -Q 2 2 .From the original proof, we conclude via the Talagrand transport inequality that:EQUATIONFrom Lemma 11, and using the fact that r(m) ≥ 0, we note that:E[Q 1 ] + 1 2 E[Q 2 1 ] ≤ 3 + mβ 2 4(m 2 -1) 2Now, let us bound ER(Q 1 ). Via a straightforward application of the taylor series, we have almost surely:EQUATIONFrom Lemma 11, we conclude that|Q 1 | ≤ β 2 + 1 2(m 2 -1) . Since ν ≤ β 2 , we must have m = 1 + 1 ν ≥ 1 + 1 β 2 . Thus, we have |Q 1 | ≤ 3β 22 almost surely. Using this in Equation ( 32), we have:EQUATIONIn the last step, we have used item 3 of Lemma 11 along with the fact that mβ 2 ≥ 1.Combining these estimates with Equation ( 31), we conclude:W 2 ⟨Z, e 1 ⟩, ⟨Z ′ , e 1 ⟩ + ⟨N, e 1 ⟩ √ 1 + ν 2 ≤ mβ 4 + m 2 β 2 2(m 2 -1) 2 exp( 3β 2 2 ) + 3 + mβ 2 2(m 2 -1) 2Now, we use the fact that m = 1 + 1 ν , which implies 1 m 2 -1 ≤ 1 (m-1) 2 ≤ ν 2 . Thus, we have:EQUATIONUsing the fact that β 2 ≥ ν, we have:3ν 4 +β 2 ν 3 2 ≤ 2β 2 ν 3 . Thus, 2(1 + ν)β 2 ν 3 ≤ 2β 2 ν 3 + 2β 4 ν 3 ≤ 2(1 + ν)(β 2 ν 2 + β 4 ν 3 ). Plugging this into Equation ( 27), we conclude the result.In this section, we will prove Theorem 2 after developing some key results regarding OLMC. Recall that in this case b(x, τ ) = -∇F (x) for some F :R d → R, G α K = α K I and Γ α K = 2α K I. Let N t := K-1 j=0 H t,j. Throughout this section, we assume that ∇F is L-Lipschitz. We will assume that T is a power of 2 in this entire section, which useful to apply Lemma 26. The results hold for any T by considering the closest power of 2 above T instead.Recall B tK+i from Theorem 1. Instantiating this for Overdamped Langevin Dynamics, we conclude that for any t, i ∈ N ∪ {0} and 0 ≤ i ≤ K -1, we must have:EQUATIONHere, we have used the fact that ∇F is L-Lipschitz. Now, consider β tK+i in Theorem 1. For any p ≥ 1, t, i ∈ N ∪ {0} and 0 ≤ i ≤ K -1, we have:EQUATIONIn order to apply Theorem 1, we will now proceed to bound the quantities in the RHS of Equations ( 35) and ( 36)The following lemma gives us almost sure control over quantities of interest. We refer to Section I.1 for the proof.Lemma 12. Suppose the stepsize αL < 1. Let M t,k := 2α K sup 0≤i≤k-1 i j=0 Z tK+j . Then the following hold almost surely for every 0≤ k ≤ K 1. sup 0≤i≤k-1 ∥ XtK+i -XtK ∥ ≤ α∥∇F ( XtK )∥ + M t,k 2. sup 0≤i≤k ∥ XtK+i -XtK+i ∥ ≤ αLN t sup i≤k-1 ∥ XtK+i -XtK ∥ 3. E[M p t,K ] ≤ C(p)(αd) p 2We will now prove the following growth estimate for the trajectory XtK . We refer to Section I.2 for its proof. Lemma 13. Let p ≥ 1 be fixed. There exists a large enough constant Cp which depends only on p such that whenever s -t ≤ 1 αL Cp , we have:sup t≤h≤s E∥ XhK -XtK ∥ p | XtK 1 p ≤ 3α(s -t)∥∇F ( XtK )∥ + C p αd(s -t)We apply Lemma 26 along with Lemma 13 to conclude the following result which is proved in Section I.4. Lemma 14. There exists a constant c p such that whenever αL ≤ c p , we have:EQUATIONWhile Lemma 14 controlledT t=1 E∥∇F ( XtK )∥ 2p in terms of T t=1 ∥∇F ( XtK )∥ 2, in the Lemma below we control T t=1 ∥∇F ( XtK )∥ 2 . We refer to Section I.5 for its proof. Lemma 15. Suppose p ≥ 1 be arbitrary. There exists c p > 0 small enough such that whenever αL < c p , we have for some constant C p > 0 depending only on p:EQUATIONWe combine the results of Lemma 14 and Lemma 15 to conclude the following result:Lemma 16. Given p ≥ 1 arbitrary, there exists a constant c p > 0 depending only p such that whenever αL < c p and α 3p-1 L 3p-1 T p-1 < c p , we must have:T t=1 E∥∇F ( XtK )∥ 2p ≤ C p L p d p T (1 + (αLT ) p-1 ) + C p L p-1 α E|(F (X 0 ) -F ( XKT )) + | p + 1 Lemma 17. 1. E∥B tK+i ∥ 2 ≤ CL 4 α 5 K E∥∇F ( XtK )∥ 2 + CL 4 α 4 K d 2. Eβ 2p tK+i ≤ C p L 2p α 3p K p E∥∇F ( XtK )∥ 2p + L 2p α 2p K p d pProof.1. Using Equation ( 35), we have:EQUATION2. We use Equation ( 36) and proceed as in item 1.For p ≥ 1, we define ∆ (p) := E[(F (X 0 ) -F (x * )) p ]. Let ≲ denote ≤ up to a universal positive multiplicative constant on the RHS. We now combine Lemma 17, Lemma 15 with Theorem 1 (taking r = 7) to conclude the following bound:KL Law((X P t ) 0≤t≤T ) Law((X Kt ) 0≤t≤T ) ≲ L 4 α 4 (∆ (1) + 1) + L 5 α 5 K(∆ (2) + 1) + L 8 α 8 K 2 (∆ (3) + 1) + L 14 α 14 K 3 (∆ (5) + 1) + L 20 α 20 K 7 (∆ (7) EQUATIONIn this section, we will prove Theorem 3 after developing some key results regarding ULMC. We will assume that T is a power of 2 in this entire section, which useful to apply Lemma 26. The results hold for any T by considering the closest power of 2 above T instead. Recall that in the case of ULMC, we have:A h := I d 1 γ (1 -e -γh )I d 0 e -γh I d , G h := 1 γ (h -1 γ (1 -e -γh ))I d 0 1 γ (1 -e -γh )I d 0 b(X t ) := -∇F (U t ) 0 Γ 2 h := 2 γ h -2 γ (1 -e -γh ) + 1 2γ (1 -e -2γh ) I d 1 γ (1 -2e -γh + e -2γh )I d 1 γ (1 -2e -γh + e -2γh )I d (1 -e -2γh )I dThroughout this section, we assume that ∇F (•) is L-Lipshitz. We let N t := K-1 i=0 H t,i . We let Π be the projector to the first d-dimensions in the standard basis. That is, Π := I d 0 0 0 .The 2d dimensional space can be decomposed into position (sub-space spanned by the first d standard basis vectors) and velocity (sub-space spanned by the last d standard basis vectors) sub-spaces. We use the convention that X ∈ R 2d is such that X = U V where U ∈ R d is the position and V ∈ R d is the velocity. Throughout this section, we implicity let Xτ = Ũτ Ṽτ . With some abuse of notation, we let U = ΠX and V = (I -Π)X.The following lemma collects some useful bounds on A h , G h and Γ h , and is proved in Section I.6 Lemma 18.G ⊺ h Γ -2 h G h ⪯ C h exp(2γh) γ I d 0 0 0 For all 0 ≤ j ≤ K -1, ∥ΠA jα K G α K ∥ ≤ 3α 2 2K exp( γα K )Analogous to the analysis of Overdamped Langevin Dynamics, we consider the following dynamics.Using the scaling relations in Section 2.1, we write down the iteration of PS(A, G, Γ, b, α, K) for Underdamped Langevin Dynamics as:EQUATIONWhere, α∆ t :=K-1 i=0 KH i A α(K-i-1) K G α K b( XtK+i ) -b( XtK ) and Γ α Ȳt := K-1 j=0 A α(K-1-j) K Γ α K Z tK+j .Applying the triangle inequality, we conclude the following lemma. Lemma 19. Suppose that ∥ • ∥ any is any semi-norm over R d . Then the following hold almost surely for every 0≤ k ≤ K 1. ∥ XtK+i -XtK ∥ any ≤ ∥(A iα K -I) XtK ∥ any + ∥G iα K b( XtK )∥ any + ∥ i-1 j=0 A α(i-1-j) K Γ α K Z tK+j ∥ any 2. ∥ XtK+i -XtK+i ∥ any ≤ i-1 j=0 KH j ∥A (i-1-j)α K G α K (b( XtK+i ) -b( XtK ))∥ any Define ψ t := ŨKt + ṼKt γ .The following Lemma, proved in Section I.7, gives the time evolution of ψ t . Lemma 20.ψ t+1 -ψ t = - α γ ∇F ( ŨtK ) - K-1 i=0 H i α γ ∇F ( ÛtK+i ) -∇F ( ŨtK ) + ΨtWhere Ψt ∼ N (0, 2α γ I d ) is independent of XtK but not necessarily indepependent of ÛtK+i for i > 0.Lemma 21. Consider the seminorm given by ∥x∥EQUATION2.EQUATION3. For any p ≥ 1, we have:EM p t,K ≤ C exp( pγα 2 )γ p 2 α 3p 2 (d + log K) p 2We refer to Section I.8 for the proof. Analogous to Section E, for any p ≥ 1, we defineS 2p (V ) := T -1 t=0 ∥ ṼtK ∥ 2p , S 2p (∇F ) := T -1 t=0 ∥∇F ( ŨtK )∥ 2p .The following lemma, proved in Section I.9, bounds the moments of these quantities. Lemma 22. There exists a constant c 0 such that whenever γα < 1, αL γ < c 0 . Then the following relationships hold:1.EQUATION[E(S 2 (∇F )) p ]EQUATIONSuppose additionally that γ ≥ C 0 √ L for some large enough universal constant C 0 . Then, the bounds above imply the following inequalities:1.EQUATION[E(S 2 (∇F )) p ]EQUATIONThe following lemma gives a growth bound for Overdamped Langevin dynamics. We give the proof in Section I.10. Lemma 23. Let s, t ∈ N ∪ {0} such that s > t. There exists a constant c p > 0 such that whenever αL(s-t) γ ≤ c p , αγ(s -t) ≤ c p and α 2 L(s -t) ≤ c p , the following statements hold:EQUATION)EQUATIONWe combine Lemma 23 with Lemma 26 to conclude: Lemma 24. Let p ≥ 1 be given. There exists c p > 0 such that for any N which satisfies the following conditions:1. N is an integer power of 2 and N ≤ T . The following satements hold:Now, consider the following term in the statement of Theorem 1:B tK+i := Γ -1 α K G α K [b( XtK+i ) -b( XtK+i )]Using Lemma 18, and under the conditions of Theorem 4 holding with p = 1, we have:EQUATIONIn the second step, we have used item 2 of Lemma 19 along with the bounds in Lemma 18.Applying item 1 from Lemma 21 and Theorem 4 we conclude:EQUATIONIn the first step we have used the fact that N t is independent of sup 0≤i≤K-1 ∥ ÛtK+i -ŨtK ∥ 2 and the second step follows from the fact that EN 2 t ≤ 2. In the third step we have used item 1 from Lemma 21. In the fourth step, we have used item 3 of Lemma 21. In the final step we have used Theorem 4.Now consider β tK+i = ∥KΓ -1 α K G α K [b( XtK+i ) -b( XtK )]∥.Using Lemma 18 and item 1 of Lemma 21, we have:EQUATIONTherefore, we note that for any p ≥ 1, under the assumptions of Theorem 4, we must have:EQUATIONFor p ≥ 1, we define∆ (p) = E(F (U 0 + V0 γ ) -F (x * )) p + E∥V 0 ∥ 2p + 1.By LHS ≲ RHS we denote LHS ≤ C.RHS for some universal positive constant C. We now apply Theorem 1 (with r = 7) along with the bounds in Equations ( 57) and (59):EQUATIONWe now prove the results concerning convergence under the assumption that π ⋆ satisfies the Logarithmic Sobolev Inequalities (LSI). We first define the LSI by following the discussion in [46] and refer to this work for further details. Definition 1. We say that a measure π ⋆ over R d satsifies λ-LSI for some λ > 0, if for every smooth g : R d → R such that E X∼π ⋆ g 2 (X) < ∞, the following inequality is satisfied:E π ⋆ g 2 log g 2 -E π ⋆ g 2 log E π ⋆ g 2 ≤ 2 λ E π ⋆ ∥∇g∥ 2We note that whenever the density π ⋆ (x) = e -F (x) and F is λ strongly convex, then π ⋆ satisfies the λ-LSI. It is well known that λ-LSI implies the Poincare inequality (PI) with parameter λ defined below: Definition 2. We say that a measure π ⋆ over R d satsifies λ-PI for some λ > 0, if for every smooth g : R d → R such that E X∼π ⋆ g 2 (X) < ∞, the following inequality is satisfied:E π ⋆ g 2 -(E π ⋆ g) 2 ≤ 1 λ E π ⋆ ∥∇g∥ 2We note the following useful lemma: Lemma 25. Let π ⋆ (x) ∝ exp(-F (x)) satisfy λ-PI or λ-LSI and let F be L-smooth. Then, L ≥ λ.Proof. Let X, X ′ ∼ π ⋆ be i.i.d. Applying integration by parts, we have: E∇F (X) = 0 and E⟨X, ∇F (X)⟩ = d. Note that LSI implies PI and under λ-PI, we must have for any unit norm vector v ∈ R d :var(⟨X, v⟩) ≤ 1 λLet Σ be the covariance of X (it exists due to the assumption of PI). Therefore,EQUATIONIn the second step we have used the fact that whenever F is L-smooth, then F (x) -F (y) ≤ ⟨∇F (y), x -y⟩ + L 2 ∥x -y∥ 2 for all x, y ∈ R d .We now state [45, Theorem 1] (adapting to our notation) which gives convergence guarantees for OLMC under the assumption of λ-LSI. Theorem 5. Suppose π * (x) ∝ exp(-F (x)) satisfies the λ-LSI and that F is L-smooth. Consider OLMC with set size η which satsifies 0 < η ≤ λ 2L 2 . Then, LMC satisfies:KL Law(X N ) π ⋆ ≤ e -ληN KL Law(X 0 ) π ⋆ + 8dL 2 η λ .Where X 0 , . . . , X N are the iterates of OLMC.Whenever ϵ ≤ 1, we take α K = λϵ 2 16dL 2 . T = 1 αλ log 2KL Law(X0) π ⋆ ϵ 2. With the suggested initialization in [45] i.e X 0 ∼ N (x, I L ) for any stationary pointx of F we have KL Law(X 0 ) π ⋆ ≤ Õ(d). Now, take α = c 0 min ϵ √ λ L 3 2 d 3 4, λ , T = Õ( C1 λα ) andK = C2 αL √ d .Let X 0 , . . . , X T K be the iterates of OLMC with step-size α K with the parameters as chosen above. Applying Theorem 5, we conclude: KL Law(X T K ) π ⋆ ≤ ϵ 2 8 . Let X P 0 , . . . , X P T be the iterates of PS(A, G, Γ, α, K) with A, G, Γ corresponding to OLMC. Applying Theorem 3 (with the explicit lower order terms given in Equation ( 40)), we conclude that:KL Law(X P T ) Law(X T K ) ≤ ϵ 2 8. Now, by Pinsker's inequality and the triangle inequality for TV distance, we have:TV(Law(X P T ), π ⋆ ) ≤ 2KL Law(X P T ) Law(X T K ) + 2KL Law(X T K ) π ⋆ ≤ ϵ (62)This proves the result in Table 1 for OLMC.We now consider the analogous result for ULMC below. The following result is a re-statement of [53, Theorem 7] . We note that C LSI in this work is 1 λ in our work. Theorem 6. Let X 0 , . . . , X N be the iterates of ULMC with step-size η. We make the following assumptions:1. F is L-smooth and that π ⋆ satisfies λ-LSI.2. Suppose x be any stationary point of F . Initialize X 0 such that U 0 , V 0 are independent with U 0 ∼ N (x, ζI) (ζ as given in [53] ) and V 0 ∼ N (0, I).3. F (x) -F (x * ) = Õ(d) and E U,V ∼π ⋆ ∥U ∥ ≤ Õ(d). Let ϵ ≤ 1. Then, we let η = Θ( ϵ √ λ L √ d ), c √ L ≤ γ ≤ C √ L, T = Θ( L 3 2 √ d λ 3 2 ϵ) then,TV(Law(U N ), π ⋆ ) ≤ ϵLet X 0 , . . . , X T K be the iterates of ULMC with step-size α K . We take:EQUATION),EQUATION) and N = Θ( √ L λα ). We now consider X P 0 , . . . , X P T to be the iterates of PS(A, G, Γ, α, K) with A, G, Γ corresponding to ULMC, such that X P 0 and F satisfy the same assumptions in Theorem 6. Applying Theorem 3 along with Pinsker's inequality, we conclude that:TV(Law(X (P ) T ), Law(X T K )) ≤ ϵ 2 .Now, using Theorem 6, along with the triangle inequality for TV, we conclude:TV(Law(X (P ) T ), π ⋆ ) ≤ ϵThis proves the bounds for ULMC in Table 1 .We state the following technical lemma which will be proved in Section I.3. Lemma 26. Let a 1 , . . . , a T be random vectors in R d . Consider the partitioning {1, . . . , T } = ∪ ⌈T /N ⌉ k=1 T k where T k := {(k -1)N + 1, . . . , min(kN, T )}. Assume that N divides T and T ≥ N . Then, we have:T j=1 E∥a j ∥ 2p ≤ 2 p N T N k=1 j∈T k j ′ ∈T k j ′ >j E∥a j -a j ′ ∥ 2p + 2 N p-1 E( T j=1 ∥a j ∥ 2 ) pLemma 27. Let Z 1 , . . . , Z T be i.i.d. standard Gaussian random variables, and a sequence of random vectors g 1 , . . . , g T such that g t is independent of Z t , . . . , Z T . Then,E| T s=1 ⟨g s , Z s ⟩| p ≤ C p E( T s=1 ∥g s ∥ 2 ) p Proof. Let λ ∈ R. Consider the random variable M t = exp( t s=1 λ⟨g s , Z s ⟩ -λ 2 ∥gs∥ 2 2). M t is a super martingale and hence, we have:E[M T ] ≤ 1.Applying the Chernoff bound, we conclude that for anly λ, λ > 0 we must have:P(| T s=1 ⟨g s , Z s ⟩| > λ 2 T s=1 ∥g s ∥ 2 + h λ ) ≤ 2 exp(-h) . =⇒ P(| T s=1 ⟨g s , Z s ⟩| p > C p λ p ( T s=1 ∥g s ∥ 2 ) p + h) ≤ 2 exp(-C p λh 1 p )Now, for any variable X, P(X + > h) = P(X > h) whenever h ≥ 0. Now using the fact that:EX ≤ EX + = ∞ 0 P(X + > h)dh, and taking X = | T s=1 ⟨g s , Z s ⟩| p -C p λ p ( T s=1 ∥g s ∥ 2 ) p : E| T s=1 ⟨g s , Z s ⟩| p ≤ C p λ p E( T s=1 ∥g s ∥ 2 ) p + ∞ 0 2 exp(-c p λh 1 p )dh ≤ C p λ p E( T s=1 ∥g s ∥ 2 ) p + C p λ pWe conclude the result by choosing λ 2p = 1 E( T s=1 ∥gs∥ 2 ) p . Lemma 28. Let K be any positive integer. Consider G ∼ Bin(K, 1 K ). Then for any p ≥ 1, we must have: EG p ≤ C p for some constant depending only on p.Proof. The proof is simple forK = 1. Assume K ≥ 2. For any 0 ≤ n ≤ K, let p n := P(G = n). Then, n!p n = (1 -1 K ) K-n K! K n (K -n)! ≤ (1 -1 K ) K-1 ≤ 2 eIn the second step we have used the inequality (1 -1 n ) n ≤ 1 e and the fact that K ≥ 2. Note that 1 en! is the probability that the standard poisson random variable is n (denote this by p * n ). Therefore, we must have: p n ≤ 2p * n . Thus, we must have:EG p ≤ n∈N 2n p p * n ≤ C p I Proofs of Technical Lemmas I.1 Proof of Lemma 12Proof.1. This follows by applying the triangle inequality to the definition of XtK+i .2. Consider:EQUATION3. This follows by a direct application of Doob's inequality.Proof. We will use b() and ∇F () interchangeably in the proof. Consider the updateX(t+1)K = XtK + αb( XtK ) + α∆ t + √ 2α Ȳt .Where Ȳt = 1√ K K-1 i=0 Z tK+i and ∆ t = K-1 i=0 H i (b( XtK+i ) -b( XtK )) Therefore, we have for any s > t:EQUATIONIn this proof only, for any random variable U , we letEQUATIONp . Using the triangle inequality for M p , we conclude:EQUATIONIn the second line, we have used the fact that N h is independent of sup 0≤i≤K-1 ∥ XhK+i -XhK ∥ and the fact that M p (N t ) ≤ C p for some constant C p (Lemma 28 ). In the third step we have applied item 1 from Lemma 12.Putting this back in Equation 65, we conclude:EQUATIONTherefore, we must have:EQUATIONNow, supposing s -t ≤ 1 CpαL for some large enough constant Cp which depends only on p. This ensures that αL(s -t)(1 + αLC p ) ≤ 1 2 in the equation above. Thus, we have:sup t≤h≤s M p ( XhK -XtK ) ≤ 3αL(s -t) ∥b( XtK )∥ L + C p dα(s -t) I.3 Proof of Lemma 26 Proof. āk := 1 |T k |j∈T k a j Consider the following quantity:EQUATIONNow, by Jensen's inequality, we must have:∥ā k ∥ 2 ≤ 1 |T k |j∈T k ∥a j ∥ 2 . Therefore, we conclude:EQUATIONNow, considerEQUATIONUsing this with Equation (70), we conclude the statement of the lemma.Proof. We now apply Lemma 26 with a t := b( X(t-1)K ). Whenever αL ≤ c p for some small enough constant c p , we can take N to be the largest integer power of 2 lower than ⌈ 1 αL C2p ⌉ for a large enough constant C2p such that the result of Lemma 13 with s -t ≤ N . Applying Lemma 13, we conclude:EQUATIONApplying Lemma 26EQUATIONNote that whenever the constant C2p defining N above is large enough, this impliesEQUATIONProof. Since ∇F is L Lipschitz, we have:EQUATIONAlso note that: X(t+1)K = XtK -α∇F ( XtK ) + α∆ t + √ 2α Ȳt .Where Ȳt = 1√ K K-1 i=0 Z tK+i and ∆ t = - K-1i=0 H t,i (∇F ( XtK+i ) -∇F ( XtK )) Thus, summing Equation (75) from t = 0 to t = T -1, we conclude:EQUATIONIn the last step, we have used the fact that 2|⟨∇F ( XtK ), ∆ t ⟩| ≤ ∥∇F ( XtK )∥ 2 + ∥∆ t ∥ 2 and the assumption that αL ≤ c p for some small enough constant c p .Therefore, we conclude:EQUATIONThe properties of Gaussians show that:E( T -1 t=0 ∥ Ȳt ∥ 2 ) p ≤ C p T p d p .By Lemma 27 and the AM-GM inequality, we conclude that for any κ > 0 arbitrary, we have:EQUATIONIn the first step, we have used Jensen's inequality. In the second step we have used the fact that ∇F is L-Lipshcitz and that N t is independent of XtK+i for 0 ≤ i ≤ K -1. In the last step, we have used the fact that E(N t ) 2p ≤ C p for some constant C p (Lemma 28).By Lemma 12, we have:E sup 0≤i≤k ∥ XtK+i -XtK ∥ 2p ≤ C p α 2p E∥∇F ( XtK )∥ 2p + C p α p d pPlugging all of these bounds, and taking κ in Equation (78) to be large enough, we conclude:EQUATIONI.6 Proof of Lemma 18Proof. For the sake of clarity, definex = 1 -exp(-γh). ρ = (1-x) 2 γ 1 2 γ (h- 2(1-x) γ + (1-x 2 ) 2γ )(1-x 2 )Algebraic manipulations show that:ρ = (1 -x) 3 2(1 + x) γh -(1 -x) -(1-x) 2 2 (81) Now note that γh = -log(1 -(1 -x)) = ∞ i=1 (1-x) i i. Therefore, we conclude: γh -(1 -x) -(1-x) 2 2 ≥ (1-x) 3 3 + (1-x) 4 4. Therefore,EQUATIONNow, we note that for scalars A, B > 0 and ρ ∈ [0, 1], we have:A 2 I d ρABI d ρABI d B 2 I d ⪰ A 2 (1 -ρ)I d 0 0 B 2 (1 -ρ)I dThus, we conclude from the above computations that:EQUATION)We have used the fact that:2 γ h - 2 γ (1 -exp(-γh)) + 1 2γ (1 -exp(-2γh)) ≥ 2 3γ 2 (1 -exp(-γh)) 3Using the fact that for PSD matrices A, B A ⪯ B implies B -1 ⪯ A -1 , we have:EQUATIONThe second inequality follows easily by elementary manipulations.Proof. Explicit computations show that:ψ t+1 = ψ t - α γ ∇F ( ŨtK ) + α I d I d γ ∆ t + I d I d γ Γ α ỸtNow, note that I dI d γ A h = I I d γ and I d I d γ G h = h γ I d 0 .Therefore, from the definition of ∆ t , we conclude that:α I d I d γ ∆ t = - K-1 i=0 H i α γ ∇F ( ÛtK+i ) -∇F ( ŨtK )A straight forward calculation shows that: Ψt := I dI d γ Γ α Ỹt ∼ M(0, 2h γ I d ).Proof.1. Note thatÛtK+i -ŨtK = Π(A iα K -I) XtK + ΠG iα K b( XtK ) + i-1 j=0 ΠA α(i-1-j) K Γ α K Z tK+jWe bound each of the terms separately.∥(A iα K -I) XtK ∥ Π ≤ iα K ∥ ṼtK ∥ ≤ α∥ ṼtK ∥ ∥G iα K b( XtK )∥ Π ≤ α 2 i 2 2K 2 ∥∇F ( ŨtK )∥ ≤ α 2 2 ∥∇F ( ŨtK )∥Plugging these inequalities gives the result.2. Since ∇F is L-Lipschitz, we have:EQUATIONUsing Lemma 20 to evaluate ψ t+1 -ψ t , we conclude:EQUATIONIn the second step, we have use the fact that ∇F is L-Lipschitz along with the Cauchy-Schwarz inequality. In the third step, we have used the AM-GM inequality which states that for any a, b ≥ 0, we have 2ab ≤ a 2 + b 2Similarly, we note that:EQUATIONWe have again used the AM-GM inequality in the second step above. Convexity of x → x 2 implies:EQUATIONPlugging these upper bounds into Equation (85), we have:EQUATION3. Using the scaling relations (Section 2.1), it is easy to show that for any i such that 0≤ i ≤ K -1, i-1 j=0 ΠA α(i-1-j) K Γ α K Z tK+j is a Gaussian with covariance matrix Σ i such that Σ i ⪯ 2 exp(γα)γα 3 3 I d . Applying Gaussian concentration for i-1 j=0 ΠA α(i-1-j) K Γ α KZ tK+j along with the union bound over all 0 ≤ i ≤ K -1, we conclude:P(M t,K > C exp( γα 2 ) γα 3 ( √ d + β + log K)) ≤ exp(-β 2 /2)We conclude the moment bounds by integrating tail probabilities.Proof.1. Let g := e -γα . Then, we have:EQUATIONNow, consider:EQUATIONIn the second step, we have use the fact that| 2g(1-g) γ ⟨∇F ( ŨtK ), ṼtK ⟩| ≤ g(1 - g) ∥∇F ( ŨtK )∥ 2 γ 2 + g(1 -g)∥ ṼtK ∥ 2 .In the third step, we have used the fact that:2g⟨ ṼtK , (I -Π)α∆ t ⟩ ≤ ( √ g -g)∥ ṼtK ∥ 2 + ∥(I-Π)α∆t∥ 2 1- √ g , -2(1-g) γ ⟨∇F ( ŨtK ), (I - Π)α∆ t ⟩ ≤ 1-g γ 2 ∥∇F ( ŨtK )∥ 2 + (1 -g)∥(I -Π)α∆∥ 2 and ∥(I -Π)(α∆ t + Γ α Ỹt )∥ 2 ≤ 2∥(I -Π)(α∆ t )∥ 2 + 2∥(I -Π)Γ α Ỹt ∥ 2 . With extention of notation, let S 2 ((I-Π)α∆) := T -1 t=0 ∥(I-Π)∆ t ∥ 2 , S 2 ((I-Π)Γ α Ỹ ) := T -1 t=0 ∥(I -Π)Γ α Ỹt ∥ 2 .Using Equation (91), we conclude:EQUATIONNote that (I -Π)Γ α Ỹt ∼ N (0, ( E| T -1 t=0 ⟨ ṼtK , (I -Π)Γ α Ỹt ⟩| p ≤ C p (γα) p 2 E(S 2 (V )) p E| T -1 t=0 ⟨∇F ( ŨtK ), (I -Π)Γ α Ỹt ⟩| p ≤ C p (γα) p 2 E(S 2 (∇F )) pBy L smoothness of F , we have:∥(I -Π)(α∆ t )∥ ≤ LαN t sup 0≤i≤K ∥ ÛtK+i -ŨtK ∥.Using the result in Lemma 21, we conclude:EQUATIONIn the second step, we have used jensen's inequality to show that ( 1T t a 2 t ) p ≤ 1 T t a 2p t .In the third step, we have used the fact that N t is independent of XtK and that EN 2p t ≤ C p for some constant C p (Lemma 28).We now use the fact that for any two random variables X, Y such that E∥X∥p , E∥Y ∥ p < ∞, [E∥X +Y ∥ p ] 1 p ≤ [E∥X∥ p ] 1 p +[E∥Y ∥ p ] 1 p. Using the bounds established above and applying them to Equation (91), we conclude:EQUATIONHere, we have used that fact that 1 -√ g ≥ Cγα whenever γα < 1. We now use the AM-GM inequality to show thatC p [E(S 2 (V )) p ] 1 2p √ γα ≤ [E(S 2 (V )) p ] 1 p 10 + C ′ p γα .Similarly, we have:C p α γ [E(S 2 (∇F )) p ] 1 2p ≤ [E(S 2 (∇F )) p ] 1 p 2γ 2 + C ′ p γα .Applying this to the RHS of Equation ( 94) and re-arranging, we conclude the statement of the lemma.2. Using similar methods as in item 1 above, we apply Lemma 27 and collect the following moment bounds:E| T -1 t=0 ⟨∇F (ψ t ) -∇F ( ŨtK ), Ψt ⟩| p ≤ C p L p α p 2 γ 3p 2 E(S 2 (V )) p E| T -1 t=0 ⟨∇F ( ŨtK ), Ψt ⟩| p ≤ C p α p 2 γ p 2 E(S 2 (∇F )) pBy a calculation similar to that in item 1, we have:EQUATIONSumming item 2 of Lemma 21 over t = 0 to t = T -1 and applying the triangle inequality for p-th moments, we have:EQUATIONIn the last step, we have use the fact that γα < 1 and Lα γ < c 0 . Applying AM-GM inequality toC p γ α [E(S 2 (∇F )) p ] 1 2p and C p L 1 γα [E(S 2 (V )) p ] 1 2psimilar to item 1, we conclude the result.Proof.XsK-XtK = (A α(s-t) -I) XtK + s-1 h=t A α(s-h-1) G α b( XhK ) + α∆ h + Γ α ỸhIn this proof only, for any random variable W and any projection operator P over R 2d we letEQUATIONp . We use the convention that M p (W ) = (E∥W ∥ p ) 1 p . Using the triangle inequality for M p , we conclude:M p ( XsK -XtK , Π) ≤ M p ((A α(s-t) -I) XtK , Π) + s-1 h=t M p (A α(s-h-1) G α b( XhK ), Π) + s-1 h=t M p (A α(s-h-1) α∆ h , Π) + M p ( s-1 h=t A α(s-h-1) Γ α Ỹh , Π) (97)Similarly, we have:EQUATIONBy scaling relations given in Section 2.1, we have:s-1 h=t ΠA α(s-h-1) Γ α Ỹh ∼ N (0, Π(Γ 2 (s-t)α )Π) and s-1 h=t (I -Π)A α(s-h-1) Γ α Ỹh ∼ N (0, (I -Π)(Γ 2 (s-t)α )(I -Π)). Notice that Π(Γ 2 α(s-t) )Π ≲ 4α(s-t) γ 0 0 0 and (I -Π)(Γ 2 α(s-t) )(I -Π) ≲ 0 0 0 2γα(s -t)Therefore, we conclude:M p ( s-1 h=t A α(s-h-1) Γ α Ỹh , Π) ≤ C p dα(s -t) γ M p ( s-1 h=t A α(s-h-1) Γ α Ỹh , I -Π) ≤ C p dαγ(s -t) Notice that ∥Π(A α(s-t) -I) XtK ∥ ≤ α(s -t)∥ ṼtK ∥ and ∥(I -Π)(A α(s-t) -I) XtK ∥ ≤ αγ(s - t)∥ ṼtK ∥. This implies: M p ((A α(s-t) -I) XtK , Π) ≤ α(s -t)M p ( ṼtK ) M p ((A α(s-t) -I) XtK , I -Π) ≤ αγ(s -t)M p ( ṼtK ) Notice that: ∥ΠA α(s-h-1) G α b( XhK )∥ ≤ α γ ∥∇F ( ŨhK )∥ and ∥(I -Π)A α(s-h-1) G α b( XhK )∥ ≤ α∥∇F ( ŨhK )∥ . This implies: M p (A α(s-h-1) G α b( XhK ), Π) ≤ α γ M p (∇F ( ŨhK )) ≤ αL γ M p ( XhK -XtK , Π)+ α γ M p (∇F ( ŨtK )) M p (A α(s-h-1) G α b( XhK ), I-Π) ≤ αM p (∇F ( ŨhK )) ≤ αLM p ( XhK -XtK , Π)+αM p (∇F ( ŨtK )) Notice that: ∥ΠA α(s-h-1) α∆ h ∥ ≤ αLN h γ sup 0≤i≤K-1 ∥ ÛhK+i -ŨhK ∥ and ∥(I - Π)A α(s-h-1) α∆ h ∥ ≤ αLN h sup 0≤i≤K-1 ∥ ÛhK+i -ŨhK ∥.Applying item 1 in Lemma 21, we conclude: In the first step, we have used the triangle inequality for M p along with the fact that N h is independent of M h,K and XhK and the fact that its p-th moment is bounded by a constant C p . In the second step, we have controlled the norm of ∇F ( ŨhK ) with that of ∇F ( ŨtK ) using Lipschitzness. In the third step, we have invoked item 3 of Lemma 21 to bound the moments of M h,K . In the fourth step, we have controlled the norm ṼhK in terms of XtK and XtK -XhK .EQUATIONSimilarly, we have: Applying the estimates derived above to Equation (97), along with the assumption that αL(s-t) γ ≤ c p , αγ(s -t) ≤ c p and α 2 L(s -t) ≤ c p for some small enough constant c p which depends only on p, we conclude: • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).* This work was done when SK was a student researcher at Google Research.The prior works considered the compexity for Wasserstein distance W2(output, target) ≲ ϵ whereas we consider TV ≤ ϵ. This is a standard comparison in the sampling literature[53].The time convention in the DDPM literature is reverse: Xt-1 = atXt + bt∇ log pτ n-1-t (Xt) + σtZt. 


Decision: This paper considers Poisson Midpoint Method, which is an alternative to Randomized Midpoint Method. Authors prove that the proposed scheme can achieve quadratic speed up of LMC. Authors demonstrate the performance via numerical studies.

This paper was reviewed by three reviewers with the following Scores/Confidence: 5/3, 7/3, 5/4. I think the paper is studying an interesting topic and the results are relevant to NeurIPS community. The following concerns were brought up by the reviewers:

- Theorem statements are quite convoluted and should be improved.
    
- Proof sketch can be more user friendly for the broader audience.
  
- Poor presentation should be improved in a few places. See reviewers' comments for details.

- The condition on initialization is certainly not restrictive but should be clarified.

Authors should carefully go over reviewers' suggestions and address any remaining concerns in their final revision. Based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper to the NeurIPS 2024 program.

Reviews: review: paper_summary: paper_summary: Based on the Randomized Midpoint Method (RMM) proposed for simulating the Langevin Dynamics numerically, the authors proposed a variant called the Poisson Midpoint Method (PMM). Theoretical analysis is performed to bound the deviation between the continuous dynamics and the discretized dynamics. Moreover, when PMM is applied to simulate the Langevin Dynamics, the authors have shown that PMM can accelerate both the Overdamped Langevin Monte Carlo (OLMC) and the Underdamped Langevin Monte Carlo (ULMC) under additional assumptions like isoperimetry. Furthermore, numerical experiments on image generation with DDPM are conducted to illustrate the effectiveness of the proposed method.
strength_weakness: strength_weakness: The presentation of this paper can be potentially improved from a few aspects. Firstly, it seems a bit weird to use the phrase "lower order terms" (see for instance, line 167-168, 172-173, 205-206) in mathematical equations and derivations. One possible alternative is to replace it via the big O notation. Secondly, the authors might consider clarifying the differences between Langevin Dynamics and the backward SDE in diffusion model. One important thing that probably needs to stated clearly is that PMM is a method for simulating general time-dependent SDE (like the reversed SDE in diffusion models), but this paper only provides theoretical analysis for the case when PMM is applied to simulate Langevin Dynamics. Furthermore, it seems that the authors didn't discuss potential topics for future work.
questions: questions: One major concern of the reviewer is that one assumption in the paper seems to be too strong. Specifically, Assumption 2 in the paper bounds the 14-th moment of the distance between the starting point and the global minima. To the best of the reviewer's knowledge, such assumption seems to be uncommon in sampling literature. Would it be possible for the authors to elaborate a bit more on why this assumption is essential here? It seems that this assumption essentially restricts the starting point to be a warm start, which makes the sampling task easier.
limitations: limitations: A few typos (the list is probably incomplete) in the current version of the manuscript are also listed below:

1. From line 179-180, the two phrases "to OLMC" and "to ULMC" appeared twice. The two sentences probably need to revised.
2. In line 206, the subtitle should be "Sketch for Theorem 3". Moreover, the following sentence should be "This is similar to Theorem 2".
paper_summary: paper_summary: In this work, the authors draw inspiration from the randomized midpoint method for Langevin dynamics and design a new method called the poisson midpoint method. They establish guarantees for this method relative to the EM discretisation with thinning, and discuss theoretical and practical applications to sampling.
strength_weakness: strength_weakness: The writing of the paper can be definitely improved. It appears a bit rushed in multiple key areas, and several definitions and procedures are stated without much consideration towards the readability of those portions. Since there is a new method being proposed, it would help if intuition for this method is proposed, not just for the idea, but also the design of the method itself. Additionally, when certain complex steps (relative to other methods described in the paper like LMC) are mentioned (see the Algorithm and Computational complexity paragraph on Page 5), it would help to explain each (or a subset of them) in layman's terms to connect them to method (the two equations (7) and (8) on the previous page). The theoretical result in Section 3, specifically theorem 1 has inconsistent terminology ($B$ is the Brownian motion, $\mathsf{S}(...)$ and $\mathsf{PS}$ are lacking a $b$, $B$, $b$, $\beta$, the role of $r$ and the randomness of the expectation is unclear. In the next theorems, the authors state "Lower Order Terms" instead of a small-O notation would have seemed more appropriate. The paragraph on quadratic speedup is also quite hastily presented in my opinion, because it is still not clear to me how the speedup is represented.
questions: questions: 1. What is $p_{\tau}$ in Eq. (4)?
2. As mentioned in Page 2, $A_\alpha$, $G_\alpha$, $\Gamma_\alpha$ are time independent, and so DDPM doesn't fall under the model in (1) right? How is DDPM relevant to the scheme (1) in that case?
3. The scaling relations part of Section 2.1 could perhaps be placed later on? It's difficult to understand its purpose since it seems disconnected from the rest of the main draft.
4. In Line 109, "speed up" is usually a computational term. Is this alluding to the "speed up" discussed later in Section 3 (in which case it might help to reference it here)?
5. Since RLMC is also relatively complex (in comparison to a standard forward EM discretization), it would help to remark on how this is implemented, even if briefly. As written, it is essentially a combination of two discretizations: one forward, and one backward, and the backward discretization is the harder one to implement.
6. The section of text between Lines 118--123 is very hard to follow. Since it is essential to the intuition, I would recommend that it be reworded suitably to be more formal and less colloquial. What is $\mathrm{law}(Z + \epsilon)$ and $\mathrm{law}(Z + \epsilon Z')$?
7. I feel like the claim of "replacing K steps with one step" is not quite true, since eqs (7), (8) are each performing K steps each right?
8. In eq (7), could the authors perhaps give some intuition about the summation on the RHS, and what the cheap interpolation is actually performing? A diagram would help greatly, but it is understandable if this is difficult to provide.
9. In (8), what is the purpose of H_{t,i}? Why don't the $\Gamma$, $A$, $G$ terms not have a "i" scaling like in the cheap interpolation?
10. Why does $I_t$ have $N$ elements, and what is $N$ here? Initial conditions / conventions for the recursion in step 2 would be better off being stated apriori. What is $\mathrm{Corr}$ in Step 4?
11. Where is the computational complexity discussed in the paragraph on the top of page 5?
12. How is there no dependence on the $\{X_{t}\}$ sequence in Eq. 10?
13. Why is this strange dependence on 14th moment in Assumption 2 present?
limitations: limitations: Yes.
paper_summary: paper_summary: The paper proposes a novel, multi-step discretization method for SDEs, with particular applications to the Langevin diffusions (overdamped and underdamped) and the denoising diffusion model SDE. The rates obtained are better than any other known guarantees in KL divergence.
strength_weakness: strength_weakness: The algorithm as described in the paper is somewhat cumbersome to read. Perhaps some simpler exposition or visualization may be helpful.

The rates obtained are worse than those of the randomized midpoint discretizations (noting that these latter rates are only valid in Wasserstein distance).
questions: questions: Although the notation in Section 2 is necessary to capture all the cases considered in this work, it is a bit cumbersome and makes the results difficult to parse. Perhaps for the proof sketch, only the OLMC case could be presented.

The final rates have some somewhat strange dependence on the key parameters; for instance, the dependence on dimension is $d^{5/12}$ for the underdamped PLMC. Do the authors have any intuition whether the rates obtained in their analysis are close to optimal for this discretization, and whether it would be possible to improve it with slight algorithmic modifications?

Is the method extendible to the case where the tail growth of the potential is of order $|x|^\alpha$, for $\alpha < 2$? Or the case of weak smoothness?

Minor:

59: stochastic approximations LMC -> "stochastic approximations of LMC"

70: denotes identity -> denotes the identity

108 "Motive" -> "motivate"

206: Theorem 1; I believe this is meant to refer to Theorem 3.

Many equations/inline formulas are missing punctuation. For instance Eq. 6, Steps 3-5 (ll 145-149).
limitations: limitations: The authors have adequately addressed the limitations of their work.

Comments: title: Thank you
Thank you for your rebuttal
Response to the authors' rebuttal
Cannot see individual rebuttals
comment: Thank you very much!
I thank the authors for their response. In particular, given that they claim to be able to extend this analysis even to the weakly smooth and Poincare cases, I believe this framework has enough power and novelty to be accepted. Hence, I have raised my score.

I was mistaken regarding the comparison with randomized midpoint; the dimension dependence is worse (which is still mysterious to me) but the inverse accuracy dependence is indeed better as noted by the authors. I thank the authors for their correction.
I have no further questions, and I feel that the contributions of the paper is very interesting, albeit poorly presented. Improvements to the presentations as suggested by multiple reviewers would certainly make this a good paper. I maintain by score on the basis of this.
Dear authors,

Thank you so much for your detailed response, which have addressed most of my concerns. By reading other reviewers' comments, I agree with them that the presentation of the paper can be improved from many aspects. It would help improve the quality of the manuscript if the authors could follow the suggestions to revise the paper accordingly.

Best regards,

Reviewer F8sy
As per the email from the PC, the official rebuttal becomes visible to the reviewers on August 7th. We have submitted the rebuttals, but they are currently invisible to you. Apologies about the confusion -- we did not realize that this comment would be visible earlier.
Hello authors, I'm unable to see the individual rebuttal (at least those pertaining to my review). Could you please check if the visibility of the rebuttal has been set correctly?