Improved Techniques for Training Consistency Models

Thank you for your thoughtful feedback. We address your questions or concerns below.

Q1: I am wondering how scalable consistency models are with the proposed modifications. Can one train, for instance, text-to-image consistency models? Or how about training on higher-resolution images?

A: We anticipate that improved consistency models will scale easily to larger datasets. Indeed, all generative models that perform well at the scale of 64x64 ImageNet have been successfully used for large-scale text-to-image generation and larger resolution images.

Q2: It would be interesting to validate that consistency models can also be successfully trained on non-image data (e.g. audio, graphs, 3D, video, etc).

A: We agree with the reviewer and plan to leave this as future work.

Q3: When the paper defines the ground truth consistency function in Section 3.2, is there a typo in the ground truth consistency function?

A: Thank you for pointing this out. We have corrected the typo in our revision.

Q4: What are $s_0$ and $s_1$ in Section 3.4?

A: $s_0$ and $s_1$ control the minimum and maximum number of discretization steps. We have added their definitions to the draft.

Q5: In Section 3.4, the authors use an exponential discretization curriculum, which leads to the visualization in Fig. 3(a). Why is this exponential curriculum not included in the visualizations and ablations in Figs. 3(b) and 3(c)?

A: In our latest revision, we have made the exponential discretization curriculum our default recommendation. We have also thoroughly compared all discretization curriculums in Figures 3b and 3c.

Thank you for your thoughtful feedback. We address your questions or concerns below.

Q1: This paper functions primarily as a technical exploration, compiling successful practices in training diffusion models. While it includes numerous ablation studies, it lacks sufficient theoretical backing to fully support these practices.

A: We want to emphasize that this work focuses on consistency models, not diffusion models. Although theoretical backing would be beneficial, we don't believe it exists for all practical techniques, particularly those in deep learning.

Q2: Whether this training would be effective when the generation process is not the reverse of a diffusion process but a more general corruption process, such as those described in [1] and [2].

A: We believe consistency training can be easily generalized to flow matching, rectified flows, and improved Poisson flow. However, this is not the primary focus of this work.

Q3: Could the author elaborate more on why higher discretization steps can reduce bias but increase variance?

A: Higher discretization steps decrease bias, as consistency training theoretically holds in the limit of $N\to\infty$. The observation that higher discretization steps also cause larger variance is primarily empirical. This is because models trained with larger discretization steps typically converge slower than those trained with smaller ones (as shown in Figure 3d in [1])

Reference:

[1] Song, Y., Dhariwal, P., Chen, M. and Sutskever, I., 2023. Consistency models.

Thank you for your thoughtful feedback. Please find our response to your questions below.

Q1: Clarification on the theoretical analysis in Section 3.2.

A: There are two competing theories justifying consistency training in previous work. As our analysis in Section 3.2 demonstrates, Argument (ii) is the correct one, requiring a zero EMA for the teacher network.

While Equation (6) holds regardless of the relationship between $\theta$ and $\theta^-$ (as predicted by Argument (i)), it doesn't necessarily validate consistency training. This is because the training dynamics are defined by the gradient of the loss function due to the stop gradient operator in the consistency training objective, not the loss function itself. In fact, the gradient of Equation (6) is always zero, offering no training signal as $N\to\infty$.

To avoid degenerate gradients, the consistency training objective's gradient must be scaled by $1/\Delta \sigma$. This is demonstrated in Equation (7) and supported by Argument (ii). In this scenario, the gradient is finite only when $\theta=\theta^-$, which further validates Argument (ii) and justifies our decision to set the teacher EMA decay rate to zero.

We would like to emphasize that even if the consistency training objective doesn't exist, it's acceptable as long as its gradient is well-defined. Both conditions can occur simultaneously because the loss function includes the stop gradient operator.

Q2: Do the authors have any insight as to why the Pseudo-Huber loss could bring this much improvement?

A: As Figure 5(b) in Appendix B demonstrates, Pseudo-Huber loss leads to smaller variance in Adam updates, which leads to faster convergence in training. This is likely because Pseudo-Huber losses are less affected by outlier gradients in training data compared to l2.

Q3: Do the authors have insights on why insensitive time embeddings are helpful or in this case crucial?

A: As mentioned in Section 3.1, we believe it is essential that noise embeddings are sufficiently sensitive to minute differences to offer training signals, yet too much sensitivity can lead to training instability.

Q4: How to feed noise level to positional embeddings.

A: We followed EDM by feeding $0.25 \log \sigma$ to the positional embedding layer.

Q5: No definitions of $s_0$ and $s_1$, and the typo in Section 3.2

A: Thank you for the suggestions! Indeed, $s_0$ and $s_1$ represent the start and end discretization steps, and we have updated our draft with their definitions. We have also corrected the typo in Section 3.2.

I thank the authors for their responses and revision. I hold my original assessment of accepting the paper.

Thank you for your thoughtful feedback. We address your questions or concerns below.

Q1: Pseudo-Huber loss is not convincing compared to LPIPS

A: We respectfully disagree with the reviewer. One major goal of this work is to eliminate LPIPS. Unlike LPIPS, Pseudo-Huber losses are not restricted to the image domain. Additionally, we have identified two other weaknesses of LPIPS in introduction: First, there could be potential bias in evaluation since the same ImageNet dataset trains both LPIPS and the Inception network in FID, which is the predominant metric for image quality. As analyzed in Kynkäänniemi et al. (2023), improvements of FIDs can come from accidental leakage of ImageNet features from LPIPS, causing inflated FID scores. Second, learned metrics require pre-training auxiliary networks for feature extraction. Training with these metrics requires backpropagating through extra neural networks, which increases the demand for compute.

Q2: Apart from the significant change of removing EMA, the other enhancements resemble engineering optimizations rather than foundational advances. Their relevance to broader applications is questionable, especially since the empirical evaluation is limited to smaller datasets like CIFAR-10 and ImageNet-64x64. Expanding experiments to include higher-resolution images on ImageNet or more varied datasets such as COCO or LAION could substantiate the model's versatility.

A: We agree with the reviewer. This paper provides both foundational advances and engineering optimizations, not merely limited to the former. We plan to undertake large-scale consistency training on datasets such as COCO or LAION in future work.

Q3: The training speed for the consistency model is comparatively slow, especially when compared with models like GANs.

A: We respectfully disagree with the reviewer. Given the same model size, consistency training is faster than GANs. This is because GANs requires backpropagation through both the generator and discriminator, whereas consistency training only requires backpropagation through the student network.

Q4: The applicability of consistency training appears confined to training unguided models.

A: We leave generalization to guided models to future work.

Q5: Does the current consistency training outperform consistency distillation when applying the enhancements presented in this paper to the latter?

A: Yes. Consistency training surpasses consistency distillation, even when provided with the same enhancements.

Q6: Given the apparent benefits of scaling up the network architecture for larger datasets (e.g., ImageNet as shown in Table 3), is there a risk that one(few)-step methods will be inherently limited by network capacity? In other words, can these methods directly approximate the ODE solution with a network of the same capacity as the original diffusion model, or would alternative approaches like RectFlow [1] provide a more viable solution?

A: We anticipate that one-step models will be larger than multi-step ones. Consistency training can easily be generalized to rectified flows and may benefit from a similar reflow procedure.

Thank you for the reminder. We apologize for the delay in responding to the reviewers. Below, we address your questions.

Q1: On a scale of 1 to 10, with the difficulty of reproducing DDPM on a 1D toy data being 1 and that of reproducing GAN being 5, how would you rate the difficulty of reproducing your 'improved consistency model'?

A: In this work, we do not have any results on 1D toy data. If you are referring to the difficulty of implementing consistency training for 1D datasets, it is quite easy. We would rank its difficulty level on par with implementing DDPMs. However, we observe that consistency models might not perform well on one-dimensional toy datasets; they are typically much more effective for high-dimensional datasets.

Q2: How robust is the 'improved consistency model' across different training runs? Are there any "implementation-level" tricks that haven't been disclosed to the reader at this time?

A:Improved consistency training is as robust as training diffusion models. No additional tricks, other than those mentioned in this paper, are needed.

Q3: Is there a plan to release the code, and if so, do you have a tentative date in mind?

A: We plan to release the code upon paper publication.

Thank you for your reply to our questions.

However, we observe that consistency models might not perform well on one-dimensional toy datasets; they are typically much more effective for high-dimensional datasets.

1D toy data, such as a mixture of three Gaussian, and 2D toy data, such as a mixture of 8 Gaussian, double moon, and spiral, are often valuable for understanding how well a generative model captures data distribution. It would be insightful to know the reasons behind the suboptimal performance of consistency models on these toy datasets. Could you provide insights into why they might not perform well? For example, are they prone to collapsing into one or a few modes, or do they struggle to distinguish between different density modes?

Additionally, I am curious about the performance of consistency models when the number of reverse steps exceeds 2. Would the performance improve, or would it decline as the number of reverse steps increases beyond 2? Any observations or insights into this aspect would be appreciated.