Truncated Consistency Models

TL;DR

Truncated Consistency Models (TCM) address a key limitation of standard consistency models, which must balance denoising at early PF ODE times against generation at late times. By training on a truncated time range $[t',T]$ and employing a two-stage procedure with a boundary condition from a full-range model, TCM reallocates capacity toward generation and achieves stronger one-step and two-step generation with smaller networks. The authors formalize a boundary-consistent parameterization, decompose the training loss into a boundary loss and a consistency loss, and validate improvements on CIFAR-10 and ImageNet $64\times64$ against state-of-the-art CM baselines, along with comprehensive ablations. These results suggest a practical path to faster, more stable diffusion-based generation with reduced model size and improved sample quality.

Abstract

Consistency models have recently been introduced to accelerate sampling from diffusion models by directly predicting the solution (i.e., data) of the probability flow ODE (PF ODE) from initial noise. However, the training of consistency models requires learning to map all intermediate points along PF ODE trajectories to their corresponding endpoints. This task is much more challenging than the ultimate objective of one-step generation, which only concerns the PF ODE's noise-to-data mapping. We empirically find that this training paradigm limits the one-step generation performance of consistency models. To address this issue, we generalize consistency training to the truncated time range, which allows the model to ignore denoising tasks at earlier time steps and focus its capacity on generation. We propose a new parameterization of the consistency function and a two-stage training procedure that prevents the truncated-time training from collapsing to a trivial solution. Experiments on CIFAR-10 and ImageNet $64\times64$ datasets show that our method achieves better one-step and two-step FIDs than the state-of-the-art consistency models such as iCT-deep, using more than 2$\times$ smaller networks. Project page: https://truncated-cm.github.io/

Figures (11)

• Figure 1: (a) Two-stage training of TCM. In Stage 1, a standard consistency model is trained to provide both the boundary condition and initialization for TCM training in Stage 2. TCM focuses on learning in the $[t',T]$ range, discarding denoising tasks at earlier times and allocating network capacity toward generation-like tasks at later times. (b) Sample quality (FID, lower is better) of the two training stages. TCM (Stage 2) improves over standard consistency training (Stage 1) across datasets. Additionally, standard consistency training shows instability on challenging datasets like ImageNet 64x64, where the model could diverge during training.
• Figure 2: Evolution of the denoising FID (dFID$_t$) during standard consistency training for different $t$, where $0 < t \le 80$ follows the EDM noise schedule karras2022elucidating. The model gradually sacrifices its denoising capability at smaller times ($t<1.0$) to trade for the improved generation quality at $t=80$ as training proceeds.
• Figure 3: Denoising FID (dFID) for continuation of standard consistency training at later iterations (Stage 1) and TCM model (Stage 2) at various $t$s on CIFAR-10 during the course of training. For TCM, we set the dividing time $t'=1$. We can see, in the second stage, the dFID exhibits a dramatic increase at times below the dividing time $t'$, while the dFID at times above $t'$ and FID at $t=T$ continue to improve. Notably, the rate of dFID in the truncated stage increase at earlier times is significantly faster compared to standard consistency training, suggesting a more efficient "forgetting" of the denoising tasks.
• Figure 4: (a) Comparison of Student-$t$ distributions with different standard deviations $\sigma$ and degree of freedom $\nu$. (b) FID evolution on CIFAR-10 for different $\sigma$ and $\nu$. $w_b=0.1, \rho=0.25$, $t'=1$, and a batch size of 128 are used for all plots. (c) Effect of $\rho$ and $w_b$ on the FID on CIFAR-10. We use a batch size of 128. $t'$ is set to 1.
• Figure 5: Uncurated one-step (top) and two-step (bottom) generated samples from TCM (EDM) on CIFAR-10 and TCM (EDM2-XL) on ImageNet 64$\times$64, respectively.