The Discretization Complexity Analysis of Consistency Models under Variance Exploding Forward Process

Thank you for your valuable comments and suggestions. We provide our response to each question below. Please click the PDF to view our revision paper.

Weakness 1: The notation.

Thanks again for the helpful comments. We have reorganized the notation of this paper in a paragraph (including the definition of discretization complexity and the notation for diffusion and consistency models). Please read line 298-314 in the revision of our paper (highlighted in green).

Weakness 2: The table.

We have polished the caption of Table 1 and put it above the table. Please read line 108-113 in the revision of our paper (highlighted in green).

Weakness 3 (a): The explanation and proof of Assumption 4

In this part, by calculating the exact form of $\nabla \log q\_t(\cdot)$ and $\nabla^2 \log q\_t(\cdot)$}, we show the reason why Assumption 4 matches the true order of $\nabla \boldsymbol{f}^{\text{ex}}$. Let $\mathbf{m}\_t(X\_t):=\mathbb{E}\_{q\_{0 \mid t}\left(\cdot \mid X\_t\right)}\left[X\_0\right]$ and $\boldsymbol{\Sigma}\_t(X\_t):=\operatorname{Cov}\_{q\_{0 \mid t}\left(\cdot \mid \mathbf{x}_t\right)}\left(X_0\right)$ be the posterior mean and variance. We have the following results.

Lemma 1. [The VESDE version of Lemma 5 [1]] Considering the VESDE forward process, for all $t\ge 0$ and $\forall X\_t\in \mathbb{R}^d$, we have that

It is clear that $\mathbf{m}\_t$ directly maps the noised data to the clean target data distribution, which is exactly $\boldsymbol{f}^{\text{ex}}$ in our work. Then, the ground-truth score function can be parameterized as $\nabla \log q\_{T-t^{\prime}}(Y,T-t^{\prime})= (\boldsymbol{f}^{\text{ex}}(Y,t^{\prime})-Y)/\sigma\_{T-t^{\prime}}^2$. Using the above parametrization and the second result of Lemma 1, we know that

With the above equality, we know that $\nabla \boldsymbol{f}^{\text{ex}}(Y,t^{\prime})=\boldsymbol{\Sigma}\_{T-t^{\prime}}/\sigma\_{T-t^{\prime}}^2$. With the bounded support assumption, we know that $\\|\boldsymbol{\Sigma}\_{T-t^{\prime}}\\|_{op}\leq R^2$. Hence, Assumption 4 matches the true order of $\nabla \boldsymbol{f}^{\text{ex}}$.

We have now added the discussion of the Lipschitz constant in line 349-375 in the revision of our paper (highlighted in green).

Weakness 3 (b): The explanation and discussion of Assumption 2&3.

Since the form of Assumption 3 is the same compared with the consistency model objective function $\mathcal{L}\_{\mathrm{CD}}^K$, Assumption 3 means the loss in the training phase is small. Furthermore, the $\left(t_{k+1}^{\prime}-t_k^{\prime}\right)^2$ term matches the $h^2$ order of the discretization error of one-step PFODE process. Hence, our Assumption 3 is reasonable.

For Assumption 2, this term is contained in the score-matching objective function. Hence, this assumption requires the score-matching loss to be small enough in the training phase.

We note that the training process in the application usually guarantees the loss (objective function) converges to a small term. Hence, Assumptions 2 and 3 are realistic assumptions.

Weakness 4: Refrain from evaluating their own work.

Thanks again for the helpful comments. We have polished the presentation of Remark 2 (line 382, highlighted in green).

Thank you for your valuable comments and suggestions.

As a beginning, we first emphasize our goal in this theoretical work. Our work aims to achieve the competitive discretization complexity compared with diffusion models under the mild assumption and realistic setting (VESDE forward process and EDM stepsize), which is satisfied and used in applications. Then, we say that "we explain the great performance of consistency models from the theoretical perspective". In this work, as shown in Table 1, we achieve $1/\epsilon_{W\_2}^4$ discretization complexity, which has the same order as the SOTA complexity result of diffusion models (In our rebuttal, we discuss each assumption in detail.).

We provide our response to each question below. Please click the PDF to view our revision paper.

Weakness 1& Q1: Empirical evidence.

Thanks again for your valuable comments on the empirical evidence. One core contribution of our work is to analyze consistency models with great performance in the application (VESDE forward process and EDM stepsize). Hence, empirical works with this setting can be our evidence [1] [2].

Weakness 2& Q2 (Assumption): The approximated score and consistency function.

For Assumption 2, this term is contained in the score-matching objective function. Hence, this assumption requires the score-matching loss to be small enough in the training phase.

For Assumption 3, since the form of Assumption 3 is the same as the consistency model objective function $\mathcal{L}\_{\mathrm{CD}}^K$, it means the loss in the training phase of the consistency function is small. Furthermore, the $\left(t_{k+1}^{\prime}-t_k^{\prime}\right)^2$ term matches the $h^2$ order of the discretization error of one-step PFODE process.

We note that the training process in the application usually guarantees the loss (objective function) converges a small term. Hence, Assumptions 2 and 3 are realistic assumptions.

Weakness 3& Q3: The application to the high-dimension dataset.

In the application, when dealing with high-dimension data, empirical works use the latent consistency models, which learn the consistency function in the latent space [3]. There are many works [4] [5] that have used LCM in 3D generation and can be viewed as empirical evidence of our work.

Let the latent space dimension be $d^{\prime}$, which is significantly smaller than $d$. Our analysis still holds for latent $d^{\prime}$ and avoids the dependence of $d$.

Weakness 4& Q4 & Q5: The discussion of EDM parameter.

Thanks for the comments on the non-uniform stepsize. We note that the EDM stepsize can cover most of current stepsize with different decay patterns (with different $a$). For example,

• When $a=1$, the EDM stepsize becomes the uniform stepsize and achieves $1/\epsilon_{W\_2}^{5.5}$, which is worse than the results presented in Theorem 1, which also matches the empirical observation (Figure 13 of [6]).
• When $a\ge 1$, the stepsize becomes a decay stepsize and a larger $a$ means a larger decay ration. [6] show that $a=7$ (Figure 13 of [6]) achieves great performance, and our analysis also can provide the discretization complexity under this choice (our revision paper (line 399-400)).
• When $a\rightarrow +\infty$, the EDM stepsize becomes the theoretical friendly exponential-decay stepsize and a $1/\epsilon\_{W\_2}^4$ result.

For the adaptive stepsize, while it has been used in empirical work [7], it has not been analyzed from the theoretical perspective, and our EDM stepsize covers almost all stepsize choices (since the previous works focus on two extremes: the uniform and exponential-decay). We leave the analysis for the adaptive stepsize as an interesting future work.

Weakness 4& Q4 (Assumption): The discussion on Lipschitz constant.

For the Lipschitz constant of consistency function, as shown in line 349-375 in the revision of our paper (highlighted in green), we prove that the Lipschitz constant of consistency function matches the true order. Hence, Assumption 4 is reasonable.

Weakness 5 & Q6 (Assumption): The bounded support assumption and unbounded distribution.

We note that the bounded support assumption is naturally satisfied by the image dataset, which indicates this assumption allows multimodal distributions. Furthermore, this assumption allows the blow-up phenomenon (a widely observed phenomenon in application) of the score function at the end of the reverse process ($t^{\prime}=T-\delta$ or $t=\delta$). Hence, this assumption is always viewed as a weak and realistic assumption in theoretical works [8].

Though the bounded support assumption is natural, we still consider some special unbounded distribution and show that $L_{f,0}$ also has $1/\sigma_T^2$ order at the beginning of the reverse process (this is crucial for our analysis, line 691 of our revision paper), which indicates our analysis still hold for these unbounded distributions. We also give some intuition as to why it is possible for this special distribution to achieve a better discretization complexity by using the property of these special distributions.

• We first consider a Gaussian distribution $X_0\sim \mathcal{N}(\mu,\Sigma)$, where the posterior mean $\mathbf{m}\_t(X\_t):=\mathbb{E}\_{q\_{0 \mid t}\left(\cdot \mid X\_t\right)}\left[X\_0\right]$ (can be viewed as $\boldsymbol{f}^{\text{ex}}$ in our work) has an analytical form

It is clear that $\nabla_{X_t}\mathbb{E}[X_0|X_t]=\left( \Sigma+\sigma_t^2 I_d\right)^{-1} \Sigma$ (the $\nabla \boldsymbol{f}^{\text{ex}}$ term), which has the order $1/\sigma_T^2$ at the beginning of the reverse process ($t^{\prime}=0$ or $t=T$). When at the end of the reverse process, since the existence of $\Sigma$ in $\left( \Sigma+\sigma_t^2 I_d\right)^{-1}$, the score will not blow up, which avoids the $1/\delta$ dependence in Theorem 1 and leads to a better discretization complexity.

• For the multimodal distribution such as a Gaussian mixture, the score function is highly nonlinear. Hence, similar to [9], we consider a special 2-mode Gaussian mixture distribution $X_0 \sim 1/2*\mathcal{N}(\mu, \sigma^2I_d)+1/2*\mathcal{N}(-\mu, \sigma^2I_d)$. Under this assumption, the ground-truth score has the following form (Appendix A.2 of [9], we transform it from VP to VESDE setting)

Then, we know that

At the beginning of the reverse process ($t=T$), the $\tanh$ can be approximated by the linear part (Taylor's theorem). It is clear the first part of $\nabla^2 \log q_t(X_t)$ has at least $1/\sigma_T^4$ order. With the discussion of Lemma 1 (in our revision paper, line 349-375, highlighted in green), $\nabla \boldsymbol{f}^{\text{ex}}$ has $1/\sigma_T^2$ order. When at the end of the reverse process, similar to the Gaussian distribution, the score will not blow up, leading to a better discretization complexity.

Weakness 6 & Q7: The consistency distillation paradigm.

As discussed in Remark 1 of our work, no existing results analyze the consistency training paradigm (which does not require the pre-trained score function) under a realistic setting with great performance. Hence, we left it as an interesting future work to explain the empirical success of the consistency training paradigm from the theoretical perspective under the setting in the application.

Q8&Q11: The error analysis.

The goal of this work is to control each error term to guarantee $W_2\left(\boldsymbol{f}_{\boldsymbol{\theta}, 0 \sharp \mathcal{N}}\left(\mathbf{0}, T^2 I_d\right), q_0\right)\leq \epsilon\_{W_2}$. We require each term to be strictly smaller than $\epsilon\_{W_2}$ and then obtain an upper bound for the discretization complexity, which is a standard analysis process in the complexity analysis area. Hence, our analysis does not misestimate the discretization complexity (steps). In other words, if all assumptions are satisfied, using discretization steps $K$ provided in our work, $W\_2\left(\boldsymbol{f}\_{\boldsymbol{\theta}, 0 \sharp \mathcal{N}}\left(\mathbf{0}, T^2 I_d\right), q_0\right)$ must smaller than $\epsilon\_{W_2}$.

However, it is an interesting future work to analyze the lower bound of the discretization complexity and show our upper bound result is nearly optimal.

Dear reviewer,

Thanks again for your comments on adaptive stepsize and EDM stepsize. In response to Weakness 4&Q4 &5, we show that EDM stepsize covers most of the current stepsize with different decay patterns. In this part, we further discuss the difference between adaptive stepsize and EDM stepsize and show that the current adaptive stepsize relies on a trained NN [7]. On the contrary, EDM stepsize is widely used in empirical works [6] [10] and does not rely on a trained NN.

For the first point, to determine the adaptive stepsize for different prompts, [7] train an additional NN by using the reinforcement learning method, which would introduce additional training costs. As shown in Section 4 of [6] ([6] use EDM stepsize with $a=7$), it is an open question if adaptive solvers can be a net win over a well-tuned fixed schedule in sampling diffusion models. Currently, the empirical work still adopts EDM stepsize and achieves great performance [10].

[10] Karras, T., Aittala, M., Lehtinen, J., Hellsten, J., Aila, T., & Laine, S. (2024). Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 24174-24184).

Thank you for your valuable comments and suggestions. We provide our response to each question below. Please click the PDF to view our revision paper.

Weakness 1: Empirical evidence.

Thanks again for your valuable comments on the empirical evidence. One core contribution of our work is to analyze consistency models with great performance (VESDE forward process and EDM stepsize). Hence, empirical works with this setting can be our evidence [1] [2].

Question 1: The definition of great performance.

The goal of our work is to achieve the competitive discretization complexity compared with diffusion models under the mild assumption and realistic setting (VESDE forward process and EDM stepsize), which is satisfied and used in applications. Then, we say that "we explain the great performance of consistency models from the theoretical perspective". In this work, as shown in Table 1, we achieve $1/\epsilon_{W\_2}^4$ discretization complexity, which has the same order as the SOTA complexity result of diffusion models.

In the following part, we discuss each assumption.

• Assumption 1: The bounded support assumption is naturally satisfied by the image dataset and is widely used in theoretical analysis. Furthermore, the bounded support assumption allows the blow-up phenomenon (a widely observed phenomenon in application [3]) of the score function at the end of the reverse process ($t^{\prime}=T-\delta$ or $t=\delta$).

• Assumption 2&3: For Assumption 2, this term is contained in the score-matching objective function. Hence, this assumption requires the score-matching loss to be small enough in the training phase.

  For Assumption 3, since the form of Assumption 3 is the same as the consistency model objective function $\mathcal{L}\_{\mathrm{CD}}^K$, it means the loss in the training phase is small, which is usually satisfied in the application. Furthermore, the $\left(t_{k+1}^{\prime}-t_k^{\prime}\right)^2$ term matches the $h^2$ order of the discretization error of one-step PFODE process.

  We note that the training process in the application usually guarantees the loss (objective function) converges a small term. Hence, Assumptions 2 and 3 are realistic assumptions.

• For Assumption 4, as shown in line 349-375 in the revision of our paper (highlighted in green), the Lipschitz constant of the approximated consistency function matches the true order and is reasonable.

[1] Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. (2023). Consistency models. arXiv preprint arXiv:2303.01469.

[2] Kim, D., Lai, C. H., Liao, W. H., Murata, N., Takida, Y., Uesaka, T., ... & Ermon, S. (2023). Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279.

[3] Kim, D., Shin, S., Song, K., Kang, W., & Moon, I. C. (2021). Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. arXiv preprint arXiv:2106.05527.

Thank you for your valuable comments and suggestions. We provide our response to each question below. Please click the PDF to view our revision paper.

Weakness 1(b). The order of the Lipschitz constant of consistency function.

In this part, by calculating the exact form of $\nabla \log q\_t(\cdot)$ and $\nabla^2 \log q\_t(\cdot)$}, we show the reason why Assumption 4 matches the true order of $\nabla \boldsymbol{f}^{\text{ex}}$ and would not lead to a meaningless result.

Let $\mathbf{m}\_t(X\_t):=\mathbb{E}\_{q\_{0 \mid t}\left(\cdot \mid X\_t\right)}\left[X\_0\right]$ and $\boldsymbol{\Sigma}\_t(X\_t):=\operatorname{Cov}\_{q\_{0 \mid t}\left(\cdot \mid \mathbf{x}_t\right)}\left(X_0\right)$ be the posterior mean and variance. We have the following results.

Lemma 1. [The VESDE version of Lemma 5 [1]]Considering VESDE forward process, for all $t\ge 0$ and $\forall X\_t\in \mathbb{R}^d$, we have that

It is clear that $\mathbf{m}\_t$ directly maps the noised data to the clean target data distribution, which is exactly $\boldsymbol{f}^{\text{ex}}$ in our work. Then, the ground-truth score function can be parameterized as $\nabla \log q\_{T-t^{\prime}}(Y,T-t^{\prime})= (\boldsymbol{f}^{\text{ex}}(Y,t^{\prime})-Y)/\sigma\_{T-t^{\prime}}^2$. Using the above parametrization and the second result of Lemma 1, we know that

With the above equality, we know that $\nabla \boldsymbol{f}^{\text{ex}}(Y,t^{\prime})=\boldsymbol{\Sigma}\_{T-t^{\prime}}/\sigma\_{T-t^{\prime}}^2$. With the bounded support assumption, we know that $\\|\boldsymbol{\Sigma}\_{T-t^{\prime}}\\|_{op}\leq R^2$. Hence, Assumption 4 matches the true order of $\nabla \boldsymbol{f}^{\text{ex}}$.

We have now added the discussion of the Lipschitz constant in line 349-375 in the revision of our paper (highlighted in green).

Weakness 1 (c). The order of special Gaussian and Gaussian Mixture cases. (Unbounded distribution)

Since the bounded support assumption is naturally satisfied by the image dataset, the $R^2/\sigma_t^2$ order for $\nabla \boldsymbol{f}^{\text{ex}}$ holds for a wide range of distribution. Furthermore, the bounded support assumption allows the blow-up phenomenon (a widely observed phenomenon in application [2]) of the score function at the end of the reverse process ($t^{\prime}=T-\delta$ or $t=\delta$).

Though the bounded support assumption is natural, we still consider some special unbounded distribution and show that $L_{f,0}$ also has $1/\sigma_T^2$ order at the beginning of the reverse process (this is crucial for our analysis, line 691 of our revision paper). We also give some intuition as to why it is possible for this special distribution to achieve a better discretization complexity by using the property of these special distributions.

• We first consider a Gaussian distribution $X_0\sim \mathcal{N}(\mu,\Sigma)$, where the posterior mean $\mathbf{m}\_t(X\_t):=\mathbb{E}\_{q\_{0 \mid t}\left(\cdot \mid X\_t\right)}\left[X\_0\right]$ (can be viewed as $\boldsymbol{f}^{\text{ex}}$ in our work) has an analytical form

​ It is clear that $\nabla_{X_t}\mathbb{E}[X_0|X_t]=\left( \Sigma+\sigma_t^2 I_d\right)^{-1} \Sigma$ (the $\nabla \boldsymbol{f}^{\text{ex}}$ term), which has the order $1/\sigma_T^2$ at the beginning of the reverse process ($t^{\prime}=0$ or $t=T$). When at the end of the reverse process, since the existence of $\Sigma$ in $\left( \Sigma+\sigma_t^2 I_d\right)^{-1}$, the score will not blow up, which avoids the $1/\delta$ dependence in Theorem 1 and leads to a better discretization complexity.

• For the multimodal distribution such as a Gaussian mixture, the score function is highly nonlinear. Hence, similar to [3], we consider a special 2-mode Gaussian mixture distribution $X_0 \sim 1/2*\mathcal{N}(\mu, \sigma^2I_d)+1/2*\mathcal{N}(-\mu, \sigma^2I_d)$. Under this assumption, the ground-truth score has the following form (Appendix A.2 of [3], we transform it from VP to VESDE setting) $$\nabla \log q_t(X_t)=\tanh \left(\frac{\mu^{\top} X_t}{\sigma_t^2+\sigma^2}\right) \frac{\mu}{\sigma_t^2+\sigma^2}-\frac{X_t}{\sigma_t^2+\sigma^2}.$$ Then, we know that $$\nabla^2 \log q_t(X_t)=\left(1-\tanh^2 \left(\frac{\mu^{\top} X_t}{\sigma_t^2+\sigma^2}\right)\right)\frac{\mu^{\top}\mu I_d}{(\sigma_t^2+\sigma^2)^2}-\frac{I_d}{\sigma_t^2}.$$ At the beginning of the reverse process ($t=T$), the $tanh$ can be approximated by the linear part (Taylor's theorem). It is clear the first part of $\nabla^2 \log q_t(X_t)$ has at least $1/\sigma_T^4$ order. With the discussion of Lemma 1, $\nabla \boldsymbol{f}^{\text{ex}}$ has $1/\sigma_T^2$ order. When at the end of the reverse process, similar to the Gaussian distribution, the score will not blow up, leading to a better discretization complexity.

Thanks for your quick and detailed response. Since $\mathbb{E}\left[X_0|X_t\right]$ is posterior mean, the $\boldsymbol{f}^{\text{ex}}(X_t,t)= \mathbb{E}\left[X_0|X_t\right]$ can generate the ground-truth mean for a large $T$. Now, let us consider $q_0=\mathcal{N}(\mu,1)$ (by similar calculation with the above discussion, Lemma 1 still hold). Then, we have that

It is clear that $X_t-\mu\sim \mathcal{N}(0,1+\sigma_t^2)$. We note that $(1+\sigma_t^2)^{-1}$ coefficient is the core to generate true mean. More specifically, $\left(1+\sigma_t^2\right)^{-1}\left(X_t-\mu\right)\sim \mathcal{N}(0,(1+\sigma_t^2)^{-1})$, which would goes to $0$ with a large enough $T$ (In our work, we also choose a large enough $T$) and does not introduce a large influence on the first $\mu$ term.

We also provide the solution of posterior covariance

which is independent of $X_t$ and $\nabla_{X\_t} Cov\left[X_0|X_t\right] = 0$.

Though the good results of posterior covariance (under Gaussian distribution), we note that when considering a more general $q_0$, the posterior covariance maybe depends on $X_t$. However, as shown in [4], they directly view $\mathbb{E}\left[X_0|X_t\right]$ as the ground truth denoised autoencoder (DAE) (the consistency function in our work) and empirical consistency models [5] also follow this setting, we say that $\mathbb{E}\left[X_0|X_t\right]$ is the consistency function in our work.

Thanks again for your vauleabl comments. We will add the above discussion in our revision paper.