Reverse Transition Kernel: A Flexible Framework to Accelerate Diffusion Inference

Thanks for your positive comments and helpful suggestions. We will response to your concerns point by point in the following.

W1: (1). Writing Issues, (2). The paper suggests numerical experiments in the abstract, but only includes the experiments in the appendix. (3). Should the paper clarify the use of the  notation \tilde{O}? Minor: the paper uses "closed" as opposed to "close".

AW1: (1). We will revise them in the next version. (2). We will move the experiments part to the main body of this paper. (3). We clarified the notation $\tilde{O}$ in the notations paragraph of Section 2, i.e., Preliminaries. We will highlight it in the next version. Minor: We will revise them in the next version.

Q1: 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?

AQ1: Although there are some similarities in the use of ULD to accelerate the inference process of diffusion models, our work is orthogonal to [Dockhorn et al.] mentioned in your comments, and the motivations of these two papers are different. Specifically, [Dockhorn et al.] expected to change the way of modeling the probability flow, while our work aimed to utilize well-trained scores to accelerate the inference process. Compared with the reverse SDE, although we use a different Markov chain to generate the data, we do not require another training process as [Dockhorn et al.]. Under these conditions, we think it is good to combine our work, such as applying RTK in CLD-based probability flow, rather than comparison.

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

AQ2: This is a good observation. Actually, in order to achieve the same KL/TV convergence as SDE-based methods in theory, existing ODE-based methods require either an SDE corrector (DPOM and DPUM [Chen et al.]) or some strict assumptions (a small estimation error of energy Hessian [Li et al.]), which are shown in our Table. 1. To the best of our knowledge, whether the pure ODE-based generator can achieve the same theoretical results as those for the SDE-based generator, under mild assumptions, is still an open question. We have known that (1) stochasticity is the only difference between the implementations of ODE and SDE-based methods and (2) the contraction of KL divergence highly depends on the diffusion term from a Fokker-Planck equation perspective. Therefore, it can be inferred that stochasticity could ameliorate possible imperfections in the score estimation.

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

We would like to clarify that, to remove the influence of the score estimation error on the convergence, we conduct the experiments on the synthetic data with exact scores, which can be explicitly calculated since the target distribution is a Gaussian mixture. Besides, we tune the hyper-parameters only under small NFE settings, which may require a relatively large step size. Such a large step size may magnify the differences between the generated samples. In the latest supplementary materials for the rebuttal, we added some new experiments. We especially compare RTK-based methods and DDPM on the real-world dataset MNIST, which demonstrates that RTK-based methods are better than DDPM even with score estimation errors in the same setting.

Thanks for your positive comments and helpful suggestions. We will response to your concerns point by point in the following.

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

A good question! We should note that the main idea of RTK techniques is to replace the discretization of SDE/ODE with a sampling subproblem. Therefore, this framework can also be applied to some of the latest generative models, especially for flow matching [Lipman et al.]. Instead of defining the forward process as an SDE in diffusion models, flow matching defines the forward process as a distribution conditioning on the initial data distribution. Actually, for their optimal transport implementation, the forward process can be converted to the following SDE $d x_t^\to = \frac{x^\to_t}{t-1}d t + \sqrt{\frac{2t}{1-t}}d B_t$ equivalently. According to the exponential integrator, we can first calculate the forward transition kernel and then deduce the reverse transition kernel. By decomposing the reverse process of flow matching with a serious of strongly log-concave sampling subproblem and choosing proper inner samplers, I believe it is possible to accelerate the flow matching inference.

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

A2: In this paper, we mainly focus on the theory (the RTK framework) about the convergence of diffusion inference process, and highlight how to design algorithms (RTK-ULD and RTK-MALA) to achieve faster convergence rate in theory. In the meanwhile, we have also provided some experimental results for the Gaussian mixture targets shown in Appendix. A, which justify the effectiveness and efficiency of the RTK based methods.

We also added some new experiments in our latest supplementary material of rebuttal. These experiments are conducted on another two Mixtures of Gaussians (MoG) and the MNIST dataset. Besides, we add an additional related approach, i.e., [Song et al.], for comparison. All experimental results validate that the RTK-based inference process outperforms DDPM in different settings. More details about the supplementary experiments can be found in the common response.

We will add this additional context to our paper in the next version.

[Lipman et al.] Flow Matching for Generative Modeling.

[Song et al.] Generative Modeling by Estimating Gradients of the Data Distribution.

Thanks for your comments. We will response to your question in the following.

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

We have conducted additional experiments on both synthetic and real-world data. Additionally, we have included comparisons with an existing method and evaluated a new numerical metric. The detailed results of these experiments are included in the supplementary material of the rebuttal.

For synthetic data, we conducted additional experiments on spiral-shaped and chessboard-shaped Mixture of Gaussians (MoG), as shown in Figures 1 and 2. Our RTK-based methods achieved significantly better marginal accuracy with a small number of function evaluations (NFE). Comparing Annealed Langevin Dynamics (ALD) [Song et al.] with our RTK-based methods, we found that ALD has much lower marginal accuracy, particularly when NFE is small. Furthermore, using the Wasserstein Distance metric (Figure 3), we found that our RTK-based methods outperformed DDPM and ALD, especially with a small NFE.

For real-world data, we conducted experiments on the MNIST dataset, as shown in Figure 4. In Figure 4(a), we show that, with the same NFE, the RTK-based methods achieve better FID scores than DDPM, especially when NFE is small. Figures 4(b-d) further illustrate that the RTK-based methods generate images of higher quality than DDPM when NFE is small. We hope these supplementary experiments will address your concerns.

We will include these additional results and try more real-world dataset in the revision.

W2: Weakness. 2: The 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*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.

Actually, the gradient complexities (or the number of function evaluations) we analyzed in Theorem 4.1 and Theorem 4.4 are indeed calculated by $K\cdot S$. For Theorem 4.1, we provide the detailed chose of $K$ and $S$ at Line. 945 — Line. 947. For Theorem 4.4, we provide the $K$ and $S$ at Line. 301 — Line. 302 and Line. 304 — Line. 305 in the main body of this paper.

We also provided the number of function evaluations (NFE) in our experiments for the Gaussian mixture targets. According to the top-right figure in Figure 1, we compare the marginal accuracy [Mangoubi et al.] of different inference algorithms, whose x-axis explicitly denotes NFE. Therefore, we think the comparison is fair, and we will move our experiments to the main text in revision.

Furthermore, we also provide some comparisons in FID-NFE for both synthetic and real-world data which can be found in Figure. 3 (calculated by Wasserstein 2 distance) and Figure. 4 in our latest supplementary material of rebuttal. Both the FID-NFE graph and the case study illustrate that the RTK-based methods generate images of higher quality than DDPM when NFE is small. We hope these additional experiments will solve your concerns.

Minor comments:

A: We will revise these typos in our next version.

[Mangoubi et al.] Dimensionally Tight Bounds for Second-Order Hamiltonian Monte Carlo

Thanks for your positive comments and helpful suggestions. We will response to your concerns point by point in the following.

W1: The principal justification for this strategy is theoretical, and it is not yet clear whether it is practically beneficial.

AW1: In this paper, we mainly focus on the theory (RTK framework) about the convergence of the diffusion inference process and highlight how to design algorithms (RTK-ULD and RTK-MALA) to achieve a faster convergence rate in theory. In the meantime, we have also provided some experimental results for the Gaussian mixture targets shown in the Appendix. A, which justifies the effectiveness and efficiency of the RTK-based methods. Besides, we also add some new experiments in our latest supplementary material of rebuttal. These experiments are conducted on another two Mixture of Gaussians (MoG) and the MNIST dataset, which validates RTK-based inference process outperforms DDPM in different settings. More details about the supplementary experiments can be found in the common response. We will add them to our paper in the next version.

W2: Some of the English language was confusing at times, for example referring to "Gaussian process" when I believe "Gaussian approximation" (or similar) was intended.

AW2: We will revise them in the next version.

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

AQ1: In summary, learning the score function and solving the reverse SDE are separate subproblems of diffusion modeling. Specifically, solving the reverse SDE is the bottleneck for the inference process. For the practitioners, if they would like to generate data based on pre-trained models, the efficiency of generation only depend on the inference process , i.e., solving the SDE with well-trained scores. On the other hand, if one would like to develop a diffusion model for some new datasets, learning score in the training process will be the main focus.

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

AQ2: That is a good question. Our RTK framework definitely allows for plug-and-play in terms of the sampling method. This is one of the core advantages of the RTK framework.

As for the conditions of a new sampler that need to be satisfied, we require it to converge to any strong log-concave target distribution with arbitrary error tolerance for KL divergence or TV distance. To guarantee the final convergence, we first set the step size $\eta$ and the iteration number $K$ of outer loops as Corollary C.3 and C.5, i.e., $\eta = \frac{1}{2}\cdot \log \frac{2L+1}{2L},\quad K = 4L\cdot\log \frac{(1+L^2)d+\|\nabla f_*(0)\|^2}{\epsilon^2},$ which is independent of the choice of inner samplers. Then, if the inner sampler can guarantee the TV distance convergence, we require it to converge to the target distribution with an error tolerance at most TV\left(\hat{p}\{(k+1)\eta|k\eta}(\cdot|\hat{x}), p^\gets{(k+1)\eta|k\eta}(\cdot|\hat{x})\right)\le \frac{\epsilon}{K}\le \frac{\epsilon}{4L}\cdot \left(\log \frac{(1+L^2)d+\|\nabla f_*(0)\|^2}{\epsilon^2}\right)^{-1} which is shown in our Corollary C.3. A similar result can be found in Corollary C.5 if the inner sampler can achieve the convergence of KL divergence.

Thank you very much for your careful review. We will answer your concerns one by one in the following.

W1: The convergence results in section 4 are built upon a particular choice in the step size...

AW1: Sorry for the confusion. The particular step size choice in Section 4 is to simplify our theory. Any step size smaller than such a choice will be feasible, making the target distribution of the sampling subproblem strongly log-concave.

For general step size $\eta$ in RTK-ULD, the gradient complexity should be

$\tilde{O}\left( \underbrace{C_{\mathrm{LSI}}(\eta)^{3/2} \cdot \left(L+\frac{e^{-2\eta}}{1-e^{-2\eta}}\right)^{3/2}\cdot d^{1/2}\epsilon^{-1}\eta^{-1/2}}\_{S} \cdot \underbrace{\eta^{-1}\cdot \log \frac{(1+L^2)d+\\|\nabla f\**(0)\\|^2}{\epsilon^2}}\_{K} \right),$ which can be obtained with similar proof techniques in Theorem 4.4.

Note that the LSI constants of the subproblems depend on $\eta$. If the step size is large ($\Theta(1)$-level), the constant $C_{\mathrm{LSI}}$ will blow up and grow exponentially wrt $L\eta$, i.e., $C_{\mathrm{LSI}} = \exp(O(L\eta))$ (Proposition 1 of [Ma et al.]). If the step size is smaller than $\ln (1+1/2L)$, the target distribution of inner loops will be log-concave, and we have $C_{\mathrm{LSI}}(\eta) = \left(-L + \frac{e^{-2\eta}}{1-e^{-2\eta}}\right)^{-1}.$ Omitting the log-terms, the gradient complexity will be $\tilde{O}\left(\underbrace{ \left(\frac{L+\frac{e^{-2\eta}}{1-e^{-2\eta}}}{-L+\frac{e^{-2\eta}}{1-e^{-2\eta}}}\right)^{3/2}\cdot d^{1/2}\epsilon^{-1}\eta^{-1/2}}\_{S} \cdot \underbrace{\eta^{-1}}\_{K} \right).$ Therefore, plugging our setting $\eta=1/2\cdot \ln (1+1/2L)$, we can obtain the complexity shown in Theorem 4.4.

When $\eta$ is extremely small ($\Theta(\epsilon)$-level), we can use a one-step update for the inner sampler. Under these conditions, we have $S=1$, and the complexity will only be dependent on $K$. We defer the explanation in Question 3. In the next version, we will add this theorem to demonstrate the gradient complexity for general cases.

Q1: How the step size is derived? Is this a unique choice in some sense?

AQ1: Sorry for the confusion. The step size is derived from Line 188-189. Specifically, we requested the Hessian matrix of the RTK to be $L$-positive-definite, $-\nabla^2 \log p_{(K-k-1)\eta|(K-k)\eta}(x) = \nabla^2 f_{(K-k-1)\eta}(x)+ \frac{e^{-2\eta}}{1-e^{-2\eta}}\cdot I \succeq -L\cdot I + \frac{e^{-2\eta}}{1-e^{-2\eta}}\cdot I \succeq L\cdot I.$ where the first inequality follows Assumption 1 (L-smoothness of $\log p_t$). With this inequality, we have $\eta\le 1/2\cdot \ln(1+1/2L)$. Thus, we can provide the upper bound of the step size rather than the unique one. Some results shown in the para after alg 2 are also derived by the Hessian matrix. We will make it clear in the next version.

Q2: What is a practical way of choosing the step size for general target data like real-world image datasets?

AQ2: Since data distribution smoothness is usually unknown, we tune the step sizes with grid search. Practically, we first set the NFE as $M$. Then, we tune the outer iteration number $K$. For any $K$, the outer step size is set as $\eta = 1.0/K$, and the inner iteration number is $S = M/K$. Finally, we tune the inner step size $\tau$.

Q3: What are the challenges in complexity analysis compared to the cited results [8, 9, 21] in Table 1? ... can you show the method discussed in this paper yields a more general technique in terms of complexity analysis?

In summary, compared with [8, 9, 21] in Tab 1, the technical challenges of this paper are as below

• For the outer complexity analysis, how can we break away from the paradigm (used in [8, 9, 21]) of using the Girsanov Theorem to analyze the discretized reverse process? As a solution, we calculate the close form of the reverse transition kernel (RTK), which considers the inference as a series sampling subproblem.
• For the inner complexity analysis, how can we derive the gradient complexity when sampling algorithms can only access an approximate drift (score estimation $s_\theta$) or energy differences rather than the exact ones ($\nabla\ln p_t$)?
• For the total complexity analysis, how can we control the error propagation of a series of sampling sub-problems whose target distributions are RTKs? As a solution, we introduce chain rules of $f$-divergence.

For the generality, DDPM is a special case of our analysis when $\eta$ is chosen at an $O(\epsilon)$-level. When we choose DDPM to generate data, the ideal and practical reverse transition kernels are shown as follows p^\gets(x|x^\prime) \propto \exp\\left(-f_{k\eta}(x) - \frac{\\|e^\eta x^\prime - x\\|^2}{2(e^{2\eta}-1)}\right). $\hat{p}(x|x^\prime)\propto \exp\left(-\frac{\\|e^\eta x^\prime - 2(1-e^{\eta})s\_{\theta,(k+1)\eta}(x^\prime)- x\\|^2}{2(e^{2\eta}-1)}\right).$ with Lemma 3.1 and Line 123-132 in Section 2. We suppose step size is small, e.g., $e^{2\eta} - 1 = \epsilon$, and have $-\nabla_x^2 \log p(x|x^\prime) = \nabla\_{x}^2 f_{k\eta}(x) + \epsilon^{-1}\cdot I \succeq (2\epsilon)^{-1}\cdot I.$ This means the RTK is an $1/2\epsilon$-strongly log-concave distribution. With LSI, we have $\mathbb{E}\*{x^\prime}\left[ KL\left(\hat{p}(\cdot|x^\prime)\|p^\gets(\cdot|x^\prime)\right)\right]\le 2\epsilon\cdot \int \hat{p}(x,x^\prime)\cdot \left\|\nabla \log \frac{\hat{p}(x|x^\prime)}{p^\gets(x|x^\prime)}\right\|^2 dx \approx 2\epsilon \cdot \int \hat{p}(x,x^\prime)\cdot \left\|\nabla \ln p\_{k\eta} - s_{\theta,(k+1)\eta}(x^\prime)\right\|^2 dx,$ where RHS of this inequality is the same as the discretization error controlled in [Chen et al.], which can be further relaxed by Eq. (5.6) of [Chen et al.]. Thus, the method discussed yields a more general technique in terms of complexity analysis.

[Ma et al.] Sampling Can Be Faster Than Optimization.

[Chen et al.] Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.