Title: Reverse Transition Kernel: A Flexible Framework to Accelerate Diffusion Inference
Venue: NeurIPS 2024
Accepted: yes
Keywords: diffusion inference

reverse transition kernel
Abstract: To generate data from trained diffusion models, most inference algorithms, such as DDPM, DDIM, and other variants, rely on discretizing the reverse SDEs or their equivalent ODEs. In this paper, we view such approaches as decomposing the entire denoising diffusion process into several segments, each corresponding to a reverse transition kernel (RTK) sampling subproblem. Specifically, DDPM uses a Gaussian approximation for the RTK, resulting in low per-subproblem complexity but requiring a large number of segments (i.e., subproblems), which is conjectured to be inefficient. To address this, we develop a general RTK framework that enables a more balanced subproblem decomposition, resulting in $\tilde O(1)$ subproblems, each with strongly log-concave targets. We then propose leveraging two fast sampling algorithms, the Metropolis-Adjusted Langevin Algorithm (MALA) and Underdamped Langevin Dynamics (ULD), for solving these strongly log-concave subproblems. This gives rise to the RTK-MALA and RTK-ULD algorithms for diffusion inference. In theory, we further develop the convergence guarantees for  RTK-MALA and RTK-ULD in total variation (TV) distance: RTK-ULD can achieve $\epsilon$ target error within $\tilde{\mathcal O}(d^{1/2}\epsilon^{-1})$ under mild conditions, and  RTK-MALA enjoys a $\mathcal{O}(d^{2}\log(d/\epsilon))$   convergence rate under slightly stricter conditions. These theoretical results surpass the state-of-the-art convergence rates for diffusion inference and are well supported by numerical experiments.

Reverse Transition Kernel: A Flexible Framework to Accelerate Diffusion Inference

Abstract

To generate data from trained diffusion models, most inference algorithms, such as DDPM [17], DDIM [31], and other variants, rely on discretizing the reverse SDEs or their equivalent ODEs. In this paper, we view such approaches as decomposing the entire denoising diffusion process into several segments, each corresponding to a reverse transition kernel (RTK) sampling subproblem. Specifically, DDPM uses a Gaussian approximation for the RTK, resulting in low per-subproblem complexity but requiring a large number of segments (i.e., subproblems), which is conjectured to be inefficient. To address this, we develop a general RTK framework that enables a more balanced subproblem decomposition, resulting in ˜ O (1) subproblems, each with strongly log-concave targets. We then propose leveraging two fast sampling algorithms, the Metropolis-Adjusted Langevin Algorithm (MALA) and Underdamped Langevin Dynamics (ULD), for solving these strongly log-concave subproblems. This gives rise to the RTK-MALA and RTK-ULD algorithms for diffusion inference. In theory, we further develop the convergence guarantees for RTK-MALA and RTK-ULD in total variation (TV) distance: RTK-ULD can achieve ϵ target error within ˜ O ( d 1 / 2 ϵ -1 ) under mild conditions, and RTKMALAenjoys a O ( d 2 log( d/ϵ )) convergence rate under slightly stricter conditions. These theoretical results surpass the state-of-the-art convergence rates for diffusion inference and are well supported by numerical experiments.

1 Introduction

Generative models have become a core task in modern machine learning, where the neural networks are employed to learn the underlying distribution from training examples and generate new data points. Among various generative models, denoising diffusions have produced state-of-the-art performance in many domains, including image and text generation [14, 2, 28, 29], text-to-speech synthesis [27], and scientific tasks [34, 37, 5]. The fundamental idea involves incrementally adding noise and gradually transform the data distribution to a prior distribution that is easier to sample from, e.g., Gaussian distribution. Then, diffusion models parameterize and learn the score of the noised distributions to progressively denoise samples from priors and recover the data distribution [36, 32].

Under this paradigm, generating data in denoising diffusion models involves solving a series of sampling subproblems, i.e., generating samples from the distribution after one-step denoising. To this end, DDPM [17], one of the most popular sampling methods in diffusion models, has been developed for this purpose. DDPM uses the Gaussian transition with carefully designed mean and covariance to approximate the solutions to these sampling subproblems. By sequentially stacking

Hanze Dong

UIUC

a series of Gaussian transitions, DDPM successfully generates high-quality samples that follow the data distribution. The empirical success of DDPM has immediately triggered various followup work [33, 25], aiming to accelerate the inference process and improve the generation quality. Alongside rapid empirical research on diffusion models and DDPM-like sampling algorithms [19, 39], theoretical studies have emerged to analyze the convergence and sampling error of DDPM. In particular, [21, 24, 8, 7, 3, 9] have established polynomial convergence bounds, in terms of dimension d and target sampling error ϵ , for the generation process under various assumptions. Previous work usually assume the score estimation error to construct the sampling analysis [8, 15, 3, 19]. A typical bound under minimal data assumptions on the score of the data distribution is provided by [8, 3], which establishes an ˜ O ( dϵ -2 ) score estimation guarantees to sample from data distribution within ϵ -sampling error in the Total Variation (TV) distance.

In essence, the denoising diffusion process can be approached through various decompositions of sampling subproblems, where the overall complexity depends on the number of these subproblems multiplied by the complexity of solving each one. Within this framework, DDPM can be regarded as a specific solver for the denoising diffusion process that heavily prioritizes the simplicity of subproblems over their quantity. In particular, it adopts simple one-step Gaussian approximations for the subproblems, with O (1) computation complexity, but needs to deal with a relatively large number-approximately O ( dϵ -2 ) -of target subproblems to ensure the cumulative sampling error is bounded by ϵ in TV distance. This imbalance raises the question of whether the DDPM-like approaches stand as the most efficient algorithm, considering the extensive potential subproblem decompositions of the denoising diffusion process. We therefore aim to:

accelerate the inference of diffusion models via a more balanced subproblem decomposition in the denoising process.

In this work, we propose a novel framework called reverse transition kernel (RTK) to achieve exactly that. Our approach considers a generalized subproblem decomposition of the denoising process, where the difficulty of each sampling subproblem and the total number of subproblems are determined by the step size parameter η . Unlike DDPM, which requires setting η = ϵ 2 , resulting in approximately ˜ O (1 /η ) = ˜ O (1 /ϵ 2 ) subproblems, our framework allows η to be feasible in a broader range. Furthermore, we demonstrate that a more balanced subproblem decomposition can be attained by carefully selecting η = Θ(1) as a constant, resulting in approximately ˜ O (1) sampling subproblems, with each target distribution being strongly log-concave. This nice property further enables us to efficiently solve the sampling subproblems using well-established acceleration techniques, such as Metropolis Hasting step and underdamped discretization, without encountering many subproblems. Consequently, our proposed framework facilitates the design of provably faster algorithms than DDPM for performing diffusion inference. Our contributions can be summarized as follows.

We propose a flexible framework that enhances the efficiency of diffusion inference by balancing the quantity and hardness of RTK sampling subproblems used to segment the entire denoising diffusion process. Specifically, we demonstrate that with a carefully designed decomposition, the number of sampling subproblems can be reduced to approximately ˜ O (1) , while ensuring that all RTK targets exhibit strong log-concavity. This capability allows us to seamlessly integrate a range of well-established sampling acceleration techniques, thereby enabling highly efficient algorithms for diffusion inference.

Building upon the developed framework, we implement the RTK using the Metropolis-Adjusted Langevin Algorithm (MALA), making it the first attempt to adapt this highly accurate sampler for diffusion inference. Under slightly stricter assumptions on the estimation errors of the energy difference and score function, we demonstrate that RTK-MALA can achieve linear convergence with respect to the sampling error ϵ , specifically O (log(1 /ϵ )) , which significantly outperforms the ˜ O (1 /ϵ 2 ) convergence rate of DDPM [8, 3]. Additionally, we consider the practical diffusion model where only the score function is accessible and develop a score-only RTK-MALA algorithm. We further prove that the score-only RTK-MALA algorithm can achieve an error ϵ with a complexity of ˜ O ( ϵ -2 / ( u -1) · 2 u ) , where u can be an arbitrarily large constant, provided the energy function satisfies the u -th order smoothness condition.

We further implement Underdamped Langevin Dynamics (ULD) within the RTK framework. The resulting RTK-ULD algorithm achieves a state-of-the-art complexity of ˜ O ( d 1 / 2 ϵ -1 ) for both d

and ϵ dependence under some mild assumptions as [13]. Compared with the ˜ O ( dϵ -2 ) complexity guarantee for DDPM, it improves the complexity with an ˜ O ( d 1 / 2 ϵ -1 ) factor. This result also matches the state-of-the-art convergence rate of the ODE-based methods [9], though those methods require Lipschitz conditions for both the ground truth score function and the score neural network.

2 Preliminaries

In this section, we first introduce the notations used in subsequent sections. Then, we present several distinct Markov processes to demonstrate the procedures for adding noise to existing data and generating new data. Besides, we specify the assumptions required for the target distribution in our algorithms and analysis.

Notations. We say a complexity h : R → R to be h ( n ) = O ( n k ) or h ( n ) = ˜ O ( n k ) if the complexity satisfies h ( n ) ≤ c · n k or h ( n ) ≤ c · n k [log( n )] k ′ for absolute contant c, k and k ′ . We use the lowercase bold symbol x to denote a random vector, and the lowercase italicized bold symbol x represents a fixed vector. The standard Euclidean norm is denoted by ∥ · ∥ . The data distribution is presented as p ∗ ∝ exp( -f ∗ ) . Besides, we define two Markov processes R d , i.e.,

In the above notations, T presents the mixing time required for the data distribution to converge to specific priors, K denotes the iteration number of the generation process, and η signifies the corresponding step size. Further details of the two processes are provided below.

Adding noise to data with the forward process. The first Markov process { x t } corresponds to generating progressively noised data from p ∗ . In most denoising diffusion models, { x t } is an Ornstein-Uhlenbeck (OU) process shown as follows

If we denote underlying distribution of x t as p t ∝ exp( -f t ) meaning f 0 = f ∗ , the forward OU process provides an analytic form of the transition kernel, i.e.,

for any t ′ ≥ t , where p t ′ ,t denotes the joint distribution of ( x t ′ , x t ) . According to the Fokker-Planck equation, we know the stationary distribution for SDE. (1) is the standard Gaussian distribution.

Denoising generation with a reverse SDE. Various theoretical works [21, 24, 8, 7, 3] based on DDPM [17] consider the generation process of diffusion models as the reverse process of SDE. (1) denoted as { x ← t } . According to the Doob's h -Transform, the reverse SDE, i.e., { x ← t } , follows from

whose underlying distribution p ← t satisfies p T -t = p ← t . Similar to the definition of transition kernel shown in Eq. (2), we define p ← t ′ | t ( x ′ | x ) = p ← t ′ ,t ( x ′ , x ) /p ← t ( x ) for any t ′ ≥ t ≥ 0 and name it as reverse transition kernel (RTK).

To implement SDE. (3), diffusion models approximate the score function ∇ ln p t with a parameterized neural network model, denoted by s θ ,t , where θ denotes the network parameters. Then, SDE. (3) can be practically implemented by

with a standard Gaussian initialization, x 0 ∼ N ( 0 , I ) . Eq. (4) has the following closed solution

which is exactly the DDPM algorithm.

DDPM approximately samples the reverse transition kernel. DDPM can also be viewed as an approximated sampler for RTK, i.e., p ← t ′ | t ( x ′ | x ) for some t ′ > t . In particular, the update of DDPM at the iteration k applies the Gaussian-type transition, i.e.,

to approximate the distribution p ← ( k +1) η | kη ( ·| x ) [17]. Specifically, by the chain rule of KL divergence, the gap between the data distribution p ∗ and the generated distribution p T satisfies

where K = T/η is the total number of iterations. For DDPM, to guarantee a small sampling error, we need to use a small step size η to ensure that p ( k +1) η | kη is sufficiently close to p ← ( k +1) η | kη . Then, the required iteration numbers K = T/η will be large and dominate the computational complexity. In Chen et al. [8, 10], it was shown that one needs to set η = ˜ O ( ϵ 2 ) to achieve ϵ sampling error in TV distance (assuming no score estimation error) and the total complexity is K = ˜ O (1 /ϵ 2 ) .

Intuition for General Reverse Transition Kernel. As previously mentioned, DDPM approximately solves RTK sampling subproblems using a small step size η . While this allows for efficient one-step implementation, it necessitates a large number of RTK sampling problems. This naturally creates a trade-off between the quantity of RTK sampling problems and the complexity of solving them. To address this, one can consider a larger step size η , which results in a relatively more challenging RTK sampling target p ← ( k +1) η | kη and a reduced number of sampling problems ( K = T/η ). By examining a general choice for the step size η , the generation process of diffusion models can be depicted through a comprehensive framework of reverse transition kernels, which will be explored in depth in the following section. This framework enables the design of various decompositions for RTK sampling problems and algorithms to solve them, resulting in an extensive family of generation algorithms for diffusion models (that encompasses DDPM). Consequently, this also offers the potential to develop faster algorithms with lower computational complexities, e.g., applying fast sampling algorithms for sampling the RTK, i.e., p ← ( k +1) η | kη with a reasonably large η .

General Assumptions. Similar to the analysis of DDPM [8, 7], we make the following assumptions on the data distribution p ∗ that will be utilized in the theory.

[A1] For all t ≥ 0 , the score ∇ ln p t is L -Lipschitz.

[A2] The second moment of the target distribution p ∗ is upper bounded, i.e., E p ∗ [ ∥·∥ 2 ] = m 2 2 .

Assumption [A1] is standard one in diffusion literature and has been used in many prior works [4, 10, 21, 9]. Moreover, we do not require the isoperimetric conditions, e.g., the establishment of the log-Sobolev inequality and the Poincaré inequality for the data distribution p ∗ as [21], and the convex conditions for the energy function f ∗ as [4]. Therefore, our assumptions cover a wide range of highly non-log-concave data distributions. We emphasize that Assumption [A1] may be relaxed only to assume the target distribution is smooth rather than the entire OU process, based on the technique in [7] (see rough calculations in their Lemmas 12 and 14). We do not include this additional relaxation in this paper to clarify our analysis. Assumption [A2] is one of the weakest assumptions being adopted for the analysis of posterior sampling.

3 General Framework of Reverse Transition Kernel

This section introduces the general framework of Reverse Transition Kernel (RTK). As mentioned in the previous section, the framework is built upon the general setup of segmentation: each segment has length η ; within each segment, we generate samples according to the RTK target distributions. Then, the entire generation process in diffusion models can be considered as the combination of a series of sampling subproblems. In particular, the inference process via RTK is displayed in Alg. 1.

The implementation of RTK framework. We begin with a new Markov process { ˆ x kη } k =0 , 1 ,...,K satisfying Kη = T , where the number of segments K , segment length η , and length of the entire process T correspond to the definition in Section 2. Consider the Markov process { ˆ x kη } as the

Algorithm 1 INFERENCE WITH REVERSE TRANSITION KERNEL (RTK)

1: Input: Initial particle ˆ x 0 sampled from the standard Gaussian distribution, Iteration number K , Step size η , required convergence accuracy ϵ ;

2: for k = 0 to K -1 do

3: Draw sample ˆ x ( k +1) η with MCMCs from ˆ p ( k +1) η | kη ( ·| ˆ x kη ) which approximates the groundtruth reverse transition kernel, i.e.,

4: return ˆ x K .

generation process of diffusion models with underlying distributions { ˆ p kη } , we require ˆ p 0 = N ( 0 , I ) and ˆ p Kη ≈ p ∗ , which is similar to the Markov process { x ← kη } . In order to make the underlying distribution of output particles close to the data distribution, we can generate ˆ x kη with Alg. 1, which is equivalent to the following steps:

Initialize ˆ x 0 with an easy-to-sample distribution, e.g., N ( 0 , I ) , which is close to p Kη .

Update particles by drawing samples from ˆ p ( k +1) η | kη ( ·| ˆ x kη ) , which satisfies

Under these conditions, if ˆ p kη ( z ) ≈ p ( K -k ) η ( z ) , then we have

for any k ∈ { 0 , 1 , . . . , K } . This means we can implement the generation of diffusion models by solving a series of sampling subproblems with target distributions p ← ( k +1) η | kη ( ·| ˆ x kη ) .

The close form of reverse transition kernels. To implement Alg. 1, the most critical problem is determining the analytic form of RTK p ← t ′ | t ( x ′ | x ) for and t ′ ≥ t ≥ 0 which is shown in the following lemma whose proof is deferred to Appendix B.

Lemma 3.1. Suppose a Markov process { x t } with SDE. 1, then for any t ′ > t , we have

The first critical property shown in this Lemma is that RTK p t | t ′ is a perturbation of p t with a l 2 regularization. This means if the score of p t , i.e., ∇ f t , can be well-estimated, the score of RTK, i.e., ∇ log p t | t ′ can also be approximated with high accuracy. Moreover, in the diffusion model, ∇ f t = ∇ log p t is exactly the score function at time t , which is approximated by the score network function s θ,t ( x ) , then

which can be directly calculated with a single query of s θ,t ( x ) . The second critical property of RTK is that we can control the spectral information of its score by tuning the gap between t ′ and t . Specifically, considering the target distribution, i.e., p ( K -k -1) η | ( K -k ) η for the k -th transition, the Hessian matrix of its energy function satisfies

According to Assumption [A1] , the Hessian ∇ 2 f ( K -k -1) η ( x ) = -∇ 2 log p ( K -k -1) η can be lower bounded by -L I , which implies that RTK p ( K -k -1) η | ( K -k ) η will be L -strongly log-concave and 3 L -smooth when the step size is set η = 1 / 2 · log(1 + 1 / 2 L ) . This further implies that the targets

of all subsampling problems in Alg. 1 will be strongly log-concave, which can be sampled very efficiently by various posterior sampling algorithms.

Sufficient conditions for the convergence. According to Pinsker's inequality and Eq. (7), we can obtain the following lemma that establishes the general error decomposition for Alg. 1.

Lemma 3.2. For Alg 1, we have

for any K ∈ N + and η ∈ R + .

It is worth noting that the choice of η represents a trade-off between the number of subproblems divided throughout the entire process and the difficulty of solving these subproblems. By considering the choice η = 1 / 2 · log(1+1 / 2 L ) , we can observe two points: (1) the sampling subproblems in Alg. 1 tend to be simple, as all RTK targets, presented in Lemma 3.1, can be provably strongly log-concave; (2) the total number of subproblems is K = T/η = ˜ O (1) , which is not large. Conversely, when considering a larger η that satisfies η ≫ log(1 + 1 /L ) , the RTK target will no longer be guaranteed to be log-concave, resulting in high computational complexity, potentially even exponential in d , when solving the corresponding sampling subproblems. On the other hand, if a much smaller step size η = o (1) is considered, the target distribution of the sampling subproblems can be easily solved, even with a one-step Gaussian transition. However, this will increase the total number of sampling subproblems, potentially leading to higher computational complexity.

Therefore, we will consider the setup η = 1 / 2 · log(1+1 / 2 L ) in the remaining part of this paper. Now, the remaining task, which will be discussed in the next section, would be designing and analyzing the sampling algorithms for implementing all iterations of Alg. 1, i.e., solving the subproblems of RTK.

4 Implementation of RTK inner loops

In this section, we outline the implementation of Step 3 in the RTK algorithm, which aims to solve the sampling subproblems with strong log-concave targets, i.e., p ← ( k +1) η | kη ( ·| ˆ x kη ) ∝ exp( -g ) . Specifically, we employ two MCMC algorithms, i.e., the Metropolis-adjusted Langevin algorithm (MALA) and underdamped Langevin dynamics (ULD). For each algorithm, we will first introduce the detailed implementation, combined with some explanation about notations and settings to describe the inner sampling process. After that, we will provide general convergence results and discuss them in several theoretical or practical settings. Besides, we will also compare our complexity results with the previous ones when achieving the convergence of TV distance to show that the RTK framework indeed obtains a better complexity by balancing the number and complexity of sampling subproblems.

RTK-MALA. Alg. 2 presents a solution employing MALA for the inner loop. When it is used to solve the k -th sampling subproblem of Alg. 1, x 0 is equal to ˆ x kη defined in Section 3 and used to initialize particles iterating in Alg. 2. In Alg. 2, we consider the process { z s } S s =0 whose underlying distribution is denoted as { µ s } S s =0 . Although we expect µ S to be close to the target distribution p ← ( k +1) η | kη ( ·| x 0 ) , in real practice, the output particles z S can only approximately follow p ← ( k +1) η | kη ( ·| x 0 ) due to inevitable errors. Therefore, these errors should be explained in order to conduct a meaningful complexity analysis of the implementable algorithm. Specifically, Alg. 2 introduces two intrinsic errors:

[E1] Estimation error of the score function: we assume a score estimator, e.g., a well-trained diffusion model, s θ , which can approximate the score function with an ϵ score error, i.e., ∥ s θ,t ( z ) -∇ log p t ( z ) ∥ ≤ ϵ score for all z ∈ R d and t ∈ [0 , T ] .

[E2] Estimation error of the energy function difference: we assume an energy difference estimator r which can approximate energy difference with an ϵ energy error, i.e., | r t ( z ′ , z ) + log p t ( z ′ ) -log p t ( z ) | ≤ ϵ energy for all z , z ′ ∈ R d .

Under these settings, we provide a general convergence theorem for Alg. 2. To clearly convey the convergence properties, we only show an informal version.

Algorithm 2 MALA/PROJECTED MALA FOR RTK INFERENCE

1: Input: Returned particle of the previous iteration x 0 , current iteration number k , inner iteration number S , inner step size τ , required convergence accuracy ϵ ;

2: Draw the initial particle z 0 from

3: for s = 0 to S -1 do

4: Draw a sample ˜ z s from the Gaussian distribution N ( z s -τ · s θ ( z s ) , 2 τ I ) ;

5: if z s +1 ̸∈ B ( z s , r ) ∩ B ( 0 , R ) then

6:

z s +1 = z s ; ▷ This condition step is only activated for Projected MALA.

7: continue ;

8: Calculate the accept rate as

9: Update the particle z s +1 = ˜ z s with probability a , otherwise z s +1 = z s .

10: return z S ;

Theorem 4.1 (Informal version of Theorem C.17) . Under Assumption [A1] -[A2] , for Alg. 1, we choose

and implement Step 3 of Alg. 1 with Alg. 2. Suppose the score [E1] , energy [E2] estimation errors and the inner step size τ satisfy

and the hyperparameters, i.e., R and r , are chosen properly. We have

where ρ is the Cheeger constant of a truncated inner target distribution exp( -g ( z )) 1 [ z ∈ B ( 0 , R )] and Z denotes the maximal l 2 norm of particles appearing in outer loops (Alg. 1).

It should be noted that the choice of η choice ensures the L strong log-concavity of target distribution exp( -g ( z )) , which means its Cheeger constant is also L . Although the Cheeger constant ρ in the second term of Eq. 9 corresponding to truncated exp( -g ( z )) should also be near L intuitively, current techniques can only provide a loose lower bound at an O ( √ L/d ) -level (proven in Corollary C.8). While in both cases above, the Cheeger constant is independent with ϵ . Combining this fact with an ϵ -independent choice of inner step sizes τ , the second term of Eq. 9 will converge linearly with respect to ϵ . As for the diameter Z of particles used to upper bound τ , though it may be unbounded in the standard implementation of Alg. 2, Lemma C.18 can upper bound it with ˜ O ( L 3 / 2 ( d + m 2 2 ) ρ -1 ) under the projected version of Alg. 2.

Additionally, to require the final sampling error to satisfy TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) , Eq. 9 shows that the score and energy difference estimation errors should be ϵ -dependent and sufficiently small, where ϵ score corresponding to the training loss can be well-controlled. However, obtaining a highly accurate energy difference estimation (requiring a small ϵ energy ) is hard with only diffusion models. To solve this problem, we can introduce a neural network energy estimator similar to [38] to construct r ( z ′ , z , t ) , which induces the following complexity describing the calls of the score estimation.

Corollary 4.2 (Informal version of Corollary C.19) . Suppose the estimation errors of score and energy difference satisfy

Algorithm 3 ULD FOR RTK INFERENCE

1: Input: Returned particle of the previous iteration x 0 , current iteration number k , inner iteration number S , inner step size τ , velocity diffusion coefficient γ , required convergence accuracy ϵ ;

2: Initialize the particle and velocity pair, i.e., (ˆ z 0 , ˆ v 0 ) with a Gaussian type product measure, i.e., N ( 0 , e 2 η -1) ⊗N ( 0 , I ) ;

3: for t = s to S -1 do

4: Draw noise sample pair ( ξ z s , ξ v s ) from a Gaussian type distribution.

6: ˆ v s +1 = e -γτ ˆ v s -γ -1 (1 -e -γτ ) s θ (ˆ z s ) + ξ v t

5: ˆ z s +1 = ˆ z s + γ -1 (1 -e -γτ )ˆ v s -γ -1 ( τ -γ -1 (1 -e -γτ )) s θ (ˆ z s ) + ξ z s

7: return z S ;

If we implement Alg. 1 with the projected version of Alg. 2 with the same hyperparameter settings as Theorem 4.1, it has TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) with an O ( L 4 ρ -2 · ( d + m 2 2 ) 2 Z 2 · log( d/ϵ )) complexity.

Considering the loose bound for both ρ and Z , the complexity will be at most ˜ O ( L 5 ( d + m 2 2 ) 6 ) which is the first linear convergence w.r.t. ϵ result for the diffusion inference process.

Score-only RTK-MALA. However, the parametric energy function may not always exist in real practice. We consider a more practical case where only the score estimation is accessible. In this case, we will make use of estimated score functions to approximate the energy difference, leading to the score-only RTK-MALA algorithm. In particular, recall that the energy difference function takes the following form:

Since the quadratic term can be obtained exactly, we only need to estimate the energy difference. Then let f ( z ) = -log p ( K -k -1) η ( z ) and denote h ( t ) = f (( z ′ -z ) · t + z ) , the energy difference g ( z ′ ) -g ( z ) can be reformulated as

where we perform the standard Taylor expansion at the point t = 0 . Then, we only need the derives of h i (0) , which can be estimated using only the score function. For instance, the h (1) ( t ) can be estimated with score estimations:

Moreover, regarding the high-order derivatives, we can recursively perform the approximation: ˜ h ( i +1) (0) = ( ˜ h ( i ) (∆ t ) -˜ h ( i ) (0)) / ∆ t . Consider performing the approximation up to u -order derivatives, we can get the approximation of the energy difference:

Then, the following corollary states the complexity of the score-only RTK-MALA algorithm.

Corollary 4.3. Suppose the estimation errors of the score satisfies ϵ score ≪ ρϵ/ ( Ld 1 / 2 ) , and the log-likelihood function of p t has a bounded u -order derivative, e.g., ∥ ∥ ∇ ( u ) f ( z ) ∥ ∥ ≤ L , we have a nonparametric estimation for log-likelihood to make we have TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) with a complexity shown as follows

This result implies that if the energy function is infinite-order Lipschitz, we can nearly achieve any polynomial order convergence w.r.t. ϵ with the non-parametric energy difference estimation.

RTK-ULD. Alg. 3 presents a solution employing ULD for the inner loop, which can accelerate the convergence of the inner loop due to the better discretization of the ULD algorithm. When it is used

to solve the k -th sampling subproblem of Alg. 1, x 0 is equal to ˆ x kη defined in Section 3 and used to initialize particles iterating in Alg. 2. Besides, the underlying distribution of noise sample pair is

In Alg. 3, we consider the process { (ˆ z s , ˆ v s ) } S s =0 whose underlying distribution is denoted as { ˆ π s } S s =0 . We expect the z -marginal distribution of ˆ π S to be close to the target distribution presented in Eq. 8. Then, the complexity of RTK-ULD to achieve the convergence of TV distance is provided as follows, and the detailed proof is deferred to Theorem D.6. Besides, we compare our theoretical results with the previous in Table 1.

Theorem 4.4. Under Assumptions [A1] -[A2] and [E1] , for Alg. 1, we choose

and implement Step 3 of Alg. 1 with projected Alg. 3. For the k -th run of Alg. 3, we require Gaussian-type initialization and high-accurate score estimation, i.e.,

If we set the hyperparameters of inner loops as follows. the step size and the iteration number as

It can achieve TV(ˆ p Kη , p ∗ ) ≲ ϵ with an ˜ O ( L 2 d 1 / 2 ϵ -1 ) gradient complexity.

5 Conclusion and Limitation

This paper presents an analysis of a modified version of diffusion models. Instead of focusing on the discretization of the reverse SDE, we propose a general RTK framework that can produce a large class of algorithms for diffusion inference, which is formulated as solving a sequence of RTK sampling subproblems. Given this framework, we develop two algorithms called RTK-MALA and RTK-ULD, which leverage MALA and ULD to solve the RTK sampling subproblems. We develop theoretical guarantees for these two algorithms under certain conditions on the score estimation, and demonstrate their faster convergence rate than prior works. Numerical experiments support our theory.

We would also like to point out several limitations and future work. One potential limitation of this work is the lack of large-scale experiments. The main focus of this paper is the theoretical

understanding and rigorous analysis of the diffusion process. Implementing large-scale experiments requires GPU resources and practitioner support, which can be an interesting direction for future work. Besides, though we provided a score-only RTK-MALA algorithm, the ˜ O (1 /ϵ ) convergence rate can only be achieved by the RTK-MALA algorithm (Alg. 2). However, this faster algorithm requires a direct approximation of the energy difference, which is not accessible in the existing pretrained diffusion model. Developing practical energy difference approximation algorithms and incorporating them with Alg. 2 for diffusion inference are also very interesting future directions.

References

Jason M Altschuler and Sinho Chewi. Faster high-accuracy log-concave sampling via algorithmic warm starts. In 2023 IEEE 64th Annual Symposium on Foundations of Computer Science (FOCS) , pages 2169-2176. IEEE, 2023.

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d -linear convergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum? id=r5njV3BsuD .

Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto-encoders and Langevin sampling. arXiv preprint arXiv:2002.00107 , 2020.

Nicholas M. Boffi and Eric Vanden-Eijnden. Probability flow solution of the Fokker-Planck equation, 2023.

Peter Buser. A note on the isoperimetric constant. In Annales scientifiques de l'École normale supérieure , volume 15, pages 213-230, 1982.

Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning , pages 4735-4763. PMLR, 2023.

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.

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.

Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for a proximal algorithm for sampling. In Conference on Learning Theory , pages 2984-3014. PMLR, 2022.

Xiang Cheng and Peter Bartlett. Convergence of langevin mcmc in kl-divergence. In Algorithmic Learning Theory , pages 186-211. PMLR, 2018.

Sinho Chewi. Log-Concave Sampling . 2024.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

Hanze Dong, Xi Wang, LIN Yong, and Tong Zhang. Particle-based variational inference with preconditioned functional gradient flow. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=6OphWWAE3cS .

Raaz Dwivedi, Yuansi Chen, Martin J Wainwright, and Bin Yu. Log-concave sampling: Metropolis-Hastings algorithms are fast. Journal of Machine Learning Research , 20(183):1-42, 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Xunpeng Huang, Hanze Dong, Yifan Hao, Yian Ma, and Tong Zhang. Monte Carlo sampling without isoperimetry: A reverse diffusion approach, 2023.

Xunpeng Huang, Hanze Dong, Yifan HAO, Yian Ma, and Tong Zhang. Reverse diffusion monte carlo. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=kIPEyMSdFV .

Holden Lee, Andrej Risteski, and Rong Ge. Beyond log-concavity: Provable guarantees for sampling multi-modal distributions using simulated tempering Langevin Monte Carlo. Advances in neural information processing systems , 31, 2018.

Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. arXiv preprint arXiv:2206.06227 , 2022.

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.

Yin Tat Lee and Santosh Srinivas Vempala. Eldan's stochastic localization and the KLS hyperplane conjecture: an improved lower bound for expansion. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) , pages 998-1007. IEEE, 2017.

Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards non-asymptotic convergence for diffusion-based generative models. In The Twelfth International Conference on Learning Representations , 2023.

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.

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.

Vadim Popov, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Gradtts: A diffusion probabilistic model for text-to-speech. In International Conference on Machine Learning , pages 8599-8608. PMLR, 2021.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical textconditional image generation with clip latents. arxiv 2022. arXiv preprint arXiv:2204.06125 , 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems , 35:36479-36494, 2022.

Ohad Shamir. A variant of azuma's inequality for martingales with subgaussian tails. arXiv preprint arXiv:1110.2392 , 2011.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2020.

Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motifscaffolding problem. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=6TxBxqNME1Y .

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

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023.

Xingyu Xu and Yuejie Chi. Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction. arXiv preprint arXiv:2403.17042 , 2024.

Dinghuai Zhang, Ricky T. Q. Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=OIsahq1UYC .

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.

Difan Zou, Pan Xu, and Quanquan Gu. Faster convergence of stochastic gradient langevin dynamics for non-log-concave sampling. In Uncertainty in Artificial Intelligence , pages 11521162. PMLR, 2021.

Contents

A Numerical Experiments

In this section, we conduct experiments when the target distribution p ∗ is a Mixture of Gaussian (MoG) and compare RTK-based methods with traditional DDPM. Specifically, we are considering a forward process from an MoG distribution to a normal distribution in the following

where K is the number of Gaussian components, µ k and σ 2 k are the means and variances of the Gaussian components, respectively. The solution of the SDE follows

Since x 0 and ξ are both sampled from Gaussian distributions, their linear combination x t also forms a Gaussian distribution, i.e.,

Then, we have

We can also calculate the score of x t , i.e.,

We consider a MoG consisting of 12 Gaussian distributions, each with 10 dimensions, as shown in Fig. 2 (f). The means of the 12 Gaussian distributions are uniformly distributed along the circumference of a circle with a radius of one in the first and second dimensions, while the remaining dimensions are centered at the origin. Each component of the mixture has an equal probability and a variance of 0.007 across all dimensions.

We evaluate Alg. 1 with unadjust Langevin algorithm (ULA), which leads to RTK-ULA, Alg 2, 3 implementations, and DDPM under the same Number of Function Evaluations (NFE). Specifically, while DDPM models x η across a sequence of η timesteps spanning from 0 to T in increments of 0 . 001 × T (i.e., [0 , 0 . 001 T, 0 . 002 T, . . . , T ] ), we execute Alg. 1, 2, and 3 at fewer timesteps within x [0 , 0 . 2 T, 0 . 4 T, 0 . 6 T, 0 . 8 T ] , and we distribute the NFE uniformly to these timesteps for MCMC. The experiments are taken on a single NVIDIA GeForce RTX 4090 GPU. We evaluate the sampling quality using marginal accuracy, i.e.,

where ˆ p i ( x ) is the empirical marginal distribution of the i -th dimension obtained from the sampled data, p i ( x ) is the true marginal distribution of the i -th dimension, and d is the total number of dimensions.

ULA

ULD

1.00

(a)

(@

Ground Truth

(c)

Ground Truth

ULA

ULD

Marginal Accuracy

0.95

DDPM

0.90

ULA

ULD

0.85

MALA_ES

MALA

0.80

0.75

100

200

300

400

500

0.4

0.2

0.0

0.2

0.4

0.4

0.2

0.0

0.2

0.4

Number of Function Evaluations

Dimension 2

Dimension 2

DDPM

MALA_ES

MALA

(e)

n punoug

(f)

pno

(g)

no

0.030

DDPM

MALA_ES

MALA

0.025

0.010 -

0.005

0.000

0.4

0.2

0.0

0.2

0.4

0.4

0.2

0.0

0.2

0.4

0.4

0.2

0.0

0.2

0.4

Dimension 2

Dimension 2

Dimension 2

Fig. 1 (a) shows the marginal accuracies of our RTK sampling algorithms and DDPM along NFE. We observe that all algorithms using RTK converge quickly. Among all RTK algorithms, RTK-MALA achieves the highest marginal accuracy. Score-only RTK-MALA is worse than RTK-MALA since the estimated energy contains errors, yet it is still slightly better than RTK-ULD. Along all RTK algorithms, RTK-ULA demonstrates the lowest performance in terms of marginal accuracy, but it still outperforms DDPM with a large margin especially when NFE is small.

Fig. 1 (b-f) shows the histograms of sampled MoG by DDPM and RTK-based methods. We observe that DDPM cannot reconstruct the local structure of MoG. ULA can roughly reconstruct the MoG structure, but it is still weak in complex regions, specifically around the peaks and valleys. In contrast, RTK-ULD, score-only RTK-MALA, and RTK-MALA can reconstruct more fine-grained structures in complex regions.

Fig. 2 (a-e) shows the clusters sampled by DDPM and RTK-based methods. We observe that DDPM fails to accurately reconstruct the ground truth distribution. In contrast, all methods based on RTK can generate distributions that closely approximate the ground truth. Additionally, RTK-MALA shows superior performance in accurately reconstructing the distribution in regions of low probability.

We perform other experiments on various MoG settings.Figure 3 and 4 shows the experiments on a spiral-shaped and chessboard-shaped MoG as the ground truth distribution, respectively. In these experiments, we also compared Annealed Langevin Dynamics (ALD) [] with our RTK-based method. In these figures, we observe that compared to DDPM and ALD, our RTK-based methods achieve significantly better marginal accuracy when NFE is small. Besides, we find that our RTK-based methods have a much better estimation on low-probability region. Furthermore, in Figure 5, we evaluated the methods using the Wasserstein Distance metric, which corresponds to Fréchet Inception Distance. Our results indicate that our RTK-based methods have a lower Wasserstein Distance compared to DDPM and ALD, especially when NFE is small.

Furthermore, we conducted experiments on the MNIST dataset, as shown in Figure 6. We first trained a score model following the typical variance-preserving noise schedule and then compared different sampling methods using the Fréchet Inception Distance (FID) evaluation criterion. Figure 6 (a) shows that compared with DDPM, our RTK-based methods achieve better FID scores than DDPM,

(a)

DDPM

(b)

ULA

(c)

ULD

1.5

1.5

1.5

1.0

1.0

1.0

0.5

0.5

0.5

Dimension 2

0.0

0.0-

0.0-

0.5

0.5

0.5

1.0

1.0

1.0 -

1.5

-1.5

-1.5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.5

-1.0

0.5

0.0

0.5

1.0

1.5

Dimension 1

Dimension 1

Dimension 1

(d)

MALA_ES

(e)

MALA

(f

Ground Truth Samples

1.5

1.5

1.5

1.0

1.0

1.0

0.5

0.5

0.5

Dimension 2

Dimension 2

Dimension 2

0.0

00

0.0 +

0.5

0.5

0.5

1.0

1.0

1.0

1.5

1.5

1.5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.5

-1.0

0.5

0.0

0.5

1.0

1.5

Dimension 1

Dimension 1

Dimension 1

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

ALD

1.00

0.95

Marginal Accuracy

0.90

0.85

0.80

DDPM

MLA_ES

ALD

0.75

ULA

ULD

0.70

MALA_ES

0.65

MALA

100

200

300

400

500

NumberofFunctionEvaluations

particularly when NFE is small. Figure 6 (b) shows that our RTK-based methods generate images of higher quality than DDPM when NFE is small.

Overall, these numerical experiments demonstrate the benefit of the RTK framework for developing faster algorithms than DDPM in diffusion inference. Besides, experimental results also well support our theory, showing that RTK-MALA achieves faster convergence than RTK-ULA and RTK-ULD, even with estimated energy difference via score functions.

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

4LD

1.00

0.95

Marginal Accuracy

0.90

0.85

0.80

DDPM

MLA_IS

0.75

ALD

ULA

0.70

ULD

MALA_ES

0.65

MALA

0.60

100

200

300

400

500

NumberofFunctionEvaluations

(a)

(b)

0.35

DDPM

DDPM

WassersteinDistance

ALD

WassersteinDistance

0.16

ALD

0.30

ULA

ULA

ULD

ULD

MALA_ES

0.14

MALA_ES

0.25

MALA

MALA

0.12

0.20

0.10

0.15

0.08

100

200

300

400

500

100

200

300

400

500

NumberofFunctionEvaluations

NumberofFunctionEvaluations

B Inference process with reverse transition kernel framework

Proof of Lemma 3.1. According to Bayes theorem, the following equation should be validated for any x ∈ R d and t ′ > t ,

To simplify the notation, we suppose the normalizing constant of p t , i.e.,

Besides, the forward OU process, i.e., SDE. 1, has a closed transition kernel, i.e.,

Then, we have

(a)

(b)

(c)

(d)

3

100

DDPM

ULD

MALA_ES

08

3

FID

60

40

9

8

20

20

Number ofFunction Evaluations

40

60

80

100

Plugging this equation into Eq. 10, and we have

Moreover, when we plug the reverse transition kernel

into the previous equation and have

Hence, the proof is completed.

Lemma B.1 (Chain rule of TV) . Consider four random variables, x , z , ˜ x , ˜ z , whose underlying distributions are denoted as p x , p z , q x , q z . Suppose p x,z and q x,z denotes the densities of joint distributions of ( x , z ) and (˜ x , ˜ z ) , which we write in terms of the conditionals and marginals as

then we have

Besides, we have

Proof. According to the definition of the total variation distance, we have

With a similar technique, we have

Hence, the first inequality of this Lemma is proved. Then, for the second inequality, we have

Hence, the proof is completed.

Lemma B.2. For Alg 1, we have

for any K ∈ + and η ∈ +

Proof. For any k ∈ { 0 , 1 , . . . , K -1 } , let ˆ p ( k +1) η,kη and p ← ( k +1) η,kη denote the joint distribution of (ˆ x ( k +1) η , ˆ x kη ) and ( x ← ( k +1) η , x ← kη ) , which we write in term of the conditionals and marginals as

Under this condition, we have

where the inequalities follow from Lemma B.1. By using the inequality recursively, we have

where p ∞ denotes the stationary distribution of the forward process. In this analysis, p ∞ is the standard since the forward SDE. 1, whose negative log density is 1 -strongly convex and also satisfies LSI with constant 1 due to Lemma E.9.

For Term 1. we have

where the first inequality follows from Pinsker's inequality, the second one follows from Lemma E.1, and the last one follows from Lemma E.2. It should be noted that the smoothness of p 0 required in Lemma E.2 is given by [A1] .

Plugging this inequality into Eq. 11, we have

Hence, the proof is completed.

Corollary B.3. For Alg 1, if we set

and

we have the total variation distance between the underlying distribution of Alg 1 output and the data distribution p ∗ will satisfy TV(ˆ p Kη , p ∗ ) ≤ 2 ϵ .

Proof. According to Lemma B.2, we have

for any K ∈ N + and η ∈ R + . To achieve the upper bound TV( p ∞ , p Kη ) ≤ ϵ , we only require

For Term 2. For any x ∈ R d , the formulation of p ← k +1 | k ( ·| ˆ x ) is

whose negative log Hessian satisfies

Note that the last inequality follows from [A1] . In this condition, if we require

then we have

To simplify the following analysis, we choose η to its upper bound, and we know for all k ∈ { 0 , 1 , . . . , K -1 } , the conditional density p ← k +1 | k ( x | ˆ x ) is strongly-log concave, and its score is 3 L -Lipschitz. Besides, combining Eq. 12 and the choice of η , we require

which can be achieved by then we have

where the latter equation implies

when we suppose L ≥ 1 without loss of generality. In this condition, if there is a uniform upper bound for all conditional probability approximation, i.e.,

then we can find Term 2 in Eq. 11 will be upper bounded by ϵ . Hence, the proof is completed.

Lemma B.4 (Chain rule of KL) . Consider four random variables, x , z , ˜ x , ˜ z , whose underlying distributions are denoted as p x , p z , q x , q z . Suppose p x,z and q x,z denotes the densities of joint distributions of ( x , z ) and (˜ x , ˜ z ) , which we write in terms of the conditionals and marginals as

Proof. According to the formulation of KL divergence, we have

where the last inequality follows from the fact

With a similar technique, it can be obtained that

Hence, the proof is completed.

Proof of Lemma 3.2. This Lemma uses nearly the same techniques as those in Lemma B.2, while it may have a better smoothness dependency in convergence since the chain rule of KL divergence. Hence, we will omit several steps overlapped in Lemma B.2.

For any k ∈ { 0 , 1 , . . . , K -1 } , let ˆ p ( k +1) η,kη and p ← ( k +1) η,kη denote the joint distribution of (ˆ x ( k +1) η , ˆ x kη ) and ( x ← ( k +1) η , x ← kη ) , which we write in term of the conditionals and marginals as

Besides, we consider a reference Markov process { ˜ x k } whose initial marginal distribution and transition kernels satisfy

Under these conditions, we have

Since { ˜ x k } and { x ← k } share the same transition kernel, then we have

where the first inequality follows from the chain rule of TV distance, the second inequality follows from Pinsker's inequality, and the last inequality follows from Lemma E.2. Besides, we have

where the first inequality follows from Pinsker's inequality, the second and the third inequalities follow from Lemma B.4. By using this inequality recursively, we have

where the equation follows from the definition of process { ˜ x k } . Therefore, combining Eq. 14 with Eq. 13, the proof is completed.

Corollary B.5. For Alg 1, if we set

and

we have the total variation distance between the underlying distribution of Alg 1 output and the data distribution p ∗ will satisfy TV(ˆ p Kη , p ∗ ) ≤ 2 ϵ .

Proof. According to Lemma 3.2, we have

To achieve the upper bound Term 1 ≤ ϵ , we only require

we know for all k ∈ { 0 , 1 , . . . , K -1 } , the conditional density p ← k +1 | k ( x | ˆ x ) is strongly-log concave, and its score is 3 L -Lipschitz. In this condition, we require

when we suppose L ≥ 1 without loss of generality. Then, to achieve Term 2 ≤ ϵ , the sufficient condition is to require a uniform upper bound for all conditional probability approximation, i.e.,

Hence, the proof is completed.

Remark 1. To achieve the TV error tolerance shown in Corollary B.3, .i.e.,

it requires the KL divergence error to satisfy

Compared with the results shown in Corollary B.5, this result requires a higher accuracy with an O ( L ) factor, which is not acceptable sometimes.

Lemma B.6. Suppose Assumption [A1] -[A2] hold, the choice of η keeps the same as that in Corollary B.5, and the second moment of the underlying distribution of ˆ x kη is M k , then we have

Proof. Considering the second moment of ˆ x ( k +1) η , we have

For Term 2, by choosing

Then, we focus on the innermost integration, suppose ˆ γ y ( · , · ) as the optimal coupling between ˆ p ( k +1) η | kη ( ·| y ) and p ← ( k +1) η | kη ( ·| y ) . Then, we have

Since p ← ( k +1) η | kη is strongly log-concave, i.e.,

the distribution p ← ( k +1) η | kη also satisfies 1 /L log-Sobolev inequality due to Lemma E.9. By Talagrand's inequality, we have

Plugging Eq 18 and Eq 19 into Eq 17, we have

To upper bound the innermost integration, we suppose the optimal coupling between p ( K -k -1) η and p ← ( k +1) η | kη ( ·| y ) is γ y ( · , · ) . Then it has

Since p ← ( k +1) η | kη satisfies LSI with constant 1 /L . By Talagrand's inequality and LSI, we have

where the last inequality follows from the choice of η = 1 / 2 · log(2 L + 1) / 2 L and the fact E p ( K -k -1) η [ ∥ x ∥ 2 ] ≤ ( d + m 2 2 ) obtained by Lemma E.7. Plugging this results into Eq. 20, we have

C Implement RTK inference with MALA

In this section, we consider introducing a MALA variant to sample from p ← k +1 | k ( z | x 0 ) . To simplify the notation, we set

and consider k and x 0 to be fixed. Besides, we set

According to Corollary B.5 and Corollary B.3, when we choose

the log density g will be L -strongly log-concave and 3 L -smooth. With the following two approximations,

We left the approximation level here and determined when we needed the detailed analysis. we can use the following Algorithm to replace Line 3 of Alg. 1.

In this section, we introduce several notations about three transition kernels presenting the standard, the projected, and the ideally projected implementation of Alg. 2.

Standard implementation of Alg. 2. According to Step 4, the transition distribution satisfies

Considering a 1 / 2 -lazy version of the update, we set

the transition kernel for the standard implementation of Alg, 2 will be

Projected implementation of Alg. 2. According to Step 4, the transition distribution satisfies

Considering the projection operation, i.e., Step 5 in Alg 1, if we suppose the feasible set

the transition distribution becomes

Hence, a 1 / 2 -lazy version of the transition distribution becomes

the transition kernel for the projected implementation of Alg, 2 will be

with a density function with a density function

Ideally projected implementation of Alg. 2. In this condition, we know the accurate g ( z ) -g ( z ′ ) and ∇ g ( z ) . In this condition, the ULA step will provide

Considering the projection operation, i.e., Step 5 in Alg 1, the transition distribution becomes

Hence, a 1 / 2 -lazy version of the transition distribution becomes

the transition kernel for the accurate projected update will be

Lemma C.1. Suppose we have with a density function

which means

then the target distribution of the Inner MALA, i.e., p ← ( z | x 0 ) will be L -strongly log-concave and 3 L -smooth for any given x 0 .

Proof. Consider the energy function g ( z ) of p ← ( z | x 0 ) , we have

whose Hessian matrix satisfies

Under these conditions, if we have

For the analysis convenience, we set

that is to say g ( z ) is L -strongly convex and 3 L -smooth.

C.1 Control the error from the projected transition kernel

Here, we consider the marginal distribution of { z s } and { ˜ z s } to be the random process when Alg. 2 is implemented by the standard and projected version, respectively. The underlying distributions of these two processes are denoted as z s ∼ µ s and ˜ z s ∼ ˜ µ s , and we would like to upper bound TV( µ S , ˜ µ S ) for any given x 0 .

Rewrite the formulation of z S , we have

̸

where 1 ( · ) is the indicator function. In this condition, for any set A , we have

̸

which means

̸

̸

Therefore, the total variation distance between µ S and ˆ µ S can be upper bounded with

̸

̸

Hence, to require TV( µ S , ˜ µ S ) ≤ ϵ/ 4 a sufficient condition is to consider Pr[ z S = ˜ z S ] . The next step is to show that, in Alg. 2, the projected version generates the same outputs as that of the standard version with probability at least 1 -ϵ/ 4 . It suffices to show that with probability at least 1 -ϵ/ 4 , projected MALA will accept all S iterates. In this condition, let { z 1 , z 2 , . . . , z S } be the iterates generated by the standard MALA (without the projection step), our goal is to prove that with probability at least 1 -ϵ/ 4 all z s stay inside the region B ( 0 , R ) and ∥ z s -z s -1 ∥ ≤ r for all s ≤ S . That means we need to prove the following two facts

With probability at least 1 -ϵ/ 8 , all iterates stay inside the region B ( 0 , R ) .

With probability at least 1 -ϵ/ 8 , ∥ x s -x s -1 ∥ ≤ r for all s ≤ S .

Lemma C.2. Let µ S and ˜ µ S be distributions of the outputs of standard and projected implementation of Alg. 2. For any ϵ ∈ (0 , 1) , we set

where z ∗ is denoted as the global optimum of the energy function, i.e., g , defined in Eq. 22. Suppose P( ∥ z 0 ∥ ≥ R/ 2) ≤ ϵ/ 4 and set

then we have

̸

Proof. We borrow the proof techniques provided in Lemma 6.1 of [41] to control the TVD gap between the standard and the projected implementation of Alg. 2.

̸

Particles stay inside B ( 0 , R ) . We first consider the expectation of ∥ z s +1 ∥ 2 when z s is given, and have

where the second equation follows from Eq. 28, the forth equation follows from Eq. 26 and the fifth equation follows from Eq. 27 and Eq. 25. Note that q ( z ′ | z s ) is a Gaussian-type distribution whose mean and variance are z s -τ · s θ ( z s ) and 2 τ respectively. It means

Suppose z ∗ is the global optimum of the function g due to Lemma C.1, we have

where the first inequality follows from the combination of L -strong convexity of g and Lemma E.3 , the second inequality follows from the 3 L -smoothness of g The strong convexity and the smoothness of g follow from Lemma C.1.

Combining Eq. 34, Eq. 35 and Eq. 36, we have

By requiring ϵ score ≤ τ ≤ L/ (2 + 72 L 2 ) < 1 , we have

Suppose a radio R satisfies

To prove ∥ z s ∥ ≤ R for all s ≤ S , we only need to consider z s satisfying ∥ z s ∥ ≥ 4 √ ∥∇ g ( 0 ) ∥ 2 /L 2 + d/L , otherwise ∥ z s ∥ ≤ R/ 2 ≤ R naturally holds. Then, by the concavity of the function log( · ) , for any ∥ z s ∥ ≥ R/ 2 , we have

Consider the random variable

obtained by the transition kernel Eq. 24, Note that ∥ ξ ∥ is the square root of a χ ( d ) random variable, which is subgaussian and satisfies

for any t ≥ 0 . Under these conditions, requiring

we have

In Eq. 40, the first inequality follows from the definition of transition kernel T z s shown in Eq. 28 and the second inequality follows from

According to the fact

where the second inequality follows from the smoothness of g , and the last inequality follows from Eq. 39 and ∥ z s ∥ ≤ R , we have

which implies the last inequality of Eq. 40 for all t ≥ 0 . Furthermore, suppose ∥ z s ∥ ≥ R/ 2 , it follows that

Therefore, we have log( ∥ z s +1 ∥ 2 ) -log( ∥ z s ∥ 2 ) is also a sub-Gaussian random variable and satisfies

We consider any subsequence among { z k } S k =1 , with all iterates, except the first one, staying outside the region B ( 0 , R/ 2) . Denote such subsequence by { y s } S ′ s =0 where ∥ y 0 ∥ ≤ R/ 2 and S ′ ≤ S . Then, we know y s and y s +1 satisfy Eq. 38 and Eq. 42 for all s ≥ 1 . Under these conditions, by requiring

∥ z 0 ∥ ≤ R/ 2 with a probability at least 1 -ϵ/ 16 , we only need to prove all points in { y s } S ′ s =0 will stay inside the region B ( 0 , R ) with probability at least 1 -ϵ/ 16 .

Then, set E s to be the event that it implies

which satisfies E s -1 ⊆ E s . Besides, suppose the filtration F s = { y 0 , y 1 , . . . , y s } , the sequence

is a super-martingale, and the martingale difference has a subgaussian tail, i.e., for any t ≥ 0 ,

where the first inequality is established when

Under these conditions, suppose

which follows from the fact ( a -b ) 2 ≥ a 2 / 4 -b 2 / 3 for all a, b ∈ R . Then, for any u ≥ 0 , we have

which implies that the martingale difference is subgaussian. Then by Theorem 2 in [30], for any s , we have

with the probability at least 1 -ϵ ′ conditioned on E s -1 . Taking the union bound over all s = 1 , 2 , . . . , S ′ ( S ′ ≤ S ) and set ϵ = 16 ϵ ′ S ′ , we have with probability at least 1 -ϵ/ 16 , for all s = 1 , 2 , . . . , S ′ , it holds

By requiring

we have log( ∥ y s ∥ 2 ) ≤ log( R 2 ) , which is equivalent to ∥ y s ∥ ≤ R . Combining with the fact that with probability at least 1 -ϵ/ 16 the initial point y 0 stays inside B ( 0 , R/ 2) , we can conclude that with probability at least 1 -ϵ/ 8 all iterates stay inside the region B ( 0 , R ) .

The difference between z s +1 and z s is smaller than r . In this paragraph, we aim to prove ∥ z s +1 -z s ∥ ≤ r for all s ≤ S . Similar to the previous techniques, we consider

According to the transition kernel Eq. 28, it has

where the second inequality follows from the triangle inequality, and the last inequality follows from Eq. 41 when the choice of τ satisfies Eq. 39. Under these conditions, by choosing

Eq. 45 becomes

which means

Taking union bound over all iterates, we know all particles satisfy the local condition, i.e., ∥ z s +1 -z s ∥ ≤ r with the probability at least 1 -ϵ/ 8 . Hence, the proof is completed.

C.2 Control the error from the approximation of score and energy

Lemma C.3. Under Assumption [A1] -[A2] , we set

For any ϵ ∈ (0 , 1) , we set

Suppose it has

we have

for any set A ⊆ B (0 , R ) and point z ∈ B (0 , R ) .

Proof. Note that the Markov process defined by ˜ T z ( · ) and ˜ T ∗ , z ( · ) are 1 / 2 -lazy. We prove the lemma by considering two cases: z ̸∈ A and z ∈ A .

When z ̸∈ A , we have

Similarly, we have

In this condition, we consider

which means

First, we try to control Term 1, which can be achieved by investigating ˜ q ( z ′ | z ) / ˜ q ∗ ( z ′ | z ) as follows.

In this condition, we have

It means

where the last inequality follows from the fact z ′ ∈ B ( z , r ) ∩ B ( 0 , R ) / { z } , z ∈ B ( 0 , R ) and

According to the definition of R and r shown in Lemma C.2, we choose

Under this condition, we require

then we have

and

with the definition of Term 1 shown in Eq. 46.

Then, we try to control Term 2 of Eq. 46 and have

According to Eq. 49, it has

then we can upper and lower bounding Term 2 by investigating ˜ a z ( z ′ ) / ˜ a ∗ , z ( z ′ ) as follows

In this condition, for any 0 < δ ≤ 1 , we first consider two cases. When

we have

Besides, when

we have

Then, we start to consider finding the range of ln(Term 2 . 1) as follows

where the last inequality follows from Eq. 48. Besides, similar to Eq. 47, we have

which means

where the last inequality follows from the fact z ′ ∈ B ( z , r ) ∩ B ( 0 , R ) / { z } and

Combining this result with Eq. 48, we have

Plugging this result into Eq. 54, it has

By requiring ϵ energy ≤ 0 . 1 , we have

Combining this result with Eq. 52 and Eq. 53, we have

which implies

Plugging Eq. 55 and Eq. 51 into Eq. 50, we have

In this condition, combining Eq. 56, Eq. 49 with Eq. 46, we have

Hence, we complete the proof for z ̸∈ A .

When z ∈ A , suppose there exist some r ′ satisfying

We can split A into AΩ ′ z and Ω ′ z . Note that by our results in the first case, we have

Then for the set Ω ′ z , we have

where the last inequality follows from Eq. 57 and the property of 1 / 2 lazy, i.e.,

In this condition, we have

Hence, we complete the proof for z ∈ A .

Corollary C.4. Under the same conditions as shown in Lemma C.3, if we require

then we have for any set A ⊆ B (0 , R ) and point z ∈ B (0 , R ) .

C.3 Control the error from Inner MALA to its stationary

In this section, we denote the ideally projected implementation of Alg. 2 whose Markov process, transition kernel, and particles' underlying distributions are denoted as { ˜ z ∗ ,s } S s =0 , Eq. 33, and ˜ µ ∗ ,s respectively. According to [41], we know the stationary distribution of the time-reversible process { ˜ z ∗ ,s } S s =0 is

Here, we denote Ω = B ( 0 , R ) and

In the following analysis, we default

Under this condition, the smoothness of g is 3 L and the strong convexity constant is L .

we aim to build the connection between the underlying distribution of the output particles obtained by projected Alg 2, i.e., ˜ µ S , and the stationary distribution ˜ µ ∗ though the process { ˜ z ∗ ,s } S s =0 . Since the ideally projected implementation of Alg. 2 is similar to standard MALA except for the projection, we prove its convergence through its conductance properties, which can be deduced by the Cheeger isoperimetric inequality of ˜ µ ∗ .

Under these conditions, we organize this subsection in the following three steps:

Find the Cheeger isoperimetric inequality of ˜ µ ∗ .

Find the conductance properties of ˜ T ∗ .

Build the connection between ˜ µ S and ˜ µ ∗ through the process { ˜ z ∗ ,s } S s =0 .

C.3.1 The Cheeger isoperimetric inequality of ˜ µ ∗

Definition 1 (Definition 2.5.9 in [12]) . A probability measure µ defined on a Polish space ( X , dis) satisfies a Cheeger isoperimetric inequality with constant ρ > 0 if for all Borel set A ⊆ X , it has

Lemma C.5 (Theorem 2.5.14 in [12]) . Let µ ∈ P 1 ( X ) and let Ch > 0 . The following are equivalent.

1. µ satisfies a Cheeger isoperimetric inequality with constant Ch .

For all Lipschitz f : X → R , it holds that

Remark 2. For a general non-log-concave distribution, a tight bound on the Cheeger constant can hardly be provided. However, considering the Cheeger isoperimetric inequality is stronger than the Poincaré inequality, [6] lower bound the Cheeger constant ρ with Ω( d 1 / 2 c P ) where c P is the Poincaré constant of ˜ µ ∗ . The lower bound of c P can be generally obtained by the Bakry-Emery criterion and achieve exp( -˜ O ( d )) . While for target distributions with better properties, ρ can usually be much better. When the target distribution is a mixture of strongly log-concave distributions, the lower bound of ρ can achieve 1 / poly( d ) by [20]. For log-concave distributions, [23] proved that ρ = Ω(1 / (Tr(Σ 2 )) 1 / 4 ) , where Σ is the covariance matrix of the distribution ˜ µ ∗ . When the target distribution is m -strongly log-concave, based on [16], ρ can even achieve Ω( √ L ) . In the following, we will prove that the Cheeger constant can be independent of x 0 .

Lemma C.6. Suppose µ ∗ and ˜ µ ∗ are defined as Eq. 22 and Eq. 58, respectively, where R in ˜ µ ∗ is chosen as that in Lemma C.2. For any ϵ ∈ (0 , 1) , we have

Proof. Suppose µ ∗ ∝ exp( -g ) and ˜ µ ∗ are the original and truncated target distributions of the inner loops. Following from Lemma C.13, it has

when ˜ µ ∗ is deduced by the R shown in Lemma C.2. Under these conditions, supposing Ω = B ( 0 , R ) , then we have

Suppose

then the first term of RHS of Eq. 60 satisfies

and the second term satisfies

Combining all these things, we have

where we suppose ϵ ≤ 1 without loss of generality. Hence, the proof is completed.

Lemma C.7. Suppose µ ∗ , ˜ µ ∗ and ϵ are under the same settings as those in Lemma C.6, the variance of ˜ µ ∗ can be upper bounded by 2 d/L .

Proof. According to the fact that µ ∗ is a L -strongly log-concave distribution defined on R d with the mean v m , which satisfies

following from Lemma E.8. Suppose

where R shown in Lemma C.2, then the variance bound can be reformulated as which implies

Note that the last inequality follows from Lemma C.6. Besides, suppose the mean of ˜ µ ∗ is v ˜ m , then we have

Combining Eq. 61 and Eq. 62, the variance of ˜ µ ∗ satisfies

Hence, the proof is completed.

Corollary C.8. For each truncated target distribution defined as Eq. 58, their Cheeger constant can be lower bounded by ρ = Ω( √ L/d ) .

Proof. It can be easily found that ˜ µ ∗ is log-concave distribution, which means their Cheeger constant can be upper bounded by ρ = Ω(1 / (Tr(Σ)) 1 / 2 ) , where Σ is the covariance matrix of the distribution ˜ µ ∗ . Under these conditions, we have

where the last inequality follows from Lemma C.7. Hence, ρ = Ω( √ L/d ) and the proof is completed.

C.3.2 The conductance properties of ˜ T ∗

We prove the conductance properties of ˜ T ∗ , z with the following lemma.

Lemma C.9 (Lemma 13 in [22]) . Let ˜ T ∗ , z be a be a time-reversible Markov chain on Ω with stationary distribution ˜ µ ∗ . Fix any ∆ > 0 , suppose for any z , z ′ ∈ Ω with ∥ z -z ′ ∥ ≤ ∆ we have TV ( ˜ T ∗ , z , ˜ T ∗ , z ′ ) ≤ 0 . 99 , then the conductance of ˜ T ∗ , z satisfies ϕ ≥ Cρ ∆ for some absolute constant C , where ρ is the Cheeger constant of ˜ µ ∗ .

In order to apply Lemma C.9, we have known the Cheeger constant of ˜ µ ∗ is ρ . We only need to verify the corresponding condition, i.e., proving that as long as ∥ z -z ′ ∥ ≤ ∆ , we have TV ( ˜ T ∗ , z , ˜ T ∗ , z ′ ) ≤ 0 . 99 for some ∆ . Recalling Eq. 33, we have

where the second inequality follows from Eq. 31 and the last inequality follows from Eq. 30. Then the rest will be proving the upper bound of TV ( ˜ T ∗ , z , ˜ T ∗ , z ′ ) , and we state another two useful lemmas as follows.

Lemma C.10 (Lemma B.6 in [41]) . For any two points z , z ′ ∈ R d , it holds that

Proof. This lemma can be easily obtained by plugging the smoothness of g , i.e., 3 L , into Lemma B.6 in [41].

Corollary C.11 (Variant of Lemma 6.5 in [41]) . Under Assumption [A1] -[A2] , we set

If we set

there exist absolute constants c 0 , such that ϕ ≥ c 0 ρ τ where ρ is the Cheeger constant of the distribution ˜ µ ∗ .

Proof. By the definition of total variation distance, there exists a set A ⊆ Ω satisfying

Due to the closed form of ˜ T ∗ , z shown in Eq. 63, we have

Under this condition, we have

Upper bound Term 1 . We first consider to lower bound ˜ a ∗ , ˆ z (˜ z ) in the following. According to Eq. 32, we have which means

Since Term 1 . 2 and Term 1 . 3 are grouped to more easily apply the strong convexity and smoothness of g (Lemma C.1), it has

Besides, by requiring τ ≤ 1 / 3 L , we have

Therefore,

and

where the last inequality follows from

Under these conditions, we have

where the last inequality follows from the Markov inequality shown in the following

Then, by choosing

it has 6 Lτ 2 (3 LR + ∥∇ g ( 0 ) ∥ ) 2 ≤ 1 / 40 . Besides by choosing

where the last inequality follows from the range of R shown in Lemma C.2, it has 18 Ldτ ≤ 1 / 40 . Under these conditions, considering Eq. 64, we have

Then, combining the step size choices of Eq. 65, Eq. 66, and Lemma C.2, since the requirement can be achieved by

the range of τ can be determined.

Upper bound Term 2 . In This part, we use similar techniques as those shown in Lemma 6.5 of [41]. According to the triangle inequality, we have

Then, we upper bound Term 2 . 1 as follows

According to the definition, ˜ q ∗ , z ( · ) is Gaussian distribution with mean z -τ ∇ g ( z ) and covariance matrix 2 τ I , thus we have

Then, we start to lower bound

Then, we require

where the latter condition can be easily covered by the choice in Eq. 68 when d ≥ 3 without loss of generality. Under this condition, we have

Since we have

by the smoothness, it has

Plugging Eq. 72 and Eq. 71 into Eq. 70, we have

where the last inequality follows from the choice of r shown in Lemma C.3, i.e.,

Under these conditions, we have

Then combine the above results and apply Lemma C.10, assume τ ≤ 1 / (3 L ) , we have

Plugging the above into Eq. 69, we have

where the second inequality follows from Eq. 67.

After upper bounding Term 1 and Term 2 , we have

where the last inequality can be established by requiring ∥ z -z ′ ∥ ≤ √ 2 τ . Combining Lemma C.9, the conductance of ˜ µ ∗ satisfies √

Hence, the proof is completed.

The connection between ˜ µ S and ˜ µ ∗ . With the conductance of truncated target distribution, we are able to find the convergence of the projected implementation of Alg. 2. Besides, the gap between the truncated target ˜ µ ∗ and the true target µ ∗ can be upper bounded by controlling R while such an R will be dominated by the range of R shown in Lemma C.2. In this section, we will omit several details since many of them have been proven in [41].

Lemma C.12 (Lemma 6.4 in [41]) . Let ˜ µ S be distributions of the outputs of the projected implementation of Alg. 2. Under Assumption [A1] -[A2] , if the transition kernel ˜ T z ( · ) is δ -close to ˜ T ∗ , z with δ ≤ min { 1 -√ 2 / 2 , ϕ/ 16 } ( ϕ denotes the conductance of ˜ µ ∗ ), then for any λ -warm start initial distribution with respect to ˜ µ ∗ , it holds that

Lemma C.13 (Lemma 6.6 in [41]) . For any ϵ ∈ (0 , 1) , set R to make it satisfy

and ˜ µ ∗ be the truncated target distribution of µ ∗ . Then the total variation distance between µ ∗ and ˜ µ ∗ can be upper bounded by TV(˜ µ ∗ , µ ∗ ) ≤ ϵ/ 4 .

C.4 Main Theorems of InnerMALA implementation

Lemma C.14. Under Assumption [A1] -[A2] , we can upper bound G = ∥∇ g ( 0 ) ∥ as

Furthermore, we can reformulate R as

to make it satisfy the requirement shown in Lemma C.3. Then, the range of inner step sizes, i.e., τ , will satisfy

where the absolute constant C τ = 2 -4 · 3 -8 · 7 -2 .

Proof. To make the bound more explicit, we control R and G in our previous analysis. For G = ∥∇ g ( 0 ) ∥ , according to Eq. 22, we have

which means

Besides, we should note f ( K -k -1) η is the smooth (Assumption [A1] ) energy function of p ( K -k -1) η denoting the underlying distribution of time ( K -k -1) η in the forward OU process. Then, we have

where the first inequality follows from Lemma E.6, and the third inequality follows from Lemma E.7. Under these conditions, we have

Then, for R defined as

we can choose R to be the upper bound of RHS. Considering

then we choose

After determining R , the choice of τ can be relaxed to

where the absolute constant C τ = 2 -4 · 3 -8 · 7 -2 , since we have

Hence, the proof is completed.

Theorem C.15. Under Assumption [A1] -[A2] , for any ϵ ∈ (0 , 1) , let ˜ µ ∗ ( z ) ∝ exp( -g ( z )) 1 [ z ∈ B ( 0 , R )] be the truncated target distribution in B ( 0 , R ) with

r in Alg. 2 satisfies

and ρ be the Cheeger constant of ˜ µ ∗ . Suppose ˜ µ 0 ( {∥ x ∥ ≥ R/ 2 } ) ≤ ϵ/ 16 , the step size satisfy

the score and energy estimation errors satisfy

then for any λ -warm start with respect to µ ∗ the output of both standard and projected implementation of Alg. 2 satisfies

Proof. We characterize the condition on the step size τ . Combining Lemma C.2 and Corollary C.11, it requires the range of τ to satisfy

Under this condition, we have

which implies

Then, we have

which matches the requirement of Lemma C.3. Under this condition, if we require

it makes

and satisfies the requirements shown in Lemma C.12.

Then, we are able to put the results of these lemmas together to establish the convergence of Alg. 2. Note that if µ 0 is a λ -warm start to µ ∗ , it must be a λ -warm start to ˜ µ ∗ since ˜ µ ∗ ( A ) ≥ µ ∗ ( A ) for all A ∈ Ω . Combining Lemma C.2, Lemma C.12 and Lemma C.13, we have

After combining this result with the choice of parameters shown in Lemma C.14, the proof is completed.

Lemma C.16. Under the same assumptions and hyperparameter settings made in Theorem C.15, we use Gaussian-type initialization

If we set the iteration number as which means

the standard and projected implementation of Alg. 2 can achieve

Proof. We reformulate the target distribution µ ∗ and the initial distribution µ 0 as follows

Under this condition, we have

Due to Assumption [A1] , we have

and

Since the function ϕ ( z ) is strongly log-concave, it satisfies

due to Lemma E.3 and Lemma E.4. Under these conditions, we have

Besides, we have

which implies

Plugging Eq. 76, Eq. 77 and Eq. 78 into Eq. 75, we have

Due to the strong convexity of ϕ , it has

and where the inequality follows from Eq. 73. Combining with the fact

Eq. 79 can be relaxed to

which is independent on ∥ x 0 ∥ . Then, In order to ensure the convergence of the total variation distance is smaller than ϵ , it suffices to choose τ and S such that

where the last two inequalities follow from Theorem C.15. Hence, the proof is completed.

Theorem C.17. Under Assumption [A1] -[A2] , for Alg. 1, we choose

and implement Step 3 of Alg. 1 with projected Alg. 2. For the k -th run of Alg. 2, we use Gaussian-type initialization

If we set the hyperparameters as shown in Lemma C.16, it can achieve

TV(ˆ p Kη , p ∗ ) ≤ ϵ + ˜ O ( Ld 1 / 2 ρ -1 ϵ score ) + O (ˆ τ -1 / 2 · L 2 ( d 1 / 2 + m 2 + Z ) ρ -1 ϵ energy with a gradient complexity as follows

for any ˆ τ ∈ (0 , 1) where Z denotes the maximal l 2 norm of particles appearing in outer loops (Alg. 1).

Proof. According to Lemma B.3, we know that under the choice

it requires to run Alg. 2 for K times where

For each run of Alg. 2, we require the total variation error to achieve

Combining with Lemma C.16, we consider a step size

where τ ′ ∈ (0 , 1) , to solve the k -th inner sampling subproblem. Then, the maximum iteration number will be

This means that with the total gradient complexity

where Z denotes the maximal l 2 norm of particles appearing in outer loops (Alg. 1), we can obtain

Hence, the proof is completed.

Lemma C.18. Suppose we implement Alg. 2 with its projected version, we have

where Z denotes the maximal l 2 norm of particles appearing in outer loops (Alg. 1)

Proof. Suppose we implement Alg. 2 with its projected version, where each update will be projected to a ball with a ratio r shown in Lemma C.2. Under these conditions, we have

For each i ∈ { 1 , 2 , . . . K } , we have

Then, we have

Hence, the proof is completed.

C.5 Control the error from Energy Estimation

Corollary C.19. Suppose the diffusion model s θ satisfies

and another parameterized model ˆ l ˆ θ ( x , t ) is used to estimate the log-likelihood of p t ( x ) satisfying

If we implement Alg. 1 with the projected version of Alg. 2, it has

with the following gradient complexity

Proof. Since we have highly accurate scores and energy estimation, we can construct s θ and r θ ′ (shown in Eq. 23) for the k -th inner loop as follows

Under these conditions, we have

Plugging these results into Theorem C.17 and setting ˆ τ = 1 / 2 , we have

with the following gradient complexity

Corollary C.20. Suppose the score estimation is extremely small, i.e.,

and the log-likelihood function of p t has a bounded 3 -order derivative, e.g.,

we have a non-parametric estimation for log-likelihood to make we have TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) with

gradient calls.

Proof. Combining the Alg. 2 and the definition of ϵ energy shown in Lemma C.4, we actually require to control

for any s ∈ [0 , S -1] . Then, we start to construct r θ (˜ z s , z s ) . Since we have

we should only estimate the difference of the energy function f ( K -k -1) η which will be presented as f for abbreviation. Besides, we define the following function

which means

Under the high-order smoothness condition, i.e.,

where ∥ · ∥ denotes the nuclear norm, then we have

It means we need to approximate h ( i ) with high accuracy.

For i = 1 , the ground truth h (1) (0) is

we can approximate it numerically as

since we have score approximation. Then it has

Then, for i = 2 , we obtain the ground truth h (2) (0) by

which means

If we use the differential to approximate h (2) (0) , i.e.,

we find the error term will be

If we use smoothness to relax the integration term, we have

which means

Combining Eq. 80, Eq. 81 and Eq. 82, we have

which means the final energy estimation error will be

Considering ϵ score is extremely small (compared with the output performance error tolerance ϵ ), we can choose t depending on ϵ score , e.g., t = √ ϵ score , to make Term 1 and Term 2 in Eq. 83 diminish. Under this condition, the term Lr 3 / 6 will dominate RHS of Eq. 83. Besides, we have

then we have

where the last equation follows from the choice of τ shown in Theorem C.15. Then, plugging this result into Theorem C.17 and considering ϵ score ≪ ϵ , we have

with a gradient complexity as follows

Then, by choosing

we have TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) with

Hence, the proof is completed.

Remark 3. If we consider more high-order smooth, i.e.,

with similar techniques shown in Corollary C.20, we can have the following bound, i.e.,

when ϵ score is extremely small. Under this condition, since it has

we have

Then, plugging this result into Theorem C.17 and considering ϵ score ≪ ϵ , we have

where we suppose L ≥ 1 in the last equation without loss of generality. Then, by supposing

we have TV(ˆ p Kη , p ∗ ) ≤ ˜ O ( ϵ ) with

where the last 2 u appears since the estimation of high-order derivatives requires an exponentially increasing call of score estimations.

D Implement RTK inference with ULD

In this section, we consider introducing a ULD to sample from p ← k +1 | k ( z | x 0 ) . To simplify the notation, we set

and consider k and x 0 to be fixed. Besides, we set

According to Corollary B.5 and Corollary B.3, when we choose

the log density g will be L -strongly log-concave and 3 L -smooth.

For the underdamped Langevin dynamics, we utilize a form similar to that shown in [40], i.e.,

with a little abuse of notation for t ∈ [ sτ, ( s +1) τ ) . We denote the underlying distribution of (ˆ z t , ˆ v t ) as ˆ π t , and the exact continuous SDE

has the underlying distribution ( z t , v t ) ∼ π t . The stationary distribution of the continuous version is defined as

where the z -marginal of π ← ( ·| x 0 ) is p ← ( ·| x 0 ) which is the desired target distribution of inner loops. Therefore, by taking a small step size for the discretization and a large number of iterations, ULD will yield an approximate sample from p ← ( ·| x 0 ) . Besides, in the analysis of ULD, we usually consider an alternate system of coordinates

their distributions of the continuous time iterates π M t and the target in these alternate coordinates π M , respectively. Besides, we need to define log-Sobolev inequality as follows

Definition 2 (Log-Sobolev Inequality) . The target distribution p ∗ satisfies the following inequality

with a constant C LSI for all smooth function g : R d → R satisfying E p ∗ [ g 2 ] < ∞ .

Remark 4. Log-Sobolev inequality is a milder condition than strong log-concavity. Suppose p satisfies m -strongly log-concavity, it satisfies 1 /m LSI, which is proved in Lemma E.9.

Definition 3 (Poincaré Inequality) . The target distribution p satisfies the following inequality

with a constant C PI for all smooth function g : R d → R satisfying E p ∗ [ g 2 ] < ∞ .

In the following, we mainly follow the idea of proof shown in [40], which provides the convergence of KL divergence for ULD, to control the error from the sampling subproblems.

Lemma D.1 (Proposition 14 in [40]) . Let π M t denote the law of the continuous-time underdamped Langevin diffusion with γ = c √ 3 L for c ≥ √ 2 in the ( ϕ, ψ ) coordinates. Suppose the initial distribution π 0 has a log-Sobolev (LSI) constant (in the altered coordinates) C LSI ( π M 0 ) , then { π M t } satisfies LSI with a constant that can be uniformly upper bounded by

Lemma D.2 (Adapted from Proposition 1 of [26]) . Consider the following Lyapunov functional

For targets π ← ∝ exp( -g ) which are 3 L -smooth and satisfy LSI with constant 1 /L , let γ = 2 √ 6 L . Then the law π t of ULD satisfies

Lemma D.3 (Variant of Lemma 4.8 in [1]) . Let ˆ π t denote the law of SDE. 85 and π t denote the law of the continuous time underdamped Langevin diffusion with the same initialization, i.e., ˆ π 0 = π 0 . If γ ≍ √ L and the step size τ satisfies

then we have

Proof. The main difference of this discretization analysis is whether the score ∇ log p t can be exactly obtained or only be approximated by s θ . Therefore, in this proof, we will omit various steps the same as those shown in [1].

We consider the following difference

From Girsanov's theorem, we obtain immediately using Itô's formula

According to Corollary 20 of [40], we have

where the last inequality can be established by requiring

since exp( u ) ≤ 1 + 2 u for any u ∈ [0 , 1] .

With similar techniques utilized in Lemma 4.8 of [1], we know that if

it holds that

Furthermore, for

it has

Therefore, we have

Then, still with similar techniques utilized in Lemma 4.8 of [1], we have

In summary, we have

and the proof is completed.

Corollary D.4. Under the same assumptions and hyperparameter settings made in Lemma D.3. If the step size τ and the score estimation error ϵ score satisfies

Then we have χ 2 (ˆ π T ∥ π T ) ≲ ϵ 2 .

Proof. We can easily obtain this result by plugging the choice of τ and ϵ into Lemma D.3. Noted that we suppose L ≥ 1 without loss of generality.

Theorem D.5 (Variant of Theorem 6 in [40]) . Under Assumption [A1] -[A2] , for any ϵ ∈ (0 , 1) , we require Gaussian-type initialization and high-accurate score estimation, i.e.,

If we set the step size and the iteration number as

the marginal distribution of output particles ˆ p T will satisfy KL ( ˆ p T ∥ ∥ p ← ( ·| x 0 ) ) ≤ O ( ϵ 2 ) .

Proof. Consider the underlying distribution of the twisted coordinates ( ϕ, ψ ) for SDE. 85, the decomposition of the KL using Cauchy-Schwarz:

Using LSI of the iterations via Lemma D.1, we have

Then, we start to upper bound the relative Fisher information. Since π M = M # π ← ( ·| x 0 ) , then

and similarly for ∇ log π M T . This yields the expression

According to the definition of M , we have

For any c 0 > 0 and

we have

The determinant is

for c 0 > 0 sufficiently small, which means that

Therefore, Eq. 88 becomes

According to Lemma D.2, the decay of the Fisher information requires us to set

which yields KL ( π M T ∥ ∥ π M ) ≤ ϵ 2 . Besides, we can easily have

According to the definition of LSI, we also have

Recall as well that this requires γ ≍ 3 L in SDE. 85. For the remaining KL ( ˆ π M T ∥ ∥ π M T ) and χ 2 ( ˆ π M T ∥ π M T ) in Eq. 87, we invoke Lemma D.3 with the value T = Sτ specified and desired accuracy ϵ , , which consequently yields

Under this condition, we start to consider the initialization error. Suppose we have π 0 = N ( 0 , e 2 η -1) ⊗N ( 0 , I ) , which implies

Following the η setting, i.e., which yields

where the inequality follows from Eq. 73. Therefore, combining Eq. 91, Eq. 90 and Eq. 89, we have

which implies

In this condition, the score estimation error is required to be

Hence, the proof is completed.

Theorem D.6. Under Assumption [A1] -[A2] , for Alg. 1, we choose

and implement Step 3 of Alg. 1 with projected Alg. 3. For the k -th run of Alg. 3, we require Gaussian-type initialization and high-accurate score estimation, i.e.,

If we set the hyperparameters as shown in Lemma D.5, it can achieve TV(ˆ p Kη , p ∗ ) ≲ ϵ with an ˜ O ( L 2 d 1 / 2 ϵ -1 ) gradient complexity.

Proof. According to Corollary B.5, we know that under the choice

it requires to run Alg. 3 for K times where

For each run of Alg. 3, we require the KL divergence error to achieve

Combining with Theorem D.5, we consider a step size

then the iteration number will be

For an expectation perspective, we have

where the last inequality follows from Lemma B.6. This means that with the total gradient complexity

Hence, the proof is completed.

E Auxiliary Lemmas

Lemma E.1 (Theorem 4 in [35]) . Suppose p ∝ exp( -f ) defined on R d satisfies LSI with constant µ > 0 . Along the Langevin dynamics, i.e.,

Lemma E.2. Suppose p ∝ exp( -f ) defined on R d satisfies LSI with constant µ > 0 where f is L-smooth, i.e.,

If p 0 is the standard Gaussian distribution defined on R d , then we have

Proof. According to the definition of LSI, we have

where the third inequality follows from the L -smoothness of f ∗ and the last equation establishes since E p 0 [ ∥ x ∥ 2 ] = d is for the standard Gaussian distribution p 0 in R d .

Lemma E.3 (Variant of Lemma B.1 in [41]) . Suppose f : R d → R is a m -strongly convex function and satisfies L -smooth. Then, we have

where x ∗ is the global optimum of the function f .

Proof. According to the definition of strongly convex, the function f satisfies

Besides, we have

Combining the above two inequalities, the proof is completed.

Lemma E.4 (Lemma A.1 in [41]) . Suppose a function f satisfy

where x t ∼ p t , then it has then we have

Lemma E.5 (Lemma 1 in [18]) . Consider the Ornstein-Uhlenbeck forward process

and denote the underlying distribution of the particle x t as p t . Then, the score function can be rewritten as

Lemma E.6 (Lemma 11 in [35]) . Assume p ∝ exp( -f ) and the energy function f is L -smooth. Then

Lemma E.7 (Lemma 10 in [8]) . Suppose that Assumption [A1] -[A2] hold. Let { x t } t ∈ [0 ,T ] denote the forward process, i.e., Eq. 1, for all t ≥ 0 ,

Lemma E.8. Suppose q is a distribution which satisfies LSI with constant µ , then its variance satisfies

Proof. It is known that LSI implies Poincaré inequality with the same constant, i.e., µ , which means if for all smooth function g : R d → R ,

In this condition, we suppose b = E q [ x ] , and have the following equation

where g i ( x ) is defined as g i ( x ) := ⟨ x , e i ⟩ and e i is a one-hot vector ( the i -th element of e i is 1 others are 0 ). Combining this equation and Poincaré inequality, for each i , we have

Hence, the proof is completed.

Lemma E.9 (Variant of Lemma 10 in [11]) . Suppose -log p ∗ is m -strongly convex function, for any distribution with density function p , we have

By choosing p ( x ) = g 2 ( x ) p ∗ ( x ) / E p ∗ [ g 2 ( x ) ] for the test function g : R d → R and E p ∗ [ g 2 ( x ) ] < ∞ , we have

which implies p ∗ satisfies 1 /m -log-Sobolev inequality.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

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

Answer: [Yes]

Justification: We have clear claim that we improve the diffusion inference by RTK framework.

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 discussed the limitation in section 5.

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: See Section 4 A1, A2, E1, E2, E3 for Assumptions. Proof are provided in the Appendix.

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: See our Appendix.

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: [Yes]

Justification: We have detailed information in the Appendix to reproduce the emprical results.

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: See the Appendix for more information.

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: [NA]

Justification: We just have illustrative empirical result about the trend and generated samples. Not claims for accuracy or error from the empiral side.

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: See appendix.

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

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: It is a theoretical paper. No societal impact is visible in a short term.

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: No real data in this paper.

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: [NA]

Justification: No real data are used.

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: Only synthetic data are used with detailed instructions.

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

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

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:\nDiffusion models rely on the use of the reverse SDE (or ODE) for sample generation, but the discrete-time simulation of the reverse SDE, when regarded as solving a sequence of subproblems, is potentially less efficient compared to the usage of better MCMC samplers for solving each subproblem. This paper suggests that given a fixed computational budget, employing MALA and ULD to simulate samples from a de-noised distribution allows for higher efficiency. The paper then analyzes the theoretical properties of 2 RTK samplers and showcase a more efficient use of computational expenses when compared to SDE based samplers.\n\nSTRENGTHS:\n- The paper provides a compelling argument of accelerating the inference of reverse SDE, which seems more interesting than previous attempts to solve the same problem using specialized SDE solvers. \n- The paper also presents a valid argument in taking bigger step sizes in order to take advantage of Lipchitz continuity. \n- The paper can help bridge between diffusion modeling and MCMC, which can lead to interesting future work on further improving the sampling of de-noised distributions.\n\nWEAKNESSES:\nI like this paper overall -- the main weaknesses are mostly about how it is written, as there exist small mistakes or confounding statements. \n\n- The paper's usage of the word \"Gaussian process\" (e.g., in L. 125) seems confusing. I would assume Gaussian process mostly refers to infinite-dimensional Gaussian distribution in the space of functions, but finite discretization of the reverse SDEs simply yields a sequence of Gaussian conditional distributions. \n- The paper suggests numerical experiments in the abstract, but only includes the experiments in the appendix. \n- Should the paper clarify the use of the $ \\tilde{\\mathcal{O}} $ notation? It does not seem like a standard notation and has already been used in the abstract. \n- Minor: the paper uses \"closed\" as opposed to \"close\", e.g., in Algorithm 1 and Line 172, 176.\n\nQUESTIONS:\n- The use of an underdamped dynamical system in diffusion modeling is not entirely a new idea: For instance, Dockhorn et al. suggest using a critically dampened Langevin diffusion as the forward process. While a critically dampened LD encodes a different evolution of marginal distributions, solving one discretization step of the backward SDE would be quite similar to the use of ULD in this paper. What do the authors think about the pros and cons of underdamping the forward process directly vs. your approach of only modifying the inference of the backward SDE? \n- Another advantage of SDE or RTK samplers, when compared the deterministic simulation using probability flow ODEs, is the addition of stochasticity that ameliorates possible imperfections in the score estimation. I think it is reasonable to assume that the adoption of RTK as a generator might further improve this issue. \n- The experiment figures in the appendix often show the failures of DDPM, but I'm not sure why there exists a dramatic difference between DDPM and others in figure 1 and 2. I had previously assumed that DDPM should be able to sample from the same data distribution, even though it might be computationally more expensive compared to the methods proposed. Can the authors provide some clarification about why the generated samples are so different (I would assume it's either inaccurate estimates of the scores, or the step size being too large)?\n\n[1] Dockhorn T, Vahdat A, Kreis K. Score-Based Generative Modeling with Critically-Damped Langevin Diffusion. In: International Conference on Learning Representations [Internet]. 2021.\n\nLIMITATIONS:\n- The paper suggests in the final section several limitations of the current paper. I agree with the characterization of limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper presents a framework for accelerating the inference process in denoising diffusion probabilistic models (DDPMs) by optimizing the balance between the number and complexity of sampling subproblems.

STRENGTHS:
- The proposed Reverse Transition Kernel (RTK) framework allows for a more efficient decomposition of the denoising diffusion process, enabling the use of advanced sampling algorithms like the Metropolis-Adjusted Langevin Algorithm (MALA) and Underdamped Langevin Dynamics (ULD). 
- By incorporating MALA and ULD, the paper leverages well-established acceleration techniques. This integration not only improves the convergence rates but also provides robust theoretical guarantees.
- The authors provide theoretical convergence guarantees for these algorithms surpass current state-of-the-art rates for diffusion inference, supported by numerical experiments.

WEAKNESSES:
- While the paper focuses on diffusion models, it raises the question of whether the RTK framework can be applied to other generative models or machine learning tasks.
- The paper lacks sufficient numerical evaluation to thoroughly demonstrate the effectiveness of the proposed method. Additionally, there is a lack of detailed comparison to other related approaches, making it difficult to fully assess the advantages and limitations of the RTK framework in practical settings.

QUESTIONS:
See above.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscript is on accelerating or improving the inference in diffusion models. This denoising process corresponds to the discretization of an ODE or SDE. In this work, the authors view this process as multiple reverse transition kernel sampling subproblems. They introduce a general framework for the reverse transition kernel, and leverage two sampling algorithms to solve it. Theoretical results on the convergence of both sampling methods are presented.\n\nSTRENGTHS:\nThe method seems novel and founded on a strong theoretical analysis. The authors find the analytic solution of the RTK for specific SDEs which enables practical algorithms. \n\nThe overall presentation of the manuscript is good, and the theoretical analysis improves on existing ones.\n\nWEAKNESSES:\nThe manuscript lacks of experiments and especially comparison with existing inference methods. Theoretically, these methods seem to improve on existing one, but it would be great to show it in practice as well, e.g. comparing FIDs of a pretrained model for different inference methods. \n\nThe discretization of an ODE or SDE is replaced with $K$ subproblems consisting of sampling from the reverse transition probabilities. Within each of these subsampling problems, we need to simulate a chain of size $S$, requiring at least $S$ calls of the score network. My understanding is that you would need to call the score network at least $K\\times S$ times. I wonder if the authors could show experiments comparing the number of function evaluations (or computation time) for different methods and their scores (e.g. FID of generated images). I believe some of this analysis on synthetic datasets is presented in the appendix. It would be great to improve these experiments and move them to the main body of the text; it would help convey that for the same NFE the accuracy is greater, so the proposed method indeed accelerates inference without decreasing the performance.\n\nQUESTIONS:\n### Minor comments\n\n- potential typo on line 118 \"To implement SDE.\"\n- wrong punctuation at the end of eq. 5. \n- Alg. 3 \"of the previous iteration $x_0$\", shouldn't it be $x_{k-1}$?\n\nLIMITATIONS:\nThe authors address the limitations of their work in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper takes a fresh perspective on simulating the backward SDE in diffusion modelling, proposing to replace the usual Euler-Maruyama discretisation (i.e. based on Gaussian approximations) with a more accurate approximation (based on Metropolis--Hastings and related techniques).  The approach is explored in theoretical detail, and it is argued that in principle fewer time steps are required to simulate from the backward SDE.\n\nSTRENGTHS:\n- Casting the solution of the backward SDE as a sequence of \"simple\" sampling problems is an interesting approach that opens the door to new methods and it is conceivable that these could deliver practical benefit.\n\n- The theoretical analysis appears to be extremely detailed, taking into account both the usual SDE discretisation errors that occur when approximating the backward SDE and the sampling errors that are incurred using Metropolis--Hastings and related techniques.\n\nWEAKNESSES:\n- The principal justification for this strategy is theoretical, and it is not yet clear whether it is practically beneficial.\n\n- Some of the English language was confusing at times, for example referring to \"Gaussian process\" when I believe \"Gaussian approximation\" (or similar) was intended.\n\nQUESTIONS:\n- Is solving the backward SDE really a computational bottleneck in diffusion modelling?  I had imagined that learning the score function was actually the main bottleneck.\n\n- Do the author(s) results allow for plug-and-play in terms of the sampling method used?  For instance, if I design a new sampling method for use in this context (e.g. some new version of Metropolis--Hastings) what conditions would my method need to satisfy in order to inherit the theoretical guarantees that have been established?  A clear statement could help other researchers to engage with these results.\n\nLIMITATIONS:\nThe main limitation is the lack of empirical support, which the author(s) acknowledge in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper discusses an acceleration method for the inference of the diffusion models. It considers the denoising diffusion process as a sequence of reverse transition kernel (RTK) sampling subproblems. Then, the paper integrate the RTK subproblems to the Metropolis-Adjusted Langevin Algorithm (MALA) and Underdamped Langevin Dynamics (ULD), for solving strongly log-concave subproblems.\n\nSTRENGTHS:\nThis paper presents a new modification of diffusion models, where it views DDPM as approximately solving RTK sampling subproblems using a small step size. Based on this, the authors propose two efficient algorithms for diffusion inference and demonstrate the acceleration theoretically and empirically through a mixed Gaussian example.\n\nWEAKNESSES:\nThe choice of step size is important as there is a trade-off between the log-concavity of RTK subproblems and the number of subproblems. Either too small or too large of the step size results in a high computational complexity. The convergence results in section 4 are built upon a particular choice in the step size, which depends on the Lipschitz constant associated with the target data. The proposed method will gain more clarity and generality, and be beneficial to algorithm design if a concrete computational complexity in terms of the step size can be analyzed. See questions below for more details.\n\nQUESTIONS:\n1. There is a certain step size choice used in this paper (for example after Lemma 3.1 and Theorem 4.1), and the choices are to ensure the strong long-concavity of target distribution. How is this derived? Is this a unique choice in some sense? There is a paragraph after algorithm 2 commenting on the step size choice, I'm wondering if this result can be made more concrete and precise.\n\n2. What is a practical way of choosing the step size for general target data like real-world image datasets?\n\n3. What are the challenges in complexity analysis compared to the cited results [8, 9, 21] in Table 1? As DDPM can be seen as solving RTK sampling subproblems using a small step size, can you show the method discussed in this paper yields a more general technique in terms of complexity analysis?\n\nLIMITATIONS:\nThe limitations have been discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}