Improved Convergence Rate for Diffusion Probabilistic Models

Thanks a lot for your detailed and positive comments on our paper! We sincerely appreciate the time you took to review our work. Below, we address your points in detail:

1. Technical overview to explain the intuition

Thanks a lot for your explicit summary. The main improvement of our work over prior works is based on Lemma 2, which provides a simple framework for transferring the sampling task to discretizing the ODE (16) with a small $\ell_2$ error. Once this reduction is established, we leverage the randomized midpoint method, as in prior works, to achieve a more efficient discretization of the ODE and attain the desired bounds.

We will incorporate a detailed technical overview in the revised manuscript, to explain the intuition following your summary.

2. Discussion of non-Lipschitz setting

We will add a remark to explain the difficulty in removing the smoothness condition as follows: The extension to non-Lipschitz settings is significantly more challenging due to the critical role the smoothness assumption plays in our analysis. Previous analysis on diffusion models focus primarily on stepwise error propagation, analyzing how errors at one step affect the next. In contrast, our randomized design requires tracking error propagation across multiple steps simultaneously, which necessitates uniform control over these errors. Without smoothness, only a high-probability bound that depends on the dimension $d$ can be established, which is hard to use compared to the uniform bound and may result in worse bound since it scales with $d$.

3. Modification for Lemma 3 compared to Gupta et al.

We are sorry for missing this discussion regarding Lemma 3. While similar bounds are indeed presented in Gupta et al. (Appendix A.1), they are derived under an exponential integrator formulation. In our work, we derive Lemma 3 within the context of our framework and notations, tailored specifically to our method.

We will include the following remark in the revised manuscript: Similar bounds as Lemma 3 are also presented in Gupta et al. (Appendix A.1) under an exponential integrator formulation. Our derivation adapts these ideas to the context of the probability flow ODE discretization in this work.

Thanks a lot for your positive comments! We sincerely appreciate your time for reviewing our work. Below, we address your questions and suggestions in detail:

1. Core intuition and proof overview

The overarching framework of this paper is as follows. First, by injecting noise into the probability flow ODE (16), we transfer the sampling task to one of discretizing the ODE (16) with small $\ell_2$ error (see Lemma 2). Then we make use of the techniques proposed by Shen & Lee, (2019) and Gupta et al. (2024), utilizing randomized midpoints to improve efficiency in estimating the ODE solution (see Lemma 3). The final algorithm is designed based on the discretization form of ODE (16), where estimates of $x_{\tau_{k, i}}$ are iteratively computed using this form and are inserted into the subsequent estimates. In addition, the main sources of error include: discretization error controlled in Lemma 3, estimation error arising from the approximation of $x_{\tau_{k, i}}$, which is bounded in Lemma 4, and the initialization and early stopping errors, which are controlled in Lemma 5 using techniques similar to prior works.

To enhance clarity, we will include a detailed explanation for the core intuition and a proof overview in the revised manuscript.

2. Extension to log-concave sampling

Thanks a lot for your suggestion regarding log-concave sampling. As our result relies only on the smoothness condition, the total variation (TV) bound $d^{1/3}\epsilon^{-2/3}$ in Theorem 1 holds true directly for the log-concave case. Hense, our TV bound improves upon prior works, such as Theorem 1.3 in Gupta et al. (2024), by at least $d^{1/12}$, without requiring the stongly-log-concave condition.

We will add a detailed discussion on this extension in the revised manuscript.

3. Discussion on the non-Lipschitz setting

Extending our techniques to the more general case is more challenging due to the absence of the smoothness assumption. In most previous analyses of diffusion models, the primary focus is on the stepwise influence of error: how errors propagate from one step to the next. However, obtaining the benefit of the random design in our method requires tracking error propagation across multiple steps simultaneously. This multi-step dependency critically depends on the smoothness condition, which allows uniform control over error propagation. Without smoothness, the analysis can only guarantee some high probability bound that depends on the dimension $d$. This bound is hard to use compared to the uniform bound, and may lead to a worse bound since it scales with $d$.

We will add a remark in the manuscript to clarify the difficulties of extending our approach to the non-Lipschitz setting and the limitations of current techniques in this regard.

We sincerely appreciate your time for reviewing our paper. Below, we address your comments point by point:

1. Practical implementation details.

In practice, the neural network evaluations of score functions are typically conducted on batches (denoted as $B$) of data. This enables our algorithm to be implemented with at least $B$ parallelized processes without incurring communication overhead. As you said, the primary bottleneck in this setup is GPU memory capacity, which scales near linearly in $B = d^{1/3}\epsilon^{-2/3}$ in our Theorem 2. Moreover, we believe that the communication costs among multiple GPUs are negligible compared to the computation costs, as it is common practice to train and infer large models across many GPUs.

For example, using a single NVIDIA V100 GPU (32GB), it is feasible to set a batch size of over $50$ for latent diffusion model sampling on datasets such as FFHQ and LSUN (church/bedroom) [1]. This implies that the score estimates $s(x_t, t)$ for more than $50$ different $(x_t, t)$ pairs can be computed simultaneously, resulting in an approximately $50$-fold acceleration. This performance is sufficient to handle lower-accuracy scenarios where $\epsilon > 0.3$, ensuring that $d^{1/3}\epsilon^{-2/3} < B = 50$ for $d < 12288$.

To address the limitations of this parallel algorithm, we will include the following remark in the revised manuscript: "The implementation of this parallel algorithm assumes that the GPU memory is capable of supporting score estimations for a large batch of data simultaneously. Parallelizing across multiple GPUs introduces additional communication overhead, which may impact efficiency."

[1] Rombach, Robin, et al. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

2. Data assumption

Our framework imposes two assumptions on data: bounded second order moment and smoothness, which we believe encompass many practical applications. However, we acknowledge that exceptions exist:

The assumption of bounded second-order moments can accommodate many heavy-tailed distributions as long as their density decays no slower than $1/x^3$. This assumption fails only for distributions with density scales at the order of $1/x^c$ for $1<c<3$.

The smoothness assumption excludes non-smooth (e.g., discrete) distributions. It remains valid for certain multimodal distributions. For instance, if the data are generated from a Gaussian Mixture Model (GMM), $X_0 \sim \sum_i \pi_i\mathcal{N}(\mu_i, 1)$, then the score function remains smooth.

We will include the following remark in the revised manuscript to explain these assumptions: ''Our assumptions of bounded second-order moments and smoothness hold for many real-world distributions. However, there exist exceptions including extremely heavy-tailed or non-smooth distributions. For example, densities decaying slower than $1/x^3$ or purely discrete data may not align with our theoretical framework.''

We are grateful for your time for reviewing our paper. Below, we provide point-by-point responses to your comments and questions:

1. Smoothness condition

We agree that the requirement for a smoothness condition on the score function and its estimation is a limitation of our theoretical framework. However, we want to argue that our result provides improved convergence compared with prior results under general conditions in many cases, such as $L < 3381$ for $\epsilon = 1$ and $L < 72845$ for $\epsilon = 0.1$, as discussed in Section 1.2 and 3.2. Addressing the challenge of achieving sublinear dependency on $d$ without imposing the smoothness assumption needs much more extra effort, and will be left as our future research.

2. More explanations for algorithm and proof in Lemma 3

The core of the algorithm, as described in Lemma 2, involves discretizing the ODE (16) to generate $\widehat{X}_{k+1}$ from $\widehat{X}_k$. Given an initialization point $x\_{\tau\_{k,0}}=x_k$, the term $\theta(\tau; x_k) := \frac{s\_{\tau}^{\star}(x\_{\tau})}{2(1-\tau)^{3/2}}$ depends only on $\tau$. Solving (16) is equivalent to calculating the integral of $\theta(\tau; x_k)$. Consequently, it is natural for the algorithm to approximate the integral by discretizing $\tau$ and estimating $\theta(\tau; x_k)$ at some discrete time points. The intuition behind Lemma 3 lies in controlling the error between the true value $x\_{\tau\_{k, n}}$ and its estimation obtained through this discretization as defined in (21).

As for the proof of Lemma 3, we decompose the error into four terms, each of which is analyzed separately. The first two terms correspond to errors at the intial and end points, the second term contains the discretization error for all middle points, and the last term includes estimation errors. The second term is the main term and its analysis leverages the randomized design.

We will add a detailed explanation for Lemma 3 in the revised manuscript.

3. Connection with random RK method

Thanks a lot for pointing out the connection with random designed RK method. We will cite and compare with [1] appropriately. Here our method is termed ''randomized midpoint'' following the literature: it aligns with its initial introduction in [2] for log-concave sampling, and subsequent adaptation for diffusion models by [3].

4. Extension to SDE-type sampler

This is a quite interesting direction. The benefit of the randomized design in our work relies heavily on the deterministic nature of ODE process (16). It is still unclear for us how to adapt this approach to deal with the inherent stochasticity in SDEs while maintaining similar improvements. This will be left as our future work.

[1] Huang, Daniel Zhengyu, Jiaoyang Huang, and Zhengjiang Lin, Convergence analysis of probability flow ODE for score-based generative models, 2024.

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

[3] Shivam Gupta, Linda Cai, and Sitan Chen. Faster diffusion-based sampling with randomized midpoints: Sequential and parallel, arXiv:2406.00924, 2024.