Improving Consistency Models with Generator-Induced Flows

Authors: Thibaut Issenhuth, Sangchul Lee, Ludovic Dos Santos, Jean-Yves Franceschi, Chansoo Kim, Alain Rakotomamonjy
Primary area: generative models
Keywords: generative models, diffusion models, flow models, consistency models

Abstract:
Consistency models imitate the multi-step sampling of score-based diffusion in a single forward pass of a neural network.
They can be learned in two ways: consistency distillation and consistency training. The former relies on the true velocity field of the corresponding differential equation, approximated by a pre-trained neural network.
In contrast, the latter uses a single-sample Monte Carlo estimate of this velocity field. The related estimation error induces a discrepancy between consistency distillation and training that, we show, still holds in the continuous-time limit. To alleviate this issue, we propose a novel flow that transports noisy data towards their corresponding outputs derived from the currently trained model - as a proxy of the true flow. Our empirical findings demonstrate that this approach mitigates the previously identified discrepancy.
Furthermore, we present theoretical and empirical evidence indicating that our generator-induced flow surpasses dedicated optimal transport-based consistency models in effectively reducing the noise-data transport cost.
Consequently, our method not only accelerates consistency training convergence but also enhances its overall performance.

Decision: reject
Reviewer scores: 8, 5, 3, 8

Reviews and discussion:
<thread>
creation date:
1730619912409

thread_type:
official_review

<thread_object>
creator: reviewer_lcwl

<content>
summary:
The paper gave a negative answer to the  resorbent in the continuous-time limit of the discrepancy between consistency distillation
and training. Further, they had a new disovery that the an induced flow inside the generator, which could be used to train a consistent model.
They conducted great work on verifying their discovery in theory.

strengths:
A very nice work on observation and anlysis  that showed some very interesting insight on CM. It remove an unclear fact about the CMs, which will be very useful for later discussions.

weaknesses:
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.

questions:
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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730881741435

thread_type:
official_review

<thread_object>
creator: reviewer_obag

<content>
summary:
This paper studies the estimation error in the single-sample Monte Carlo estimates in the consistency training. The authors show that in the continuous-time limit, there exists a discrepancy between consistency training and consistency distillation (consistency distillation uses a pre-trained model to sample adjacent points instead of a single-sample Monte Carlo estimate. To remedy this issue, the authors propose generator-induced flows, which introduce a correlation between noise and the data at $t=0$. Experiments on image datasets show that the coupling leads to improved performance.

strengths:
- The paper studies a new way of sampling the "clean" data in consistency training. The clean data $\hat{x}_0$ are online predictions of the current neural network, given $x_t \sim p_t$ , where $x_t=x_0+t z, x_0 \sim p_0, z \sim N(0,I)$. In this way, it introduces a "generator-induced coupling" (GC) between $\hat{x}_0$ and $z$

- The paper studies the distribution shift between independent coupling in the original consistency training and the GC. The authors propose a mix-training strategy to overcome this issue.

- Experimental results show that the mix-training strategy gives best results across image datasets.

weaknesses:
- The paper has severe flaws in the main theorem (theorem 1) and the algorithm. I kindly request the authors to address them carefully or correct me if I'm wrong:

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" (please refer to the proof in the original consistency model paper). The idea is very similar to how the denoising score-matching objective works: the posterior of the "single-sample estimate" gives a ground-truth score, and CT is theoretically sound in the continuous-time limit. 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?

2. As the theorem is incorrect, so is Algorithm 1. Think of a counter-example in Algorithm 1: the current online model $f_\theta$ only maps all the intermediate points $x_t$ to a certain class, say "dog" class. As a result, the model will only learn to generate the doggy images since the algorithm replaces the clean data with the generated one by $f_\theta$. This scenario is very likely, as the initial $f_\theta$ likely drops many modes in the data distribution when $t$ is relatively large --- once it drops the modes, these lost modes will never be recovered.

3. The observations in Fig.4a and Fig.4b can be easily explained by the point 2 above. The authors utilize a hacky way (mix-training) to fix this.


- Despite the flaws discussed above, the proposed algorithm seems to improve over baseline on several image datasets. I think there are two possible reasons:

1. The current FID is in a poor regime, and the baselines are not well-tuned. For example, the iCT model can achieve a FID < 3 on CIFAR-10, instead of >7 in the paper. These poor FIDs make the results look unliable. In addition, could the authors test the results on ImageNet-64? This dataset is more representative than the datasets used in the paper.

2. 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. 


Due to the current status, I'm leaning toward rejecting the paper unless the authors clarify some of the critical points above.

questions:
Please see the weakness section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731384284953

thread_type:
official_review

<thread_object>
creator: reviewer_hfie

<content>
summary:
This paper, "Improving Consistency Models with Generator-Induced Flow," studies the coupling mechanism used in training consistency models. Consistency models are a recently developed class of generative models that learn to perform the multi-step sampling of score-based diffusion in a single forward pass.Instead of relying on sampling trajectories from a pre-trained diffusion model, consistency training uses stochastic interpolation between the data distribution and noise distribution, known as independent coupling (IC). This approach shows a discrepancy from consistency distillation. The authors analyze the differences between consistency training under IC and consistency distillation. To address the limitations of IC, the paper proposes a new type of flow called generator-induced flow, which relies on generator-induced coupling (GC) for training consistency models. This proposed generator-induced coupling accelerates convergence and enhances model performance compared to the base IC method and optimal transport approaches.

**Theoretical Results**

- This paper characterizes the discrepancy between consistency distillation and consistency training in values and gradients using different couplings. 
- Under mild conditions, the loss values differ between consistency distillation and consistency training for distance functions other than the standard L2 loss. 
- Though the limiting gradient in the continuous time coincides with the L2 metric, the gradient of the limiting loss values still differs due to the error term when using the independent coupling (IC).
- Authors show that the proposed GC could reduce the discrepancy between consistency training and consistency distillation.

**Algorithmic Improvements**

- The authors discuss potential couplings that could reduce the gap with the velocity field.
- The paper introduces generator-induced coupling (GC) that capitalizes consistency models to approximate the velocity field.

**Empirical Performance**

- The proposed GC adds an affordable cost to training while reducing the transport cost on CIFAR10 compared with batched optimal transport and independent coupling.
- The mixed strategy using IC and GC achieves improved FID compared with IC and batched OT baselines.

strengths:
- The theorectical analysis investigates the discrepancy between consistency distillation and consistency training. This analysis is novel for understanding consistency models.
-  The proposed generator-induced coupling (GC) is cost-effective compared to other methods like optimal transport (OT).
- This paper is well-written. The logical flow from the consistency model definition to the theoretical results on the gradient discrepancy, followed by the generator-induced coupling (GC) for reduced transport cost.
- The analysis and experiments on the transport cost induced by different couplings are helpful for understanding.
- Key insights are clear and discussed in detail. For example, theorem 1 characterizes the discrepancies between consistency distillation and consistency training under IC. This is a clear motivation for the proposed method. The authors also discuss the implication of this analysis, showing how coupling could lead to different transport costs for consistency models.

weaknesses:
- The paper acknowledges that consistency models trained exclusively on GC can suffer from distribution shift issues. Specifically, the marginal distribution shift between IC and GC during intermediate timesteps affects performance. Using a pre-trained consistency model through independent coupling could avoid divergence points to an inherent limitation but largely increase the computational cost. 

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

- 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. Including ReFlow [1] as a baseline coupling for comparison would strengthen the analysis. Additionally, conducting comparisons with efficient benchmarks like ECMs [1] for larger model sizes and comparing them with the standard iCT-IC would provide a better evaluation.

[1] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003. \
[2] Geng, Z., Pokle, A., Luo, W., Lin, J., & Kolter, J. Z. (2024). Consistency Models Made Easy. arXiv preprint arXiv:2406.14548.

While I acknowledge that this paper presents interesting insights into analyzing consistency models, the proposed generator-induced coupling is less attractive in practice and may suffer from distribution shifts. Therefore, I am not fully convinced of its practical impact at this stage.

questions:
Some of these questions have been already mentioned in the weaknesses.

It would be helpful if the authors could discuss the following questions.

- Why does the distribution shift occur specifically at intermediate timesteps between IC and GC trajectories? 
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731434757600

thread_type:
official_review

<thread_object>
creator: reviewer_lmst

<content>
summary:
This paper identifies a previously overlooked discrepancy between training and distillation in consistency models and introduces a novel approach called generator-induced flows to address it. This method improves model convergence speed and performance. Additionally, to tackle the distribution shift issue between independent coupling (IC) and generator-induced coupling (GC), the paper proposes a mixed training strategy, resulting in efficient and high-quality image generation.

strengths:
- This paper is well-structured and easy to read. Each section clearly states its claims and conclusions, making it easy to understand.
- The Generator-Induced Coupling introduced by the authors is very straightforward and serves as a convincing method for resolving the discrepancy. Additionally, the paper's contribution is significant, as it points out aspects of the discrepancy that have been overlooked in previous studies.
- The proposed method surpasses the state-of-the-art iCT, demonstrating high practicality and effectiveness. Therefore, I believe this paper should be accepted.

weaknesses:
The main difficulty of the proposed method in this study lies in the dilemma that, while aiming to reduce the discrepancy by predicting endpoints from the model, the model itself is still being trained and thus has weak performance as an endpoint predictor, leading to distribution shifts. To address this, the authors propose mixing IC and GC during training, but this solution is a compromise supported mainly by empirical evidence, and introduces a new challenge in determining the appropriate mixing ratio as a tunable parameter.

questions:
- Could you please explain more about the "weak $g_\phi$ 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?
- The performance of varying $\mu$ 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731942473291

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global answer

comment:
We thank the reviewers for their positive and constructive feedback on our work. We provided initial responses, along with a revision of our submission, to initiate the discussion which we look forward to pursue.

We responded individually to each reviewer. Additionally, we would like to share to all the following information.

**Paper revision.** 

Changes w.r.t. the initial submission are highlighted in blue.

1. As requested by **Reviewer obag**, we added a definition of $\dot{x}_t$ to remove the potential confusion between $\dot{x}_t$ and $v_t$. *(Line 127)*
2. We added $\tilde{\mathcal{R}}\_{\text{batch-OT}}$ in *Figure 2*. It further confirms the relevance of our analytical framework. Indeed, we see that $\tilde{\mathcal{R}}\_{\text{GC}} < \tilde{\mathcal{R}}\_{\text{batch-OT}} < \tilde{\mathcal{R}}\_{\text{IC}}$, which is the ordering that we observe on almost all datasets for final performance model (see Table 4). 
3. As suggested by **Reviewer LCWL**, we added an high-level intuitive explanation of the proof of Lemma 1. *(Line 1000)*
4. As requested by **Reviewer LmST**, we added some details about the "weak partially-trained endpoint predictor". *(Line 476)*
5. Study for $\mu$ on ImageNet (**Reviewers LmST & hFie**): we added **a new experiment** in Appendix Figure 8. On CIFAR-10, the baseline iCT-IC is improved for $\mu={0.3,0.5,0.7}$. On ImageNet, it is improved for $\mu={0.3,0.5}$. This shows the stability of our performance w.r.t. to this new hyperparameter. *(Line 1103)* 
6. Description of the rationale for the selection of $\mu=0.5$ (**Reviewers LmST & hFie**) in the Appendix. *(Line 1133)*

**Statement on experiments.** As noted by some reviewers, some of our baselines may appear weaker compared to those shown in the literature. This is primarily due to our limited access to high-end GPUs, which has constrained our ability to fully reproduce results from SOTA models as described in their original papers. We hope that this context will be taken in consideration, as we have fully complied to the scientific comparison methodology.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732201817465

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely thank the reviewers once again for their constructive feedback, which has been invaluable in improving and clarifying the paper. In particular, we addressed key concerns, such as the potential confusion between $\dot{x}_t$ and $v_t$ raised by **Reviewer obag**, as well as the remarks related to our empirical setting highlighted by **Reviewer hFie**, as detailed in our previous responses. We would be happy to provide further clarification on this or any other points if needed. Please let us know if there are additional questions or concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733086543380

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
New experiment in ECM setting

comment:
Dear Reviewers,

We would like to thank **Reviewer hFie** for their valuable suggestion. Following their advice, we conducted an **additional experiment** using the efficient ECM setting described in Geng et al. [1]. This configuration enables high performance with approximately 1 GPU-hour of fine-tuning (8k training steps) 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1738837090803

thread_type:
public_comment

<thread_object>
creator: ~thibaut_issenhuth1

<content>
title:
Response to Meta Review

comment:
We would like to thank the reviewers and area chair for their detailed feedback and for acknowledging the promise of our work. Following the paper's rejection, we would like to respectfully address and dispute the three main concerns raised in the meta-review.

### 1. Discrepancy Between Theory and Practice

Reviewers LmSt and hFie regretted the empirical need to rely on a mixing strategy between IC and GC for stability, which is not an issue as we showed in our initial answer. In the area chair's assessment, mixing introduces a discrepancy betwen theory and practice, casting doubts on our claims. To address simultaneously both concerns, **we make clear in this message and in the new version of our paper ([https://arxiv.org/abs/2406.09570](https://arxiv.org/abs/2406.09570)) that mixing, on the contrary, stems naturally from our theory**.

The theoretical motivation for GC in our paper assumes access to an optimal generator on IC trajectories *(Theorem 2, Lemma 1 and its corollaries)*. However, when training with GC alone, the assumption of an ideal generator is broken. Consequently, our theory does not claim GC-only should succeed: **the divergence issue of GC alone does not raise concerns on the validity of our theoretical claims**. 

Relying on an optimal, i.e. pre-trained in practice, IC generator is impractical and limits GC-based generator’s performance, as demonstrated in Section 5.2. To comply with the *theoretical requirement* of having an IC-generator with great performance, we propose a mixed coupling that combines training on trajectories from GC and IC. It is rooted in both practical constraints (one phase training) and theoretical considerations (optimal generator on IC). As a result, **the mixing strategy is motivated by our theoretical analysis**.

### 2. Scale of Experimental Validation

While we acknowledge that our experimental setup does not match the scale of some recent high-resource studies -- stated as an important concern by the area chair and Reviewer obag only --, we maintain that **it is sufficient to validate our claims** regarding convergence speed and improved performance.
- Despite constrained resources, we provide meaningful insights into the training dynamics of consistency models.
- We do so within reasonable FID values -- our IC baseline operating within the ranges of the original consistency model of [Song et al. (2023)](https://proceedings.mlr.press/v202/song23a.html) on CIFAR-10 -- and on four standard datasets -- more than several other consistency models: [Song & Dhariwal (2024)](https://openreview.net/forum?id=WNzy9bRDvG), [Lu & Song (2024)](https://arxiv.org/pdf/2410.11081), [Geng et al. (2024)](https://arxiv.org/pdf/2406.14548).
- Our experiments are methodologically sound within the investigated scope, reimplementing baselines for fair comparison, using identical model sizes, batch sizes, and training steps for all.

Regarding our constrained resources, the area chair's analogy with "social science experiment made on 10 persons" is **inappropriate**. In this analogy, claims are not verifiable because of lacking statistical significance. In our case, our claims are explicitly validated within the reasonable and robust experimental framework we present. We believe that this is valuable and can be of interest for the community to replicate on a larger scale and build upon.

### 3. Metrics 

The area chair questioned our choice of comparison metrics. We acknowledge the known limitations of metrics like FID, KID, and IS, as highlighted in recent critiques. However, **these metrics remain standard for generative model evaluation** unlike their alternatives, and their widespread use enables comparisons with prior work. Some of the most impactful and recently published diffusion/consistency papers are based on FID, *e.g.* [Karras et al. (2022)](https://openreview.net/forum?id=k7FuTOWMOc7), [Song et al. (2023)](https://proceedings.mlr.press/v202/song23a.html), as well as [Lu & Song (2024)](https://arxiv.org/pdf/2410.11081) and [Geng et al. (2024)](https://arxiv.org/pdf/2406.14548) at the present ICLR edition. This choice of metrics should not be ground for rejection.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the coupling mechanisms used to train consistency models and proposes a new generator-induced coupling (GC) which is meant to address the limitations of independent couplings (IC). In practice, using GC alone leads to distribution shifts, so a mixed coupling involving GC and IC is recommended.

Reviewers noted discrepancies between the experimental results and practical recommendations, compared to the motivation and theoretical discussions given by the authors. In particular, the need to combine GC with IC (a method for which the authors point out various flaws and set out to fix) is unsatisfying. It raises concerns as to the validity of the author’s claims. Also, the small scale and non-comprehensive nature of the experiments leaves room for doubt as to the actual benefits of the GC method. The models trained by the authors achieve FID values that are well above those reported in previous works, which might imply that the models are not well-converged.

In my view, the experimental results and practical recommendations do not align with the motivation and theoretical discussions given by the authors, which raises concerns as to the validity of those claims. The small scale and non-comprehensive nature of the experiments leaves room for doubt as to the actual benefits of the GC method.

Other concerns from the reviewers were addressed, such as some concerns with the theoretical results. Generally the reviewers appreciated the ideas in the work. Overall, this work is promising, and I believe that based on the feedback given by reviewers, it can be improved to the quality required for acceptance in a comparable venue. However, those revisions are extensive enough to require further peer-review, so I am not recommending acceptance at ICLR.