PaperlayerView exact revision on arXiv

Source-linked AI summary

Consistency Models

Yang Song, Prafulla Dhariwal, Mark Chen, Ilya Sutskever

arXiv:2303.01469v2cs.LGcs.CVstat.ML

TL;DR

Diffusion models are powerful but slow because sampling is iterative and computationally costly. The paper introduces consistency models, which directly map noise to data in one step while retaining multistep sampling and zero-shot editing. They outperform existing diffusion-distillation methods in few-step generation and form a standalone generative-model family when trained without diffusion models.

  • Problem

    Diffusion models’ iterative sampling typically requires 10–2000 times more compute than single-step models, causing slow inference and limiting real-time applications.

  • Method

    Consistency models learn self-consistent mappings from probability-flow ODE trajectory points to their origins and can be trained by distilling diffusion models or in isolation.

  • Results

    Consistency models outperform existing diffusion-distillation methods across datasets and achieve one-step FIDs of 3.55 on CIFAR-10 and 6.20 on ImageNet 64 × 64.

  • Takeaways & Limitations

    The models provide one-step and few-step generation plus zero-shot editing, while standalone training establishes them as an independent generative-model class.

  • Takeaways & Limitations

    The multistep sampler assumes FID is unimodal in each next time point, and better time-point-selection strategies are left for future work.

Abstract

from arXiv · show

Diffusion models have significantly advanced the fields of image, audio, and video generation, but they depend on an iterative sampling process that causes slow generation. To overcome this limitation, we propose consistency models, a new family of models that generate high quality samples by directly mapping noise to data. They support fast one-step generation by design, while still allowing multistep sampling to trade compute for sample quality. They also support zero-shot data editing, such as image inpainting, colorization, and super-resolution, without requiring explicit training on these tasks. Consistency models can be trained either by distilling pre-trained diffusion models, or as standalone generative models altogether. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one- and few-step sampling, achieving the new state-of-the-art FID of 3.55 on CIFAR-10 and 6.20 on ImageNet 64x64 for one-step generation. When trained in isolation, consistency models become a new family of generative models that can outperform existing one-step, non-adversarial generative models on standard benchmarks such as CIFAR-10, ImageNet 64x64 and LSUN 256x256.

1. Introduction

Diffusion models generate high-quality data across multiple fields but require costly iterative sampling. Consistency models address this bottleneck by mapping noise to data in one step while retaining multistep quality trade-offs and zero-shot editing.

  • Diffusion models have achieved major success in image, audio, and video generation.
  • 10–2000 times more compute is typically required for diffusion sampling than for single-step generative models, limiting inference speed and real-time applications.The iterative process can improve sample quality and supports zero-shot inverse-problem solving.
  • Consistency models map points on probability-flow ODE trajectories to their starting points, enforcing self-consistent outputs.Random noise endpoints can therefore be converted into data samples with one network evaluation.
  • Consistency models retain multistep sampling, allowing compute to be exchanged for sample quality and enabling zero-shot data editing.
  • Consistency models use either diffusion distillation or standalone training, and neither approach requires adversarial training.The methods permit flexible neural-network architectures.
  • 3.55 and 6.20 are the reported one-step FIDs for CIFAR-10 and ImageNet 64 × 64, respectively, with consistency models outperforming existing diffusion-distillation methods.The models also achieve 2.93 and 4.70 for two-step generation on those datasets.

2. Diffusion Models

Continuous-time diffusion models use stochastic dynamics and an equivalent probability-flow ODE to connect data and tractable noise. Sampling reverses this process numerically, but repeated score-model evaluations remain computationally expensive.

  • Diffusion models progressively perturb data into noise through Gaussian perturbations and generate samples through sequential denoising.
  • The probability-flow ODE has solution trajectories whose samples at time t follow the corresponding diffusion distributions.
  • The score function ∇log p_t(x) enables the diffusion model formulation as a score-based generative model.
  • The terminal diffusion distribution is typically chosen close to a tractable Gaussian distribution π(x).
  • Sampling initializes an empirical probability-flow ODE from Gaussian noise and solves it backward with a numerical ODE solver to obtain an approximate data sample.The solver commonly stops at a small positive time ϵ for numerical stability.
  • More than 10 score-model evaluation steps are still needed for competitive samples, while many distillation methods require expensive pre-collection of diffusion-model samples.

3. Consistency Models

Consistency models are designed around single-step generation while preserving iterative quality–compute trade-offs and zero-shot editing. They estimate a self-consistent mapping from probability-flow trajectories to their origins and can be trained through distillation or isolation.

  • Consistency Models: Consistency models support single-step generation at their core while allowing iterative generation for quality–compute trade-offs and zero-shot data editing.
  • Definition: A consistency function maps any point (x_t, t) on a probability-flow ODE trajectory to its endpoint x_ϵ.
  • Definition: Self-consistency requires points from the same probability-flow trajectory to produce identical outputs.
  • Parameterization: The boundary condition f(x_ϵ, ϵ) = x_ϵ makes the consistency function an identity at ϵ and is required for successful training.
  • Parameterization: Consistency models can be parameterized with differentiable skip connections that preserve differentiability at t = ϵ.
  • Sampling: One-step sampling evaluates f_θ(x̂_T, T) from Gaussian noise, while repeated denoising and noise injection can improve sample quality.
  • Sampling: The multistep sampler chooses time points greedily with ternary search, assuming FID is unimodal in each next time point.This assumption holds empirically in the experiments, while better strategies remain future work.
  • Zero-Shot Data Editing: Zero-shot consistency-model editing includes interpolation, denoising, inpainting, colorization, super-resolution, and stroke-guided editing without explicit task training.

4. Training Consistency Models via Distillation

Consistency distillation trains a consistency model by matching outputs on adjacent points along probability-flow ODE trajectories generated from a pre-trained score model. The resulting framework uses one-step ODE updates and EMA-stabilized online and target networks, with theoretical accuracy improving as solver step size decreases.

  • Consistency Distillation: A pre-trained score model and numerical ODE solver generate adjacent trajectory points for distilling a consistency model.The solver estimates an earlier point from a later point using one discretization step.
  • Limitations: The work considers only one-step ODE solvers and leaves generalization to multistep solvers as future work.A broader solver result is also left for future work in the theoretical discussion.
  • Consistency Distillation: The distillation loss minimizes output differences between the consistency model evaluated at adjacent points on the same trajectory.Training samples data points, noise levels, and neighboring trajectory states, then compares the corresponding model outputs.
  • Optimization: Training uses stochastic gradient descent for the online network and an exponential moving average update for the target network.The EMA target and stop-gradient update are reported to stabilize training and improve final performance.
  • Theory: Under Lipschitz and bounded-local-error assumptions, sufficiently small ODE solver steps can make the estimated consistency model arbitrarily accurate.The theoretical justification uses asymptotic analysis of the solver discretization and consistency loss.
  • Limitations: Continuous-time distillation extensions avoid specifying discrete time steps but require Jacobian-vector products and forward-mode automatic differentiation.These implementation requirements may not be well-supported in some deep learning frameworks.

5. Training Consistency Models in Isolation

Consistency models can also be trained independently of pre-trained diffusion models by replacing the score model with an unbiased score estimator derived from noisy data. The resulting consistency-training objective supports asymptotic analysis, while progressive schedules trade bias and variance during training.

  • Standalone Training: Consistency models trained without pre-trained diffusion models form an independent family of generative models.This distinguishes consistency training from existing diffusion distillation techniques.
  • Standalone Training: An unbiased estimator uses a clean data point and its Gaussian-noised version to estimate the score as −(x_t − x)/t^2.This estimator can replace the pre-trained diffusion model for Euler-based consistency training in the N → ∞ limit.
  • Objective: The consistency-training loss depends only on the online and target consistency networks and is agnostic to diffusion-model parameters.The loss is described as dominating the remainder term as the discretization becomes finer.
  • Optimization: Progressively increasing the number of time steps during training moves from lower-variance, higher-bias objectives toward higher-variance, lower-bias objectives.The schedule is paired with a changing EMA decay rate for best practical performance.
  • Limitations: Continuous-time consistency training removes the discrete objective’s bias and schedule functions but requires forward-mode automatic differentiation.The continuous-time formulation is established when the target network is stop-gradient from the online network.

6. Experiments

Experiments evaluate consistency distillation and consistency training across image datasets, training choices, sampling settings, and image-editing tasks. The results show strong few-step generation and zero-shot editing capabilities.

  • Experiments compare consistency models on CIFAR-10, ImageNet 64 × 64, and LSUN 256 × 256 using FID, Inception Score, and Precision.
  • Training consistency models: LPIPS, Heun’s solver, and N = 18 are the best reported configuration choices for consistency distillation on CIFAR-10.Heun’s solver uniformly outperforms Euler’s solver at the same N, while sufficiently large N values make performance insensitive to N.
  • Few-step image generation: Consistency distillation outperforms progressive distillation across datasets and sampling steps, except for single-step generation on Bedroom 256 × 256.
  • Training consistency models: Adaptive schedules for N and µ significantly improve consistency-training convergence speed and sample quality.Smaller N converges faster but produces worse samples, whereas larger N converges more slowly but produces better samples after convergence.
  • Few-step image generation: Consistency training outperforms VAEs and normalizing flows on CIFAR-10 and achieves quality comparable to one-step progressive-distillation samples without distillation.Samples from CT and EDM generated from the same initial noise share significant structural similarity.
  • Zero-shot image editing: Consistency models support zero-shot colorization, super-resolution, inpainting, interpolation, denoising, and stroke-guided image generation.These capabilities are demonstrated without explicit training on the corresponding editing tasks.

7. Conclusion

The conclusion presents consistency models as generative models designed for one-step and few-step generation. It reports advantages over diffusion distillation methods and most existing single-step non-adversarial models, while noting connections to methods in other fields.

  • Consistency models are specifically designed to support one-step and few-step generation.
  • Consistency distillation outshines existing diffusion-model distillation techniques on multiple image benchmarks and with small sampling iteration counts.
  • Standalone consistency models generate better samples than existing single-step generation models except GANs.
  • Consistency models share similarities with techniques from deep Q-learning and momentum-based contrastive learning, suggesting prospects for cross-field exchange of ideas and methods.

Appendices

The appendices define notation and develop the error analysis underlying consistency-model training. The proof uses Lipschitz consistency models, bounded ODE-solver error, and recursive error control.

  • A.1. Notations: The notation defines fθ(x, t) as a parameterized consistency model and f(x, t; ϕ) as the empirical probability-flow ODE consistency function.
  • A.1. Notations: The appendix specifies random sampling conventions for data points, discretization indices, and noisy states.
  • Error analysis: Theorem 1 assumes a Lipschitz consistency model and uniformly bounded local error from the numerical ODE solver.
  • Error analysis: The proof initializes the error at the first time point and derives a recursion relation for later error vectors.
  • Error analysis: The probability-flow ODE consistency function satisfies the initial-point parameterization used to establish the base case of the induction.
  • Error analysis: The derivation concludes by applying induction to the recursive error bound.

A.3. Consistency Training

Theoretical results justify consistency-training objectives, including continuous-time extensions and direct training without pre-trained diffusion models, under stated regularity assumptions. These extensions introduce practical and metric-specific limitations.

  • Theoretical Assumptions: Theoretical guarantees assume smooth models, bounded derivatives and weighting, bounded score-function moments, and an exact pre-trained score model.
  • Continuous-Time Extensions: Consistency training can be extended to infinitely many time steps under suitable conditions.
  • Practical Limitations: The continuous-time theory assumes an Euler solver, while analogous guarantees for more general solvers remain future work.
  • Practical Limitations: Continuous-time objectives require Jacobian-vector products and forward-mode automatic differentiation, which may be poorly supported in some frameworks.
  • Consistency Training: The pseudo-objective can be estimated without a pre-trained diffusion model, enabling direct consistency training.

B.3. Experimental Verifications

Experiments on CIFAR-10 compare metrics, solvers, objectives, and discretization choices for consistency distillation and training. LPIPS and stop-gradient continuous-time distillation perform particularly well, while continuous-time training requires stabilization.

  • Experimental Setup: The experiments compare consistency objectives using ℓ1, ℓ2, and LPIPS metrics across discrete- and continuous-time formulations.
  • Consistency Distillation: The stop-gradient continuous-time distillation objective outperforms the non-stop-gradient version for both LPIPS and ℓ2 metrics.
  • Consistency Distillation: LPIPS performs best across the evaluated distillation approaches.
  • Consistency Training: Continuous-time consistency training is initialized from a pre-trained EDM model to stabilize training, whereas discrete-time CT works with random initialization.
  • Consistency Training: LPIPS improves continuous-time consistency-training performance and continuous-time CT outperforms discrete-time CT with the same metric.

C. Additional Experimental Details

The experiments reuse established diffusion architectures and modify EDM skip connections to satisfy the consistency boundary condition. Training uses scheduled discretization and EMA settings across several datasets and distillation regimes.

  • Architectures: The experiments use NCSN++ for CIFAR-10 and architectures from Dhariwal and Nichol for ImageNet 64 × 64 and LSUN 256 × 256.
  • Architectures: Consistency models use EDM architectures with modified skip connections to satisfy the boundary condition.
  • Training Schedules: Training schedules progressively change discretization steps and EMA decay during optimization.
  • Training Setup: Distillation experiments use pre-trained EDM weights, while consistency-training models are randomly initialized in the general experimental setup.
  • Training Setup: Experiments use Rectified Adam without learning-rate decay, warm-up, or weight decay, together with EMA updates.

D. Additional Results on Zero-Shot Image Editing

Consistency models support zero-shot image editing through modified multistep sampling, covering colorization, super-resolution, inpainting, interpolation, denoising, and stroke-guided generation. The procedure injects reference-image information through masks and invertible transformations.

  • Editing Applications: The model demonstrates zero-shot colorization, super-resolution, inpainting, interpolation, denoising, and stroke-guided image generation.
  • Editing Procedure: Most editing tasks use a small modification of multistep sampling that repeatedly applies the consistency model and restores masked reference information.
  • Inpainting: Inpainting treats missing pixels as the masked region and uses the identity transformation.
  • Colorization: Colorization reformulates the task as inpainting after transforming grayscale and color information into a decoupled space.
  • Super-Resolution: Super-resolution uses patch-based downsampling together with an orthogonal transformation to separate conditional information.
  • Denoising: A single consistency model can denoise Gaussian-perturbed images whenever the noise scale lies in [ϵ, T].

E. Additional Samples from Consistency Models

The section presents additional consistency-model samples across CIFAR-10, ImageNet, LSUN Bedroom, and LSUN Cat, including editing and generation examples. It also shows colorization, super-resolution, inpainting, interpolation, denoising, and stroke-guided generation samples.

  • Additional samples cover CIFAR-10, ImageNet 64 ˆ 64, LSUN Bedroom 256 ˆ 256, and LSUN Cat 256 ˆ 256.
  • Colorization examples place grayscale images beside consistency-model outputs and ground-truth images.
  • Super-resolution examples compare 32 ˆ 32 inputs with 256 ˆ 256 consistency-model outputs and 256 ˆ 256 ground truth.
  • Inpainting, interpolation, denoising, and stroke-guided generation are illustrated through additional consistency-model examples.
  • Uncurated samples are shown for CIFAR-10, ImageNet 64 ˆ 64, LSUN Bedroom 256 ˆ 256, and LSUN Cat 256 ˆ 256, with corresponding samples sharing the same initial noise.
Loading 2303.01469v2…