PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

Cheng Lu, Yang Song

arXiv:2410.11081v2cs.LGstat.ML

TL;DR

Existing consistency models face discretization and timestep-scheduling issues, while continuous-time variants have suffered from unstable training. The paper introduces TrigFlow and targeted parameterization, architecture, and objective improvements, reaching 1.5 billion parameters and narrowing the FID gap with state-of-the-art diffusion models to within 10% using two-step generation.

  • Problem

    Most existing CMs use discretized timesteps that introduce hyperparameters and discretization errors, while continuous-time CMs have faced training instability.

  • Method

    The paper unifies EDM and Flow Matching with TrigFlow and stabilizes continuous-time CM training through parameterization, architectural, and objective improvements.

  • Results

    1.5 billion parameters is reached on ImageNet 512×512, while two-step sCMs narrow the FID gap with state-of-the-art diffusion models to within 10%.

  • Takeaways & Limitations

    The resulting sCMs scale predictably across datasets and model sizes and outperform other few-step sampling approaches at large scales.

  • Takeaways & Limitations

    The method can still produce noticeable artifacts, retains a small two-step gap to state-of-the-art diffusion models, and requires further validation for high-resolution or fine-detail generation.

Abstract

from arXiv · show

Consistency models (CMs) are a powerful class of diffusion-based generative models optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretization errors. While continuous-time formulations can mitigate these issues, their success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffusion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous-time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512x512. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR-10, 1.48 on ImageNet 64x64, and 1.88 on ImageNet 512x512, narrowing the gap in FID scores with the best existing diffusion models to within 10%.

1 INTRODUCTION

Consistency models enable few-step generation but discrete-time training introduces timestep-grid and discretization issues, while continuous-time training has been unstable. The paper proposes a simplified, stabilized framework and scales continuous-time CMs to large models with strong two-step results.

  • Diffusion models achieve strong generative results but often require dozens to hundreds of sampling steps per sample.
  • Consistency models avoid synthetic-dataset supervision and adversarial training, and can be trained from scratch with consistency training.
  • Discrete-time CMs introduce discretization errors and require careful timestep scheduling, whereas continuous-time CMs face training instability.
  • TrigFlow unifies EDM and Flow Matching, while architectural and objective changes target the root causes of continuous-time CM instability.
  • 1.5 billion parameters is the scale reached by sCMs on ImageNet 512×512, while two-step generation narrows the FID gap with state-of-the-art diffusion models to within 10%.
  • Continuous-time CMs improve sample quality as adjacent timestep gaps shrink, and sCMs produce more diverse samples with better guidance compatibility than VSD.

2 PRELIMINARIES

Diffusion models learn to reverse a progressively noised data distribution, while consistency models learn direct noisy-to-clean mappings along the probability-flow ODE. Continuous-time training removes solver-based discretization errors but has suffered from severe optimization instability.

  • 2.1 DIFFUSION MODELS: Diffusion models perturb clean data x0 with Gaussian noise to form xt = αtx0 + σtz, with larger t indicating greater noise.
  • 2.1 DIFFUSION MODELS: EDM uses αt = 1 and σt = t, and parameterizes the diffusion model with manually designed coefficients that equalize objective variance across timesteps.
  • 2.2 CONSISTENCY MODELS: A consistency model maps noisy xt directly to clean x0 in one step by following the probability-flow ODE trajectory and must satisfy fθ(x, 0) ≡ x.
  • 2.2 CONSISTENCY MODELS: Discrete-time CMs enforce consistency between adjacent finite timesteps, but their results depend on ∆t, annealing schedules, and numerical ODE-solver errors.
  • 2.2 CONSISTENCY MODELS: As ∆t approaches zero, the discrete-time objective’s gradient converges to a tangent-based continuous-time expression along the probability-flow ODE.
  • 2.2 CONSISTENCY MODELS: Continuous-time CMs avoid ODE solvers and discretization errors, but prior work reports severe optimization instability that limits empirical performance and adoption.
  • 2.2 CONSISTENCY MODELS: Consistency distillation uses a pretrained diffusion model’s probability-flow ODE, whereas consistency training learns CMs from scratch.

3 SIMPLIFYING CONTINUOUS-TIME CONSISTENCY MODELS

TrigFlow simplifies diffusion and consistency-model formulations by combining EDM properties with a Flow Matching special case. The resulting parameterization supports simple probability-flow and CM expressions while motivating stabilized continuous-time training.

  • Previous EDM-based CM coefficients improve training efficiency but have complex relationships with t and σd that complicate theoretical analysis.
  • TrigFlow sets cskip(t) = cos(t), cout(t) = −σd sin(t), and cin(t) ≡ 1/σd while retaining EDM properties.
  • TrigFlow is a special case of Flow Matching and v-prediction, combining the advantages of both formulations with simpler diffusion, ODE, objective, and CM expressions.
  • Under TrigFlow, the noisy sample is xt = cos(t)x0 + sin(t)z for t ∈ [0, π/2], and the prior has distribution N(0, σd^2I).
  • The TrigFlow diffusion model uses Fθ(xt/σd, cnoise(t)), with cnoise(t) transforming time to facilitate conditioning.
  • The TrigFlow CM is parameterized as a single-step solution of the probability-flow ODE using a first-order ODE solver while enforcing the boundary condition.

4 STABILIZING CONTINUOUS-TIME CONSISTENCY MODELS

The paper stabilizes continuous-time consistency-model training by targeting unstable parameterization, time embeddings, network normalization, and gradient objectives. These changes improve training stability and enable continuous-time models to outperform discrete-time variants across discretization settings.

  • Root causes of instability: Continuous-time CM instability originates primarily from the time-derivative term involving sin(t)∂tFθ−, while other tangent components are relatively stable.The analysis decomposes the tangent function and identifies the time derivative as the unstable component.
  • Parameterization: Setting cnoise(t) = t avoids the EDM time transformation’s blow-up as t approaches π/2.Under EDM, sin(t) · ∂tcnoise(t) = 1/cos(t), which becomes unbounded near π/2.
  • Time embeddings: Positional time embeddings replace high-scale Fourier embeddings because their derivatives are more stable for both noisy samples and time.Large Fourier scales produce larger, more oscillatory derivatives and cause instability, whereas TrigFlow with positional embeddings remains stable.
  • Network architecture: Adaptive double normalization applies pixel normalization to adaptive scale and bias terms, retaining diffusion-model expressivity while removing observed CM-training instability.It modifies AdaGN as y = norm(x)⊙pnorm(s(t))+pnorm(b(t)).
  • Training objectives: Tangent normalization, clipping, adaptive weighting, and tangent warmup explicitly control gradient variance and spikes during CM training.Adaptive weighting minimizes loss variance across time steps, while warmup ramps the tangent coefficient from 0 to 1 over the first 10k iterations.
  • Empirical comparison: Continuous-time CMs significantly outperform discrete-time CMs across all tested discretization counts, while discrete-time quality degrades after N > 1024.Increasing N initially reduces discretization error, but excessive N introduces numerical-precision issues.

5 SCALING UP CONTINUOUS-TIME CONSISTENCY MODELS

The paper scales continuous-time consistency models using numerical and architectural improvements, evaluating consistency training and distillation across datasets, resolutions, and model sizes. The resulting sCMs achieve strong few-step quality and predictable scaling, with sCD approaching teacher diffusion performance.

  • Comparison with VSD: sCD produces more diverse samples and is more compatible with guidance than VSD, which struggles at higher guidance levels.The comparison uses EDM2-M models trained on ImageNet 512×512 across varying guidance scales.
  • Training improvements: sCM training combines improved numerical precision, JVP computation, and large-scale attention support for continuous-time models.Tangent rearrangement alleviates FP16 overflow, while a Flash Attention-style algorithm computes attention and its JVP in one forward pass.
  • Experimental setup: sCT and sCD are trained on CIFAR-10, ImageNet 64×64, and ImageNet 512×512 using FID and diffusion-model settings.The experiments use consistency training and distillation, with classifier-free guidance applied during sCD distillation but not tested for sCT.
  • Compute: Less than 20% of teacher training compute produces sCD two-step samples comparable to the teacher, with high-quality samples after 20k finetuning iterations.The effective compute per sCD training iteration is approximately twice that of the teacher model.
  • Benchmarks: 2.06 FID on CIFAR-10, 1.48 on ImageNet 64×64, and 1.88 on ImageNet 512×512 bring two-step sCM within 10% of teacher diffusion FIDs.The corresponding teacher FIDs are 2.01, 1.33, and 1.73.
  • Scaling study: 1.5 billion parameters mark the largest continuous-time consistency models trained, while both sCT and sCD improve as model FLOPs increase.sCD scales predictably across model sizes and its two-step quality becomes comparable to the teacher with more sampling steps.

6 CONCLUSION

The proposed formulations, architectures, and objectives simplify and stabilize continuous-time consistency-model training, enabling scaling to 1.5 billion parameters. Two-step generation narrows the teacher-model FID gap to within 10%.

  • 6 CONCLUSION: The method enables smooth scaling to 1.5 billion parameters on ImageNet 512×512 and predictable sample-quality improvements across datasets and model sizes.Ablations confirm the effectiveness of TrigFlow, tangent normalization, and adaptive weighting.
  • 6 CONCLUSION: Two-step generation narrows the FID gap with the teacher model to within 10% while outperforming other few-step sampling approaches at large scales.The comparison is against state-of-the-art diffusion models requiring significantly more sampling steps.

DISCUSSIONS AND LIMITATIONS

The discussion identifies limitations in latent-space consistency training, computational comparisons, image artifacts, architecture compatibility, and dependence on pretrained diffusion models.

  • Latent-space limitation: sCT is less effective than sCD in latent spaces, especially across ImageNet 512×512 model scales.The authors attribute this to higher CT variance and hypothesize that better encoder-decoder mappings could improve conditioning.
  • Computation costs: Forward-mode JVP computation gives continuous-time consistency models training compute and memory requirements similar to discrete-time counterparts.The comparison is with discrete-time models performing two forward passes per iteration.
  • Limitations: The method can produce noticeable ImageNet artifacts, retains a small two-step gap to state-of-the-art diffusion models, and may not suit some diffusion-model architectures.FID also does not capture all semantic details, and sCD relies heavily on a pretrained diffusion model.
  • Additional material: The appendix provides derivations, experimental details, training algorithms, and a discussion of TrigFlow and adaptive variational score distillation.The detailed sCM algorithm covers both sCT and sCD.

A TRAINING ALGORITHM OF SCM

Algorithm 1 defines sCM training for consistency training and distillation, initializing from data and a pretrained diffusion model and applying tangent stabilization before optimization.

  • Algorithm overview: The algorithm distinguishes consistency training sCT from consistency distillation sCD.Algorithm 1 is presented as the detailed training procedure for simplified and stabilized continuous-time consistency models.
  • Inputs and initialization: Training initializes model parameters from a pretrained diffusion model and uses data, noise scale, weighting, learning-rate, proposal, and warmup inputs.The algorithm also includes a constant c and warmup iteration H.
  • Optimization steps: The procedure applies tangent warmup, JVP rearrangement, tangent normalization, and adaptive weighting before gradient updates until convergence.Tangent normalization divides by the tangent norm plus c, while adaptive weighting is included in the training loss.

B.1 DERIVATIONS

The derivations show that unit-variance parameterization makes diffusion formulations and noise schedules equivalent up to training weighting. TrigFlow expresses this unified view with trigonometric schedules while preserving compatible sampling procedures.

  • The framework separates diffusion models into parameterization, noise schedule, training weighting, and timestep proposal distribution.
  • Under the unit variance principle, all four components can reproduce previous diffusion-model formulations within one general framework.
  • Different noise schedules are equivalent after rescaling, with their remaining influence represented through the training weighting.
  • Flow matching’s optimal transport path can be converted to other schedules; its empirical advantage is attributed to weighting and sampling choices rather than path straightness itself.
  • TrigFlow uses trigonometric transformations to make the per-time-step ℓ2 loss independent of noise-schedule choice while retaining schedule effects in the integral weighting.
  • TrigFlow unifies EDM, flow matching, and velocity prediction, and supports DDIM, DPM-Solver, and DPM-Solver++ samplers.

C.1 DERIVATIONS

The derivations formulate VSD in TrigFlow and connect its generator, diffused distributions, and score functions. An adaptive weighting variant uses learned weighting together with a log-normal timestep proposal.

  • VSD trains a generator to map Gaussian noise to the data distribution by matching diffused generator and data distributions.
  • The TrigFlow forward process uses xt = cos(t)x0 + sin(t)z and defines a corresponding velocity vt = cos(t)z − sin(t)x0.
  • The ground-truth diffusion model is related to the conditional velocity E[vt|xt], which is connected to the score function of the diffused data distribution.
  • Adaptive VSD replaces manually selected weighting with a learned weighting network and uses a log-normal proposal distribution for timesteps.

D ADAPTIVE WEIGHTING FOR DIFFUSION MODELS, CONSISTENCY MODELS AND VARIATIONAL SCORE DISTILLATION

The section recasts diffusion, consistency-model, and VSD objectives into a common neural-network–target form. Adaptive weighting is used to reduce loss variance across time, with related improvements extended to discrete-time consistency distillation.

  • Diffusion, consistency-model, and VSD objectives can be rewritten as inner products between a neural network and a same-dimensional target function.
  • Adaptive weighting trains a weighting network to estimate loss norms and balance variance across different timesteps.
  • For consistency models, prior weighting λ(t) = 1/(σd tan(t)) reduces target variance and is empirically critical for performance.
  • The adaptive approach reduces manual weighting choices and significantly lowers tuning complexity for diffusion models, consistency models, and VSD.
  • The same adaptive weighting, tangent normalization, and tangent warmup principles are applied to discrete-time consistency distillation.
  • Discrete-time sCM training uses tangent warmup, JVP rearrangement, tangent normalization, and adaptive weighting in its optimization loop.

E.2 EXPERIMENTS OF DISCRETE-TIME SCM

The section compares discrete- and continuous-time consistency models while describing training settings and a Flash Attention-style Jacobian-vector product implementation. Continuous-time models outperform discrete-time variants across discretization levels, and the proposed sCT converges faster than ECT on CIFAR-10.

  • Comparison of continuous- and discrete-time CMs: Continuous-time CMs significantly outperform discrete-time CMs across all tested discretization-step counts.Increasing discrete-time steps initially reduces discretization error, but performance degrades beyond N > 1024 because of numerical precision issues.
  • Comparison with ECT: sCT significantly outperforms ECT in 1-step CIFAR-10 FID throughout training.The comparison uses the same DDPM++ architecture, dropout rate, batch size, and other hyperparameters as ECT, with sCT-specific parameterization and settings.
  • Flash Attention JVP: A recursive Flash Attention-style algorithm computes both attention outputs and their Jacobian-vector products in one loop.The method accesses attention-matrix memory once and avoids saving intermediate activations, reducing GPU memory use.
  • Experimental settings: The experiments train teacher diffusion models on CIFAR-10, ImageNet 64×64, and ImageNet 512×512 using TrigFlow parameterization, positional time embeddings, and adaptive double normalization.The ImageNet 512×512 pipeline uses bicubic resizing, center cropping, disabled horizontal flips, and latent encoding through stable diffusion VAE2.
  • Experimental settings: The sampling setup initializes consistency models from teacher EMA parameters and uses an intermediate t = 1.1 for 2-step sampling.The starting time is chosen to match EDM and EDM2, with σmax = 80.
  • Training stability: The proposal distribution uses small Pmean values near clean data to improve training stability and final performance.The stated rationale is to reduce training error near t = 0 and thereby reduce accumulated errors.

H ADDITIONAL SAMPLES

The additional-sample figures present uncurated ImageNet 512×512 outputs from sCD-XXL using either one or two sampling steps. They cover multiple displayed sample groups without reporting quantitative comparisons.

  • 1-step samples: Figures 8, 10, and 12 show uncurated 1-step samples from sCD-XXL trained on ImageNet 512×512.
  • 2-step samples: Figures 9, 11, and 13 show uncurated 2-step samples from sCD-XXL trained on ImageNet 512×512.
Loading 2410.11081v2…