Source-linked AI summary
Improved Techniques for Training Consistency Models
Yang Song, Prafulla Dhariwal
TL;DR
Consistency models seek high-quality one-step generation without adversarial training, while prior methods rely on diffusion distillation and learned metrics with limitations. This paper improves direct consistency training through theoretical and objective-design changes, achieving strong one- and two-step FIDs and surpassing distillation results.
Problem
Prior consistency methods rely on diffusion distillation and LPIPS, which add training overhead, limit quality to the diffusion model, and can bias FID evaluation.
Method
The paper improves direct consistency training by removing EMA from the teacher, replacing LPIPS with Pseudo-Huber loss, and introducing new noise and discretization schedules.
Results
2.51 and 3.25 FID on CIFAR-10 and ImageNet 64 × 64, respectively, are achieved in one step, while two-step generation reaches 2.24 and 2.77 and surpasses distillation.
Takeaways & Limitations
Consistency training can function as an independent generative-model family with sample quality comparable to leading diffusion models and GANs.
Takeaways & Limitations
The comparison excludes FastGAN- and StyleGAN-XL-based methods because their ImageNet-pre-trained discriminator feature extractors can inflate FIDs.
Abstract
from arXiv · showhide
Consistency models are a nascent family of generative models that can sample high quality data in one step without the need for adversarial training. Current consistency models achieve optimal sample quality by distilling from pre-trained diffusion models and employing learned metrics such as LPIPS. However, distillation limits the quality of consistency models to that of the pre-trained diffusion model, and LPIPS causes undesirable bias in evaluation. To tackle these challenges, we present improved techniques for consistency training, where consistency models learn directly from data without distillation. We delve into the theory behind consistency training and identify a previously overlooked flaw, which we address by eliminating Exponential Moving Average from the teacher consistency model. To replace learned metrics like LPIPS, we adopt Pseudo-Huber losses from robust statistics. Additionally, we introduce a lognormal noise schedule for the consistency training objective, and propose to double total discretization steps every set number of training iterations. Combined with better hyperparameter tuning, these modifications enable consistency models to achieve FID scores of 2.51 and 3.25 on CIFAR-10 and ImageNet $64\times 64$ respectively in a single sampling step. These scores mark a 3.5$\times$ and 4$\times$ improvement compared to prior consistency training approaches. Through two-step sampling, we further reduce FID scores to 2.24 and 2.77 on these two datasets, surpassing those obtained via distillation in both one-step and two-step settings, while narrowing the gap between consistency models and other state-of-the-art generative models.
1 INTRODUCTION
Consistency models target high-quality one-step generation without adversarial optimization or many sampling steps. This work improves consistency training so models learn directly from data, avoid LPIPS, and achieve stronger sample quality than prior approaches.
- Consistency models generate high-quality samples with a single network evaluation, without adversarial optimization or numerous sampling steps.
- Consistency training learns directly from data, whereas consistency distillation requires a pre-trained diffusion model and limits quality to that model.
- LPIPS can bias FID through shared ImageNet features and increases compute by requiring auxiliary-network pre-training and backpropagation.
- The proposed techniques remove LPIPS dependence by eliminating EMA from the teacher, adopting Pseudo-Huber losses, and changing discretization and noise schedules.
- 2.51 and 3.25 FID on CIFAR-10 and ImageNet 64 × 64, respectively, are achieved in one sampling step, improving over previous CT by 3.5× and 4×.
- 2.24 and 2.77 FID on CIFAR-10 and ImageNet 64 × 64, respectively, are achieved with two-step generation and surpass consistency-distillation scores in both settings.
2 CONSISTENCY MODELS
Consistency models approximate the probability-flow ODE’s consistency function, mapping noisy samples toward clean data. Training discretizes noise levels and uses student–teacher consistency objectives, enabling one-step and multistep generation.
- The probability-flow ODE defines a bijective mapping from a noisy sample x_σ to a near-clean sample x_σmin, called the consistency function.
- A consistency model f_θ is a neural network trained to approximate the consistency function while satisfying the boundary condition f_θ(x, σmin) = x.
- The ODE is discretized into noise levels σmin = σ1 < σ2 < ¨ ¨ ¨ < σN = σmax, with ρ = 7 in the stated schedule.
- Consistency training compares student and EMA teacher outputs across adjacent noise levels using a weighted metric objective.
- As N increases, optimization approaches the true consistency function, while the training curriculum progressively increases N and adjusts the EMA decay rate.
- Two-step generation often improves one-step sample quality, although additional sampling steps provide diminishing benefits.
3 IMPROVED TECHNIQUES FOR CONSISTENCY TRAINING
The paper improves consistency training through revised hyperparameters, teacher updates, metrics, discretization, and noise scheduling, while correcting a theoretical flaw in prior CT analysis. These changes improve sample quality without learned metrics and support stronger CT performance.
- Weighting functions, noise embeddings, and dropout: Proper weighting, noise embeddings, and dropout substantially improve CT with the squared ℓ2 metric.The selected weighting function decreases with noise level, Fourier scale is 0.02, and dropout is 0.3.
- Weighting functions, noise embeddings, and dropout: Excessively sensitive noise embeddings can cause continuous-time CT to diverge, while a Fourier scale of 0.02 matches positional-embedding sensitivity most closely.
- Removing EMA for the teacher network: Teacher EMA should be removed in CT: theory shows nonzero teacher decay can eliminate the data-dependent training signal, while zero decay aligns the gradient with supervised denoising.Empirically, discarding teacher EMA improves CT sample quality across LPIPS and squared ℓ2 metrics.
4 PUTTING IT TOGETHER
The combined iCT techniques are evaluated on CIFAR-10 and ImageNet 64 × 64 against competing generative models using sample-quality metrics and one- or two-step generation. iCT surpasses prior diffusion distillation approaches while avoiding diffusion-model training.
- The experiments benchmark unconditional CIFAR-10 and class-conditional ImageNet 64 × 64 samples against methods from the literature.FastGAN- and StyleGAN-XL-based methods are excluded because their ImageNet-pretrained feature extractors can inflate FIDs.
- iCT-deep achieves ImageNet 64 × 64 FIDs of 3.25 with one step and 2.77 with two steps.The table reports these results for class-conditional samples.
- iCT methods surpass previous diffusion distillation approaches in both one-step and two-step generation on CIFAR-10 and ImageNet 64 × 64.The comparison uses FID, Inception score, and Precision/Recall.
- 2.51 and 3.25 are the one-step FIDs of iCT-deep on CIFAR-10 and ImageNet 64 × 64, respectively.These results are reported alongside comparisons with diffusion models and GANs.
- 2.24 is the two-step CIFAR-10 FID achieved by iCT-deep, matching the quality of Score SDE while requiring far fewer sampling steps.Score SDE reaches FID 2.20 with 2000 sampling steps.
5 CONCLUSION
The improved consistency-training techniques address prior limitations by surpassing consistency distillation without LPIPS. They produce strong FID results on CIFAR-10 and ImageNet 64 × 64 and compare favorably with diffusion and GAN methods.
- Improved consistency training surpasses consistency distillation without relying on LPIPS.The approach removes EMA from teacher networks, replaces LPIPS with Pseudo-Huber losses, and changes discretization and noise schedules.
- The resulting consistency models achieve unprecedented FID scores on CIFAR-10 and ImageNet 64 × 64.The conclusion attributes the results to the combined improved techniques.
- The reported results outpace previous consistency-training methods and challenge the sample quality of leading diffusion models and GANs.
A PROOFS
The proofs analyze when consistency-training and consistency-model objectives coincide, using a point-mass data distribution and uniform weighting with the squared ℓ2 metric.
- Under a point-mass data distribution, the perturbed distribution is Gaussian and the consistency-training and consistency-model objectives can be written explicitly.
- The constructed noisy states satisfy ξ + σ_i z, showing that the consistency-training target equals the corresponding consistency-model state.
- The consistency-model and consistency-training objectives are exactly equal in this setting.
- The consistency model is defined using the minimum noise level σ_min.
- As the number of discretization steps tends to infinity, the step-size difference Δσ tends to zero, completing the limiting argument.
B ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS
Additional details specify architectures, optimization settings, metric behavior, and sampling configurations for iCT experiments. The accompanying figures provide uncurated samples and analyze metric-induced update variance.
- iCT experiments use NCSN++ on CIFAR-10 and ADM on ImageNet 64 × 64, with iCT-deep doubling the number of residual blocks per resolution.
- All models use RAdam with learning rate 0.0001, while training lasts 400,000 iterations on CIFAR-10 and 800,000 on ImageNet 64 × 64.
- Pseudo-Huber losses smoothly interpolate between ℓ1 and squared ℓ2 metrics, and their parameter updates have lower variance than squared-ℓ2 updates.
- Two-step sampling uses intermediate noise levels of 0.821 for CIFAR-10 and 1.526 for ImageNet 64 × 64 with iCT.For iCT-deep, the corresponding values are 0.661 and 0.973.