Improved Discretization Complexity Analysis of Consistency Models: Variance Exploding Forward Process and Decay Discretization Scheme

Thank you for your valuable comments and suggestions. We provide our response to each question below.

Weakness 1: The guidance on the design of better consistency models.

This paper makes the first step to elucidate the design space of consistency models under the different diffusion processes and reveals their different advantages and disadvantages, which will heavily influence the discretization complexity and is fundamental in designing the consistency models. For the VP-based consistency models, the early stopping parameter $\delta$ has order $\epsilon_{W_2}^2$, which is worse than the VE-based consistency models ($\delta$ has order $\epsilon_{W_2}$) and is the source of large discretization complexity. However, the VE-based consistency models also have their disadvantage: the polynomial diffusion time $T$, which is much larger than $T=\log(1/\epsilon)$ for the VP-based consistency models and introduces additional $\epsilon_{W_2}$ dependence.

Hence, from the discretization perspective, a better consistency model should enjoy a Logarithmic $T$ and a $\delta$ with order $\epsilon_{W_2}$, which would lead to better complexity results. We note that the rectified flow-based one-step models have this potential. We will add the above discussion to our next version and view the design of better consistency models as important future work.

Question 1: The choice of $a$.

Our results show that with a larger $a$, the discretization complexity is better than the uniform discretization scheme ($a=1$).This phenomenon is also observed in the empirical work EDM [2], which observed that when $1\leq a\leq 7$, a larger $a$ will help the diffusion models to achieve better performance (Figure 13 (c) of [2]). When $a$ is larger than $7$, the improvement is not significant. Consistency models follow the choice of $a$ in EDM. In our Theorem 4.7, we also show that with $a=7$, the discretization complexity has order $1/\epsilon_{W_2}^{23/7}$, which is close to the $1/\epsilon_{W_2}^{3}$ of exponential decay stepsize.

[1] Lyu, Junlong, Zhitang Chen, and Shoubo Feng. "Sampling is as easy as keeping the consistency: convergence guarantee for Consistency Models." In Forty-first International Conference on Machine Learning. 2024.

[2] Karras, Tero, Miika Aittala, Timo Aila, and Samuli Laine. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.

Thank you for your valuable comments and suggestions. We provide our response to each question below.

Weakness & Suggestion: The analysis for consistency training and continuous time consistency models.

As shown by the professional reviewer, the consistency training and continuous-time consistency models are both important part of of consistency models. In this part, we discuss some possible method to obtain the discretization complexity for this method.

Consistency Training. If we can not obtain the pre-trained score function, we can construct an empirical score by using $n$ samples from the target data distribution $\\{X_{0,i}\\}\_{i=1}^n$

which has an explicit formulation, needs no additional training ([1] also use this formula) and converges to the ground-truth score function with a rate $n^{-1/d}$. Hence, we can replace the pretrained score function $s_{\phi}$ in eq. (4) with $s_{\mathrm{emp}}$. Then, the $\epsilon_{\text{score}}$ becomes $n^{-1/d}$ and achieve the guarantee for consistency models without a pre-trained score function under the VE process and EDM stepsize.

We note that though this results does not rely on the pretrained score function, it use the reverse PFODE process of diffusion. On the contrary, the consistency training paradigm only use the forward diffusion process. Hence, the above result is not the discretization complexity of consistency training paradigm. To achieve this goal, one possible way is to use similar method with [1], which use $s\_{\mathrm{emp}}$ to construct a baseline consistency function (a bridge between the target data distribution and the consistency function learned by the consistency training paradigm) instead of directly using it in the training objective function. However, as shown in our Remark 4.10, the construction of the baseline consistency function run $M$-step PFODE instead of one-step PFODE (used in application), which leads a large discretization complexity $1/\epsilon_{W_1}^{10}$. Since this result is significant larger than our $1/\epsilon\_{W_2}^3$ results, we left the discretization complexity analysis (compareable with the CD paradigm) for the CT paradigm under the setting used in application.

Continuous-time consistency models. Since continuous time model use $\frac{\mathrm{d} \boldsymbol{f}\_{\theta^{-}}(X\_t, t)}{\mathrm{d} t}$ instead of $\boldsymbol{f}\_{\theta^{-}}(X\_{t-\Delta t}, t-\Delta t))$ ($\Delta t$ is $h_{k+1}-h\_{k}$ in our work. Here we use the uniform stepsize for convenience), there are not well-defined discretization complexity $K=T/\Delta t$ for continuous-time models. However, due to the absence of $\Delta t$, the training process of continuous time models is less stable than discrete time consistency models, which is the core problem for continuous time models. Recently, [2] make a great effort to stabilize the training process of continuous time models.

Thanks again for the comments on a broader area of consistency models and we will add the above discussion in our next version.

[1] Dou, Zehao, Minshuo Chen, Mengdi Wang, and Zhuoran Yang. "Theory of consistency diffusion models: Distribution estimation meets fast sampling." In Forty-first International Conference on Machine Learning. 2024.

[2] Lu, Cheng, and Yang Song. "Simplifying, stabilizing and scaling continuous-time consistency models." arXiv preprint arXiv:2410.11081 (2024).

Thank you for your valuable comments and suggestions. We provide our response to each question below.

Q1: Theoretical Claims: The discussion on $L_{f,0}$ assumption and remove it.

Following the suggestion of the professional reviewer, we consider the $L_{f,0}$ in 2-mode GMM in the following Suggestion 1 and show $L\_{f,0}$ has the order $1/T$, which is necessary for a $W\_2$ guarantee. In this part, we mainly discuss how to remove this assumption. We prove that if considering a weaker $W\_1$ guarantee, we can remove $L\_{f,0}=O(R/T)$ assumption and achieve a $L_f^{1+1/a}/\epsilon\_{W_1}$ result:

When considering $W_1$ guarantee, the first term of line 615 (appendix B) becomes $L_fW\_1(\mathcal{N}(0,T^2I_d),q_T)$ (Using uniform $L_f$ instead of $R/T$). Different from the $W_2$ distance, the $W_1$ distance can be bounded by weight TV distance (Case 6.16 [1]):

where the second inequality follows the fact of [2]. Hence, we do not require $L_{f,0}=O(R/T)$. The other proof is exactly the same with $W_2$ distance. To guarantee $L_fW\_1(\mathcal{N}(0,T^2I_d),q_T)$ smaller than $\epsilon_{W_1}$, we require $T\ge L_fR^2/\epsilon\_{W_1}$, which is the source of additional $L_f^{1/a}$. We will add the above result in the next version.

Q2: Experimental Analysis: the calculation of $L_{f,0}$ in simulation experiments.

Since $L_{f,0}$ can be obtained by calculate the F-norm of $\nabla_{Y_0}\boldsymbol{f}(Y_0,0)$, we calculate the following equal to approximate it

where $Y_{t'}\sim q\_{T-t'}$ (sample $1000$ times and take average) and $\Delta Y = 0.01$.

Suggestion 1: The Lipschitz constant (Mixture of Gaussian).

We sincerely thanks again for the comments. We consider 2-mode GMM $X_0 \sim 1/2N(\mu, \sigma^2I_d)+1/2N(-\mu, \sigma^2I_d)$. The score has the following form (Appendix A.2 of [3], we transform it from VP to VE)

Since $f^{\mathrm{ex}}(Y_0,0)$ the associate backward mapping of the following PFODE (in the following part, we ignore the superscript of $t'$):

we need to solve it to obtain $f^{\mathrm{ex}}(Y_0,0)$. Since the score is highly nonlinear, it is hard to obtain a closed-form solution. There are two choices to overcome this hardness. The first choice is to do simulation experiments to simulate the solution (Appendix F).

The second choice is to add some assumptions on the target data to simplify the above ODE. We assume $\mu$ is smaller enough to guarantee $\tanh \left(\frac{\mu^{\top} Y_t}{(T-t)^2+\sigma^2}\right)$ can be approximated by $\frac{\mu^{\top} Y_t}{(T-t)^2+\sigma^2}$, which simplify PFODE to a linear ODE (in fact, the distribution gradually closes to Gaussian)

which have the following solution

The above results indicate

Taking the derivative of $Y_0$, we know that the $L_{f, 0}$ have order $1 / T$. This result also matches our intuition that $Y_0$ has a large variance (order $T^2$ ), and we need to multiply a $1 / T$ to avoid the influence of large variance (lines 282-287).

Further Discussion on the error of linear approximation. The above part makes a linear approximation to simplify the ODE when $\mu$ is close to $0$, which introduces some small errors. Assuming $Y_0\sim q_T= 1/2N(\mu, (T^2+\sigma^2)I)+1/2N(-\mu, (T^2+\sigma^2))$. For the variance, $Y_T=Y_0C(T)$ recover $\sigma^2$. For the mean, the recover $\mu$ of the above consistency function is approximately $\mu\sqrt{\sigma^2/(\sigma^2+T^2)}$, which is smaller than $\mu$. However, since assuming $\mu$ is close to $0$, this error term is small and is possibly introduced by the nonlinear term.

We will add the above discussion in the next version.

[1] Villani, Cédric. Optimal transport: old and new. Vol. 338. Berlin: springer, 2008.

[2] Yang et al,. Leveraging Drift to Improve Sample Complexity of Variance Exploding Diffusion Models. NeurIPS 2024.

[3] Shah et al,. Learning mixtures of gaussians using the DDPM objective. NeurIPS 2023.

Thank you for your valuable comments and suggestions. We provide our response to each question below.

Weakness 1: The approximated score and consistency function error (end-to-end analysis).

In this work, we assume the pretrained score and consistency function are accurate enough to achieve the final discretization complexity. Though they are standard assumptions in the complexity analysis area, as the friendly reviewer mentioned, the end-to-end analysis is also important, and we can use the current estimation error analysis results to achieve this goal. More specifically, for the approximated score, we use the results of [1] and replace $\epsilon_{score}$ with $n\_{score}^{-2/d}$ (where $n\_{score}$ is the number of data used to train the score function). For the approximated score function, we use the result of [2] and replace $\epsilon\_{cm}$ with $n\_{cm}^{-1/2(d+5)}$. Then, we can obtain the end-to-end complexity analysis.

Weakness 2: The Results of Multi-step Sampling Algorithm.

In fact, our analysis can be extended to $N$-step sampling algorithm and can achieve nearly $L_{f}/\epsilon\_{W_2}^{3+1/a}$ (which is better than Thm. 4.7 and Coro 4.12) under the EDM stepsize. We use $3$-step sampling algorithm as an example ($\tau_1=T,\tau_2=3T/4, \tau\_3=T/2$). Under this setting, the result becomes (here we ignore $\epsilon\_{score}$, $\epsilon_{cm}$, $R,d$ and focus on the dominated term)

To guarantee the above term smaller than $\epsilon_{W_2}$, we require $\delta=\epsilon\_{W_2}$ and $K\ge \frac{L_fT^{1/a}}{\delta^{2+1/a}\epsilon\_{W_2}}=\frac{L_fT^{1/a}}{\epsilon\_{W_2}^{3+1/a}}$, which is the same with the one-step and two-step sampling algorithms. However, 3-step algorithm only require $T\ge 1/\epsilon\_{W_2}^{1/3}$, which is better than $1/\epsilon\_{W_2}$ of 1-step and $1/\epsilon\_{W_2}^{1/2}$ of 2-step. Hence, the discretization complexity for $3$-step sampling algorithm is $L_f/\epsilon\_{W_2}^{3+4/(3a)}$, which is better than $2$-step algorithm. The above steps can be extended to $N$ steps, and the influence of $T$ decreases, and finally, $T$ does not affect the discretization complexity, which leads to a $L_{f}/\epsilon\_{W_2}^{3+1/a}$ results. We will add the above discussion in our next version.

Weakness 3: The Discretization Complexity for Piecewise Discretization Scheme (beyond the EDM with single $a$).

Before providing the complexity result for the piecewise discretization scheme, we first discuss the performance of the EDM scheme in the application (Consistency models follow the choice of $a$ of EDM.). EDM [3] shows that when $1\leq a\leq 7$, a larger $a$ will help the diffusion models to achieve better performance (Figure 13 (c) of [2]), which also matches our theoretical results. However, when $a$ is larger than $7$, the improvement is not significant and even becomes worse [3]. As a result, the exponential decay stepsize is theoretically friendly and is not widely used in applications. One possible explanation is that at the end of the reverse process, diffusion models generate image details and require a small stepsize (However, the exponential decay stepsize is too large.)

Hence, we can design a two-stage discretization scheme: (a) when $t'\in [0, T-1]$, we use the exponential decay stepsize; (b) when $t'\in (T-1, T-\delta]$, we use the EDM stepsize. With this scheme, the discretization complexity becomes $L_{f}/\epsilon\_{W_2}^{3+1/a}$, which is better than Thm. 4.7 with EDM (single $a$). This result shows the improvement of the two-stage discretization scheme from the theoretical perspective, and we leave the empirical application of this scheme as an interesting future work. We will add the above discussion in our next version.

[1]Oko, Kazusato, Shunta Akiyama, and Taiji Suzuki. "Diffusion models are minimax optimal distribution estimators." In International Conference on Machine Learning, pp. 26517-26582. PMLR, 2023.

[2]Dou, Zehao, Minshuo Chen, Mengdi Wang, and Zhuoran Yang. "Theory of consistency diffusion models: Distribution estimation meets fast sampling." In Forty-first International Conference on Machine Learning. 2024.

[3]Karras, Tero, Miika Aittala, Timo Aila, and Samuli Laine. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.