Improving Consistency Models with Generator-Induced Flows

We would like to thank the reviewer for their positive assessment of our work. We appreciate your recognition of our theoretical contributions and insights regarding consistency models.

Additional datasets

(Weakness) The experiments are not extensive. To make the observation and analysis more convining, the authors may try more datesets with different iamge sizes, and more imange generation metrics as those in paper in Ho 2020, Song 2021, Karras 2022, LIu 2023 et. al.

Given our computational constraints, we hope the reviewer understands that we are unable to perform experiments at the same scale as the cited articles. Still, we highlight that our experiments were conducted on as many datasets as most of the cited papers, including the standard benchmarks of CIFAR-10 and ImageNet.

Proof intuition

(Question) It will be better to give some high-level intuitive explanation of the key ideas behind the mathematics in the proofs; e.g., Lemma 1 and its restatement

We thank the reviewer for the relevant suggestion. We agree that Lemma 1 may seem technical without explanation and that providing more intuitive insight would be beneficial. We added the following explanation in Appendix Section A.3 with the revision of our submission.

The main purpose of the lemma is to provide a more tractable representation of $c'(t)$, the time derivative of the expected transport cost. We expect $c(t)$ to decrease with $t$ because the predicted data point $f(\mathbf{x}_t, \sigma_t)$ becomes more dependent on the noise $\mathbf{z}$ as $t$ increases. However, directly analyzing $f(\mathbf{x}_t, \sigma_t) - \mathbf{z}$ is challenging because the dependence of $f(\mathbf{x}_t, \sigma_t)$ on $\mathbf{z}$ is not explicit. Therefore, the lemma aims to:

1. identify a quantity that better captures the dependence between $\mathbf{z}$ and $\mathbf{x}_t$;
2. relate $c(t)$ to this quantity.

The proof proceeds by deriving a key property of the ground-truth consistency map $f$: it satisfies the transport equation,

$$\frac{\partial f}{\partial \sigma}(\mathbf{x}, \sigma_t) \, \dot{\sigma}_t + \frac{\partial f}{\partial \mathbf{x}}(\mathbf{x}, \sigma_t) \cdot \mathbf{v}_t(\mathbf{x}) = 0.$$

This equation is equivalent to saying that the conditional expectation of the time derivative of $f(\mathbf{x}_t, \sigma_t)$ given $\mathbf{x}_t$ is zero:

$$\mathbb{E}\left[ \frac{\partial}{\partial t} f(\mathbf{x}_t, \sigma_t) \,\middle|\, \mathbf{x}_t \right] = 0.$$

By leveraging this property, we can simplify $c'(t)$ into an expression involving $\mathbf{w}_t = \mathbf{z} - \mathbb{E}[\mathbf{z} \mid \mathbf{x}_t]$, the residual between the true noise $\mathbf{z}$ and its prediction given $\mathbf{x}_t$. This residual captures the uncertainty in predicting $\mathbf{z}$ based on $\mathbf{x}_t$, allowing us to relate $c'(t)$ directly to the prediction accuracy of $f$.

Additional dataset

In addition, could the authors test the results on ImageNet-64?

We conducted our experiments on ImageNet-32, which uses the same base dataset as ImageNet-64 but with images downsampled to 32×32 resolution rather than 64×64 because of computational constraints. The only difference between these datasets is the image resolution. Note that we include results on the 64×64 resolution with the CelebA and LSUN Church datasets.

Generator purification effect

There are some other benefits of using the generator-induced clean data. For example, the online $f_\theta$ could potentially "purify" the samples, e.g., make the samples more smooth. These purified data distributions at $t=0$ are easier to learn in the CT setting. It's analogous to training a model with soft labels instead of one-hot labels. I wonder what the authors' view on this is.

We thank the reviewer for the interesting and plausible hypothesis that aligns with some of our observations. When analyzing GC-only models, we observed that as the number of timesteps increases and performance degrades, the generated images tend to become overly smooth and simplistic -- essentially generating the "center" of each mode rather than experiencing mode collapse.

However, beyond this "purification" effect, our analysis suggests that the key factors driving improved performance are the discrepancy $\mathcal{R}$ and the transport cost. This is supported by the success of iCT-OT, which achieves improved performance without any "data purification". The effect of OT is solely to modify the transport cost and, consequently, the discrepancy term $\mathcal{R}$. We validate this further in Figure 2 of the paper, where we compare the discrepancy across IC, batch-OT (in the new revision), and GC approaches. Thus, while data purification may play a role in GC's performance, our theoretical analysis provides a more comprehensive explanation that accounts for the success of both GC and OT methods.

We would like to thank the reviewer for their constructive review and for recognizing the novelty and empirical value of our improvement of consistency models.

The reviewer's main concern is the soundness of our theoretical results and of our algorithm. We provide a detailed response which should clarify the raised issues. As a consequence, we hope this also alleviates the reviewer's concerns on our experimental results, which we address in a second time. We look forward to further discussing with the reviewer to clarify any remaining concern.

Soundness

Theorem & $\dot{x}_t \neq v_t(x_t)$

1. The time derivative of the CT model (Eq.9 in theorem 1) is incorrect. Indeed, when the CT converges, $\dot{x}_t = v_t(x_t)$, as the CT will learn ground-truth ODE instead of "single-sample estimate". [...] Also, the authors keep referring to the time derivative $\dot{x}_t$ as a "single-sample estimate", which is quite confusing. Could the authors crisply state the exact mathematical formulation of $\dot{x}_t$ in the theorem?

We appreciate this important question and opportunity to clarify our paper. The raised issue originates in a confusion between $\dot{x}_t$ -- the time derivative of the random variable $x_t$ --, and $v_t(x_t)$, the time derivative of particles $x_t$ in the PF-ODE of Equation (1). Song et al. (2023), who adopt the same notations, go from distillation to training by replacing the latter by the former. Their difference explains the discrepancy outlined in our theorem. Therefore, our theorem is correct.

Following the reviewer's suggestion, we explicitly defined $\dot{x}_t$ in the submitted revision (Line 127). We apologize for the confusion in the notations. Let us clarify in this response as well the mathematical formulation and re-define the key quantities involved.

1. The random variable $x_t$ is defined on Line 104: $x_t = x_\star + \sigma_t z$, where $x_\star \sim p_\star$ (sample from data distribution), $z\sim p_z$ (sample from the Gaussian), and $\sigma_t$ is monotonically increasing for $t\in[0,T]$.
2. Therefore, the time derivative $\dot{x}\_t$ is $\dot{x}\_t = \frac{\mathrm{d}(x_\star + \sigma_t z)}{\mathrm{d}t} = \frac{\mathrm{d} \sigma_t}{\mathrm{d}t}z=\dot{\sigma}_t z$. In the EDM case where $\sigma_t=t$, this simplifies to $\dot{x}_t = z$.
3. The velocity field $v_t$ is defined as: $v_t(x) = \mathbb{E}[\dot{x}_t | x_t=x]$ on Line 127.

This shows that $\dot{x}_t$ is indeed a single-sample estimate of $v_t(x_t)$, and should not be confused with the time derivative in the ODE (Eq. 1), which is $v_t$ itself.

Algorithm failure case

2. As the theorem is incorrect, so is Algorithm 1. Think of a counter-example in Algorithm 1: if the current online model $f_\theta$ only maps all the intermediate points $x_t$ to a certain mode, then all the other modes will never be recovered.

We thank the reviewer for the interesting remark.

First, let us notice that this failure case is shared with standard consistency models. Indeed, in a standard consistency model, the performance of the model at timestep $i$ depends on the performance of the model at timestep $i-1$. If a model drops a mode during training at timestep $i-1$, the mode will not be recovered at timestep $i$, since it learns by following the output of the current model at $i-1$. Second, note that the consistency models framework, which we share, tends to avoid this thanks to 1) a parameterization enforcing the preservation of data at timestep 0; 2) the consistency loss ensuring propagation of data from timestep $i-1$ to timestep $i$.

In our specific setting of GC models, the use of the endpoint predictor could indeed be an additional source of mode collapse, given the poor performance of the endpoint predictor on IC trajectories discussed in Section 5.1. We show, however, that our simple mixing strategy prevents this issue in Section 5.3, as partial IC training can compensate any mode collapse from GC.

Experiments

Performance of baselines

The current FID is in a poor regime, and the baselines are not well-tuned.

Note that the baseline is the standard iCT-IC. The difference between the results from the iCT paper and our baseline comes from computational constraints. While the original iCT paper used greater computational resources, our experiments maintain consistent settings across all methods (iCT-IC, iCT-GC, iCT-OT) for fair comparison, using identical model sizes, batch sizes, and training steps -- adapted to our computational constraints. Even with these constraints, our complete experimental suite required approximately 100 days of computation on A100 40GB GPUs. Matching the computational scale of the original iCT paper is unfortunately beyond our current resources.

Dear Reviewer obag,

We have carefully addressed your concerns in our initial response. We would be grateful if you could acknowledge our response and participate in the discussion to share your thoughts on the points we addressed in the rebuttal. We are happy to provide further clarifications or additional details to resolve any remaining concerns. Thanks again for your involvement.

We would like to thank the reviewer for their constructive assessment of our work. We appreciate their recognition of the provided theoretical and empirical insights on consistency models and our improvements. We address the raised weaknesses and questions below.

Hyperparameter sensitivity to circumvent distribution shifts

Specifically, the marginal distribution shift between IC and GC during intermediate timesteps affects performance. [...]

While the mixed IC-GC strategy helps mitigate this, the need for careful balancing complicates its application.

Fortunately, we found that balancing the mixed IC-GC is easy. In this paper, we have already conducted a hyperparameter grid search on mixing ratio $\mu$ using the CIFAR-10 dataset. Given the bell-shaped relationship observed between $\mu$ and FID, we opted to retain the best performing value identified on CIFAR-10, $\mu=0.5$, for all subsequent experiments, including those on other datasets, without further tuning. Importantly, even without an exhaustive hyperparameter search, our method consistently outperforms baseline approaches.

In the new paper version, we included a new experiment exploring $\mu$ on ImageNet-32 to further assess its generalizability (cf Figure 8 in Appendix). It consistently outperforms the base model for $\mu={0.3,0.5}$. We also detailed the previous explanation for our choice of $\mu$ in Appendix Section C.

Performance of baselines

The baseline used in the current draft is a relatively weak benchmark modified from iCT-IC. The mixed strategy of GC and IC improves upon this baseline only by a small margin. It is unclear whether these improvements would diminish when compared to the standard iCT-IC or iCT-OT.

Note that the baseline is the standard iCT-IC. The difference between the results from the iCT paper and our baseline comes from computational constraints. While the original iCT paper used greater computational resources, our experiments maintain consistent settings across all methods (iCT-IC, iCT-GC, iCT-OT) for fair comparison, using identical model sizes, batch sizes, and training steps -- adapted to our computational constraints. Even with these constraints, our complete experimental suite required approximately 100 days of computation on A100 40GB GPUs. Matching the computational scale of the original iCT paper is unfortunately beyond our current resources.

Questions

Why does the distribution shift occur specifically at intermediate timesteps between IC and GC trajectories?

This occurs because of the properties of the generator-induced trajectories, discussed around Line 292 and in Equation (16). On small timesteps $t$, the parameterization of the consistency model implies that $p(\tilde{x}_t) \approx p(x_t)$, and particularly $p(\tilde{x}_0) = p(x_0)$. On large timesteps, the signal-to-noise ratio is negligible, thus the distribution becomes close to the noise distribution, hence $p(\tilde{x}_T) \approx p(x_T) \approx p(\sigma_T z)$. This explains why the distribution shift mostly happens on intermediate timesteps.

When generator-induced coupling (GC) improves convergence speed, the model reaches saturation early during training. Is there a deeper analysis of why this phenomenon occurs?

Before anything else, let us remind that the saturation is not observed when using mixing, while retaining a higher convergence speed than the baseline IC. When using GC-only trajectories however, the model indeed reaches saturation early during training. In this setting, we believe that main limiting factor is the the distribution shift between IC and GC trajectories, as validated by experiments in Section 5.2. Our mixing strategy successfully solves this issue, as discussed in the paper and in this response.

Dear Reviewer hFie,

We have carefully addressed your concerns in our initial response. We would be grateful if you could acknowledge our response and participate in the discussion to share your thoughts on the points we addressed in the rebuttal. We are happy to provide further clarifications or additional details to resolve any remaining concerns. Thanks again for your involvement.

Dear Reviewer hFie,

We would like to thank you for your valuable suggestion. Following your advice, we conducted an additional experiment using the efficient ECM setting described in Geng et al. [1]. This configuration enables high performance with only 1 GPU-hour of fine-tuning from a pre-trained denoising model.

Using this setup, we compared iCT-IC and iCT-GC ($\mu=0.7$) trained with a 1-hour fine-tuning: iCT-IC achieves a 1-step FID of 7.37, while iCT-GC ($\mu=0.7$) achieves a lower FID of 6.41. These findings highlight the advantages of generator-induced flows, particularly in scenarios where rapid convergence is important.

[1] Geng et al., Consistency Models Made Easy, arXiv:2406.14548.

We would like to thank the reviewer for their positive assessment of our work. We appreciate their recognition of our identification of the discrepancy term and the introduction of GC as a way to diminish it. We address the raised weaknesses and questions below.

Distribution shifts

(Weakness) The main difficulty of the proposed method [is that] the model itself is still being trained and thus has weak performance as an endpoint predictor, leading to distribution shifts.

We thank the reviewer for raising this important point. We want to clarify that the distribution shift is not primarily caused by the endpoint predictor's performance. Rather, it is an inherent characteristic of different probability paths ($p_t$) whether using GC or other trajectory-changing methods such as batch-OT. To illustrate this, consider Figure 1 in our paper where we use an optimal ground-truth consistency model. Even in this idealized case with perfect prediction, we observe distribution shifts (specifically, lower density around $x=0$ for GC compared to IC). Similarly, when using Optimal Transport to couple noise and data, the associated probability path differs from IC. These shifts arise from the different couplings generating different probability paths.

"Weak partially-trained [predictor] as iCT-IC" in Section 5.2

(Question) Could you please explain more about the "weak partially-trained as iCT-IC" mentioned in Section 5.2? Does this imply that the endpoint predictor is trained on both GC and IC, similar to the final proposed method?

Let us clarify the terminology in Section 5.2. We discuss two distinct models:

• (i) A standard consistency model (iCT-IC) trained for 100k steps, which we call the IC predictor.
• (ii) A standard consistency model (iCT-IC) trained for only 20k steps, which we refer to as the "weak partially-trained as iCT-IC" or weak-IC predictor.

These predictors are trained purely on IC, not on both GC and IC. Using these models, we then conduct two experiments:

• (i) Training a GC-only model using the fully-trained IC predictor (100k steps).
• (ii) Training a GC-only model using the weak-IC predictor (20k steps).

This approach allows us to compare how the endpoint predictor's performance affects the GC model's performance; cf. Finding 3 in Section 5.2:

the performance of the GC model depends on the quality of the endpoint predictor evaluated on IC trajectories.

We specified this additional information in Section 5.2 of the new revision.

Sensitivity of the mixing ratio hyperparameter

(Weakness) To address this, the authors propose mixing IC and GC during training, but this solution [...] introduces a new challenge in determining the appropriate mixing ratio as a tunable parameter.

(Question) The performance of varying values on CIFAR-10 is shown in Figure 10, but what about the performance on other datasets with different mixing rates? I would also like to know how sensitive this parameter is for other models.

Fortunately, we found that the mixing ratio $\mu$ is easy to tune as a new hyper-parameter. In this version of the paper, we have already conducted a hyperparameter grid search on $\mu$ using the CIFAR-10 dataset. Given the bell-shaped relationship observed between $\mu$ and FID, we opted to retain the best performing value identified on CIFAR-10, $\mu=0.5$, for all subsequent experiments, including those on other datasets, without further tuning. Importantly, even without an exhaustive hyperparameter search, our method consistently outperforms baseline approaches.

In the new paper version, we included a new experiment exploring $\mu$ on ImageNet-32 to further assess its generalizability (cf Figure 8 in Appendix). It consistently outperforms the base model for $\mu={0.3,0.5}$. We also detailed the previous explanation for our choice of $\mu$ in Appendix Section C.