Truncated Consistency Models

We appreciate the reviewer’s valuable comments.

The truncation principle has become a standard technique in diffusion (Balaji et al., 2022; Zheng et al., 2023 -- that could be cited in the related work) and consistency (Heek et al., 2024) models.

Thank you for those references! We have added Balaji et al., 2022; Zheng et al., 2023 to our related work section. Note that since consistency models solve a more difficult task (learning to integrate PF-ODE) than diffusion models (learning the drift of PF-ODE), they can benefit more from such a strategy but also require a specific parameterization (Eq. 5) to satisfy the boundary condition.

Heek et al. Multistep Consistency Models.

We clarify that our motivation and implementation of truncated training are distinct from Heek et al. Their models predict the noised data while our models always predict the clean data, which motivates us to design new parameterization (Eq. 5). As a result, the multistep consistency model is not able to do one-step generation while we exclusively aim for it.

I do not understand the necessity of the loss decomposition in Eq. (6).

To my understanding, it amounts to giving a specific sampling weight for timesteps close to t'.

Eq. (6) is a simple decomposition of the original loss Eq. (4), and we haven't introduced any specific weighting yet. We need this decomposition to motivate such a weighting strategy.

Additionally, it appears that $S_{t'} = t'$, so the notations may be simplified in this part of the paper.

The equality only holds in the limit of $\Delta t \to 0$. Since we assume discrete-time consistency training, the notations like $S_{t’}$ help clarify that the boundary loss vanishes only when $\Delta t$ is close to zero.

The code used by the authors for the presented experiments should be included in the supplementary material for further reproducibility.

We are still going through the internal release process, but we expect the code to be released soon.

One can notice in Figure 3 that the denoising FID increases at the beginning of stage 2. I would be interested in the authors' opinion on this phenomenon.

This behavior is intended and aligns well with our motivation--moving the capacity dedicated to denoising to generation. It shows the model effectively "forgets" the unnecessary denoising tasks.

Still in Figure 3, can the authors explain why the denoising FIDs do not start from the same value between stage 1 and stage 2 (assuming they start from the same checkpoint)?

Those two are from different settings. We have updated Fig. 3 to match their initial checkpoints. Thank you for pointing this out!

While the authors acknowledge an additional training and memory cost in Section 6, I am not sure they include the additional cost induced by the supplemental training iterations to complement the pretrained consistency model (up to 200k additional optimization steps based on the different plots).

Indeed, we have considered the per-iteration training cost only. We agree with the reviewer and have updated our limitation section accordingly. Thanks for catching this!

Section 4.2 insists on the forgetting by the truncated model on the lower noise levels. This is not an advantage or a confirmation of the model's abilities, but a byproduct of the truncated training loss.

This behavior confirms our motivation for the redistribution of the neural net capacity. That said, if it had not happened, it would have been difficult to justify the enhanced performance coming from the truncated training (if the model did not forget denoising, why would we expect the generation to be improved?).

The number of optimization steps is not specified in the hyperparameters listed in Section D.

The number of iterations is as follows:

• CIFAR-10: Stage 1: 250K iterations, Stage 2: 200K iterations
• ImageNet 64x64
  • EDM2-S: Stage 1: 150K iterations, Stage 2: 120K iterations
  • EDM2-XL: Stage 1: 100K iterations (we used the pre-trained checkpoint), Stage 2: 40K iterations

The FID evolution on CIFAR-10 and ImageNet 64x64 (EDM2-S) during training is shown in Fig. 1(b). We have reflected this in the appendix.

The paper should specify where the baseline results from Tables 1 and 2 are reported from.

We mostly borrow the baseline results from the original papers. We have updated Sec. 4.3 to clarify this.

We appreciate the reviewer’s valuable comments.

While the proposed method improves considerably on small models and few time steps, the improvement on the highest quality 2-step EDM2-XL is negligible. The applicability of the proposed method on large-scale models is questionable.

EDM2-S model (with which our method achieves significant gain) has 280M parameters and is not small for the ImageNet 64x64 dataset (its size is similar to the widely used ADM architecture). Methods like variational score distillation use similar-sized networks and still can produce good FID on this dataset. The fact that consistency models require a larger network (1.1B) and more sampling steps makes them less scalable, and solving this issue is the main objective of our paper. Also, we experimented with the text-to-image setting and found that our method works well on a large-scale setting. See the general response above.

Also, we contacted the ECT authors for more details in training ECT with EDM2-XL (which was not released at the time of the submission), but still could not reproduce its results (the training diverged), especially with a large batch size, which could be due to e.g. the training details for EDM2-XL that are not available to us. We thus ended up using the checkpoint provided by the authors. We believe that with the right configuration for large batch training, the 2-step FID of EDM2-XL could be improved further.

It is in fact surprising that the truncated training works at all with EDM2-XL, given that the untruncated training with EDM2-XL does not even converge –a preliminary indication that the truncated training can be more stable.

Why variational score distillation methods are listed in separate categories and not compared against consistency models?

This is a good point! There are mainly two reasons: 1) Variational score distillation requires training a separate model (e.g. the fake score network) to estimate the gradient of their distribution-matching loss. This requires a GAN-like alternative training procedure, which introduces additional engineering burden and computational overhead. 2) The main purpose of this paper is to understand and improve the training of consistency models. We believe that our method is a step towards closing the gap between consistency models and variational score distillation in FID.

SiD seems to get better FID with significantly fewer parameters.

SiD is a variational score distillation method. It requires training a separate model and also a GAN-like alternative training procedure, which introduces more training instabilities and computational overhead. While the performance of SiD is indeed promising, we believe that the understanding of why it works is still lacking. For instance, SiD requires a very large batch size (8K) to make it not diverge on ImageNet 64x64. Moreover, their final loss function is an extrapolation of the two different loss functions they suggested, which is not justified by theory. On the contrary, consistency models do not use an alternative training procedure, require less computational overhead, and their loss function is supported by theory.

But there are many other works on consistency distillation, such as CTM / TCD. Why not comparing against them?

To our knowledge, our model outperforms all published consistency distillation/training methods on CIFAR-10 and ImageNet 64x64. We do compare with CTM (or TCD) in Table 1 where we outperform it by a large margin. Note that we have excluded the methods that rely on GAN loss, such as CTM + GAN, for fair comparison.

How does the proposed method apply to guided sampling / text-to-image applications? How does guidance scale impact the proposed method? (Including more guided sampling results would improve the submission.)

See the general response above.

We appreciate the reviewer’s valuable comments.

The approach appears a bit over-engineered, with numerous handcrafted designs and hyperparameters.

We discuss the hyperparameters introduced by truncated training in the general response. Our method does introduce a few additional hyperparameters, but we believe this is not always bad as long as one of these conditions is met: 1) there is a good guiding principle to determine their value, 2) the hyperparameters can easily be transferred from the smaller to larger models/datasets, or 3) the results are stable with a wide range of values.

Also, when developing our method, we intentionally avoided tuning these hyperparameters for any specific model or dataset and used exactly the same values for all experiments to ensure the generalizability of our method. Hence, the additional complexity of our framework is marginal.

How should one determine the terminate point of Stage 1 training? Monitoring its progress may introduce additional complexity. The number of training iterations for Stage 1 may also represent a crucial hyperparameter that may require further ablation studies.

For both stages 1 and 2, we simply monitor the 1-step FID during training and terminate the training when it converges. Determining when to terminate a generative model training can be done in many ways (e.g., validation loss or other performance metrics) and is not specific to our paper.

The paper lacks ablation studies for certain hyperparameters relevant to the proposed method, such as $\rho, \Delta_t, \Delta_{t'}$.

The effect of $\rho$ is shown in Fig. 4(c). $\Delta_t$ (or $\Delta_{t'}$) follows the prior work ECT without modification. It is a general design choice in all consistency models that is not particularly related to the idea of truncated training.

Several studies, such as [1, 2], have proposed truncating the time interval in diffusion model training to address issues like the unbounded score problem or high Lipschitz continuity when $t \approx 0$. How does the proposed method relate to this existing literature?

We thank the reviewer for pointing out these two related works. While these references and our work both employ truncation techniques, fundamental differences exist in both motivation and implementation. These works primarily focus on stabilizing regression targets during diffusion model training, whereas our approach addresses the capacity allocation in consistency models. As they focus on diffusion models, there is no notion of boundary condition in these two works, which, however, motivates a new parameterization in TCM.

In those papers, the lower bound of the truncated time interval is close to 0 as they aim to avoid an unstable regression target around t=0. In our case, we need the lower bound of the truncated time interval to be sufficiently large to represent the point where the consistency mapping transitions from noise reduction (denoising) to sample generation.

In fact, with the standard EDM pre-conditioning, the truncation proposed by these two papers is not necessary for stable training. However, in our case, the effect of the truncated training comes from the explicit redistribution of the capacity and cannot be achieved by simply reparameterizing the neural network (e.g., EDM pre-conditioning).

How can one ensure that the trained CT model, using the proposed parameterization in Eq. (5), will ultimately be continuous at the splitting time $t'$? If one performs multi-step sampling and selects (some) sampling timesteps close to $t'$, this could potentially cause issues?

We agree that, in our parameterization, there could be a discontinuity around t' as we use the stage-1 model when $t < t'$ and the stage-2 model when $t \geq t’$. However, since we only evaluate our stage-2 model on $t\geq t’$ (to ultimately discard the stage-1 model after training), it does not happen during generation. Empirically, we obtain the state-of-the-art 2-step FID by evaluating the model at $t=t'$ in our experiments.

The paper lacks ablation studies for certain hyperparameters (refer to W3).

The hyperparameters the reviewer mentioned in W3 are $N_B$ (which is equivalent to $\rho$), $\Delta_t$, and $\Delta_{t’}$. As we discussed above, we already provide the ablation for $\rho$ in our paper (see Figure 4(c)). $\Delta_{t}$ represents “the nonnegative difference between two consecutive time steps $t$ and $s$” (i.e., $\Delta_{t}=t-s$), and it directly follows ECT’s definition (see Sec. 3.3 in [a]). $\Delta_{t’}$ is a special case of $\Delta_{t}$ when setting $t=t’$. Again, this is a general design choice for consistency models and not introduced by our method. Can the reviewer provide a reason for why including ablation studies for $\Delta_t$ would help improve this submission which is about truncated training?

Guidance on determining when to terminate Stage 1 training is missing. For instance, what level of degradation occurs if Stage 1 training is imperfect? Does this significantly impact one-step or two-step sampling?

This is a good question. As mentioned earlier, we train the Stage-1 model until its 1-step FID fully converges, because the Stage-1 model serves as a boundary condition during the truncated training stage. An imperfect Stage-1 model will propagate less accurate learning signals to the Stage-2 model.

Table. Effects of Stage-1 FIDs on Stage-2 FIDs evaluated on CIFAR-10.

Here, we see that the stage-1 model with FID=2.97 yields the final FID of 2.61. If we start with a better stage-1 model (FID=2.82), we get a better final FID of 2.54. We will consider adding this result to the appendix.

Regarding [1, 2], I believe the fundamental issue is similar: addressing the challenge of large Lipschitz constants as time approaches zero. …does not suggest a fundamental difference … particularly in comparison to [2].

We clarify that TCM is motivated from a fundamentally different perspective compared to [1,2]. To our best knowledge, truncation in [1,2] is addressing the training instability issue near zero (exploding norm of the score function [1] and the large Lipschitz constants in [2]). However, we haven’t observed training instability issue near zero when training consistency models. In fact, the consistency function (that maps any intermediate point to clean data at zero) approaches the identity function as $t \to 0$ and thus becomes much easier to learn when the time step $t$ is close to zero.

This does not suggest a fundamental difference between the literature or approaches, particularly in comparison to [2].

What [2] observes is that the time derivative of the optimal $\epsilon_\theta$ in DDPM goes to $\infty$ near $t=0$. However, this is not the case for consistency models. Consider a simple case where data distribution is $\mathcal N(0,1)$. Then, the noised distribution $p_t = \mathcal N(0,1+t^2)$. In this case, the consistency function is $f(x,t) = x/(1+t^2)$, and $\partial f/\partial t = 0$ when $t \to 0$.

We also further clarify the specific difference between TCM and [2] from the methodology perspective: [2] proposes to reduce the Lipschitz constants in the time interval $[0, t’)$ by sharing the timestep conditions in this interval. However, TCM first performs consistency training in the whole time interval (Stage 1) and then performs consistency training in the time interval $[t’, T]$ (Stage 2). As we see, when it comes to truncated training, [2] focuses on stabilizing the training in the early timesteps $t \in [0, t’)$ while TCM completely gets rid of training in the early timesteps $t \in [0, t’)$.

Additionally, the comment on [2] appears inaccurate; their results demonstrate that the best truncation time is not always "the smaller, the better" (see Fig. 8).

We apologize for causing the confusion. We did not intend to say “the smaller, the better”, instead, we meant that the lower bound of the truncated time interval in [2] is close to zero, which is mentioned in [2] multiple times: “diffusion models: they frequently exhibit the infinite Lipschitz near the zero point of timesteps.”, “our approach can effectively reduce the Lipschitz constants near t = 0 to zero”, “our approach substantially alleviates the large Lipschitz constants which occur in the intervals whose noise levels are small”, “E-TSDM can effectively mitigate the large Lipschitz constants near t = 0”. However, in TCM, $t’$ (the lower bound of our truncated time interval) should not be close to zero.

References

[a] https://arxiv.org/abs/2406.14548

What are the total training iterations and time for the first and second stages, respectively? Is the first stage exactly the same as previous consistency models?

We thank the reviewer for pointing out the missing training details. The number of iterations is as follows:

• CIFAR-10: Stage 1: 250K iterations, Stage 2: 200K iterations
• ImageNet 64x64
  • EDM2-S: Stage 1: 150K iterations, Stage 2: 120K iterations
  • EDM2-XL: Stage 1: 100K iterations (we used the pre-trained checkpoint), Stage 2: 40K iterations

The FID evolution on CIFAR-10 and ImageNet 64x64 (EDM2-S) during training is shown in Fig. 1(b). We have reflected this in the appendix.

We did not try to exactly match the training iterations of previous works and used what worked best for our models. In all cases, stage 2 models outperform stage 1 models (see Fig. 1(b)).

Why the authors only initialize the truncated training stage with ECT for EDM2-XL?

That’s a good point. We tried hard to reproduce the results of ECT for EDM2-XL (the codebase for ImageNet was not released at the time of the submission). We contacted the ECT authors for more details in training ECT with EDM2-XL, but still could not reproduce its results, especially with a large batch size, which could be due to e.g. the training details for EDM2-XL that are not available to us. We thus ended up using the checkpoint provided by the authors.

It is noteworthy that our truncated training procedure succeeds even when the standard, untruncated training (ECT with EDM2-XL) fails to converge. This observation suggests that the truncated training can be more stable. While stabilizing the original consistency model training is an important research direction, it falls outside the primary scope of our paper.

I think a limit case for the truncated training is to only learn the 1-step mapping T -> 0, i.e., limiting the truncated region to a single point. The model can be directly trained on noise-data pairs. Can this method achieve better 1-step FID?

The model can be directly trained on noise-data pairs.

Note that we do not use offline noise-data pairs in our training. That case corresponds to the knowledge distillation [2] (initialized with a pre-trained consistency model) and is not a consistency model anymore. Knowledge distillation methods not only necessitate the computationally expensive construction of offline noise-data pairs but also generally struggle to match the performance of state-of-the-art generative models, as evidenced in Table 1.

Alternatively, within the TCM framework, we can explore shifting the dividing time $t'$ closer to $T$. However, as demonstrated in Table 3, increasing $t'$ beyond a certain threshold leads to performance degradation. Empirically, we observe that $t'=1$ yields the best results, better capturing the point in the probability flow ODE where the task transitions from denoising to generation.

References

[1] Jonathan Heek, et al., Multistep Consistency Models, 2024.

[2] Eric Luhman and Troy Luhman, Knowledge distillation in iterative generative models for improved sampling speed, 2021.

We appreciate the reviewer’s valuable comments.

The overall idea and motivation are actually not new. ECT already points out the trade-off between the denoising capacity and the consistency capacity

What ECT observes is that the optimal hyperparameters differ between 1-step and 2-step generation, while what we observe is that during the training of the same model, the trade-off between denoising and generation occurs after some training iteration (Fig. 2). The difference in motivation leads to the largely different solutions: ECT simply suggests choosing the hyperparameters optimized for 1 or 2-step generation, while we explicitly amplify this trade-off (Fig. 3) via truncated training to improve 1-step generation. The degree of improvement in FID coming from the truncated training cannot be achieved by simply manipulating the existing hyperparameters.

The 2-step performance improvement is relatively marginal compared to ECT.

Our method is mainly designed for 1-step generation where the improvement over ECT is more significant. Nevertheless, on ImageNet 64x64 with EDM2-S, our 2-step FID is 2.31, which surpasses the ECT's 2.79 by a nonnegligible margin (about a 17% reduction). The larger performance gain observed on the smaller model (EDM2-S) further validates the effectiveness of our truncated training procedure, by focusing the limited capacity on generation quality rather than denoising.

In addition, we would like to point out that we contacted the ECT authors for more details in training ECT with EDM2-XL, but still could not reproduce its results, especially with a large batch size, which could be due to e.g. the training details for EDM2-XL that are not available to us. We believe that with the right configuration for large batch training, the 2-step FID of EDM2-XL could be improved further.

As suggested by ECT, it is better to use a 2-step generation with a small model instead of a 1-step generation with a large one.

We agree, but we think what is even better is to improve the one-step generation without increasing the model size. This is the main purpose of TCM. As shown in Table 2, our one-step FID outperforms all existing consistency model baselines and, with EDM2-S, even matches ECT's two-step FID (2.88 for TCM one-step FID vs 2.79 for ECT two-step FID). Also, TCM can also do multi-step sampling (by choosing t greater than the splitting time t’), so one can still target 2-step generation with a small model.

Therefore, the actual benefits provided by truncated training are limited, especially considering that it requires careful tuning of the weighting functions and timestep schedules.

We discuss the hyperparameters introduced by truncated training in the general response. We largely adopted the training hyperparameters from the ECT paper, with a few adjustments tailored specifically for our truncated training procedure. Our modifications to the ECT settings include employing a Student-t distribution for time sampling and tuning the coefficient for the boundary loss and the dividing time $t'$.

We also provide a good guiding principle to determine their values, as discussed in Section 4.4. When developing our method, we intentionally avoided tuning these hyperparameters for any specific model or dataset and used exactly the same values for all experiments to ensure the generalizability of our method. We observe that these hyperparameters can easily be transferred from smaller to larger models/datasets, and their results are stable with a wide range of values. Hence, the additional complexity of our framework is marginal.

Could the authors discuss the relation with PCM?

We reviewed PCM and found that its idea is similar to multistep consistency models [1]. PCM and multistep consistency models are designed to improve multi-step generation by predicting the intermediate points on the PF ODE. In contrast, our model always predicts the data endpoint of the PF ODE as our goal is a one-step generation. Truncated training necessitates learning the boundary condition using our new parameterization and time sampling distribution, unlike these two papers. We will add a discussion of PCM alongside the multi-step consistency models in the related work section.

We appreciate the reviewer’s valuable comments.

However, with only the first method, the training still diverges (L241-L242), which shows that the part 2 seems to be more important. It is unclear whether the first part (truncation) is even necessary.

We clarify that the truncated training and the proposed time sampling distribution are not two independent design choices but rather one cannot be separated from the other. Without the proposed time sampling distribution, as $\Delta_t$ goes to zero, the truncated training permits a trivial solution (e.g. a constant function). In preliminary trials, we also observed empirically that this design choice does not work well.

It would be great if the author could conduct more rigorous ablations: Training consistency models by the original loss (i.e. the method used in stage-1) with the improved proposal distribution by focusing more on [0, t'].

That’s a good point! Such a strategy is already discussed in ref. [1], with a similar reasoning to the reviewer: "errors from minimizing consistency losses in smaller noise levels can influence larger ones and therefore should be weighted more heavily" (note that the time sampling distribution and the weighting functions are interchangeable). As such, our current weighting functions are already designed to emphasize the losses at smaller $t$'s following prior work. However, overemphasis on smaller $t$'s even further does the exact opposite of what we want (as it increases the capacity dedicated for denoising).

Nevertheless, we run experiments to see the effect of changing the weighting function to focus more on smaller $t$'s. On CIFAR-10, from the default weighting function $1/\Delta t$, we simply introduce an additional power variable $p$ to the previous weighting function, i.e., $\frac{1}{\Delta_{t,p}} = \frac{1}{(1+8 \cdot \text{sigmoid}(-t))(1-r) t^p}$ , to adjust the emphasis on smaller $t$'s. See link for visual comparisons. The results are shown below:

These results suggest that the improvement from our method is not merely due to the increased emphasis on smaller t's but rather the explicit redistribution of the capacity.

CT has the same gradient limit as CD. However, with truncation training, this conclusion will not hold anymore for t near to t'. What is the performance of CD by the truncated method?

Truncated consistency models are also consistency models with a specific parameterization and time sampling distribution. Therefore, the equivalence between CT and CD in the limit still holds for TCM (Eq. 35 in Appendix does not assume any specific parameterization of $f$). That is, the relationship between CT and CD in TCM should be the same as that in the original consistency models (the former comes with an unbiased but high-variance gradient, while the latter has a biased but low-variance gradient). For the performance of CD in TCM, please see the general response where we train TCM with CD on the text-to-image setting.

References

[1] Yang Song, Prafulla Dhariwal, Improved Techniques for Training Consistency Models, 2023.