Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion

Dongjun Kim; Chieh-Hsin Lai; Wei-Hsiang Liao; Naoki Murata; Yuhta Takida; Toshimitsu Uesaka; Yutong He; Yuki Mitsufuji; Stefano Ermon

Source-linked AI summary

Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion

Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, Stefano Ermon

arXiv:2310.02279v3Oct 1, 2023cs.LGcs.AIcs.CVstat.ML

TL;DR

Diffusion distillation accelerates sampling but lacks a reliable quality–compute trade-off, while score-based sampling remains constrained by solver discretization. CTM learns both scores and arbitrary PF ODE trajectory jumps in one framework, achieving state-of-the-art results and enabling deterministic or stochastic sampling control.

Problem
Distillation models can degrade in quality as NFE increases and lack a clear computational-budget–fidelity trade-off, while diffusion sampling remains costly.
Method
CTM jointly learns the PF ODE score and anytime-to-anytime trajectory jumps, combining distillation with denoising score matching and adversarial losses.
Results
CTM achieves state-of-the-art FID and likelihood results on CIFAR-10 and ImageNet 64 × 64, including CIFAR-10 FID 1.73 for conditional generation.
Takeaways & Limitations
CTM supports score-based ODE/SDE sampling, direct long trajectory jumps, and γ-sampling with controllable stochasticity.
Takeaways & Limitations
CTM poses risks of generating harmful or inappropriate content, requiring content filtering and moderation mechanisms.

Abstract

from arXiv · show

Consistency Models (CM) (Song et al., 2023) accelerate score-based diffusion model sampling at the cost of sample quality but lack a natural way to trade-off quality for speed. To address this limitation, we propose Consistency Trajectory Model (CTM), a generalization encompassing CM and score-based models as special cases. CTM trains a single neural network that can -- in a single forward pass -- output scores (i.e., gradients of log-density) and enables unrestricted traversal between any initial and final time along the Probability Flow Ordinary Differential Equation (ODE) in a diffusion process. CTM enables the efficient combination of adversarial training and denoising score matching loss to enhance performance and achieves new state-of-the-art FIDs for single-step diffusion model sampling on CIFAR-10 (FID 1.73) and ImageNet at 64x64 resolution (FID 1.92). CTM also enables a new family of sampling schemes, both deterministic and stochastic, involving long jumps along the ODE solution trajectories. It consistently improves sample quality as computational budgets increase, avoiding the degradation seen in CM. Furthermore, unlike CM, CTM's access to the score function can streamline the adoption of established controllable/conditional generation methods from the diffusion community. This access also enables the computation of likelihood. The code is available at https://github.com/sony/ctm.

1 INTRODUCTION

Diffusion models improve sample quality by learning scores but remain slow, while distillation accelerates sampling yet faces quality and controllability trade-offs. CTM unifies both approaches through flexible PF ODE trajectory jumps and auxiliary training losses.

Diffusion models generate high-quality samples by learning scores but require numerous evaluations during gradual denoising.
Advanced ODE and SDE solvers reduce evaluations, but intrinsic discretization error limits further improvements.
CTM jointly learns the PF ODE score integrand and integral, enabling anytime-to-anytime jumps between infinitesimally close and distant times.
γ-sampling alternates forward and backward trajectory jumps, with γ controlling stochasticity.
CTM combines distillation, denoising score matching, and adversarial losses, achieving state-of-the-art density estimation and image generation on CIFAR-10 and ImageNet 64 × 64.

2 PRELIMINARY

Diffusion sampling solves a reverse stochastic or deterministic process, typically through numerical PF ODE integration or learned distillation. These approaches trade computation, discretization accuracy, sample quality, and sampling variance differently.

A forward diffusion process adds noise, while reverse-time SDE sampling uses the score to recover data distributions.
The deterministic probability flow ODE provides a reverse-time counterpart whose solution matches the forward process marginals.
The denoiser E[x|x_t] equals x_t + t∇log p_t(x_t), providing an alternative expression for the score.
Neural denoisers are trained with denoising score matching and then used to form an empirical PF ODE.
Score-based sampling applies time-discretized ODE solvers, whereas distillation directly estimates the PF ODE integral with a neural network.
Distillation amortizes sampling cost but can degrade with increasing NFE and produces uncontrollable sample variance in multistep sampling.

3 CTM: AN UNIFICATION OF SCORE-BASED AND DISTILLATION MODELS

CTM represents PF ODE trajectory jumps for arbitrary time pairs, unifying score-based and distillation sampling. Its soft consistency objective and auxiliary losses support flexible trajectory learning and improved student training.

3 CTM: AN UNIFICATION OF SCORE-BASED AND DISTILLATION MODELS: CTM integrates score-based solving and direct anytime-to-anytime jumps along the PF ODE trajectory.
3.1 DECODER REPRESENTATION OF CTM: CTM defines G(x_t,t,s) as the PF ODE solution from initial time t to any final time s ≤ t.
3.1 DECODER REPRESENTATION OF CTM: The neural jump parameterization enforces G_θ(x_t,t,t)=x_t automatically, converting constrained training into unconstrained optimization.
3.2 DISTILLATION LOSS: SOFT CONSISTENCY LOSS: Soft consistency matches teacher and student predictions at time s to estimate the full solution trajectory more precisely.
3.2 DISTILLATION LOSS: SOFT CONSISTENCY LOSS: Sampling the teacher over different intervals makes soft consistency span global consistency, local consistency, and CM’s target when s = 0.
3.3 AUXILIARY LOSSES FOR BETTER TRAINING OF STUDENT: Distillation, DSM, and adversarial losses provide complementary training signals for the student, while DSM improves precision for small jumps.

4 SAMPLING WITH CTM

CTM supports score-based sampling and deterministic or stochastic γ-sampling through flexible time traversal along PF ODE trajectories. Its γ parameter connects existing sampling methods while controlling stochasticity and mitigating discretization or accumulated-jump errors.

Score-based sampling: CTM evaluates scores through gθ(xt, t, t), enabling standard ODE/SDE score-based sampling alongside its trajectory jumps.CTM performs comparably to EDM with Heun’s method in high-dimensional image synthesis.
Empirical comparison: Figure 4 compares EDM, CM, and CTM sampling methods across different NFE settings on 256 × 256 AFHQ-cat.The comparison includes score-based and distillation approaches.
γ-sampling: γ-sampling unifies distillation and score-based sampling by traversing PF ODE trajectories between arbitrary time points.The sampler alternates denoising and forward-diffusion transitions until reaching time 0.
γ-sampling: γ = 1 reproduces CM’s fully stochastic multistep sampler, whereas γ = 0 yields deterministic PF ODE distillation sampling.Intermediate 0 < γ < 1 generalizes EDM’s stochastic sampler.
γ-sampling: γ-sampling variance scales proportionally with γ^2, while smaller γ values preserve closer semantic alignment in stroke-based generation.The paper notes that the best γ depends on practical usage and empirical configuration.
Error behavior: γ = 1 can accumulate overlap errors as NFE increases, whereas γ = 0 avoids this accumulation and removes score-based discretization error.The stated error behavior is O(√T + t1 + · · · + tN) for γ = 1 and O(T) for γ = 0.

5 EXPERIMENTS

Experiments evaluate CTM on CIFAR-10 and ImageNet using diffusion teachers, auxiliary losses, and multiple sampling or guidance settings. CTM achieves strong few-step results, while DSM improves large-NFE quality and soft consistency outperforms local consistency.

Experimental setup: CTM is evaluated on CIFAR-10 and ImageNet 64 × 64 with EDM and CM checkpoints serving as teacher models.The experiments combine CTM with DSM and adversarial losses using established architectures and discriminator configurations.
Quantitative results: FID 1.73 with CTM surpasses EDM’s 1.82 at 35 NFE and StyleGAN-XL’s 1.85 on conditional CIFAR-10 generation.CTM also reports FID 1.98 at NFE 1 and FID 1.63 at 2 NFEs in the cited CIFAR-10 results.
Additional capabilities: CTM’s score approximation enables negative log-likelihood evaluation and classifier-rejection sampling with an FID-IS trade-off comparable to high-NFE classifier guidance.Classifier-rejection sampling is combined with NFE 1 sampling in the cited evaluation.
Quantitative results: CTM surpasses previous non-guided generative models in ImageNet FID and reports classifier-likelihood behavior closest to validation data.The reported ImageNet results are achieved within 30K training iterations.
Loss analysis: Soft consistency outperforms local consistency and performs comparably to global consistency in CTM loss experiments.The paper attributes soft matching’s advantage to sampling the amount of teacher supervision across the full time interval.
Loss analysis: DSM preserves quality at small NFE and significantly improves large-NFE sampling when its scale does not outweigh CTM.The experiments use an adaptive DSM weight throughout the paper.
Training without pretraining: Without a pretrained diffusion model, CTM combined with DSM and GAN losses reaches FID 2.39 at NFE 1 on unconditional CIFAR-10.This is reported as comparable to performance with pretrained diffusion models.

6 CONCLUSION

CTM unifies consistency and diffusion models by learning intermediate PF ODE solutions, enabling flexible time traversal and strong few-step generation results. The paper reports state-of-the-art FID and likelihood on CIFAR-10 and ImageNet 64 × 64.

Conclusion: CTM learns intermediate PF ODE solutions to support unrestricted time traversal and combine consistency- and diffusion-model training advantages.The framework is presented as both a generative model and a unified approach to training and sampling.
Conclusion: CTM surpasses its teacher model and achieves state-of-the-art FID and likelihood for few-step diffusion sampling on CIFAR-10 and ImageNet 64 × 64.The conclusion highlights CTM’s versatility across training and sampling settings.

ETHICS STATEMENT

CTM raises content-safety risks because it can generate harmful or inappropriate media, motivating strong filtering and moderation mechanisms.

ETHICS STATEMENT: CTM can generate harmful or inappropriate content, including deepfakes, graphic violence, and offensive material.The stated mitigation is strong content filtering and moderation to prevent unethical or harmful media generation.

B.2 CONVERGENCE ANALYSIS – DISTILLATION FROM TEACHER MODELS

The convergence analysis studies CTM under discrete timesteps and establishes trajectory and density convergence under teacher, solver, and regularity assumptions. It also characterizes trajectory injectivity, density transitions, and γ-sampling behavior.

Convergence along Trajectory in a Time Discretization Setup: Discrete-timestep CTM training converts the continuous-time loss into a counterpart using timesteps t0 = 0 < t1 < ··· < tN = T.The analysis initially assumes local consistency matching and notes extension to soft matching.
Convergence along Trajectory in a Time Discretization Setup: Under uniform Lipschitzness, bounded solver truncation error, and zero CTM loss, CTM converges to the teacher transition for every initial and end time.The solver’s local truncation error is assumed to be O((∆N^t)^(p+1)) with p ≥ 1.
Convergence of Densities: If CTM targets perfectly approximate the true G-function, its induced density converges to the data distribution as N →∞.More generally, the induced density converges uniformly to the teacher-induced density under the stated assumptions.
Non-Intersecting Trajectory: An optimal CTM transition is bi-Lipschitz and injective, so distinct initial points remain distinct along the PF ODE trajectory.The bi-Lipschitz bounds are e^−Lϕ(t−s) and e^Lϕ(t−s) times the initial-distance norm.
Density Transition: Under score regularity assumptions, the optimal CTM pushes the density at time t forward to the predefined density at any earlier time s.The result is stated as ps = Tt→s♯pt and supports retrieving true densities at any time.

C.4 COMPARISON OF GAN EFFECTS IN GENERATION

Adversarial training improves CTM’s visual detail and FID, with reported benefits for both one-step and larger-NFE generation.

Visual Comparison: GAN-augmented CTM samples show finer details, improved overall shapes, and enhanced brightness, contrast, and saturation.The reported visual comparison includes the teacher, CTM without GAN, and CTM with GAN at NFE 1.
FID Comparison: GAN loss improves CTM’s FID relative to training without GAN and also benefits large-NFE samples.Table 3 compares CTM variants trained with identical hyperparameters except for the GAN loss.
Downstream Applications: CTM supports loss-based trajectory optimization for downstream tasks such as image inpainting or colorization.The proposed algorithm uses trajectory traversal and a loss-embedded corrector to control sample variance or interpret guidance as posterior sampling.

D.1 TRAINING DETAILS

Training uses EDM-style parameterization and a staged optimization procedure combining CTM, denoising score matching, and adversarial objectives.

Optimization procedure: The training loop samples data, noise, and ordered times, computes a teacher ODE solution, then updates generator and discriminator parameters.The procedure samples x0, ε, t, s, and u; constructs xt; evaluates Solver(xt, t, u; ϕ); and alternates gradient updates.
Model parameterization: EDM-style skip and output scalings initialize CTM near the teacher denoiser when most teacher parameters are copied to the student.This initialization is reported to partially explain fast convergence.
Compute and adversarial training: CIFAR-10 training uses 4×V100 GPUs, while ImageNet training uses 8×A100 GPUs and warms up the λGAN hyperparameter.On CIFAR-10, adversarial training is disabled initially and activated after 50K iterations, or 200K without a pretrained diffusion model.
Time discretization: CTM training uses N = 18 discretized timesteps for CIFAR-10, with at most 17 ODE steps per iteration.The maximum number of ODE steps is limited to prevent an iteration from taking too long.
Auxiliary objectives and stabilization: Training also uses LPIPS for LCTM, mixed time sampling for LDSM, feature-space transformations for LGAN, and EMA updates for the student parameters.The reported default EMA decay is μ = 0.999, with μ = 0.9999 also effective for GAN training.

D.2 EVALUATION DETAILS

Evaluation solves the PF ODE with RK45 for likelihood and uses dataset-specific statistics and time discretizations for sample-quality measurements.

Likelihood evaluation: Likelihood evaluation solves the PF ODE with an RK45 Dormand–Prince solver using tolerance 1e-3 and tmin = 0.002.These settings follow the practice suggested by Kim et al. (2022b).
Metric computation: CIFAR-10 FID uses Karras et al. statistics, whereas ImageNet metrics use Dhariwal and Nichol’s pre-calculated statistics.The CIFAR-10 result is summarized in Figure 15.
Metric computation: Figure 15 compares CIFAR-10 state-of-the-art results, with points closer to the origin indicating better performance.The figure caption defines the visual comparison rule rather than reporting a specific value.
Sampling schedule: For large-NFE sampling, the EDM time discretization equi-divides [0, 1] into n points before transforming them to the diffusion-time scale.The paper emphasizes that time discretization for training and sampling is a modeler’s choice.

E ADDITIONAL GENERATED SAMPLES

The additional material presents generated-sample comparisons and theoretical results concerning CTM trajectories, convergence, and their induced distributions.

Additional generated samples: Figure 16 compares identical-starting-point samples from EDM, CTM, and CM across four ImageNet classes at different FIDs and NFEs.The listed settings include EDM with FID 2.44 at NFE 79, CTM with FID 2.19 at NFE 1 and FID 1.90 at NFE 2, and CM with FID 6.20 at NFE 1.
Additional generated samples: The PF ODE provides a deterministic counterpart to the reverse-time SDE, and its solution maps samples between diffusion times.The reverse-time process and PF ODE are described in terms of marginal distributions and trajectory mappings.
Convergence and trajectories: Under the stated assumptions, the learned CTM transition converges to the teacher for arbitrary initial and terminal times as the timestep spacing tends to zero.The result is stated for a uniformly Lipschitz CTM and uses the discretized training objective.
Distributional convergence: As the discretization is refined, the CTM-induced sample distribution converges in distribution and, under additional conditions, uniformly in density to the teacher distribution.The density convergence statements rely on uniform equicontinuity and boundedness assumptions.
Convergence and trajectories: Optimal CTM trajectories do not intersect, because PF ODE solution operators are injective and bi-Lipschitz under the stated regularity assumptions.Distinct initial values remain distinct at later times, and the flow satisfies a Lipschitz bound.

F.7 PROOF OF PROPOSITION 9

The proof establishes that the reverse-time PF ODE has unique trajectories for almost every initial point and transports the terminal distribution to the diffusion marginals.

Existence and uniqueness: Under technical regularity conditions, the reverse-time PF ODE admits a unique solution on [0, T] for pT-almost every initial point.The proof invokes a proposition on ODE well-posedness after establishing regularity of the vector field.
Distributional transport: The flow generated by the reverse-time PF ODE pushes pT forward to pt and, more generally, pushes pt forward to ps for s ≤ t.This identifies the PF ODE flow with the diffusion marginals through uniqueness of the continuity equation and the ODE.