Broadening Target Distributions for Accelerated Diffusion Models via a Novel Analysis Approach

We thank the reviewer for providing the highly inspiring feedback. In the revised paper, we have made changes based on your review comments and have highlighted all revisions in blue fonts.

Q1: Eq. (6): There is no guarantee that this matrix is positive semi-definite and, therefore, can be used as a covariance matrix. I agree that for most $x_t$ this will be the case for large $T$, but it is necessary here to have a definition that works for every $x_t$.

A1: Thanks for the great question/suggestion! To proceed more rigorously, we can project the matrices $\Sigma_t$ and $\widehat{\Sigma}_t$ onto the space of positive-semi definite (PSD) matrices for those $x_t$’s where either of these two matrices is not PSD (through eigen-decomposition and zeroing-out all negative eigenvalues). Since the measure of the events containing such bad $x_t$’s decreases to zero, all theoretical results in this paper, which are derived in expectation, will not be affected. We have made a footnote there in our revision.

Q2: lines 277-289: It is generally not recommended to refer to material, such as an assumption, before it has been introduced. In this case, Assumption 3 is mentioned before it is properly stated. Similarly, within the text of Assumption 3, Assumption 4 is referenced prior to its statement. I strongly recommend that the authors adjust the flow of the text to avoid such inversions.

A2: Thank you for your suggestion! To prevent such inversions, we have moved the statement of Lemma 2 into the appendix. Note that Lemma 2 is not the main contribution of this paper. We have also moved all discussions of the assumptions to be after the assumptions are stated.

Q3: Line 287: It is unclear why the factor $(1-\alpha_t)$ is retained on the left-hand side while $O(1-\alpha_t)$ is used on the right-hand side. Rather than writing $(1-\alpha_t)g_t = O(1-\alpha_t)$, wouldn’t it be clearer to simply write $g_t = O(1)$? Additionally, should this relation hold uniformly for $\ell \geq 1$?

A3: Many thanks for the question! This common term was initially kept there for our own understanding of the original terms on both sides of the equation. Certainly, it can go away for simplicity without losing mathematical rigor. We have changed the notation for the requirement of $H_t$ and make it simpler (in Lines 298 and 935). Also, as suggested, this relation needs to hold uniformly for $\ell \geq 1$.

Q4: Line 300: What is the precise meaning of analytic in this context? I was unable to find a definition in the text, either in the main body or in the appendix.

A4: This “analytic” is in the context of function analysis, where a function is analytic means that its Taylor series converges to the functional value at each point in the domain. We have provided this piece of clarification in the footnote of Assumption 2.

Q5: Assumption 4: Same remark as above, I guess it is possible to suppress the factors $(1-\alpha_t)^{p \ell}$ on both sides of the equalities.

A5: Thank you for the suggestion! In our revised draft, we have made our Assumption 4 clearer by suppressing the factors and left the more general version (which is Assumption 5) in the Appendix. In particular, Assumption 4 is still satisfied under two cases: (1) when $Q_0$ has finite variance, and (2) when $Q_0$ is Gaussian mixture.

Q6: Assumption 4: What is $p$ in this assumption? Should the boundedness of $p$-th order derivatives hold for every $p$? Should the upper bound be independent of $p$?

A6: In Assumption 4, $p$ is the order of partial derivatives. From Lemmas 15 and 17, the $p$-th order derivatives are a function of $p$, which might not have an upper bound of a constant in $T$. Nevertheless, Assumption 4 is not as strong as it appears to be. We can verify Assumption 4 on the following two common cases: (1) when $Q_0$ has finite variance, and (2) when $Q_0$ is Gaussian mixture. Case 1 clearly covers a broad set of target distributions of practical interest, such as images, and many theoretical studies of diffusion models have been specially focused on such a distribution [1-2]. Case 2 has also been well studied for diffusion models [3-4].

Note that Assumption 4 can be verified for Gaussian mixture $Q_0$’s (Lemma 14) and for a proper choice of $\alpha_t$ (Lemma 17).

Q7: Line 342: I did not find the definition of $\widehat{P}’_0$.

A7: Thank you. $\widehat{P}’_t$ is the distribution at intermediate time $t$ using the estimated scores and Hessians. Thus, $\widehat{P}’_0$ is the distribution of the final output from the accelerated diffusion process. We added the definition of $\widehat{P}’_t$ in Section 3.1.

Dear Reviewer cL12,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Dear Reviewer cL12,

Since the author-reviewer discussion period will end soon, could you please kindly check whether our responses have appropriately addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Please also feel free to let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors

We thank the reviewer for providing the highly inspiring feedback. In the revised paper, we have made changes based on your review comments and have highlighted all revisions in blue fonts.

Q1: One weakness is that the sampler requires approximate access to higher-order score functions (namely, the Jacobian of the score). While there are examples of diffusion generative models in practice that use higher-order scores essentially by applying automatic differentiation to the first-order score (e.g. GENIE), this is not standard and introduces significant computational overhead. Additionally, these models in practice are not actually trained according to some kind of explicit Hessian matching loss, so it is unclear why one should expect the Jacobian of the score estimates to be close to the Jacobian of the scores. That said, the prior work of Li et al. on acceleration, which appeared in a previous ICLR, also works with second-order access, so this weakness is somewhat justifiable.

A1: We appreciate that the reviewer justifies such a requirement in the context of the existing work. Compared to previous works, our accelerated sampler can be applied to a much larger class of distributions. In particular, we do not require the Lipschitz-smooth condition along the entire sampling path nor require the target to have bounded support. Further, within the class of non-smooth targets, our sampler achieves a better dimensional dependency than previous works.

Numerically, our Fig. 1 demonstrates that our proposed sampling method substantially outperforms the previous design in [1]. Also, following the reviewer’s suggestion in Q6, we have added the work of GENIE [5] in the Related Works in the Appendix and have provided the paragraph to discuss that work. Our paper theoretically justifies the accelerated empirical performance of [6] in the regime when the Hessian of $\log q_t$ is well-estimated.

Q2: Also, not a weakness per se, but note that the result for second-order smoothness (Theorem 4) is technically incomparable to the other previous works getting scaling because while it does not assume smoothness along the trajectory, it needs to assume that the Hessian of the log-density is Lipschitz. And as mentioned above, this result requires second-order scores whereas the probability flow ODE paper of Chen et al. only uses first-order scores.

A2: We agree with the reviewer’s comment. Our results provide new contributions that complement existing studies by exploring different assumptions of distributions, which enriches the existing set of distributions studied in the literature. We have added this remark after Theorem 4.

Q3: Another note: while the authors are careful to clarify the chronology of the various related works, it is worth pointing out that in the last month or so (and before the ICLR deadline), there have been important developments in this literature that supersede some of the main results of this submission. For instance, Li & Yan showed that even for the standard SDE-based sampler (without second-order score), one can achieve an accelerated rate of $O(d/\epsilon)$, in fact under only a first moment bound.

A3: We appreciate the reviewer’s note, and we are also aware of these recent studies. However, we are following the ICLR guidelines that the paper does not need to discuss the research studies that appeared within the two months of the submission deadline, which clearly extends to the work posted after the ICLR deadline. Hence, we kindly ask the review to evaluate this paper based on the prior research, not the con-current and later research.

Further, even given those recent new developments, we believe that our paper still holds significant new contributions in (i) the proposed Hessian-based accelerated sampler is still novel, which is a representative line of accelerating DDPM following [2]; and (ii) our analysis technique is distinct in DDPM analysis and holds special advantage for certain target distributions.

Q4: A recent work of Gatmiry-Kelner-Lee also established bounds on higher-order scores of Gaussian mixtures, though specifically in the context of identity-covariance components. How do their bounds compare to the ones proved in this submission?

A4: The work by Gatmiry, Kelner, and Lee [4] addresses the problem for efficiently learning the score function. In comparison, our work addresses the problem “Given an accurately learned score and Hessian function, what guarantees can we obtain for sampling from the data distribution?”. Therefore, our work lies in an orthogonal direction to [4], which is not comparable in a straightforward way. This is also the reason why we have included [4] in the related work on Theory on Score Estimation in the related works.

Dear Reviewer ZY6t,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Dear Reviewer ZY6t,

Since the author-reviewer discussion period will end soon, could you please kindly check whether our responses have appropriately addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Please also feel free to let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors

We thank the reviewer for providing the highly inspiring feedback. In the revised paper, we have made changes based on your review comments and have highlighted all revisions in blue fonts.

Q1: Assumption 4 plays a crucial role in the construction of the tilting factor, requiring that the score function has derivatives of all orders and that the $p$-th moment is bounded $p \geq 1$. In my view, this condition is not as "rather soft" as suggested in the submission. The requirement for all derivatives to exist means that the score function is not merely smooth but infinitely differentiable, a level of regularity that is often absent in practical score functions, particularly for complex or non-smooth distributions. Furthermore, the requirement that the $p$-th moment of the derivatives is bounded indicates that the derivatives exist and do not grow excessively. This can be quite a restrictive requirement, especially for distributions with heavy tails or irregularities. While this assumption is essential for maintaining theoretical rigor, it may limit the applicability of the models to a narrower range of problems.

A1: Thanks for the question. We can verify Assumption 4 on the following two common cases: (1) when $Q_0$ has finite variance, and (2) when $Q_0$ is Gaussian mixture. Case 1 clearly covers a broad set of target distributions of practical interest, such as images, and many theoretical studies of diffusion models have been specially focused on such a distribution [1-2]. Case 2 has also been well studied for diffusion models [3-4].

A special note is that Case 1 in fact includes substantial complex or non-smooth distributions. To see this, even though $\log q_0$ might not be smooth, all Gaussian perturbed $\log q_t$’s are still infinitely smooth, which allow us to use early-stopping (Corollary 1 and Theorem 3). Furthermore, we do not require Assumption 4 in case that the Hessian of $\log q_0$ is smooth (Theorem 4).

We have added these explanations below Assumption 4 in the revised paper.

Q2: The proof is quite lengthy, and I did not have the opportunity to examine all the details. It would be beneficial if the author could elaborate on how the accelerated convergence is achieved using the proposed novel proof technique (specifically, in Step 2 regarding the tilting factor). Additionally, restating Tweedie's formula and clarifying its extension in this context would help enhance the clarity of the proof for readers.

A2: Many thanks for the question. We added the following paragraph to explain Step 2 in the revision, especially on how the acceleration is achieved:

For regular DDPMs, there is no control for the variance of the reverse sampling process, and thus $B_t(x_t) \equiv 0$. In this case, the dominating rate is determined by the expected values of $T_2$. With the variance correction in our accelerated sampler, the corresponding $B_t(x_t)$ enables us to cancel out the second-order Taylor term (see Lemma 11). As a result, the rate-determining term becomes the expected values of $T_3$, which decays faster. Thus, the acceleration is achieved.

We also provide the following clarification for how we extend Tweedie’s formula in Step 4:

To calculate posterior moments, we extend Tweedie's formula in a non-trivial way. Whereas the original Tweedie's formula provides an explicit expression for the posterior mean for Gaussian perturbed observations, we explicitly calculate the first six centralized posterior moments and provide the asymptotic order of all higher-order moments, drawing techniques from combinatorics.

We thank the reviewer again for your highly inspiring comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

[1] Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating Convergence of Score-Based Diffusion Models, Provably. In the Forty-first International Conference on Machine Learning, 2024.

[2] 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. 2024.

[3] Khashayar Gatmiry, Jonathan Kelner, and Holden Lee. Learning Mixtures of Gaussians Using Diffusion Models. arXiv preprint arXiv:2404.18869.

[4] Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning General Gaussian Mixtures with Efficient Score Matching. arXiv preprint arXiv:2404.18893.

Dear Reviewer gLE4,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Dear Reviewer gLE4,

Since the author-reviewer discussion period will end soon, could you please kindly check whether our responses have appropriately addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Please also feel free to let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors