Title: The Poisson Midpoint Method for Langevin Dynamics:  Provably Efficient Discretization for Diffusion Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Langevin Monte Carlo

Diffusion Models

MCMC
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.

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

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.

1 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.

1.1 Prior Work

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.2 Our Contributions

(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 ϵ error 2 .

(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.

1.3 Notation:

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 2 d . 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: ˜ X 4 corresponds to ˜ U 4 , ˜ V 4 ). In this case I represents identity matrix in R 2 d × 2 d 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 ˜ O, ˜ Ω , ˜ Θ are same as O, Ω , Θ up to poly-logarithmic factors in the problem parameters such as 1 ϵ , 1 α , K, T, d .

2 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:

Where 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.

Here B τ is the standard Brownian motion in R d . Under mild conditions on F and ¯ X 0 , 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 EulerMaruyama 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 2 d dimensions. We write X t = [ U t V t ] ⊺ ∈ R 2 d where U t ∈ R d is the position and V t ∈ R d is the velocity. Fix a damping factor γ > 0 . We take:

This choice of A h , Γ h , G h , b ( , ) in Equation (1) gives the Euler-Maruyama discretization of Underdamped Langevin Dynamics (a.k.a. the Kinetic Langevin Dynamics) studied extensively in Physics[14]):

The 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:

This also admits an equivalent characteristic ODE given below:

See [43, 5, 4] for further details. Here p τ is the probability density of e -t X ∗ + √ 1 -e -2 t 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.

2.1 Technical Notes

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.

Randomized 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:

Taking X t as the approximation to ¯ X tα , LMC approximates the integral ∫ ( t +1) α tα ∇ F ( ¯ X s ) ds with α ∇ F ( X t ) , giving a 'biased' estimator to the integral (conditioned on X t = ¯ X tα ). 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 ) + √ 2 uαZ t, 1 (notice X t + u is an approximation for ¯ X ( t + u ) α ). The RLMC update is:

Notice that √ 2 uα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).

2.2 The Poisson Midpoint Method

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 ( ˜ X tK ) 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 ˜ X tK , we define the interpolation:

Note that this interpolation is cheap since every one of ˆ ˜ X tK + 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:

We pick H t,i based on the following two options, independent of Z tK + i , ˜ X 0 :

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 := ✶ ( 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 ˜ X K ( t +1) given ˜ X Kt in one step by unrolling the recursion given in Equation (8). For the sake of clarity, we will relabel ˜ X tK 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:

We use the convention that i 0 = 0 , i N +1 = K -1 , A 0 = I and Γ 0 = 0 .

That is, the algorithm first generates the random mid-points H t,i , computes the interpolation ˆ ˜ X tK + i only when H t,i = 1 and then computes b ( ˆ ˜ X tK + i , ( tK + i ) α K ) 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 ( ˜ X tK ) 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.

3 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 ˜ X tK + i be the iterates in Equation (8) . Define random variables:

Then, for some universal constant C and any r > 1 :

We now apply Theorem 1 to the case of OLMC and ULMC under additional assumptions, with the proofs in Sections E and F respectively.

d d ∗

0 0 init

Assumption 1. F : R → R is L -smooth (i.e., ∇ F is L -Lipschitz). x is its global minimizer. Assumption 2. The initialization X is such that E ∥ X -x ∗ ∥ 14 < C 14 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 and α 3 L 3 T < c 2 then:

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 whenever γ > C 1 √ L , αγ < c 1 , T < c 2 γ L 2 α 3 .

OLMC 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 ≤ 2 KL 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.

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 ) :

Where ˜ Z tK + 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 ˜ X tK + i ). They are independent of Z tK + i conditioned on ˜ X tK + i . Note that the sequence ( ˜ Z tK + 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 ( ˜ Z tK + 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 ) . Wethen show that Z tK + i + S tK + i | ˜ Z 0: 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 ˜ Z 0: tK + i -1 where Σ t,i is the conditional covariance of S tK + i . This CLT step gives us the error of the form ∑ T -1 s =0 ∑ K -1 i =0 C E [ β 10 sK + i K 3 + β 6 sK + i K 2 + β 2 r 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 C E [ β 4 sK + i K 2 ] .

Sketch for Theorem 2 Applying Theorem 1 to OLMC , note that the term β tK i depends on how far the coarse estimate ˆ ˜ X tK + i is from the true value ˜ X tK + i . Indeed, under the smoothness assumption on F we show that: β 2 p tK + i ≤ ( L 2 αK 2 ) p sup 0 ≤ j ≤ K -1 ∥ ˆ ˜ X tK + j -˜ X tK ∥ 2 p . Thus:

Therefore, the proof reduces to bounding ∑ T -1 t =0 E ∥∇ F ( ˜ X tK ) ∥ 2 p . We observe that ˜ X ( t +1) K = ˜ X tK -α ( ∇ F ( ˜ X tK ) + ∆ 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:

The following sophisticated bound derived in this work is novel to the best of our knowledge:

Sketch for Theorem 3 This is similar Theorem 2, but requires us to bound E ∑ t ∥ ˜ V tK ∥ 2 p and E ∑ t ∥∇ F ( ˜ U tK ) ∥ 2 p . We track the decay of two different entities across time: (1) ∥ ˜ V tK ∥ 2 and (2) F ( ˜ U tK + ˜ V tK γ ) . Our proof shows via a similar taylor series based argument that PLMC does not allow either ∥ ˜ V tK ∥ 2 or F ( ˜ U tK + ˜ V tK γ ) to grow too large. Letting Ψ t := ˜ U tK + ˜ V tK γ we show (roughly):

5 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 50 k 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.

Werefer 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 DPMSolver, 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.

5.1 Results

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.

6 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.

7 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.

Dataset

vs. SDE Based Methods

vs. ODE Based Methods

CelebAHQ 256

LSUN

Churches

LSUN

Bedrooms

FFHQ 256

LatentDiffusionModel(CelebAHQ256)

LatentDiffusionModel(CelebAHQ256)

10

DDPM

DDIM ETA=0

9

Poisson Midpoint(Ours)

DPM-Solver

Clip-FID

Clip-FID

4

20

50

100

200

500

750 1000

20

50

100

200

500750

Number of Neural Network Calls

Number of Neural Network Calls

Latent Diffusion Model (LSUN Churches)

LatentDiffusion Model(LSUN Churches)

6.5

DDPM

DDIM ETA=0

Poisson Midpoint(Ours)

6.0

DPM-Solver

8

1

5.0

FID

6

5

4.0

20

50

100

200

500

750 1000

3.0

20

50

100

200

500

Number of Neural Network Calls

NumberofNeuralNetworkCalls

Latent Diffusion Model (LSUNBedrooms)

LatentDiffusion Model(LSUN Bedrooms)

DDPM

DDIM ETA=0

Poisson Midpoint(Ours)

DPM-Solver

5

5

50

100

200

5007501000

20

50

100

200

500750

Number ofNeural Network Calls

Number of Neural Network Calls

LatentDiffusionModel(FFHQ256)

Latent Diffusion Model(FFHQ 256)

4.5

DDPM

DDIM ETA=0

Poisson Midpoint(Ours)

Poisson Midpoint(Ours)

4.0

DPM-Solver

Clip-FID

Clip-FID

3.0

3

20

50

100

200

5007501000

2.0

20

50

100

200

500750

Number ofNeuralNetworkCalls

NumberofNeuralNetworkCalls

Acknowledgments and Disclosure of Funding

We thank Prateek Jain for helpful comments during the course of this research work. This work was done when SK was a student researcher at Google Research.

References

Krishna Balasubramanian, Sinho Chewi, Murat A Erdogdu, Adil Salim, and Shunshi Zhang. Towards a theory of non-log-concave sampling:first-order stationarity guarantees for langevin monte carlo. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 2896-2923. PMLR, 02-05 Jul 2022.

Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference on Learning Representations , 2022.

Kevin Burrage, PM Burrage, and Tianhai Tian. Numerical methods for strong solutions of stochastic differential equations: an overview. Proceedings of the Royal Society of London. Series A: Mathematical, Physical and Engineering Sciences , 460(2041):373-402, 2004.

Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast. Advances in Neural Information Processing Systems , 36, 2024.

Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations , 2023.

Xiang Cheng, Niladri S Chatterji, Peter L Bartlett, and Michael I Jordan. Underdamped langevin mcmc: A non-asymptotic analysis. In Conference on learning theory , pages 300-323. PMLR, 2018.

Sinho Chewi, Murat A Erdogdu, Mufan Li, Ruoqi Shen, and Shunshi Zhang. Analysis of langevin monte carlo from poincare to log-sobolev. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 1-2. PMLR, 02-05 Jul 2022.

Arnak S Dalalyan. Theoretical guarantees for approximate sampling from smooth and logconcave densities. Journal of the Royal Statistical Society Series B: Statistical Methodology , 79(3):651-676, 2017.

ARNAK S DALALYAN and LIONEL RIOU-DURAND. On sampling from a log-concave density using kinetic langevin diffusions. Bernoulli , 26(3):1956-1988, 2020.

Aniket Das, Dheeraj Nagaraj, and Anant Raj. Utilising the clt structure in stochastic gradient based sampling: Improved analysis and faster algorithms. arXiv e-prints , pages arXiv-2206, 2022.

Alain Durmus, Szymon Majewski, and Bła˙ zej Miasojedow. Analysis of langevin monte carlo via convex optimization. The Journal of Machine Learning Research , 20(1):2666-2711, 2019.

Alain Durmus and Éric Moulines. Nonasymptotic convergence analysis for the unadjusted Langevin algorithm. The Annals of Applied Probability , 27(3):1551 - 1587, 2017.

Andreas Eberle, Arnaud Guillin, and Raphael Zimmer. Couplings and quantitative contraction rates for langevin dynamics. The Annals of Probability , 47(4):1982-2010, 2019.

Albert Einstein. Über die von der molekularkinetischen theorie der wärme geforderte bewegung von in ruhenden flüssigkeiten suspendierten teilchen. Annalen der physik , 4, 1905.

Murat A Erdogdu and Rasa Hosseinzadeh. On the convergence of langevin monte carlo: The interplay between tail growth and smoothness. In Conference on Learning Theory , pages 1776-1822. PMLR, 2021.

Arun Ganesh and Kunal Talwar. Faster differentially private samplers via rényi divergence analysis of discretized langevin mcmc. Advances in Neural Information Processing Systems , 33:7222-7233, 2020.

Sivakanth Gopi, Yin Tat Lee, and Daogao Liu. Private convex optimization via exponential mechanism. In Conference on Learning Theory , pages 1948-1989. PMLR, 2022.

Ye He, Krishnakumar Balasubramanian, and Murat A Erdogdu. On the ergodicity, bias and asymptotic normality of randomized midpoint sampling method. Advances in Neural Information Processing Systems , 33:7366-7376, 2020.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Proceedings of the 31st International Conference on Neural Information Processing Systems , NIPS'17, page 6629-6640, Red Hook, NY, USA, 2017. Curran Associates Inc.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Richard Holley and Daniel Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of Statistical Physics , 46:1159-1194, 1987.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. CoRR , abs/1710.10196, 2017.

Peter E Kloeden, Eckhard Platen, Peter E Kloeden, and Eckhard Platen. Stochastic differential equations . Springer, 1992.

Tuomas Kynkäänniemi, Tero Karras, Miika Aittala, Timo Aila, and Jaakko Lehtinen. The role of imagenet classes in fréchet inception distance. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023 . OpenReview.net, 2023.

Yin Tat Lee and Santosh S Vempala. Convergence rate of riemannian hamiltonian monte carlo and faster polytope volume computation. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing , pages 1115-1121, 2018.

Xuechen Li, Yi Wu, Lester Mackey, and Murat A Erdogdu. Stochastic runge-kutta accelerates langevin monte carlo and beyond. Advances in neural information processing systems , 32, 2019.

Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778 , 2022.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpmsolver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095 , 2022.

Yi-An Ma, Niladri S Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L Bartlett, and Michael I Jordan. Is there an analog of nesterov acceleration for gradient-based mcmc? Bernoulli , 27(3), 2021.

Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient mcmc. Advances in neural information processing systems , 28, 2015.

Grigori N Milstein and Michael V Tretyakov. Stochastic numerics for mathematical physics , volume 39. Springer, 2004.

Wenlong Mou, Nicolas Flammarion, Martin J Wainwright, and Peter L Bartlett. Improved bounds for discretization of langevin diffusions: Near-optimal rates without convexity. Bernoulli , 28(3):1577-1601, 2022.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning , pages 8162-8171. PMLR, 2021.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

Ingram Olkin and Friedrich Pukelsheim. The distance between two random vectors with given dispersion matrices. Linear Algebra and its Applications , 48:257-263, 1982.

Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. On aliased resizing and surprising subtleties in gan evaluation. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 11400-11410, 2022.

Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli , pages 341-363, 1996.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10684-10695, June 2022.

Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems , 32, 2019.

Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems , 36, 2024.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Ramon Van Handel. Probability in high dimension. Technical report, PRINCETON UNIV NJ, 2014.

Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. Advances in neural information processing systems , 32, 2019.

Santosh S. Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems 32: NeurIPS 2019 , pages 8092-8104, 2019.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11) , pages 681-688, 2011.

Fisher Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. LSUN: construction of a large-scale image dataset using deep learning with humans in the loop. CoRR , abs/1506.03365, 2015.

Lu Yu, Avetik Karagulyan, and Arnak Dalalyan. Langevin monte carlo for strongly log-concave distributions: Randomized midpoint revisited. arXiv preprint arXiv:2306.08494 , 2023.

Alex Zhai. A high-dimensional clt in w _2 w 2 distance with near optimal convergence rate. Probability Theory and Related Fields , 170:821-845, 2018.

Alex Zhai. A high-dimensional clt in w _2 w 2 distance with near optimal convergence rate. Probability Theory and Related Fields , 170:821-845, 2018.

Shunshi Zhang, Sinho Chewi, Mufan Li, Krishna Balasubramanian, and Murat A Erdogdu. Improved discretization analysis for underdamped langevin monte carlo. In The Thirty Sixth Annual Conference on Learning Theory , pages 36-71. PMLR, 2023.

A Details of Empirical Evaluations

A.1 Pseudocode

The 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.

We take the dynamics to go backward in time (i.e., compute X t -1 from X t )

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:

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 ( ˜ X t , t, K, b ( , ) , OPTION):

Let ( A k , B k , C k ) ← InterpolationConstants ( t, k ) , for every k ∈ [ K ] .

If OPTION = 1 :

Define H i ∼ Ber (1 /K ) i.i.d for every i ∈ [ K -1] .

Define p = K .

If OPTION = 2 :

For every i ∈ [ K ] , Z i ∼ N (0 , I ) .

Define H i ← ✶ ( u = i ) where u ∼ Unif ( { 1 , . . . , K -1 } ) i.i.d. Define p = K -1 .

Perform the update:

where

Return ˜ X t -K .

InterpolationConstants( t, k ):

A k ← a t · a t -1 · a t -k +1 . t -k +1

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 ˆ X t = ˜ X t , making the correction term pH i B K,i ( b ( ̂ X t -i +1 , t -i +1) -b ( ˜ X t , t ) ) = 0 when i = 0 .

A.2 Number of Neural Network Calls

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 ( ˜ X t , t ) is always computed. b ( ˆ X t -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 -1 K . 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 2 N K .

A.3 FID Evaluation Details

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.

A.4 Hardware Description and Execution Time

For our experiments, we utilized a high-performance computational setup consisting of 8 NVIDIA A100-SXM4 GPUs, each with 40 GB of VRAM, an Intel Xeon CPU with 96 cores operating at 2.2 GHz, and 1.3 TiB of RAM. The experiments take about 16 hours to generate 50k images with 1000 neural network calls per image and batch size of 50. The time is proportionally lower when fewer neural network calls are used.

A.5 DDPM Variant Details

We describe 3 different variants for DDPM, based on 3 different choices of coefficients a t , b t , σ t used in Section A.1. Our observations indicate that each variant performs optimally for different datasets and for different ranges of the diffusion steps. For example, consider the DDPM variants of LSUN Churches in Figure 1b. Note that, Variant 3 outperforms Variant 2 for all steps less than or equal to 125 and vice versa for larger steps. Similar phenomenon can be observed for CelebAHQ 256 as shown in Figure 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.

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 σ 2 t = (1 -α t ) · β t +1 .

(a) CelebAHQ 256 DDPM Variants

(b) LSUN Churches DDPM Variants

(c) FFHQ 256 DDPM Variants

LatentDiffusionModel(CelebAHQ256)

Latent Diffusion Model (LSUN Churches)

80

172 -

DDPMVariant1

DDPM Variant 1

DDPM Variant 3

DDPMVariant3

Clip-FID

FID

20

50

100

200

5007501000

20

50

100

200

500

750 1000

Number of Neural Network Calls

NumberofNeural NetworkCalls

Latent Diffusion Model(FFHQ256)

80

DDPMVariant1

DDPMVariant2

DDPMVariant3

Clip-FID

98765 4

3

20

50

100

200

500

7501000

Number of Neural Network Calls

A.6 Poisson Midpoint Variant Details

Similar to DDPM, we perform our experiments for different choice of coefficients a t , b t , σ t . Variants 2 and 3 of Poisson Midpoint are defined with the same coefficients of the corresponding DDPM variants 2 and 3. Poisson Midpoint Method introduces additional noise due to the randomness present in H i . This becomes very large whenever K is large. Hence we reduce this randomness by reducing the variance of the Gaussian as suggested by [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.

B DPM-Solver Variant Details

For unconditional image generation of high-resolution images, third order (order=3) Multistep DPMSolvers 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.

C Proof of Theorem 1

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 ˜ X tK + 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 write

Where ˜ Z tK + i := Z tK + i + B tK + i + S tK + i such that 'bias' B tK + i is defined as

and the variance S tK + i is defined as:

C.2 Random Function Representation

By Equation (1), the iterates of S ( A,G, Γ , α K ) satisfy:

By Lemma 1, we can write down the interpolating process as:

Note that while Z tK + i are i.i.d. N (0 , I ) vectors, ˜ Z tK + 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 ( ˜ X tK + i ) t,i can be bounded by bounding the KL divergence between ( Z tK + i ) t,i and ( ˜ Z tK + 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,

CelebAHQ 256

We refer to [44, Lemma 4.18] for a proof.

Lemma 3 (Data Processing Inequality) . Let F : R k 1 → R k 2 be any measurable function. Let P, Q be any probability distributions over R k 1 . Then, the following inequality holds:

The 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 = ˜ X 0 . Then, there exists a measurable function F T such that the following hold almost surely.

and

By Lemma 3 and Lemma 4, we have:

In the last step, we have used the chain rule (Lemma 2).

C.3 Controlling KL Divergence via Wasserstein Distances

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,

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,

From Lemma 1, we show that ˜ Z τ = Z τ + B τ + S τ . Conditioned on ˜ Z 0: τ -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 A 1 + A 2 √ 2 ∼ N (0 , I ) . Therefore, letting Z τ, 1 , Z τ, 2 be i.i.d from N (0 , I ) independent of ˜ Z 0: τ -1 , X 0 , S τ , B τ , we can write:

In the second step, we have used the inequality in Lemma 5. In the third step we have used Lemma 6.

C.4 Bounding the Error Along the Trajectory:

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 + √ 2 S τ | ˜ Z 0: τ -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:

∥ N ∥ ≤ β almost surely, E N = 0 and E NN ⊺ = Σ .

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,

We also have the crude bound:

The 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 ω ∈ Ω .

Lemma 9 (Wasserstein Distance Between Gaussians, [36]) .

Lemma 10. Let

Define the random variable

Proof. Note that

Now, define the random variable β

Clearly, β is measurable with respect to the sigma algebra of ˜ Z 0: 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:

√ 2 S tK + i takes its values in a one-dimensional sub-space almost surely when conditioned on ˜ Z 0: Kt + i -1 , Z 0: Kt + i -1 , X 0 . It is independent of Z tK + i conditioned on ˜ Z 0: 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 ˜ Z 0: Kt + i -1 , X 0

Thus, almost surely:

First, note that by Lemma 8 and Jensen's inequality, we have:

First consider W 2 2 ( ¯ P,P ) . Conditioned on ˜ Z 0: 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:

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 ˜ Z 0: Kt + i -1 , X 0 ) and ν ≤ β 2 K to conclude:

Combining this with Equation (24), we conclude that for any r ≥ 1 :

C.5 Finishing The Proof

We combine Equations (16) and (17) with Lemma 10 to conclude the result of Theorem 1.

D Proof of Lemma 7

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 ′ ,e 1 ⟩ + ⟨ N ,e 1 ⟩ √ 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:

Let Z 1 := ⟨ Z , e 1 ⟩ and Z ′ 1 := ⟨ Z ′ , e 1 ⟩ , N 1 := ⟨ N , e 1 ⟩ . Wedefine m = 1+ 1 ν (and note throughout the proof that a value of infinity can be easily handled). We see that Z ′ 1 + N 1 √ 1+ ν = √ 1 -1 m Z ′ 1 + √ 1 -1 m N 1

Note 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 of Z ′ 1 + N 1 √ 1+ ν to that of Z 1 at the point x .

Where 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 )+2 mN 1 N ′ 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:

Proof. 1. Note that r ( m ) ≤ 1 2( m 2 -1) . By triangle inequality, we have:

A direct calculation shows the identity for E Q 1 .

Now, consider E Q 2 1 . We follow the proof of [52, Lemma 4.5], to conclude the following inequalities. Since E Q 1 ≤ 0 and 1 2( m 2 -1) -r ( m ) ≥ 0 , we have:

Note that the parameter σ i found in the original proof satisfies σ i = 1 in our proof. We now proceed with the proof of Lemma 7. Define R ( Q ) = exp( Q ) -1 -Q -Q 2 2 . From the original proof, we conclude via the Talagrand transport inequality that:

From Lemma 11, and using the fact that r ( m ) ≥ 0 , we note that:

Now, let us bound E R ( Q 1 ) . Via a straightforward application of the taylor series, we have almost surely:

From 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 β 2 2 almost surely. Using this in Equation (32), we have:

In 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:

Now, we use the fact that m = 1 + 1 ν , which implies 1 m 2 -1 ≤ 1 ( m -1) 2 ≤ ν 2 . Thus, we have:

Using 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.

E Overdamped Langevin Dynamics

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:

Here, 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:

In order to apply Theorem 1, we will now proceed to bound the quantities in the RHS of Equations (35) and (36)

E.1 Bounding the Moments

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

We will now prove the following growth estimate for the trajectory ˜ X tK . We refer to Section I.2 for its proof.

Lemma 13. Let p ≥ 1 be fixed. There exists a large enough constant ¯ C p which depends only on p such that whenever s -t ≤ 1 αL ¯ C p , we have:

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:

While Lemma 14 controlled ∑ T t =1 E ∥∇ F ( ˜ X tK ) ∥ 2 p in terms of ∑ T t =1 ∥∇ F ( ˜ X tK ) ∥ 2 , in the Lemma below we control ∑ T t =1 ∥∇ F ( ˜ X tK ) ∥ 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 :

We 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 α 3 p -1 L 3 p -1 T p -1 < c p , we must have:

Proof. 1. Using Equation (35), we have:

2. We use Equation (36) and proceed as in item 1.

E.2 Finishing The Proof

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:

F Underdamped Langevin Dynamics

In 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:

Throughout 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 2 d 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 2 d 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 τ = [ ˜ U τ ˜ V τ ] . 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.

For all 0 ≤ j ≤ K -1 ,

Analogous to the analysis of Overdamped Langevin Dynamics, we consider the following dynamics.

F.1 Bounding the Moments:

Using the scaling relations in Section 2.1, we write down the iteration of PS ( A,G, Γ , b, α, K ) for Underdamped Langevin Dynamics as:

Where, α ∆ t := ∑ K -1 i =0 KH i A α ( K -i -1) K G α K [ b ( ˆ ˜ X tK + i ) -b ( ˜ X tK ) ] and Γ α ¯ Y t := .

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.

Define ψ t := ˜ U Kt + ˜ V Kt γ . The following Lemma, proved in Section I.7, gives the time evolution of ψ t .

Lemma 20.

Where ˜ Ψ t ∼ N (0 , 2 α γ I d ) is independent of ˜ X tK but not necessarily indepependent of ˆ ˜ U tK + i for i > 0 .

Lemma 21. Consider the seminorm given by ∥ x ∥ 2 Π = x ⊺ Π x . Let M t,K := sup j ≤ K -1 ∥ ∑ i -1 j =0 A α ( i -1 -j ) K Γ α K Z tK + j ∥ Π . For any 0 ≤ i ≤ K -1 .

1.

2.

For any p ≥ 1 , we have:

We refer to Section I.8 for the proof. Analogous to Section E, for any p ≥ 1 , we define S 2 p ( V ) := ∑ T -1 t =0 ∥ ˜ V tK ∥ 2 p , S 2 p ( ∇ F ) := ∑ T -1 t =0 ∥∇ F ( ˜ U tK ) ∥ 2 p . 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.

2.

Suppose additionally that γ ≥ C 0 √ L for some large enough universal constant C 0 . Then, the bounds above imply the following inequalities:

1.

The 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:

We 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:

N is an integer power of 2 and N ≤ T .

The following satements hold:

1.

2.

Proof. We consider a t := ∇ F ( ˜ U ( t -1) K ) , b t := ˜ V ( 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 23

Applying Lemma 26, we conclude:

Similarly, we have:

Applying Lemma 26, we conclude:

Combining Lemmas 24 and 22 gives the following theorem.

Theorem 4. Fix p ≥ 1 . There exist constants C p , c p , ¯ c p > 0 such that whenever: γ ≥ C p √ L , αγ < c p , α 3 p -1 T p -1 L 2 p γ p +1 < ¯ c p , the following results hold:

F.2 Bounding the Bias:

Now, consider the following term in the statement of Theorem 1:

Using Lemma 18, and under the conditions of Theorem 4 holding with p = 1 , we have:

In 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:

In the first step we have used the fact that N t is independent of sup 0 ≤ i ≤ K -1 ∥ ˆ ˜ U tK + i -˜ U tK ∥ 2 and the second step follows from the fact that E N 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.

F.3 Bounding the Variance:

Now consider β tK + i = ∥ K Γ -1 α K G α K [ b ( ˆ ˜ X tK + i ) -b ( ˜ X tK )] ∥ . Using Lemma 18 and item 1 of Lemma 21, we have:

Therefore, we note that for any p ≥ 1 , under the assumptions of Theorem 4, we must have:

F.4 Finishing the Proof

For p ≥ 1 , we define ∆ ( p ) = E ( F ( U 0 + V 0 γ ) -F ( x ∗ )) p + E ∥ V 0 ∥ 2 p +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):

G Convergence Under Isoperimetry Assumptions

We 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:

We 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:

We 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 :

Let Σ be the covariance of X (it exists due to the assumption of PI). Therefore,

In 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 < η ≤ λ 2 L 2 . Then, LMC satisfies:

Where X 0 , . . . , X N are the iterates of OLMC.

Whenever ϵ ≤ 1 , we take α K = λϵ 2 16 dL 2 . T = 1 αλ log ( 2 KL ( Law ( X 0 ) ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ π ⋆ ) ϵ 2 ) . With the suggested initialization in [45] i.e X 0 ∼ N ( x , I L ) for any stationary point x of F we have KL ( Law ( X 0 ) ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ π ⋆ ) ≤

Let X 0 , . . . , X TK be the iterates of OLMC with step-size α K with the parameters as chosen above. Applying Theorem 5, we conclude: KL ( Law ( X TK ) ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ π ⋆ ) ≤ ϵ 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:

Now, by Pinsker's inequality and the triangle inequality for TV distance, we have:

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:

F is L -smooth and that π ⋆ satisfies λ -LSI.

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 ) .

Let X 0 , . . . , X TK be the iterates of ULMC with step-size α K . We take: α = ˜ Θ( √ ϵ L ( λ Ld ) 5 12 ) , K = ˜ Θ( 1 α λ 1 3 d 1 3 L 5 6 ) 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:

Now, using Theorem 6, along with the triangle inequality for TV, we conclude:

This proves the bounds for ULMC in Table 1.

H Some Technical Results

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:

Lemma27. 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,

Proof. Let λ ∈ R . Consider the random variable M t = exp( ∑ t s =1 λ ⟨ g s , Z s ⟩ -λ 2 ∥ g s ∥ 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:

Now, for any variable X , P ( X + > h ) = P ( X > h ) whenever h ≥ 0 . Now using the fact that: E X ≤ E X + = ∫ ∞ 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 :

We conclude the result by choosing λ 2 p = 1 E ( ∑ T s =1 ∥ g s ∥ 2 ) p .

Lemma 28. Let K be any positive integer. Consider G ∼ Bin ( K, 1 K ) . Then for any p ≥ 1 , we must have: E G p ≤ C p for some constant depending only on p .

Proof. The proof is simple for K = 1 . Assume K ≥ 2 . For any 0 ≤ n ≤ K , let p n := P ( G = n ) . Then,

In 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 ≤ 2 p ∗ n . Thus, we must have:

I Proofs of Technical Lemmas

I.1 Proof of Lemma 12

Proof. 1. This follows by applying the triangle inequality to the definition of ˆ ˜ X tK + i .

Consider:

This follows by a direct application of Doob's inequality.

I.2 Proof of Lemma 13

Proof. We will use b () and ∇ F () interchangeably in the proof. Consider the update

Therefore, we have for any s > t :

In this proof only, for any random variable U , we let M p ( U ) := ( E [ ∥ U ∥ p | ˜ X tK ]) 1 p . Using the triangle inequality for M p , we conclude:

))

Now, consider ∆ h = ∑ K -1 i =0 H i ( b ( ˆ ˜ X hK + i ) -b ( ˜ X hK )) . We must have

In the second line, we have used the fact that N h is independent of sup 0 ≤ i ≤ K -1 ∥ ˆ ˜ X hK + i -˜ X hK ∥ 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:

Therefore, we must have:

Now, supposing s -t ≤ 1 ¯ C p αL for some large enough constant ¯ C p which depends only on p . This ensures that αL ( s -t )(1 + αLC p ) ≤ 1 2 in the equation above. Thus, we have:

I.3 Proof of Lemma 26

Proof. ¯ a k := 1 |T k | ∑ j ∈T k a j Consider the following quantity:

Now, by Jensen's inequality, we must have: ∥ ¯ a k ∥ 2 ≤ 1 |T k | ∑ j ∈T k ∥ a j ∥ 2 . Therefore, we conclude:

Now, consider

Using this with Equation (70), we conclude the statement of the lemma.

I.4 Proof of Lemma 14

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 ¯ C 2 p ⌉ for a large enough constant ¯ C 2 p such that the result of Lemma 13 with s -t ≤ N . Applying Lemma 13, we conclude:

Applying Lemma 26

Note that whenever the constant ¯ C 2 p defining N above is large enough, this implies

I.5 Proof of Lemma 15

Proof. Since ∇ F is L Lipschitz, we have:

Also note that:

Thus, summing Equation (75) from t = 0 to t = T -1 , we conclude:

In the last step, we have used the fact that 2 |⟨∇ F ( ˜ X tK ) , ∆ t ⟩| ≤ ∥∇ F ( ˜ X tK ) ∥ 2 + ∥ ∆ t ∥ 2 and the assumption that αL ≤ c p for some small enough constant c p .

Therefore, we conclude:

The properties of Gaussians show that: E ( ∑ T -1 t =0 ∥ ¯ Y 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:

In 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 ˆ ˜ X tK + i for 0 ≤ i ≤ K -1 . In the last step, we have used the fact that E ( N t ) 2 p ≤ C p for some constant C p (Lemma 28).

Plugging all of these bounds, and taking κ in Equation (78) to be large enough, we conclude:

I.6 Proof of Lemma 18

Proof. For the sake of clarity, define x = 1 -exp( -γh ) . ρ = (1 -x ) 2 γ 1 √ 2 γ ( h -2(1 -x ) γ + (1 -x 2 ) 2 γ )(1 -x 2 ) Algebraic manipulations show that:

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,

Now, we note that for scalars A,B > 0 and ρ ∈ [0 , 1] , we have:

Thus, we conclude from the above computations that:

We have used the fact that:

Using the fact that for PSD matrices A,B A ⪯ B implies B -1 ⪯ A -1 , we have:

The second inequality follows easily by elementary manipulations.

I.7 Proof of Lemma 20

Proof. Explicit computations show that:

Now, note that [ I d I 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:

A straight forward calculation shows that: ˜ Ψ t := [ I d I d γ ] Γ α ˜ Y t ∼ M (0 , 2 h γ I d ) .

I.8 Proof of Lemma 21

Proof. 1. Note that

We bound each of the terms separately.

Plugging these inequalities gives the result.

Since ∇ F is L -Lipschitz, we have:

Using Lemma 20 to evaluate ψ t +1 -ψ t , we conclude:

In the second step, we have use the fact that ∇ F is L -Lipschitz along with the CauchySchwarz inequality. In the third step, we have used the AM-GM inequality which states that for any a, b ≥ 0 , we have 2 ab ≤ a 2 + b 2

Similarly, we note that:

We have again used the AM-GM inequality in the second step above. Convexity of x → x 2 implies:

Plugging these upper bounds into Equation (85), we have:

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 Γ α K Z tK + j along with the union bound over all 0 ≤ i ≤ K -1 , we conclude:

We conclude the moment bounds by integrating tail probabilities.

I.9 Proof of Lemma 22

Proof. 1. Let g := e

Now, consider:

In the second step, we have use the fact that | 2 g (1 -g ) γ ⟨∇ F ( ˜ U tK ) , ˜ V tK ⟩| ≤ g (1 -g ) ∥∇ F ( ˜ U tK ) ∥ 2 γ 2 + g (1 -g ) ∥ ˜ V tK ∥ 2 . In the third step, we have used the fact that: 2 g ⟨ ˜ V tK , ( I -Π) α ∆ t ⟩ ≤ ( √ g -g ) ∥ ˜ V tK ∥ 2 + ∥ ( I -Π) α ∆ t ∥ 2 1 - √ g , -2(1 -g ) γ ⟨∇ F ( ˜ U tK ) , ( I -Π) α ∆ t ⟩ ≤ 1 -g γ 2 ∥∇ F ( ˜ U tK ) ∥ 2 +(1 -g ) ∥ ( I -Π) α ∆ ∥ 2 and ∥ ( I -Π)( α ∆ t +Γ α ˜ Y t ) ∥ 2 ≤ 2 ∥ ( I -Π)( α ∆ t ) ∥ 2 +2 ∥ ( I -Π)Γ α ˜ Y t ∥ 2 .

With extention of notation, let S 2 (( I -Π) α ∆) := ∑ T -1 t =0 ∥ ( I -Π)∆ t ∥ 2 , S 2 (( I -Π)Γ α ˜ Y ) := ∑ T -1 t =0 ∥ ( I -Π)Γ α ˜ Y t ∥ 2 . Using Equation (91), we conclude:

Note that ( I -Π)Γ α ˜ Y t ∼ N (0 , (1 -e -2 γα ) I d ) . By properties of Gaussians, it is clear that for any p ≥ 1 , we must have E ( S 2 (( I -Π)Γ α ˜ Y ) ) p ≤ C p ( γαTd ) p . By Lemma 27, we conclude:

By L smoothness of F , we have: ∥ ( I -Π)( α ∆ t ) ∥ ≤ LαN t sup 0 ≤ i ≤ K ∥ ˆ ˜ U tK + i -˜ U tK ∥ . Using the result in Lemma 21, we conclude:

In the second step, we have used jensen's inequality to show that ( 1 T ∑ t a 2 t ) p ≤ 1 T ∑ t a 2 p t . In the third step, we have used the fact that N t is independent of ˜ X tK and that E N 2 p 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:

Here, we have used that fact that 1 - √ g ≥ Cγα whenever γα < 1 . We now use the AM-GM inequality to show that

Similarly, we have:

Applying this to the RHS of Equation (94) and re-arranging, we conclude the statement of the lemma.

Using similar methods as in item 1 above, we apply Lemma 27 and collect the following moment bounds:

By a calculation similar to that in item 1, we have:

Summing item 2 of Lemma 21 over t = 0 to t = T -1 and applying the triangle inequality for p -th moments, we have:

In the last step, we have use the fact that γα < 1 and Lα γ < c 0 . Applying AM-GM inequality to C p √ γ α [ E ( S 2 ( ∇ F )) p ] 1 2 p and C p L √ 1 γα [ E ( S 2 ( V )) p ] 1 2 p similar to item 1, we conclude the result.

I.10 Proof of Lemma 23

Proof.

In this proof only, for any random variable W and any projection operator P over R 2 d we let M p ( W ; P ) := ( E [ ∥P W ∥ p | ˜ X tK ]) 1 p . We use the convention that M p ( W ) = ( E ∥ W ∥ p ) 1 p . Using the triangle inequality for M p , we conclude:

Similarly, we have:

By scaling relations given in Section 2.1, we have: ∑ s -1 h = t Π A α ( s -h -1) Γ α ˜ Y h ∼ N (0 , Π(Γ 2 ( s -t ) α )Π) and ∑ s -1 h = t ( I -Π) A α ( s -h -1) Γ α ˜ Y 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:

Notice that ∥ Π( A α ( s -t ) -I ) ˜ X tK ∥ ≤ α ( s -t ) ∥ ˜ V tK ∥ and ∥ ( I -Π)( A α ( s -t ) -I ) ˜ X tK ∥ ≤ αγ ( s -t ) ∥ ˜ V tK ∥ . This implies:

Notice that: ∥ Π A α ( s -h -1) G α b ( ˜ X hK ) ∥ ≤ α γ ∥∇ F ( ˜ U hK ) ∥ and ∥ ( I -Π) A α ( s -h -1) G α b ( ˜ X hK ) ∥ ≤ α ∥∇ F ( ˜ U hK ) ∥ . This implies:

Notice that: ∥ Π A α ( s -h -1) α ∆ h ∥ ≤ αLN h γ sup 0 ≤ i ≤ K -1 ∥ ˆ ˜ U hK + i -˜ U hK ∥ and ∥ ( I -Π) A α ( s -h -1) α ∆ h ∥ ≤ αLN h sup 0 ≤ i ≤ K -1 ∥ ˆ ˜ U hK + i -˜ U hK ∥ . Applying item 1 in Lemma 21, we conclude:

Therefore, 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 ˜ X hK 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 ( ˜ U hK ) with that of ∇ F ( ˜ U 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 ˜ V hK in terms of ˜ X tK and ˜ X tK -˜ X hK . Similarly, 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:

Similarly considering Equation (98), we have:

From Equations (102) and (103), we conclude:

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The contents of the paper match the claims made in the abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have discussed the limitations in a separate section.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We explicitly reveal all the hyperparameters and open source models necessary to reproduce our results.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We do not have institutional approval for code release. We provide detailed descriptions of the algorithms and hyper-parameters.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: It is common in image generation literature to compute FID over 50k generated images, without reporting error bars since the fluctuations are usually negligible.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

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.

8. 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).

9. 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 ?

Answer: [Yes]

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).

10. 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).

11. 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe discretization proposed appears to be completely novel, and the analytic framework (using CLT type arguments) has not appeared before in this context.\n\nThe rates are state of the art in KL, clearly surpassing the previous best known discretizations in KL without any caveats.\n\nThe experimental evidence is thorough and appears to suggest that this discretization is indeed better than a naive one, for the diffusion model setting.\n\nWEAKNESSES:\nThe algorithm as described in the paper is somewhat cumbersome to read. Perhaps some simpler exposition or visualization may be helpful.\n\nThe rates obtained are worse than those of the randomized midpoint discretizations (noting that these latter rates are only valid in Wasserstein distance).\n\nQUESTIONS:\nAlthough 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.\n\nThe 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?\n\nIs 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?\n\nMinor:\n\n59: stochastic approximations LMC -> \"stochastic approximations of LMC\"\n\n70: denotes identity -> denotes the identity\n\n108 \"Motive\" -> \"motivate\"\n\n206: Theorem 1; I believe this is meant to refer to Theorem 3.\n\nMany equations/inline formulas are missing punctuation. For instance Eq. 6, Steps 3-5 (ll 145-149).\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n### Originality and Significance\n\nThe idea is quite interesting, and binning the unit interval to perform a randomized midpoint discretization is novel in my view. This clearly is not a straightforward innovation despite seeming like one. The theoretical statements also help characterize the benefits of this new strategy, which appear to be shown under weaker conditions than those required for sampling. Practically, this is also very appealing as the computational costs for the applications to generative modelling is lower than now-standard techniques when fixing the requirement (i.e., low FID).\n\n### Quality and Clarity\n\nThe idea is well-motivated, but I feel like the paper could be a lot more clearly presented, especially to be accessible to broader audiences. Please see the weaknesses / questions section below.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. What is $p_{\\tau}$ in Eq. (4)?\n2. 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?\n3. 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.\n4. 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)?\n5. 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.\n6. 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')$?\n7. 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?\n8. 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.\n9. 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?\n10. 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?\n11. Where is the computational complexity discussed in the paragraph on the top of page 5?\n12. How is there no dependence on the $\\{X_{t}\\}$ sequence in Eq. 10?\n13. Why is this strange dependence on 14th moment in Assumption 2 present?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBased 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.\n\nSTRENGTHS:\nThis paper is technically solid with a complete list of related work. The contributions are stated in a clear way. Also, a complete set of numerical experiments on diffusion models is included to illustrate the effectiveness of the proposed method.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nOne 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.\n\nLIMITATIONS:\nA few typos (the list is probably incomplete) in the current version of the manuscript are also listed below:\n\n1. From line 179-180, the two phrases \"to OLMC\" and \"to ULMC\" appeared twice. The two sentences probably need to revised.\n2. In line 206, the subtitle should be \"Sketch for Theorem 3\". Moreover, the following sentence should be \"This is similar to Theorem 2\".", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}