Reverse Transition Kernel: A Flexible Framework to Accelerate Diffusion Inference

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 $ε$ target error within $\tilde{\mathcal O}(d^{1/2}ε^{-1})$ under mild conditions, and RTK-MALA enjoys a $\mathcal{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.

Paper

Similar papers

Peer review

Reviewer sbHs7/10 · confidence 3/52024-07-11

Summary

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

Strengths

- 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. - The paper also presents a valid argument in taking bigger step sizes in order to take advantage of Lipchitz continuity. - 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.

Weaknesses

I like this paper overall -- the main weaknesses are mostly about how it is written, as there exist small mistakes or confounding statements. - 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. - The paper suggests numerical experiments in the abstract, but only includes the experiments in the appendix. - 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. - Minor: the paper uses "closed" as opposed to "close", e.g., in Algorithm 1 and Line 172, 176.

Questions

- 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? - 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. - 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)? [1] Dockhorn T, Vahdat A, Kreis K. Score-Based Generative Modeling with Critically-Damped Langevin Diffusion. In: International Conference on Learning Representations [Internet]. 2021.

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

- The paper suggests in the final section several limitations of the current paper. I agree with the characterization of limitations.

Reviewer Y93M7/10 · confidence 1/52024-07-13

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.

Rating

7

Confidence

1

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Reviewer kQ8M6/10 · confidence 2/52024-07-15

Summary

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

Strengths

The 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. The overall presentation of the manuscript is good, and the theoretical analysis improves on existing ones.

Weaknesses

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

Questions

### Minor comments - potential typo on line 118 "To implement SDE." - wrong punctuation at the end of eq. 5. - Alg. 3 "of the previous iteration $x_0$", shouldn't it be $x_{k-1}$?

Rating

6

Confidence

2

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors address the limitations of their work in the manuscript.

Reviewer PyUH6/10 · confidence 3/52024-07-23

Summary

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

Strengths

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

Weaknesses

- The principal justification for this strategy is theoretical, and it is not yet clear whether it is practically beneficial. - Some of the English language was confusing at times, for example referring to "Gaussian process" when I believe "Gaussian approximation" (or similar) was intended.

Questions

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

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The main limitation is the lack of empirical support, which the author(s) acknowledge in the conclusion section.

Reviewer 81an6/10 · confidence 3/52024-07-26

Summary

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

Strengths

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

Weaknesses

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

Questions

1. 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. 2. What is a practical way of choosing the step size for general target data like real-world image datasets? 3. 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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations have been discussed in the paper.

Reviewer kQ8M2024-08-11

Thank you

Thank you for the clarifications, especially on the complexity and NFE. I also appreciate the new experiments.

Reviewer Y93M2024-08-11

After rebuttal

The authors addressed my comments, and I've raised my score accordingly.

Authorsrebuttal2024-08-12

Thank you so much for the response. If you have further concerns, please feel free to contact us. We're always ready to discuss them further and work with you to improve our paper.

Reviewer 81an2024-08-11

The authors have addressed my concerns. I have raised my score accordingly.

Authorsrebuttal2024-08-12

Thank you so much for the response. If you have further concerns, please feel free to contact us. We're always ready to discuss them further and work with you to improve our paper.

Authorsrebuttal2024-08-12

Thank you so much for the response. If you have further concerns, please feel free to contact us. We're always ready to discuss them further and work with you to improve our paper.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC