Paper: Sampling is as easy as keeping the consistency: convergence guarantee for Consistency Models (ICLR-2024)

Reviewer comment (Theory):
Assumption 6 seems unmotivated: the authors attributed it to "technical reason", without further explanation.

Author response:
Sorry for our unclear expression. Indeed the special time scheduling was firstly suggested by Sitian Chen, *The probability flow ODE is provably fast*, second paragraph in section 3.2 Algorithm. The reason to choose this time scheduling is to control one term of error in Theorem 2: in equation (26), we get the upper bound with a term $\sum_{k=1}^N \frac{h_k^2}{ t_k^{1/2}}$. If we only take a naive time scheduling $h_k \equiv h$, this term now becomes $O(\frac{h^2}{\sqrt{\delta}})$, as $\frac{1}{t_k^{1/2}}$ will become larger when $k \to 0$. One would better choose a relatively smaller discretization step when $t_k$ is small to prevent the discretization error from being unnessary large. Thus a geometrically increasing step size will help a lot.

Reviewer comment (Reproducibility):
How is $\theta^-$ defined? Is it a moving average of past $\theta$'s?

Author response:
Yes, it is defined as $\theta^- = \text{stopgrad}(\mu\theta^- + (1-\mu) \theta)$. We sorry for putting the definition a little bit later than where it should be, we will modify our expression.

Reviewer comment (Theory):
Below Assumption 2, it says this paper does not assume convexity or dissipativity on $-\log p$, unlike previous works. What is the reason that the analysis presented here does not require such assumptions?

Author response:
In the begining of the theoretical researchs of Diffusion Models, not so many properties of the forward diffusion process, as well as its probability density $p_t(x)$ have been found in Stochastic, PDE anslyses. Thus they need to assume convexity or dissipativity on $\log p_0(x)$ to get prior estimations similar to our Lemma 7, Lemma 11, etc. On the other way, informally the Diffusion Models can be viewed as a Simulated Annealing process, which has better ability to bypass the difficulty from non-convexity or non-dissipativity comparing to traditional Langevin Dynamics.

Reviewer comment (Theory):
What is the significance of obtaining a TV bound, compared to a Wasserstein-2 bound? The result seems a lot more messy compared to the ones concerning W2 errors.

Author response:
Let the output distribution be $\mu(dx)$. A guarantee of TV error bound can reflect the continuity of its density, thus can further support other operations over Likelihood. However if only have a $W_2$ bound, the output distribution density may be discontinuous, and have some points with infinity density (like Dirac measure).

Reviewer comment (Writing):
Below (7), "such a mapping exists ..., and is smoothly relied on" What's smoothly relied? And what exactly are the mild conditions mentioned here?

Author response:
This can be found in many textbooks on Ordinary Differential Equations. A simple case is that when $v$ is differentiable and bounded, then the solution of this ODE uniquely exists, and the mapping through ODE path is differentiable.

Reviewer comment (Writing):
Below Assumption 5, "technique reason" -> "technical reason".

Author response:
Thanks for the correction and we have fixed the typos.

Reviewer comment (Writing):
What does the notation $p P_{OU}^s$ mean, for a Markov kernel $P_{OU}^s$?

Author response:
This is a density defined as follows:

Let $\xi \sim p$, and let the $P^s_{OU}$-corresponded Markov process be started at $\xi$. This Markov process generate another random variable $\xi'$, then we denote the density of $\xi'$ as $p P^s_{OU}$.

Reviewer comment (Novelty):
Not much technical contribution, most of the techniques has already been proposed in the literatures and the proof mostly follows Chen et al.(2023a). Also, Lemma 7 is not new, similar results have been established in [1]

Author response:
Thanks for your kindly comment. We admit that we use many techniques developed by former researchers, especially in proving Theorem 2, but we still have developed many new proofs, like Theorem 3, Corollary 5, Corollary 6 and Corollary 9. Besides, it is also important to use existing method to solve emerging new problems. We will keep in working harder to develop more new techniques.

Reviewer comment (Writing):
There are a lot of typos in the manuscript, e.g., on page 4: equation 4, the expression for $v^{\mathrm{em}}(x,t)$; on page 15: in the second equation, $W_2(\mathcal{N}(0,(1-e^{-2T})I_d), p_T)$ should be $W_2(\mathcal{N}(0,I_d), p_T)$. Please reexamine your manuscript carefully.

Author response:
Thanks for the correction. We have fixed these points, especially the Lemma 11 should corresponding to the Lemma 3 in Chen et al.(2023a).

Reviewer comment (Theory):
There are also some technique issues. For example, Lemma 11 in this paper is for the OU noise schedule, while Lemma 1 in Chen et al.(2023a) is for the variance explosion schedule.

Author response:
Thanks for the correction. We have fixed these points, especially the Lemma 11 should corresponding to the Lemma 3 in Chen et al.(2023a).

Reviewer comment (Theory):
In corollary 6, $n_k$ is taken to be a constant $\hat{n}$ for all $k$. While in Song et al. (2023), $n_k$ is suggested to be decreasing. The theoretical results do not seem to support a decreasing $n_k$, any explanation?

Author response:
Actually $n_k$ can be a decreasing sequence and the multistep error can still decrease. Let's look at the recursion inequality in the proof of Corollary 6,

$\mathcal{E}\_k \le L_f e^{-(t_{n_k}-\delta)} \mathcal{E}\_{k-1} + t_{n_k}D,$

the best choice of $t_{n_k}$ should be chosen to minimize the function

$ H(t) = L_f e^{-t - \delta} \mathcal{E}\_{k-1}   + t D$.

However, this is an transcendental equation, and we do not have an explicit evaluation of $\mathcal{E}\_{k-1}$, thus it would be impossible to calculate the best $n_k$. Besides, we can prove that, in upper-bound meaning, the best $n_k$ will finally converge to our defined $\hat n$. Let us suppose the best-choice $n_k$ finally converge to $n^\ast$, with corresponding $W_2$ error
$\mathcal{E}^\ast$, then we have

$\mathcal{E}^\ast\le L_f e^{-(t_{n^\ast}-\delta)} \mathcal{E}^\ast + t_{n^\ast}D,$

$\mathcal{E}^\ast \le \frac{t_{n^\ast}D}{1-L_f e^{-(t_{n^\ast}-\delta)}}$.

This suggests that we can directly find the extreme point $n^\ast$ of the best-choice sequence $n_k$ by 
$\min_{n} \frac{t_{n}D}{1-L_f e^{-(t_{n}-\delta)}}$, under some prior estimations on $\varepsilon_{cm}, \varepsilon_{sc}, L_f, L_s, d$ and $h$.

Note that a simple fixing $n_k \equiv n^\ast$ also make the $\mathcal{E}\_k$ converging to $\mathcal{E}^\ast$ exponentially, we thus choose to simply use a fixing $n_k$.

Reviewer comment (Theory):
The results on multistep sampling that only requires a constant lower bound of $T$ is amazing. However, it seems that $L_f$ implicitly depends on $T$, especially when the data distribution is complicated (so it takes more time to transform a simple Gaussian noise to the data distribution). Not sure how pratical this benefit can be.

Author response:
In actual it can be proven that, for the exact consistency function $f^{ex}$, given the data distribution is bounded-supported, it's Lipschitz constant function $L_f(t)$ has an upper bound irrelavent to $T$. You may look at our rebuttal to Reviewer 8oEy, where we gave a simple proof. 

But as the reviewer 8oEy pointed, the Lipschitz constant of $f_\theta$ could still grow with $T$ due to the score error. We thus propose an alternative definition of Lipschitz constant, and an additional penalty to keep the $f_\theta$ having a bounded Lipschitz constant. This can be found in our second rebuttal to Reviewer 8oEy.

Reviewer comment (Theory):
I'm not entirely certain about the reasonability of Assumption 5. See Questions in details.

Author response:
Thank the reviewer for pointing this out and we apologize for unsufficient explaination. Your concern is reasonable, and we should prove that the backward ODE would not accumulate the Lipschitz constants exponentially fast. In intuition, this should be true, at least in the ideal case, as when $t \to \infty$, $p_t(x)$ gradually converge to the normal distribution $N(0,I)$, and thus $\nabla \log p_t(x) \to -x$, the particles under the  backward ODE is nealy stopped in large $t$. In fact, we can prove the following statement for the exact backward ODE $dx_t = v(x,t) dt$,
 where $v(x,t) = -x - \nabla \log p_t(x)$, $p_t(x) = e^{dt}p(e^{t}x) \ast \mathcal{N}(0,(1-e^{-2t})  I_d) \sim \int_{\mathbb{R}^d} e^{dt} p(e^{t}u)e^{-\frac{\Vert x-u\Vert^2}{2(1-e^{-2t})}}du$ with a bounded data distribution density $p_0(x) = p_{data}(x), \text{supp } p(x) \subseteq B(0,R)$ for some $R > 0$.

***

**Property** . The velocity field $v(x,t)$ in the exact backward ODE satisfies:

$ \Vert \nabla v(x,t) \Vert_{op} \le e^{-2t} M,$
for some $M > 0$.

*Proof*.
By the proof of Lemma 7 in our work, we have proved that 

$\nabla^2 \log (\mu \ast \psi_{\sigma^2}(x)) = \frac{1}{\sigma^4} \int_{\mathbb{R}^d} (u -\mathbb{E}\_{\mu_{x,\sigma^2}}[u] )\otimes (u -\mathbb{E}\_{\mu_{x,\sigma^2}}[u] )  \mu_{x, \sigma^2}(du) - \frac{1}{\sigma^2} I_d$

where

 $\psi_{\sigma^2}(u - x) \sim e^{-\frac{\Vert x - u\Vert_2^2}{2\sigma^2}},$

$  \mu_{x, \sigma^2}(du) = \frac{e^{-\frac{\Vert x - u\Vert_2^2}{2\sigma^2} }\mu(du)}{\int_{\mathbb{R}^d} e^{-\frac{\Vert x - \tilde u\Vert_2^2}{2\sigma^2}} \mu(d\tilde u)}.$

Taking $\mu(u) = e^{dt} p_0(e^t u)$ which is bounded on a set of radius $e^{-t}R$ and $\sigma^2 = 1-e^{-2t}$, noticing that
$\nabla v(x,t) = -I_d - \nabla^2 \log (\mu \ast \psi_{\sigma^2})$, we immediately get the result. 

***

Now we can prove that the analytical consistency mapping $f^v(x,t)$ has a uniform Lipschitz constant for all $t \ge \delta$. We notice that, for two different end points $x_T^1, x_T^2$, we define 

$d x^1_t = v(x_t^1,t) dt$

$d x^2_t = v(x_t^2,t) dt$

and
$f^v(x_T^1,T) - f^v(x_T^2,T) = x^1_\delta - x^2_\delta$

with early stopping time $\delta >0$. Notice that, 

$ \frac{d}{dt} \Vert x^1_t - x^2_t \Vert^2_2 = \langle x^1_t - x^2_t, v(x_t^1 ,t) - v(x_t^2 , t)\rangle \le \max  \Vert \nabla v(x,t) \Vert_{op} \Vert x^1_t - x^2_t \Vert_2^2 = e^{-2t} M \Vert x^1_t - x^2_t \Vert^2_2,$

thus by Gronwall's Inequality, 

$ \Vert x^1_\delta - x^2_\delta \Vert^2_2 \le   \Vert x^1_T - x^2_T \Vert^2_2  \exp( \int_{\delta}^{T}e^{-2t}M) \le  \Vert x^1_T - x^2_T \Vert^2_2  \exp( \frac{M}{2}(1- e^{-2T})),$

which shows a bounded Lipschitz constant when $T \to \infty$.

When dealing with the model $f^{em}$ under the approximated score function $s_\phi$, it would be much more difficult, as we should ensure $-x - s_\phi(x,t)$ has a similar operator bound like $-x - \nabla \log p_t(x)$. This can be proved by further strengthen the assumption 3 to a $L_\infty$ norm, and multiply the right-hand-side with an exponential-decreasing term $e^{-2t_n}$,

 $\Vert s_{\phi}(x,t_n) - \nabla \log p_{t_n}(x) \Vert_2^2 \le \frac{e^{-2t_n}}{1-e^{-2t_n}}\varepsilon^2_{sc}$ for all $x$ and $t_n$.

In words, we believe that our assumption 5 is reasonable, however further analysis can be conducted to replace this assumption to more reasonable assunptions. This will be our future work.

Reviewer comment (Theory):
The only point I am concerned about is Assumption 5, which assumes the consistency model \( f_\theta \) is \( L_f \)-Lipschitz. But as far as I know, even if our model \( f_\theta \) can approximate the backward mapping \( f^v \) very well, we should still expect that \( L_f \) is of \( O(e^T) \) (which comes from the Gronwall's Inequality applied on the exact reverse ODE). And so the first term in Corollary 4 could look like \( \max (d^{1/2}, m) \), which cannot be arbitrarily small. Can you overcome such Gronwall-type of error? If not, I would doubt that the results of this paper prove the efficiency of CMs.

Author response:
Thank the reviewer for pointing this out and we apologize for unsufficient explaination. Your concern is reasonable, and we should prove that the backward ODE would not accumulate the Lipschitz constants exponentially fast. In intuition, this should be true, at least in the ideal case, as when $t \to \infty$, $p_t(x)$ gradually converge to the normal distribution $N(0,I)$, and thus $\nabla \log p_t(x) \to -x$, the particles under the  backward ODE is nealy stopped in large $t$. In fact, we can prove the following statement for the exact backward ODE $dx_t = v(x,t) dt$,
 where $v(x,t) = -x - \nabla \log p_t(x)$, $p_t(x) = e^{dt}p(e^{t}x) \ast \mathcal{N}(0,(1-e^{-2t})  I_d) \sim \int_{\mathbb{R}^d} e^{dt} p(e^{t}u)e^{-\frac{\Vert x-u\Vert^2}{2(1-e^{-2t})}}du$ with a bounded data distribution density $p_0(x) = p_{data}(x), \text{supp } p(x) \subseteq B(0,R)$ for some $R > 0$.

***

**Property** . The velocity field $v(x,t)$ in the exact backward ODE satisfies:

$ \Vert \nabla v(x,t) \Vert_{op} \le e^{-2t} M,$
for some $M > 0$.

*Proof*.
By the proof of Lemma 7 in our work, we have proved that 

$\nabla^2 \log (\mu \ast \psi_{\sigma^2}(x)) = \frac{1}{\sigma^4} \int_{\mathbb{R}^d} (u -\mathbb{E}\_{\mu_{x,\sigma^2}}[u] )\otimes (u -\mathbb{E}\_{\mu_{x,\sigma^2}}[u] )  \mu_{x, \sigma^2}(du) - \frac{1}{\sigma^2} I_d$

where

 $\psi_{\sigma^2}(u - x) \sim e^{-\frac{\Vert x - u\Vert_2^2}{2\sigma^2}},$

$  \mu_{x, \sigma^2}(du) = \frac{e^{-\frac{\Vert x - u\Vert_2^2}{2\sigma^2} }\mu(du)}{\int_{\mathbb{R}^d} e^{-\frac{\Vert x - \tilde u\Vert_2^2}{2\sigma^2}} \mu(d\tilde u)}.$

Taking $\mu(u) = e^{dt} p_0(e^t u)$ which is bounded on a set of radius $e^{-t}R$ and $\sigma^2 = 1-e^{-2t}$, noticing that
$\nabla v(x,t) = -I_d - \nabla^2 \log (\mu \ast \psi_{\sigma^2})$, we immediately get the result. 

***

Now we can prove that the analytical consistency mapping $f^v(x,t)$ has a uniform Lipschitz constant for all $t \ge \delta$. We notice that, for two different end points $x_T^1, x_T^2$, we define 

$d x^1_t = v(x_t^1,t) dt$

$d x^2_t = v(x_t^2,t) dt$

and
$f^v(x_T^1,T) - f^v(x_T^2,T) = x^1_\delta - x^2_\delta$

with early stopping time $\delta >0$. Notice that, 

$ \frac{d}{dt} \Vert x^1_t - x^2_t \Vert^2_2 = \langle x^1_t - x^2_t, v(x_t^1 ,t) - v(x_t^2 , t)\rangle \le \max  \Vert \nabla v(x,t) \Vert_{op} \Vert x^1_t - x^2_t \Vert_2^2 = e^{-2t} M \Vert x^1_t - x^2_t \Vert^2_2,$

thus by Gronwall's Inequality, 

$ \Vert x^1_\delta - x^2_\delta \Vert^2_2 \le   \Vert x^1_T - x^2_T \Vert^2_2  \exp( \int_{\delta}^{T}e^{-2t}M) \le  \Vert x^1_T - x^2_T \Vert^2_2  \exp( \frac{M}{2}(1- e^{-2T})),$

which shows a bounded Lipschitz constant when $T \to \infty$.

When dealing with the model $f^{em}$ under the approximated score function $s_\phi$, it would be much more difficult, as we should ensure $-x - s_\phi(x,t)$ has a similar operator bound like $-x - \nabla \log p_t(x)$. This can be proved by further strengthen the assumption 3 to a $L_\infty$ norm, and multiply the right-hand-side with an exponential-decreasing term $e^{-2t_n}$,

 $\Vert s_{\phi}(x,t_n) - \nabla \log p_{t_n}(x) \Vert_2^2 \le \frac{e^{-2t_n}}{1-e^{-2t_n}}\varepsilon^2_{sc}$ for all $x$ and $t_n$.

In words, we believe that our assumption 5 is reasonable, however further analysis can be conducted to replace this assumption to more reasonable assunptions. This will be our future work.

Reviewer comment (Writing):
The presentation of the paper could be improved. There are several typos in the main text -- I would encourage the authors to do another pass over the manuscript and fix them. Some examples:  
* "be further weaken" -> "be further weakened"  
* "consistency model is nature" -> "is natural"  
* same sentence, fix formatting of the parenthesis.  
* Figure 1: punctuation missing.  
* "a asymptotic analysis" -> "an asymptotic analysis"  
* "For technique reason" -> for technical reasons. It would also help to explain what these reasons are.

Author response:
Thanks for pointing this out, we have modified all the mentioned typos or misleading in the updated paper

Reviewer comment (Writing):
To improve the clarity of the paper, I think it would be better to explain what $\theta_{-}$ is the first time it is introduced.

Author response:
Thanks for pointing this out, we have modified all the mentioned typos or misleading in the updated paper

Reviewer comment (Theory):
For Assumption 4, I think that the expectation should be over $x_{t_{n+1}}$. Given a $x_{t_{n+1}}$ everything is deterministic. Does that affect the proof of Theorem 2?

Author response:
Thanks for pointing this out, we have modified all the mentioned typos in the updated paper. This would not affect the proof of Theorem 2, as in each individual term that needs to calculate expectations, $x_t$ are keeped in the same ODE path for all $t\in[\delta, T]$.

Reviewer comment (Theory):
Corollary 4 is a little confusing. It seems that we require a lower bound on $T$, however, in the text, it is presented as having a large value of $T$ is a bad thing -- which is why the authors claim that multistep consistency sampling helps. I think the way this result is presented is a little misleading since it requires the learning errors to be O(1/T).

Author response:
That's correct: as $T$ increase, the discretization error will increase; however, if $T$ is not large enough, the error introduced by replacing $p_T$ to a normal distribution would become large. This is also obtained in other works such as Theorem 2, *Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions*.

Reviewer comment (Theory):
I am a little concerned about how much of the complexity of the problem is hidden in the upper-bound assumption on the learning of the consistency model. The model is trained to solve the whole probability flow ODE. If the score is inaccurately learned, the prediction of the solution can be poor because of error propagation. It seems that the error in the learning of the score and the consistency errors are actually related.

Author response:
The definition of consistency error can be disentangle to the score matching error, as our Assumption 4 do not need $\hat x_{t_n}^\phi$ to be exact enough: in fact, given any score model $s_\phi(x,t)$, no matter how far from the exact score $\nabla \log p_t(x)$,  we can always define the consistency mapping $f^{em}$. And if one can not learn the score well enough, one still can not generate good enough samples by diffusion models, let alone consistency models (in consistency distillation way).

Reviewer comment (Theory):
I also don't understand why this special time scheduling is needed. I don't think this is the scheduling that has been used in practice. I also don't recall seeing this assumption in prior work and it would be great to understand where this is coming from.

Author response:
Sorry for the insufficient explaination. The special time scheduling was firstly suggested by Sitian Chen, *The probability flow ODE is provably fast*, second paragraph in section 3.2 Algorithm. The reason to choose this time scheduling is to control one term of error in Theorem 2: in equation (26), we get the upper bound with a term $\sum_{k=1}^N \frac{h_k^2}{ t_k^{1/2}}$. If we only take a naive time scheduling $h_k \equiv h$, this term now becomes $O(\frac{h^2}{\sqrt{\delta}})$, as $\frac{1}{t_k^{1/2}}$ will become larger when $k \to 0$. One would better choose a relatively smaller discretization step when $t_k$ is small to prevent the discretization error from being unnessary large. Thus a geometrically increasing step size will help a lot.