Paper: Improved Techniques for Training Consistency Models (ICLR-2024)

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
The new consistency models do not require LPIPS losses anymore and are thereby more general in that they are not limited to image synthesis anymore. 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).

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

Reviewer comment (Writing):
When the paper defines the ground truth consistency function in Section 3.2, is there a typo in the ground truth consistency function? Should it be $\xi$ instead of $\mu$ in the ground truth consistency function $f^{*}$ for the single data point data distribution?

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

Reviewer comment (Experiments):
Also 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)?

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
The paper’s improved training schemes do allow the consistency model to avoid relying on a pre-trained diffusion model. However, its theoretical basis still seems anchored in diffusion model principles. It would be beneficial if the authors explored the broader potential of the consistency training scheme. Specifically, 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].

Author response:
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.

Reviewer comment (Theory):
Could the author elaborate more on why higher discretization steps can reduce bias but increase variance?

Author response:
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])

Reviewer comment (Experiments):
The choice to use the pseudo-Huber loss is not entirely convincing, as its effectiveness appears sensitive to the chosen constant value, and it offers negligible improvements over the tuning-free LPIPS loss, which is known for its alignment with human perception and computational efficiency, especially in pixel-space models.

Author response:
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.

Reviewer comment (Evaluation):
The training speed for the consistency model is comparatively slow, especially when compared with models like GANs. It would be insightful if the authors could address this aspect and discuss its implications for scaling to larger datasets.

Author response:
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.

Reviewer comment (Experiments):
The applicability of consistency training appears confined to training unguided models, which significantly underperform in applications such as text-to-image generation.

Author response:
We leave generalization to guided models to future work.

Reviewer comment (Experiments):
Does the current consistency training outperform consistency distillation when applying the enhancements presented in this paper to the latter?

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

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Theory):
The theoretical analysis presented in Sec.3.2 does not seem justified. Although the main contribution of the paper appears to be empirical, clarification of the concerns regarding this analysis would still be appreciated.

Author response:
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.

Reviewer comment (Theory):
Eq.6 seems to apply regardless of the relation between $\theta$ and $\theta^-$.

Author response:
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$.

Reviewer comment (Writing):
In the current version, the meaning of "if the consistency loss either does not exist or is unsuitable" is unclear and does not seem to be well explained.

Author response:
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.

Reviewer comment (Evaluation):
There is a question about whether the authors have any insight into why the Pseudo-Huber loss yields such significant improvement, as it is not recalled to be used in any popular diffusion models. It is possible that consistency models, due to their unique features and training dynamics, benefit more from "bounded" gradients for all $t$.

Author response:
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.

Reviewer comment (Theory):
There is a question about whether the authors have insights on why insensitive time embeddings are helpful, or in this case crucial. A similar design choice is made in EDM, where $1/4\log(\sigma)$ has a rather small range between $\sigma_{\textrm{min}}$ and $\sigma_{\textrm{max}}$.

Author response:
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.

Reviewer comment (Reproducibility):
When examining the effects of different time embedding choices, positional embedding is mentioned. Given Fig.1 on the average $l_1$ distance, it is assumed that the authors directly feed in $\sigma$, rather than EDM's $1/4\log(\sigma)$. This point requires clarification.

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

Reviewer comment (Writing):
The paper does not contain explicit definitions of $s_0$ and $s_1$, which denote the start and end discretization steps.

Author response:
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.