Simplifying, Stabilizing and Scaling Continuous-time Consistency Models

Thank you for the detailed review and thoughtful feedback. Below we address specific questions.

Q1: Does JVP cost more time and memory?

A1: Not really. JVP (Jacobian-vector product) can be efficiently computed using forward-mode automatic differentiation, which requires the same memory and compute as a standard forward pass without saving intermediate activations. This is significantly cheaper than backpropagation, which relies on reverse-mode automatic differentiation. Consequently, our continuous-time consistency models require similar compute and memory to train when compared to their discrete-time counterparts, which perform two forward passes.

Q2: Detailed comparison with ECT.

A2: We have updated the paper to include detailed comparison with ECT in terms of training convergence and compute. See also our summary of updates, part C.

Q3: Why is sCT worse than sCD on ImageNet-512?

A3: We believe the higher training variance of CT is the main issue, particularly in the complex latent spaces defined by the pretrained encoder. We agree with the reviewer 2i3a’s hypothesis that the current encoder/decoder may not be optimal for consistency models. Theoretically, since the ground truth mapping in consistency models aims to transform a Gaussian distribution into a multimodal data distribution with potentially disconnected supports, its tangent can become ill-conditioned at boundary points, resulting in worse optimization dynamics. If we could develop a better encoder/decoder to create a more well-conditioned ground truth mapping in the latent space, the training of consistency models would likely become significantly easier. This is a fascinating research direction, and we plan to explore it in future work.

Q4: Ablation study for Adaptive Double Normalization.

A4: Adaptive Double Normalization is crucial for the model to converge. Without it, the default AdaGN layers lead to divergence, forcing us to terminate the runs prematurely. As a result, we were unable to conduct ablation studies for this feature in the same manner as for others, as the model simply fails to train without it. To clarify this point, we have updated Lines 304-305 in the paper accordingly.

Q5: Adaptive weighting achieves worse results in Figure 5(b).

A5: Thank you for pointing this out! The figure initially had incorrect legends—we mistakenly labeled the runs with adaptive weighting as "w/o adaptive weighting," leading to the confusion. We have since updated the paper to correct this. The revised Figure 5 now clearly shows that adaptive weighting improves upon fixed weighting functions for 2-step sampling and achieves comparable results for 1-step sampling.

Q6: In Figure 5c, do discrete-time CMs have a constant number of time steps $N$ or a timestep schedule up to the maximum number of steps? If it is the former, it seems to be a bit unfair to discrete-time CMs.

A6: We used a constant number of time steps for discrete-time CMs for two reasons:

1. Results in Figure 5c focus on consistency distillation. Previous papers [1][2] employed timestep schedules only for consistency training, not for distillation.

2. Figure 5c is designed to analyze how performance changes as $N$ varies. To ensure a cleaner comparison, we chose to fix $N$ throughout the course of training, avoiding confounding effects introduced by varying $N$ during training.

If the reviewer is curious about the performance of sCT versus discrete-time CT that uses timestep schedules, we have provided a comparison with the best-performing discrete-time CM (ECT) in our summary of updates, part C. Notably, our continuous-time CT consistently outperforms the best discrete-time counterparts.

References:

[1] Song, Yang, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. "Consistency models." arXiv preprint arXiv:2303.01469 (2023).

[2] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models." arXiv preprint arXiv:2310.14189 (2023).

Thank you for the detailed review and thoughtful feedback. Below we address specific questions.

Q1: Experimental comparison between Adaptive Double Normalization and AdaGN.

A1: As observed in [1] and confirmed by our experiments, AdaGN leads to divergence when used for CM training. It results in FIDs ranging from 200 to 300, with highly unrealistic samples. Given this significant performance gap, it is impractical to include the results of AdaGN alongside Adaptive Double Normalization in a plot like Figure 5. We have updated Lines 304-305 to clarify this point in the revision.

Q2: Ablation study for tangent linear warmup.

A2: Empirically, tangent warmup is optional. Removing it generally does not impact model performance, although gradients tend to have more spikes without it. While we do not consider tangent warmup a major technical contribution of our paper, we nonetheless included it in the main text for better reproducibility.

Q3: Adaptive weighting performs worse for 2-step sampling in Figure 5b.

A3: Thank you for pointing this out! The figure initially had incorrect legends—we mistakenly labeled the runs with adaptive weighting as "w/o adaptive weighting," leading to the confusion. We have since updated the paper to correct this. The revised Figure 5(b) now clearly shows that adaptive weighting improves upon fixed weighting functions for 2-step sampling and achieves comparable results for 1-step sampling.

Q4: Compute efficiency comparison with ECT.

A4: Please refer to our sumary of updates, part C.

Q5: Discussions related to text-to-image generation.

A5: We believe our method can be readily generalized to such models. Similar to our setup for ImageNet 512×512, these text-to-image diffusion models are also trained in the latent space, and are often designed for generating images of resolution around 512×512. We plan to explore this in future work.

Q6: Other potential candidates for time transformation.

A6: We explored various other time transformations during our initial experiments, such as those defined by trigonometric functions. Empirically, their results were very similar to the identity transformation, which we decided to use for simplicity.

Q7: Why is sCT less effective at higher resolutions?

A7: We believe the higher training variance of CT is the main issue, particularly in the complex latent spaces defined by the pretrained encoder. We agree with the reviewer 2i3a’s hypothesis that the current encoder/decoder may not be optimal for consistency models. Theoretically, since the ground truth mapping in consistency models aims to transform a Gaussian distribution into a multimodal data distribution with potentially disconnected supports, its tangent can become ill-conditioned at boundary points, resulting in worse optimization dynamics. If we could develop a better encoder/decoder to create a more well-conditioned ground truth mapping in the latent space, the training of consistency models would likely become significantly easier. This is a fascinating research direction, and we plan to explore it in future work.

Q8: Confusion about sCD-XXL vs. sCD-XL in Fig 6(b).

A8: To clarify, sCD-XXL outperforms sCD-XL under the two-step sCD setting. In Fig 6(b), we show the relative FID ratio (FID_CM / FID_DM), where FID_CM is the consistency model's FID, and FID_DM is the diffusion model's FID. What the reviewer observed is that sCD-XXL has a larger FID ratio than sCD-XL, this is because FID_DM of XXL is much smaller. Therefore, the actual FID of sCD-XXL (ratio × FID_DM) is still lower than that of sCD-XL.

References:

[1] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models." arXiv preprint arXiv:2310.14189 (2023).

Thank you for the detailed review and thoughtful feedback. Please find our responses to your questions below.

Q1: Adding the results for more recent flow models and the 2-step rectified flow.

A1: Thank you for pointing out the reference. We have added the results from [1] to Table 1. However, since [2] does not include experimental results on CIFAR10 or ImageNet, it is difficult to make a direct comparison. Nevertheless, we have cited [2] in the updated version for completeness. In addition, the "2-rectified-flow" listed in Table 1 already refers to the 2-step distillation of rectified flow.

Q2: An intuitive explanation for the loss in Eq(2), and mention a detailed citation (Remark 10) of it.

A2: Thanks for the suggestion. We’ve added the detailed citation in Line 165-166. The key to understanding Eq(2) lies in intuitively grasping the tangent $\frac{df}{dt}$, which represents the direction of instantaneous change that drives $f(t)$ toward $f(0)$. Optimizing $f$ along this direction, as described in Eq(2), ensures consistency with the boundary condition at $t=0$.

Q3: Adding some generated images with one step.

A3: Thanks for the suggestion. We’ve added uncurated samples for both 1-step and 2-step sampling in Appendix H, generated by our sCD-XXL trained on ImageNet 512. All corresponding 1-step and 2-step samples use the same random seed.

Q4: Additional elaboration regarding the implication of having unit norm.

A4: Thanks for the suggestion. We’ve added more discussions in the appendix accordingly.

Q5: Typos of $c_{skip}$ and $c_{out}$; typo in Eq(20) in appendix; typo in line 126; typo in line 122.

A5: Thanks for pointing it out. We’ve fixed them in the revision.

Q6: Are there any indications that the continuous consistency models will still face instability issues for even larger scales, in particular as compared to diffusion/flow models? This question mostly relates to the gap between sCT and sCD.

A6: In our experiments, we find that the training of sCM is stable at both small and large scales. We did not even apply any gradient clipping and the network activations are stable during the course of training. That said, we mostly focus on CNN-based architectures (with self-attention blocks), and leave the explorations of other architectures for future work.

We believe the gap between sCT and sCD on ImageNet 512x512 is mainly caused by the higher training variance of CT, particularly in the complex latent spaces defined by the pretrained encoder. We agree with reviewer 2i3a’s hypothesis that the current encoder/decoder may not be optimal for consistency models. Theoretically, since the ground truth mapping in consistency models aims to transform a Gaussian distribution into a multimodal data distribution with potentially disconnected supports, its tangent can become ill-conditioned at boundary points, resulting in worse optimization dynamics. If we could develop a better encoder/decoder to create a more well-conditioned ground truth mapping in the latent space, the training of consistency models would likely become significantly easier. This is a promising research direction, and we plan to explore it in future work.

References:

[1] Tong et al. 2024. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport.

[2] Kornilov et al 2024. Optimal Flow Matching: Learning Straight Trajectories in Just One Step.

Thank you for the detailed review and thoughtful feedback. Please find our responses to your questions below.

Q1: Why is sCT less effective than sCD on ImageNet 512?

A1: We believe the higher training variance of CT is the main issue, particularly in the complex latent spaces defined by the pretrained encoder. We agree with the reviewer’s hypothesis that the current encoder/decoder may not be optimal for consistency models. Theoretically, since the ground truth mapping in consistency models aims to transform a Gaussian distribution into a multimodal data distribution with potentially disconnected supports, its tangent can become ill-conditioned at boundary points, resulting in worse optimization dynamics. If we could develop a better encoder/decoder to create a more well-conditioned ground truth mapping in the latent space, the training of consistency models would likely become significantly easier. This is a fascinating research direction, and we plan to explore it in future work.

Q2: To what extent does this prior weighting contribute to variance reduction? Is it helpful to stabilize the learnable adaptive weighting layer?

A2: We can express the relationship between prior weighting and tangent normalization as:$\frac{w_t \cdot \|\frac{df}{dt}\|}{w_t \cdot \|\frac{df}{dt}\| + c} = \frac{\|\frac{df}{dt}\|}{\|\frac{df}{dt}\| + \frac{c}{w_t}}$, where $w_t$ is the prior weighting and $c$ is the tangent normalization constant. This shows that the prior weighting acts as a time-dependent tangent normalization with $c'_t = \frac{c}{w_t}$, effectively balancing the variance of the tangent across $x_t$ and $t$.

Empirically, we observe that $\|\frac{df}{dt}\|$ becomes quite large as $t \to T$ (i.e., the pure noise distribution). Prior weighting helps reduce variance and improves the numerics of tangent computation at these time steps. Prior weighting also provide slight benefits for the training stability of learnable adaptive weighting, but it is not strictly necessary.

Q3: Data-free distillation of continuous-time CMs.

A3: Thanks for the insightful suggestion! This is an interesting research direction and we plan to explore it in future work.

Q4: Conflicting notations for both diffusion models and consistency models.

A4: Thanks for pointing it out! We have changed the notation of diffusion models to $f^{\text{DM}}$.

Q5: More explanations about the Adaptive double normalization.

A5: Pixel normalization is quite similar to local response normalization. Since our network is a convolutional network, the shapes of the modulation layers $s(t)$ and $b(t)$ match the hidden feature $x$ (i.e., $C \times H \times W$). We normalize the $C$ values at each pixel position, effectively performing RMS normalization for each pixel.

Thank you for the detailed review and thoughtful feedback. Please find our responses to your questions below.

Q1: More Limitations.

A1: Thanks for the suggestion. Due to the page limit, we’ve included an additional section in the appendix for further discussions on limitations.

Q2: Elaborate positional embedding.

A2: The positional embedding for timesteps is a standard layer initially used in DDPM [1] and widely adopted in diffusion models. We’ve added further clarifications around Line 259-261.

Q3: Fig 3 does not add much value.

A3: Figure 3 highlights the key difference between discrete-time and continuous-time training, specifically that continuous-time training avoids discretization errors. We are happy to revise the illustration based on reviewer feedback.

Q4: Typos.

A4: Thanks for pointing it out! We’ve fixed it in the updated paper.

Q5: Is the method really fully stable?

A5: We believe so. Combining all the proposed techniques together, we found that the training of sCM is stable at both small and large scales, requiring no gradient clipping, and exhibiting smooth variation of network activations during the course of optimization. That said, we have only examined CNN architectures (with self-attention blocks) for CMs so far. We leave the exploration of other architectures, such as DiTs and SSMs, for future work.

References:

[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851.