Paper: Momentum-accelerated Diffusion Process for Faster Training and Sampling (ICLR-2024)

Reviewer comment (Theory):
The idea to connect SGD with forward process of diffusion process is interesting. However, there's a main difference between them: the gradient $g_t$ in SGD has unique relation to $x_t$ but $\epsilon_t$ is only a random variable, and is independent with $x_t$. They only share the formulation of the formula but has few connection in their meaning. I doubt under this situation, the momentum in the forward process still makes sense.

Author response:
Since we are not certain that we fully understand your question, we will answer it in two aspects, and hope to address your concern. 

**(1) SGD on our constructed dynamic function share the exactly same formulation with the forward diffusion process.** On the constructed dynamic function $f(x) = \mathbb E_{\zeta \sim \mathcal{N}(0, bI)} \frac{1}{2} \lVert x - \frac{\beta_t}{1 - \alpha_t} \zeta \rVert_2^2$, SGD first samples a minibatch data {$\zeta_k$}, and then  computes stochastic gradient  $g_t = x_t - \frac{\beta_t}{b(1 - \alpha_t)} \sum\nolimits_{k=1}^b \zeta_k= x_t - \frac{\beta_t}{1 - \alpha_t} \epsilon_t$ where we use re-parameterization trick by replacing $\frac{1}{b} \sum\nolimits_{k=1}^b \zeta_k \sim \mathcal{N}(0, I)$ with $\epsilon_t \sim \mathcal{N}(0, I)$. Finally, SGD updates parameter $x$ via the following stochastic optimization process:
$$	
    x_{t+1} =  x_t - (1 - \alpha_t) g_t  =  \alpha_t x_t + \beta_{t} \epsilon_t, 
\tag{1}
$$
where learning rate $(1 - \alpha_t) > 0$ is to align the diffusion and SGD processes. Above Eq. (1) exactly aligns with the forward diffusion process in DDPM.  

From the above steps, we can find that the stochastic  gradient $g_t = x_t - \frac{\beta_t}{b(1 - \alpha_t)} \sum\nolimits_{k=1}^b \zeta_k= x_t - \frac{\beta_t}{1 - \alpha_t} \epsilon_t$ in SGD involves the parameter $x_t$ and the noise $\beta_{t} \epsilon_t$. But the noise $\beta_{t} \epsilon_t$ is induced from the minibatch sample {$\zeta_k$}.  More importantly, the optimization works, e.g., [1,2], does not care the specific formulation of  stochastic gradient, such as whether stochastic gradient $g_t$ includes the parameter  $x_t$ and an independent noise $\beta_{t}\epsilon_t$, it only requires whether the stochastic  gradient $g_t$ is an unbiased estimation to the full gradient of an objective $f(x)$. Then these works can prove faster convergence speed of momentum SGD than vanilla SGD on  the traditional strongly convex problems.  For our  dynamic function, our Theorem 1 also shows faster convergence speed of momentum SGD than vanilla SGD. 

Besides, our momentum forward diffusion process can be applied to several popular diffusion frameworks, e.g. VP, VE, and EDM, and reduces their training cost by about 50% and inference cost by 66% with comparable image synthesis performance across several datasets. These practical results show the effectiveness of the momentum forward diffusion process. 

**(2) The roles of $g_t$, $x_t$ and $\epsilon_t$ in the SGD and forward process are the same.**

+ In SGD: $g_t$ represents the stochastic gradient used to update $x_t$. This gradient is a function of $x_t$ itself and the noise term $\epsilon_t$ due to the randomness of sampling a minibatch of data {$\zeta_k$}. The noise term $\epsilon_t$ in SGD is essential to simulate the stochastic nature of gradient estimation, please refer to **(1)** for details. 

+ In the forward process: The analogy in the diffusion process lies in how $x_t$ and $\epsilon_t$ interact to drive the process. The noise term in the diffusion process plays a similar role to that in SGD, introducing randomness that perturbs the process. In this context, $g_t$ can be seen as an analogous mechanism that influences the evolution of $x_t$ over time, although the specific dynamics and objectives differ from SGD.

The key point of similarity lies in the stochastic nature of both processes. In SGD, the stochasticity helps in escaping local minima and exploring the solution space, whereas in the diffusion process, it accelerates the transition from a data distribution to a Gaussian distribution. This parallel in their mechanisms underlines the relevance of using momentum in both contexts. Momentum, by accumulating historical updates (in SGD) or influencing the trajectory (in the diffusion process), helps in smoothing out the updates and providing a more stable and efficient trajectory towards the objective, whether it be finding the optimum of a function or effectively transitioning through the diffusion process.

[1] Christopher De Sa, https://www.cs.cornell.edu/courses/cs4787/2020sp/lectures/Lecture7.pdf

[2] Yanli Liu, Yuan Gao, Wotao Yin, An Improved Analysis of Stochastic Gradient Descent with Momentum, NeurIPS 2020.

Reviewer comment (Theory):
'following EDM, we remove the stochastic gradient noise in Eq. (2)'. Could you point out where EDM has stated this? Discarding stochasticity during forward process seems strange. If there's no stochasticity in forward process, then the equation makes no sense but only $x_T = \prod_{i=1}^{T} \alpha_t x_0$ and there's nothing for the model to learn.

Author response:
Thanks. For EDM (see arXiv.2206.00364), it claims this at least two places. 
On page 4, it states "The main purpose of this reframing is to bring into light all the independent components that often appear tangled together in previous work. In our framework, there are no implicit dependencies between the components -- any choices (within reason) for the individual formulas will, in principle, lead to a functioning model."

On page 21, it also says "Given that the idea of the probability flow ODE is to match a particular set of marginal distributions, it makes sense to treat the marginal distributions as first-class citizens and define the ODE directly based on $\sigma(t)$ and $s(t)$, eliminating the need for $f(t)$ and $g(t)$."

Please see more details in the corresponding sections. The EDM design space eliminates implicit dependencies between each component, including mean scaling $\mu_t$ and noise schedule $\sigma_t$. This independence grants us sufficient flexibility to separately design the sample mean scaling and noise accumulation. Inspired by this, we derive the expectation term of our momentum-based diffusion process, as elaborated from Eq. (9) to Eq. (13) in our paper. In this part of the derivation, the noise term is notably absent, as the change in expectation is not dependent on the noise term.

Reviewer comment (Experiments):
It is better to show more results by applying FDM into state-of-the-art diffusion model (Stable Diffusion).

Author response:
**Our experimental approach has been evaluated on several widely used datasets.** It is comparable in scale to those used in other papers such as [3-5]. Conducting experiments on these datasets required significant computational resources, specifically 8 A100 GPUs for approximately 4 days of training. This level of investment in terms of hardware and time underscores our commitment to rigorous testing and validation of our FDM. 

We acknowledge that ImageNet 256x256 represents a larger-scale and higher-resolution dataset that could further test the capability of our model. However, it is important to note that ImageNet 256x256 is not typically used as a standard dataset for pixel-level diffusion models due to the high amount of GPU resource requirements. See the experiments in VP, VE, and EDM to check.  
Moreover, even training on a dataset like ImageNet 64x64 would require substantial resources, specifically 32 A100 GPUs for about 14 days, which is beyond our current GPU resource capacity.

[3] Dockhorn, Tim, Arash Vahdat, and Karsten Kreis. Score-Based Generative Modeling with Critically-Damped Langevin Diffusion. ICLR 2022.

[4] Zhang, Qinsheng, Molei Tao, and Yongxin Chen. gDDIM: Generalized Denoising Diffusion Implicit Models. ICLR 2023.

[5] Du, Weitao, He Zhang, Tao Yang, and Yuanqi Du. A flexible diffusion model. ICML 2023.

Reviewer comment (Theory):
Some claims are not fully supported. Theorem 1 is mainly a convergence result of SGD with momentum on a quadratic function. It is unclear how this rate would imply a faster convergence speed of the modified forward diffusion process (eqn 2) despite some similarities in the iterates. This is further complicated by the fact that the noise term in eqn 2 is treated as a constant in eqn 9. Theorem 2 shows the convergence of the proposed diffusion process to a Gaussian distribution. However, it does not really quantify the rate of convergence, and its implications on the reverse sampling process is largely a heuristic.

Author response:
**Response 1:**
**The convergence rate of the optimization problem we examine is directly analogous to the convergence rate of the diffusion process towards equilibrium.**  This relationship is elucidated through two key points:

**(1) Alignment of Processes:**  According to Section 4.1, we have that, when minimizing a stochastic quadratic time-variant function $f(\mathbf{x}) = \mathbb{E}_{\zeta \sim \mathcal{N}(0, b\mathbf{I})} \frac{1}{2} \lVert \mathbf{x} - \frac{\beta_t}{1 - \alpha_t} \zeta \rVert_2^2$ with learning rate $\eta_t = 1 - \alpha_t$, the SGD optimization process aligns with the diffusion process in each iteration. This alignment establishes a clear link between the dynamics of the SGD and the diffusion process.

**(2) Alignment of Equilibrium State:** The optimal variable for minimizing $f(\mathbf{x})$ can be calculated as $\mathbf{x}^*=0$, which aligns with the expectation of equilibrium state of the diffusion process. According to Theorem 1, the expectation error in momentum SGD decreases more rapidly towards this optimal state. This implies that the momentum-based diffusion process also approaches equilibrium at an accelerated rate.


**Response 2:**
**In our derivation, we consistently treat the noise term as a variable, not a constant.** In the EDM design space, the probability flow ODE of the reverse process is defined by two distinct and independent factors: mean scaling $\mu_t$ and noise schedule $\sigma_t$. This independence grants us sufficient flexibility to separately design the sample mean scaling and noise accumulation. We begin by deriving the expectation term of our momentum-based diffusion process, as elaborated from Eq. (9) to Eq. (13) in our paper. In this part of the derivation, the noise term is notably absent, as the change in expectation is not dependent on the noise term. 

In subsequent developments of our model, we integrate the noise schedule $\sigma_t$ directly from the relevant baseline DMs as shown in Eq. (14) and Eq. (15). The key here is the independence of $\sigma_t$ from $\mu_t$, which allows us to adopt it unaltered. In practice, we found that numerical computation of the closed-form solution for accumulated noise like CLD does not yield practical benefits. This approach can be interpreted as integrating the expectation of momentum into the vanilla diffusion process of a certain DM. This integration is seamless, facilitating the incorporation of our momentum-based diffusion process with existing DMs, which enhances the existing DMs' performance while preserving their fundamental convergence property to the data distribution. This approach not only maintains the integrity of the baseline models but also enriches them with the efficiency and effectiveness of momentum-based diffusion.

**Response 3-1:**
The rate of convergence is specifically addressed in Theorem 1. In Theorem 1, we demonstrate the faster convergence of our momentum-based diffusion process compared to the vanilla diffusion approach. On the other hand, Theorem 2 focuses on establishing a clear relationship between $x_t$ and $x_0$, whose primary goal is to provide an intuitive understanding of how the momentum-based diffusion process evolves over time, especially to highlight the difference from the vanilla diffusion process. Theorems 1 and 2 together provide a holistic view of our model. While Theorem 1 provides quantitative analysis regarding the convergence speed, Theorem 2 contributes to the understanding of the state transitions in our momentum-based diffusion process.

**Response 3-2:**
**Since the reverse process is decided by its forward diffusion process, one can expect the accelerated speed as the forward one, and also enjoy the acceleration effect.**  This is because our momentum is implicitly incorporated into the score function in the reverse process. For example, $\nabla \log p_{\text{VP}}(x_t | x_0) \propto -\frac{x_t - e^{-\mathcal{B}(t)}x_0}{\beta(t)^2}$ in VP, while $\nabla \log p_{\text{FDM}}(x_t | x_0) \propto -\frac{x_t - e^{-\mathcal{B}(t)}(1 + \mathcal{B}(t))x_0}{\beta(t)^2}$ in our FDM. The only difference comes from the momentum. In these formulations, we can see the acceleration in the forward process results in the acceleration in the reverse process. Intuitively, the forward process in the momentum diffusion process has the formulation $x_{t+1} = x_t + \beta_t \epsilon + \alpha (x_{t}-x_{t-1})$, which leads to a larger distance between $x_{t}$ and $x_{t+1}$ than the one in the vanilla diffusion process. Thus, in the backward process of the momentum diffusion process, using $x_{t+1}$ to predict $x_{t}$ is also a more aggressive prediction. Moreover, our experiments on both EDM sampler (Table 3) and DPM-Solver++ (Table 4) affirm the acceleration on the reverse sampling process in practice.

Reviewer comment (Theory):
The actual algorithm (and its implementation) differs from the iterates considered in Theorem 2. Thus, the theory does not really capture the training dynamics.

Author response:
**Response 4:**
**The iteration considered in Theorem 2 is an abstraction for ease of notation and analysis.** In Theorem 2, we  set $\beta_t \equiv \beta$ to streamline our analysis. While the abstraction in Theorem 2 may not capture every nuance of the actual algorithm's implementation, it still faithfully reflects the fundamental mechanism of the diffusion process that gradually diffuses a data point into random noise. By abstracting the iteration process in this way, we aim to provide an intuitive and clear understanding of how the momentum-based diffusion process evolves, especially to highlight the difference from the vanilla diffusion process. This type of abstraction is a key in physics and mathematical theory to gain insights. 

Although this approach does not encompass all the complexities of the training dynamics found in practical implementations, it still offers valuable insights. It allows us to identify the fundamental aspects of the diffusion process and the role of momentum in it. The insights gained from this abstracted state transition in Theorem 2 are critical to guide the practical application and development of our FDM.

Reviewer comment (Novelty):
Using momentum to accelerate the optimation seems not a new idea, and the contribution is somewhat limited.

Author response:
Thank you for your valuable and insightful comments! In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.


**Question 1: Using momentum to accelerate the optimation seems not a new idea, and the contribution is somewhat limited.**

**Response 1:**
**The incorporation of momentum into the diffusion process, as presented in our work, represents a substantial advancement in this field.** While momentum is a well-known acceleration technique in optimization, only a few works focus on its application within the diffusion process, which we have detailed compared in Section 2 and Section 5.3.
Our contribution is listed as follows:

**(1) Theoretical Insights:** Our work bridges a crucial gap, establishing a clear and direct link between diffusion processes and gradient descent by constructing a dynamic function.   This connection inspires our exploration of the momentum-accelerated diffusion process. Moreover, we provide theoretical analysis to show the faster convergence rate in Theorem 1,  and establish a clear relationship between the perturbed image $\mathbf{x}_t$ and the clean image $\mathbf{x}_0$ in Theorem 2 which also shows the different convergence speed of momentum-accelerated diffusion process and the vanilla one.

**(2) Practical Insights:** We have effectively demonstrated the seamless integration of a momentum-based diffusion process across various diffusion models. To our knowledge, we are the first to successfully integrate momentum into the diffusion process without necessitating an *additional* momentum (velocity) space. Unlike other momentum-based methods that utilize an auxiliary space to simplify the computation of the perturbation kernel, our approach involves a more challenging task. We aim to directly derive the desired perturbation kernel of the momentum-based diffusion process exclusively within the sample space. To overcome this issue, we follow EDM and eliminate the implicit dependencies between each component, including mean scaling $\mu_t$ and noise schedule $\sigma_t$. This independence grants us sufficient flexibility to separately design the sample mean scaling and noise accumulation, thus allowing us to formulate the forward mean-varying process by a second-order damped-oscillation ODE (Eq. (12)) without the auxiliary velocity term.
This innovative approach offers numerous benefits such as numerical stability, memory efficiency, and wide applicability, which paves the way for more efficient and streamlined diffusion processes. 

**(3) Experimental Results:** Through extensive experiments, we verified improvements in training and sampling efficiency, along with superior image generation quality. 
Specifically, our FDM reduces the training cost of DMs by about 50\% and the inference cost by 66\% with a comparable image synthesis performance across several datasets. 
This practical application of our theoretical insights underscores the real-world impact of our research.

The additional benefit of our approach is evident in our experimental study presented in Table 7. In the experiments, both CLD and our FDM were applied using the traditional RK45 solver. However, a significant distinction was observed when a more efficient EDM sampler was employed. While CLD exhibited a notable performance decrease with the EDM sampler, our FDM maintained consistent performance.
Moreover, we have extended our experimental analysis to include gDDIM[1], which addresses oscillation/overshoot and numerical instability issues commonly found in the vanilla CLD. The results of these experiments are presented below. Despite the specific enhancements implemented in gDDIM to improve upon CLD, the experimental results demonstrate that our FDM still significantly outperforms gDDIM. This further underscores the robustness and efficiency of our FDM in comparison to previous methods.

**Table: Image synthesis performance on CIFAR-10 with 35 NFEs. gDDIM (CLD) use specifically designed gDDIM sampler, while VP-FDM use EDM sampler. The result denoted with $\dagger$ is from the original paper, with 50 NFEs.**
|    Method   | 30 Mimg | 50 Mimg | 70 Mimg | 100 Mimg |      Best      |
|:-----------:|:-------:|:-------:|:-------:|:--------:|:--------------:|
| gDDIM (CLD) |   5.15  |   2.95  |   2.76  |   2.51   | 2.25$^\dagger$ |
| VP-FDM      |   4.01  |   2.74  |   2.37  |   2.24   |      2.08      |

[1] Zhang, Qinsheng, Molei Tao, and Yongxin Chen. gDDIM: Generalized Denoising Diffusion Implicit Models. ICLR 2023.

Reviewer comment (Theory):
In Theorem 1, the author argues that the momentum SGD is faster by comparing the errors in expectation. However, does the upper bound of SGD tight? It is not safe to draw the conclusion that SGD is slower by simply showing that the upper bound is higher than momentum SGD, epsically when we are not aware whether this upper bound is tight.

Author response:
**Question 2: In Theorem 1, the author argues that the momentum SGD is faster by comparing the errors in expectation. However, does the upper bound of SGD tight? It is not safe to draw the conclusion that SGD is slower by simply showing that the upper bound is higher than momentum SGD, especially when we are not aware whether this upper bound is tight.**

**Response 2:**
**The bound for vanilla SGD, as presented in Theorem 1, is indeed strict and accurately reflects the convergence behavior.** To elucidate this, we delve into the step-by-step derivation of the expectation error, beginning with Eq. (22) in  our appendix:
$$
\begin{align}
	\lVert \mathbb{E}[x_{k+1}-x^*] \rVert
	&  \overset{(i)}{=} \lVert \prod_{j=0}^k (1 - \eta_j) (x_0 - x^*) + \mathbb{E} [\sum_{j=0}^k \prod_{i=1}^{k-j} (1 - \eta_i) \eta_j \sigma_j \epsilon_j ] \rVert  \\\\
	& \overset{(ii)}{=} \lVert  \prod_{j=0}^k (1 - \eta_j)  (x_0 - x^*) \rVert  \\\\
	& = \lVert x_0 - x^* \rVert \prod_{j=0}^k (1 - \eta_j) ,
\end{align}
$$

where $(i)$ is due to $x_i$ is uncorrelated to $\epsilon_j$ for any i,j, and $(ii)$ is due to $\mathbb{E}[\epsilon_j] = 0$. Similarly, for the forward diffusion process in DDPM, we can derive  $\lVert \mathbb{E}[\prod_{k=1}^t \alpha_k x_0] \rVert = \lVert \prod_{k=1}^t (1 - \eta_k) x_0 \rVert= \lVert x_0 - x^* \rVert \prod_{k=1}^t (1 - \eta_k) $ where we use $\alpha_t = 1 - \eta_t$ and $x^* = 0$. In this way, for both results of the momentum forward process and the vanilla one, their bounds are tight since both use "=". Our Theorem 1 can successfully show that our momentum-based diffusion process achieves faster convergence w.r.t the expectation.

Reviewer comment (Novelty):
The idea of aligning diffusion process with stochastic gradient descent, and adopting acceleration techniques in SGD is not entirely new.

Author response:
In their paper [2], Ma et al. first formulated gradient-based MCMC sampling as optimization in the probability space, utilizing various mathematical notations and tools such as Lebesgue measure, second moment, coupling, Wasserstein-2 distance, and optimal transport. Then they showed that an underdamped Langevin algorithm performs accelerated gradient descent by transforming accelerated gradient descent into its continuous dynamic. The most significant contribution in [2] is that they further proved the faster convergence rate of underdamped Langevin algorithm than the classical Langevin algorithm under some conditions, such as log-Sobolev inequality. These results are rigorously presented using both sample and velocity spaces, offering novel and profound insights into understanding gradient-based MCMC sampling and its accelerated version.
 
By comparison, there are two main differences between this work and [2] because of their different focuses and targeted readers.  

**(1)** Our new simple and intuitive understanding complements the previous deep mathematical insights of [2] and thus allowed us to develop practically efficient algorithms. This work constructs a dynamic function directly connecting the gradient descent with the diffusion forward process. It does not borrow complex mathematical notations and tools in [2]. Then, motivated by the faster convergence speed of momentum gradient descent, this work introduces momentum to accelerate diffusion forward process, and shows the faster speed. This work focuses on more practical and efficient diffusion model design for practical designers and engineers, in contrast to the deep mathematical insights targeted in [2]. 

**(2)** This work prioritizes efficiency, an aspect not extensively considered in both [2] and CLD [1]. This work eliminates the velocity term, yielding three benefits. **(a)** It stabilizes training, since using velocity and sample spaces like CLD and [2] needs numerically unstable Cholesky factorization (see CLD), while removing velocity allows to simplify the perturbation kernel and avoids such instability. **(b)** It reduces model memory cost which is important in practice, since using two spaces needs roughly $2\times$ model parameter and memory cost than using a single one, and poses challenges to integrate with SoTA diffusion models which are often of large model sizes. **(c)** It improves compatibility by allowing our diffusion process to be directly integrated into popular diffusion models like EDM and numerical solvers like dpm-solver, which primarily use the sample space. This is in contrast to using two spaces, which would necessitate substantial revisions to their formulations or might be incompatible.

In the final version which should allow us to use an extra page, we will discuss these related works for their big efforts on understanding gradient-based MCMC sampling and its acceleration version in detail. We do not discuss this now, since there is really no space to add new content.

Reviewer comment (Novelty):
In Dockhorn et al. (2021) [1], they first proposed the critically-damped Langevin diffusion, in which they mentioned that "the velocity in CLD accelerates mixing in the diffusion process, and is equivalent to momentum in gradient descent", where they refer to Ma et al. (2019) [2] for the equivalence.

Author response:
In their paper [2], Ma et al. first formulated gradient-based MCMC sampling as optimization in the probability space, utilizing various mathematical notations and tools such as Lebesgue measure, second moment, coupling, Wasserstein-2 distance, and optimal transport. Then they showed that an underdamped Langevin algorithm performs accelerated gradient descent by transforming accelerated gradient descent into its continuous dynamic. The most significant contribution in [2] is that they further proved the faster convergence rate of underdamped Langevin algorithm than the classical Langevin algorithm under some conditions, such as log-Sobolev inequality. These results are rigorously presented using both sample and velocity spaces, offering novel and profound insights into understanding gradient-based MCMC sampling and its accelerated version.
 
By comparison, there are two main differences between this work and [2] because of their different focuses and targeted readers.  

**(1)** Our new simple and intuitive understanding complements the previous deep mathematical insights of [2] and thus allowed us to develop practically efficient algorithms. This work constructs a dynamic function directly connecting the gradient descent with the diffusion forward process. It does not borrow complex mathematical notations and tools in [2]. Then, motivated by the faster convergence speed of momentum gradient descent, this work introduces momentum to accelerate diffusion forward process, and shows the faster speed. This work focuses on more practical and efficient diffusion model design for practical designers and engineers, in contrast to the deep mathematical insights targeted in [2]. 

**(2)** This work prioritizes efficiency, an aspect not extensively considered in both [2] and CLD [1]. This work eliminates the velocity term, yielding three benefits. **(a)** It stabilizes training, since using velocity and sample spaces like CLD and [2] needs numerically unstable Cholesky factorization (see CLD), while removing velocity allows to simplify the perturbation kernel and avoids such instability. **(b)** It reduces model memory cost which is important in practice, since using two spaces needs roughly $2\times$ model parameter and memory cost than using a single one, and poses challenges to integrate with SoTA diffusion models which are often of large model sizes. **(c)** It improves compatibility by allowing our diffusion process to be directly integrated into popular diffusion models like EDM and numerical solvers like dpm-solver, which primarily use the sample space. This is in contrast to using two spaces, which would necessitate substantial revisions to their formulations or might be incompatible.

In the final version which should allow us to use an extra page, we will discuss these related works for their big efforts on understanding gradient-based MCMC sampling and its acceleration version in detail. We do not discuss this now, since there is really no space to add new content.

Reviewer comment (Theory):
While removing the velocity term may stabilize training and potentially result in better performance compared with CLD, there is a question of how to compare the understanding in [2] — that critically damped Langevin equals gradient descent on probability space — with the viewpoint proposed in this paper.

Author response:
In their paper [2], Ma et al. first formulated gradient-based MCMC sampling as optimization in the probability space, utilizing various mathematical notations and tools such as Lebesgue measure, second moment, coupling, Wasserstein-2 distance, and optimal transport. Then they showed that an underdamped Langevin algorithm performs accelerated gradient descent by transforming accelerated gradient descent into its continuous dynamic. The most significant contribution in [2] is that they further proved the faster convergence rate of underdamped Langevin algorithm than the classical Langevin algorithm under some conditions, such as log-Sobolev inequality. These results are rigorously presented using both sample and velocity spaces, offering novel and profound insights into understanding gradient-based MCMC sampling and its accelerated version.
 
By comparison, there are two main differences between this work and [2] because of their different focuses and targeted readers.  

**(1)** Our new simple and intuitive understanding complements the previous deep mathematical insights of [2] and thus allowed us to develop practically efficient algorithms. This work constructs a dynamic function directly connecting the gradient descent with the diffusion forward process. It does not borrow complex mathematical notations and tools in [2]. Then, motivated by the faster convergence speed of momentum gradient descent, this work introduces momentum to accelerate diffusion forward process, and shows the faster speed. This work focuses on more practical and efficient diffusion model design for practical designers and engineers, in contrast to the deep mathematical insights targeted in [2]. 

**(2)** This work prioritizes efficiency, an aspect not extensively considered in both [2] and CLD [1]. This work eliminates the velocity term, yielding three benefits. **(a)** It stabilizes training, since using velocity and sample spaces like CLD and [2] needs numerically unstable Cholesky factorization (see CLD), while removing velocity allows to simplify the perturbation kernel and avoids such instability. **(b)** It reduces model memory cost which is important in practice, since using two spaces needs roughly $2\times$ model parameter and memory cost than using a single one, and poses challenges to integrate with SoTA diffusion models which are often of large model sizes. **(c)** It improves compatibility by allowing our diffusion process to be directly integrated into popular diffusion models like EDM and numerical solvers like dpm-solver, which primarily use the sample space. This is in contrast to using two spaces, which would necessitate substantial revisions to their formulations or might be incompatible.

In the final version which should allow us to use an extra page, we will discuss these related works for their big efforts on understanding gradient-based MCMC sampling and its acceleration version in detail. We do not discuss this now, since there is really no space to add new content.

Reviewer comment (Novelty):
The formulation of acceleration seems to be closely related to that in the critically damped Langevin diffusion. Specifically, equations (15) and (19) in this paper, which correspond to the update equation of \( x \) and its velocity, correspond to equation (28) in Dockhorn et al. (2021) (with \( v_0 = 0 \)). It seems that the techniques employed in EDM, particularly the conversion of the discrete process into a continuous one, play a significant role in enhancing the efficiency of the proposed method. The authors are asked to elaborate more on what is most critical in the superiority of their experimental results.

Author response:
**Despite some similarity, our FDM differs CLD from significant differences.**
When $v_0 = 0$, the velocity in CLD becomes $v_{\text{CLD}}(t) = -\mathcal{B}(t)\exp(-2\mathcal{B}(t)\Gamma^{-1}) x_0$ which differs from our FDM whose velocity is $\dot{x}_{\text{FDM}}(t) = -\beta(t)\mathcal{B}(t)\exp(-\mathcal{B}(t)) x_0$. These differences  are not caused by manual design for a superior velocity formulation, but stem from our distinct handling approach to the forward diffusion process. 

**(a) Velocity Diffusion:** CLD explicitly diffuses the velocity in an additional velocity space. In contrast, our FDM implicitly determines velocity through the time derivative of the sample $x_t$, which gets rid of  explicit velocity computation. 

**(b) Critically-Damped State:** In CLD, the critically-damped state aims to modify the noise injection level, seeking a balance between the Hamiltonian component and the Ornstein-Uhlenbeck process. However, in our FDM, the critically-damped state is designed to avoid oscillation in the expectation diffusion trajectory to streamline the model training and avoid overshoot during sampling.

**The superiority performance of our FDM is primarily based on our streamlined perturbation kernel and its wide applicability.**

**(1) Streamlined Perturbation Kernel:** Our approach is grounded in the EDM design space, where we follow the spirit of treating marginal distributions as first-class citizens and design our diffusion process directly based on mean scaling $\mu_t$ and noise schedule $\sigma_t$ within marginal distributions. This design allows our momentum-based diffusion process to employ a more streamlined perturbation kernel, as shown in Eq. (14) of our paper. This contrasts with the more complex formulation of the perturbation kernel in the CLD paper (Eq. (28) to Eq. (33)), making our approach more efficient to compute and more robust against numerical errors in practical applications.

**(2) Wide Applicability:** Our FDM operates solely within the sample space, with velocity implicitly determined through the time derivative of the sample $x_t$. This innovative approach enables seamless integration with advanced frameworks and numerical solvers. This compatibility not only ensures smooth integration but also has the potential to significantly enhance the performance of our model by utilizing the inherent advantages of these advanced techniques. The benefit of this approach is evident in our experimental study presented in Table 7. In the experiments, both CLD and our FDM were applied using the traditional RK45 solver. However, a significant distinction was observed when a more efficient EDM sampler was employed. While CLD exhibited a notable performance decrease with the EDM sampler, our FDM maintained consistent performance. Additionally, FDM is able to improve its efficiency by leveraging the advantages of the EDM sampler, showcasing its adaptability and effectiveness in various computational environments.

Moreover, we have extended our experimental analysis to include gDDIM[3], which addresses oscillation/overshoot and numerical instability issues commonly found in the vanilla CLD. The results of these experiments are presented below. Despite the specific enhancements implemented in gDDIM to improve upon CLD, the experimental results demonstrate that our FDM still significantly outperforms gDDIM. This further underscores the robustness and efficiency of our FDM in comparison to previous methods.

**Table: Image synthesis performance on CIFAR-10 with 35 NFEs. gDDIM (CLD) use specifically designed gDDIM sampler, while VP-FDM use EDM sampler. The result denoted with $\dagger$ is from the original paper, with 50 NFEs.**
|    Method   | 30 Mimg | 50 Mimg | 70 Mimg | 100 Mimg |      Best      |
|:-----------:|:-------:|:-------:|:-------:|:--------:|:--------------:|
| gDDIM (CLD) |   5.15  |   2.95  |   2.76  |   2.51   | 2.25$^\dagger$ |
| VP-FDM      |   4.01  |   2.74  |   2.37  |   2.24   |      2.08      |

[3] Zhang, Qinsheng, Molei Tao, and Yongxin Chen. gDDIM: Generalized Denoising Diffusion Implicit Models. ICLR 2023.

Reviewer comment (Theory):
Why does the mixing speed of the forward diffusion process, transitioning from the data distribution to a Gaussian distribution, relate to the training cost?

Author response:
**The convergence speed of the forward diffusion process is closely related to the training cost, because it directly determines the score function, while matching this score function is the training objective of the diffusion model.** If the forward process is too fast, then the reverse process is pushed to be fast and enforces the network to predict a too rapidly changed trajectory, which increases the learning challenges of the network; otherwise, the forward process is too slow, the trajectory converges to a Gaussian noise too slowly and hence slows down the learning of network. 

Fortunately, momentum can well trade off the convergence speed and the learning difficulty of the trajectory. During the early stages of diffusion (at a small $t$), the momentum in FDM accumulates historical information, thereby increasing the speed of the forward process as the model gains confidence in identifying the direction indicated by the score function. This implies that the score function at this stage is relatively easy to match. In later stages (at a large $t$), when the forward process nears convergence to the Gaussian distribution, the momentum gradually decreases, thus slowing down the forward process as it becomes more challenging to identify the direction due to the increasing noise. This indicates that the score function at this stage becomes more difficult to match. 
This adaptive adjustment in "step size" in FDM leads to more effective navigation of the diffusion trajectory, which optimizes the balance between the speed of convergence to the target Gaussian distribution and the learning challenge posed to the network. The efficacy of this approach is demonstrated by our empirical results based on VP, VE, and EDM.