PaperlayerView exact revision on arXiv

Source-linked AI summary

Consistency Models Made Easy

Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, J. Zico Kolter

arXiv:2406.14548v2cs.LGcs.CV

TL;DR

Consistency models offer fast sampling but remain costly and difficult to train. The paper introduces ECT, which progressively tightens consistency starting from pretrained diffusion models, and reports improved efficiency and quality alongside classic power-law scaling of CM resources.

  • Problem

    Consistency models provide fast generation but remain substantially more time-consuming and complex to train than similarly capable diffusion models, limiting their uptake.

  • Method

    ECT progressively tightens the consistency condition from a pretrained diffusion model through a continuous-time pretraining-plus-tuning scheme.

  • Results

    ECT uses 1/4–1/3 of the computational resources of iCT while achieving superior or surpassing sample quality across reported ImageNet and CIFAR-10 comparisons.

  • Takeaways & Limitations

    ECT enables scaling studies showing classic power-law behavior for CM model size, FLOPs, and training compute.

  • Takeaways & Limitations

    ECT requires a dataset for tuning, which can limit use when bespoke models’ training data are unavailable.

Abstract

from arXiv · show

Consistency models (CMs) offer faster sampling than traditional diffusion models, but their training is resource-intensive. For example, as of 2024, training a state-of-the-art CM on CIFAR-10 takes one week on 8 GPUs. In this work, we propose an effective scheme for training CMs that largely improves the efficiency of building such models. Specifically, by expressing CM trajectories via a particular differential equation, we argue that diffusion models can be viewed as a special case of CMs. We can thus fine-tune a consistency model starting from a pretrained diffusion model and progressively approximate the full consistency condition to stronger degrees over the training process. Our resulting method, which we term Easy Consistency Tuning (ECT), achieves vastly reduced training times while improving upon the quality of previous methods: for example, ECT achieves a 2-step FID of 2.73 on CIFAR10 within 1 hour on a single A100 GPU, matching Consistency Distillation trained for hundreds of GPU hours. Owing to this computational efficiency, we investigate the scaling laws of CMs under ECT, showing that they obey the classic power law scaling, hinting at their ability to improve efficiency and performance at larger scales. Our code (https://github.com/locuslab/ect) is publicly available, making CMs more accessible to the broader community.

1 INTRODUCTION

Consistency models enable high-quality generation in very few evaluations but are difficult and costly to train. ECT addresses this by progressively tightening consistency from pretrained diffusion models, improving efficiency, performance, and the study of CM scaling.

  • Motivation: Consistency models map noise and sampling-trajectory points to the same data point, enabling high-quality generation in 1–2 model evaluations.They are closely related to diffusion models and target faster generation.
  • Motivation: Training consistency models remains time-consuming, complex, and substantially more expensive than training similarly capable diffusion models.These costs and hyperparameter choices have limited community uptake.
  • Easy Consistency Tuning: ECT progressively tightens the consistency condition to bootstrap pretrained diffusion models into 1-step consistency models without extra frozen teachers.Its continuous-time schedule includes diffusion pretraining as a special stage.
  • Results: 1/4–1/3 of iCT’s computational resources are required by ECT’s pretraining-plus-tuning scheme, while tuning typically accounts for 10% or less of total cost.ECT also achieves superior 1-step and 2-step ImageNet quality and surpasses previous CIFAR-10 methods in 2-step quality.
  • Scaling and design: ECT reveals classic power-law scaling for model size, FLOPs, and training compute, while identifying scenarios favoring smaller few-step models over larger 1-step models.The efficient tuning stage also serves as a proxy for exploring CM design choices.
  • Contributions: The paper contributes ECT, the first CM scaling study, and improved continuous-time schedules and weighting functions.These contributions target efficiency, performance, and CM design-space exploration.

2 PRELIMINARIES

Diffusion models define noisy trajectories and probability-flow ODEs, while consistency models learn mappings from noisy states to clean data along those trajectories. Their training enforces consistency between adjacent discretized trajectory points, with parameterizations that satisfy the zero-noise boundary condition by design.

  • Diffusion Models: Diffusion models add Gaussian noise with increasing standard deviation σ(t) to data distributions over time.The noise level ranges from σmin at t=0 to σmax at t=T.
  • Diffusion Models: The reverse-time SDE generates samples from noise, and its probability-flow ODE preserves the same marginal probability densities.The ODE provides the trajectory framework used by consistency models.
  • Diffusion Models: With σ(t)=t, the probability-flow ODE uses the score function, while the denoising parameterization treats time and noise level interchangeably.The denoising function predicts clean image x0 from noisy image xt.
  • Consistency Models: Consistency models learn a function that maps noisy images back to clean images along the probability-flow ODE’s bijective data–noise mapping.The function must satisfy a boundary condition at t=0.
  • Consistency Models: The parameterization fθ(xt,t)=cskip(t)xt+cout(t)Fθ(xt,t) guarantees the boundary condition through cskip(0)=1 and cout(0)=0.Fθ is trainable, while cskip and cout are time-dependent scaling factors.
  • Training objective: CM training discretizes the probability-flow ODE into N−1 subintervals and minimizes a metric between adjacent trajectory points.Earlier formulations used LPIPS, uniform weighting, teacher-based or unbiased score estimation, and later iCT modifications such as Pseudo-Huber metrics and adaptive curricula.

3 PROBING CONSISTENCY MODELS

The paper formulates consistency training in differential form, identifies why enforcing infinitesimal consistency from scratch is unstable, and proposes ECT to progressively tighten the condition from diffusion pretraining to fast-sampling consistency models.

  • Differential consistency: Differential consistency requires the model output to remain constant along the trajectory, together with a boundary condition matching the clean image at zero noise.The differential form alone admits collapsed constant solutions, so the boundary condition is necessary.
  • Differential consistency: Finite-difference training compares predictions at two noise levels generated from the same clean image and shared perturbation direction.The training objective aligns the higher-variance prediction with a stop-gradient self-teacher at lower variance.
  • The curse of consistency: As ∆t approaches zero, each interval becomes easier but the number of intervals increases, so accumulated consistency errors can slow convergence or worsen prediction error.The relevant error bound depends on the product of the number of intervals and the maximum interval error.
  • The curse of consistency: Experiments indicate that strict infinitesimal consistency training from scratch can converge slowly because errors accumulate across many discretization intervals.Figure 1 characterizes this regime as slower and less stable as N increases.
  • Easy Consistency Tuning: ECT combines an iteration-dependent schedule with adaptive weighting that controls gradient variance and upscales gradients when prediction differences become small.The adaptive factor is designed to prevent vanishing gradients during fine-grained learning.
  • Easy Consistency Tuning: ECT starts from diffusion pretraining and gradually shrinks ∆t, treating diffusion training as loose-discretization consistency training before tightening the condition.The initialization also supplies useful self-teacher targets and helps avoid trivial solutions.
  • Easy Consistency Tuning: The resulting pretraining-plus-tuning scheme reduces computational cost relative to training consistency models from scratch and uses tuning to explore design choices for the full pipeline.The paper presents tuning insights as transferable to improve pretraining.

4 EXPERIMENTS

The experiments evaluate ECT across datasets, training schemes, compute budgets, model scales, and inference steps. ECT achieves strong few-step quality with substantially lower training compute and exhibits favorable scaling behavior.

  • Experimental setup: ECT is evaluated on CIFAR-10 and ImageNet 64×64 using FID, FDDINOv2, NFEs, training FLOPs, inference cost, and generative performance.
  • Comparison of training schemes: 2-step ECM matches Score SDE-deep in FID using 1/1000 of its inference cost and latency, and outperforms EDM with DPM-Solver-v3 at NFE=10.
  • Comparison of training schemes: 2-step ECM achieves FID 2.20 versus 2.93 for 2-step CD on CIFAR-10 while using around 1/3 of CD’s training compute.
  • Comparison of training schemes: ECT rivals iCT-deep using 1/4 of the overall training compute, 1/8 of the tuning-stage compute, and 1/2 of the model size.
  • Scaling laws: ECT follows classic power-law scaling with training compute, model size, and model FLOPs, with stronger compute-quality correlations for 1-step and 2-step generation.The correlations between log training compute and log FDDINOv2 are −0.9940 for 1-step and −0.9996 for 2-step generation, both with p-values < 10−4.
  • Scaling laws: ECT achieves better 2-step generation than state-of-the-art CMs using 33% of iCT’s overall computational budget, while its tuning stage accounts for 0.60%–1.91% of pretraining compute.The tuning stage uses 12.8M training images, with the percentage varying by model size.
  • Scaling laws: On CIFAR-10 and ImageNet 64×64, smaller 2-step models can surpass larger 1-step models, such as 498M ECM-M versus 1.1B ECM-XL.

5 RELATED WORK

Related work includes consistency models, diffusion distillation, distribution matching, and fast samplers. These approaches differ in whether they train directly, match trajectories or distributions, or reduce sampling steps without additional training.

  • Diffusion Distillation: Diffusion distillation uses a pretrained diffusion model as a teacher and trains a student to reproduce its sample quality in fewer steps.
  • Diffusion Distillation: Trajectory-matching methods learn teacher-trajectory points, including offline distillation that uses synthetic data generated from a pretrained diffusion model.
  • Diffusion Distillation: Consistency distillation uses adjacent teacher points but relies on a fixed sampler-derived schedule, which can introduce discretization errors and limits quality to the pretrained diffusion model.
  • Diffusion Distillation: Distribution-matching methods minimize divergence between data and model distributions using score distillation or adversarial losses.
  • Diffusion Distillation: Training-based acceleration methods require additional post-pretraining procedures, which can be computationally intensive.
  • Fast Samplers for Diffusion Models: Fast samplers reduce diffusion sampling steps without training by using advanced solvers to lower discretization error when simulating SDEs or ODEs.

6 CONCLUSION

The paper presents ECT as an efficient route to training consistency models and uses it to study their scaling and design space. The released code supports further prototyping, study, and deployment.

  • ECT enables state-of-the-art few-step generative capabilities at minimal tuning cost and benefits from scaling.
  • The authors release code to support future prototyping, studying, and deploying consistency models.

7 LIMITATIONS

ECT requires a dataset to tune diffusion models into consistency models, limiting use when bespoke training data are unavailable. A data-free ECD variant is discussed, but it applies only to distillation settings.

  • ECT learns the consistency condition on a given dataset through a self teacher, unlike data-free methods that use a frozen diffusion teacher.
  • ECT may be limited when bespoke models’ training data are unavailable publicly.
  • The data-free ECD variant generates synthetic data from a consistency model on the fly but applies only to distillation settings.

BROADER IMPACTS AND ETHICS STATEMENT

ECT is presented as a computationally efficient way to train high-quality generative models, with potential benefits for artists and creators. The stated impact is faster content production through reduced computational costs.

  • ECT uses a small fraction of current consistency-model training and diffusion-distillation requirements while producing state-of-the-art few-step generators.
  • Reduced computational costs may help artists and creators produce content more efficiently.

REPRODUCIBILITY STATEMENT

The paper provides reproducibility materials and details design choices for ECT’s mapping, weighting, metric, and training schedule. These choices aim to control gradient balance and progressively tighten the consistency condition.

  • The authors provide experimental settings, hyperparameters, and source code to support reproducing the reported results.
  • ECT progressively shrinks ∆t toward dt through an exponentially decreasing parameterization during training.
  • The mapping function adjusts r relative to t because prediction difficulty varies across noise levels, which can otherwise imbalance gradient flows.
  • Downscaling r/t near t = 0 balances gradient flow across noise levels and controls gradient variance.
  • The pseudo-Huber differential decomposes into adaptive weighting and squared-L2 components, motivating exploration of adaptive weighting terms.
  • ECT’s schedule differs from iCT’s discrete curriculum, which partitions [0, T] into consecutive intervals and increases their number during training.

B EXPLORING DESIGN SPACE & SCALING OF CONSISTENCY MODELS

ECT enables broad exploration of consistency-model weighting, regularization, training schedules, and scaling. The experiments find no universal weighting function, show benefits from adaptive weighting and tuned dropout, and report power-law scaling with strong few-step performance.

  • ECT’s efficiency enables low-cost exploration of weighting functions, training schedules, and regularization for consistency models.
  • There is no universal timestep weighting that outperforms all alternatives across datasets, models, and target metrics.
  • Adaptive per-sample weighting improves results on both CIFAR-10 and ImageNet 64×64 by controlling gradient variance.
  • On ImageNet 64×64, increasing dropout from 0.10 to 0.40 reduces 2-step FID from 4.53 to 3.24.
  • A dropout rate tuned for one weighting function transfers well to other weighting functions and, for ImageNet 64×64, across model sizes.
  • Within 1 hour on one A100 GPU, ECT achieves a 2-step FID of 2.73, outperforming Consistency Distillation’s 2-step FID of 2.93.
  • ECT-pretrained EDM achieves FDDINOv2 of 150.39 unconditional and 121.05 class-conditional, versus 168.17 and 145.20 for the EDM baseline.
  • The reported Pareto frontier exhibits power-law scaling, while configurations not optimized for the current compute budget can deviate from that trend.

C EXTENSION TO CONSISTENCY DISTILLATION.

The authors extend ECT to continuous-time consistency distillation, yielding Easy Consistency Distillation (ECD), and also examine a data-free variant. Continuous-time distillation improves standard CD, while ECD and ECT trade advantages according to batch size, model scale, and data availability.

  • Continuous-time Consistency Distillation: ECD replaces CD’s fixed discrete schedule with a continuous-time schedule to reduce discretization error in approximating differential consistency.The method extends ECT’s continuous-time schedule to consistency distillation.
  • Continuous-time Consistency Distillation: Continuous-time distillation improves standard CD on ImageNet 64×64 under the reported pretrained EDM2-S baseline.The comparison uses the same general consistency-distillation setting described for ImageNet 64×64.
  • ECD versus ECT: ECD can outperform ECT at smaller batch sizes because teacher models provide lower-variance score estimates than ECT’s Monte Carlo estimation.The reported advantage is associated with resource-constrained settings and smaller models.
  • ECD versus ECT: The ECD–ECT performance gap narrows at larger model scales, making computational resources and model size relevant when choosing between them.The paper states that scaling compensates for limited training budgets and variance.
  • Data-Free ECD: Data-free ECD generates synthetic x0 samples from the consistency model itself instead of sampling x0 from a dataset.The resulting self-synthetic data are used in the same ECD training step.
  • Data-Free ECD: 4.38 1-step FID and 2.77 2-step FID are achieved by data-free ECD on ImageNet 64×64, comparable to data-dependent ECD and ECT at the same budget.This supports data-free ECD as an alternative when large datasets are limited or unavailable.

D EXPERIMENTAL DETAILS

The experiments initialize ECMs from pretrained diffusion models, evaluate them with FID and FDDINOv2, and vary model, data, compute, and sampling configurations. The setup emphasizes efficient tuning, broad scaling experiments, and reproducible implementation details.

  • Model Setup: ECMs are initialized from pretrained EDM for CIFAR-10 and pretrained EDM2 models ranging from 280M to 1.1B parameters for ImageNet 64×64.The initialization differs by dataset and conditionality.
  • Computational Cost: ECT’s ImageNet 64×64 tuning stage uses 0.39% of the iCT training budget and 0.60%–1.91% of EDM2 pretraining compute, depending on model size.The exact per-model resource requirements are reported in Table 8.
  • Training Details: Experiments use RAdam on CIFAR-10 and Adam on ImageNet 64×64, with dataset-specific beta settings and learning-rate schedules.CIFAR-10 omits decay, weight decay, and warmup, while ImageNet follows EDM2’s inverse-square-root decay.
  • Training Scope: The study uses full-parameter tuning, including for models exceeding 1B parameters, while leaving parameter-efficient fine-tuning for future work.The experiments are non-adversarial to focus on understanding consistency models and avoid inflated FID.
  • Sampling Details: The reported intermediate sampling schedules originally developed for iCT also perform well with ECMs, though determining optimal schedules remains open.Two-step sampling uses intermediate t = 0.821 for CIFAR-10 and t = 1.526 for ImageNet 64×64.
  • Evaluation Metrics: FID and FDDINOv2 are computed from 50k generated images for both CIFAR-10 and ImageNet 64×64.FDDINOv2 is used because it aligns better with human evaluation according to cited prior work.
  • Scaling Experiments: Scaling experiments vary batch size, iterations, training-image budgets, model capacity, and FLOPs across CIFAR-10 and ImageNet 64×64.CIFAR-10 budgets range from 12.8M to 409.6M training images, while ImageNet models scale from 12.8M to 102.4M.

E QUALITATIVE RESULTS

The paper provides randomly generated two-step samples from ECMs trained on CIFAR-10 and ImageNet 64×64. The visualizations organize samples by class and, for CIFAR-10, compare training-data and model-scale conditions.

  • Qualitative Samples: Figures 8 and 9 show randomly generated 2-step samples from ECMs trained on CIFAR-10 and ImageNet 64×64, respectively.The samples are presented as qualitative visualizations of two-step generation.
  • CIFAR-10: Each row in Figure 8 corresponds to a different CIFAR-10 class.The figure is class-conditional.
  • ImageNet 64×64: Each row in Figure 9 corresponds to a different ImageNet 64×64 class.The figure is class-conditional.
Loading 2406.14548v2…