[Paper Review] Multistep Consistency Models

Jonathan Heek, Emiel Hoogeboom+1 more|arXiv (Cornell University)|Mar 11, 2024

Complex Systems and Decision MakingDecision Sciences3 citations

TL;DR

This paper introduces Multistep Consistency Models, a unification of consistency models and TRACT that interpolates between single-step consistency and full diffusion sampling. By training segment-wise consistency models over 2–8 steps with shared parameters, it achieves state-of-the-art FID scores (1.4 on ImageNet64, 2.1 on ImageNet128) in 8 steps, matching standard diffusion quality while retaining high sampling speed.

ABSTRACT

Diffusion models are relatively easy to train but require many steps to generate samples. Consistency models are far more difficult to train, but generate samples in a single step. In this paper we propose Multistep Consistency Models: A unification between Consistency Models (Song et al., 2023) and TRACT (Berthelot et al., 2023) that can interpolate between a consistency model and a diffusion model: a trade-off between sampling speed and sampling quality. Specifically, a 1-step consistency model is a conventional consistency model whereas a $\infty$-step consistency model is a diffusion model. Multistep Consistency Models work really well in practice. By increasing the sample budget from a single step to 2-8 steps, we can train models more easily that generate higher quality samples, while retaining much of the sampling speed benefits. Notable results are 1.4 FID on Imagenet 64 in 8 step and 2.1 FID on Imagenet128 in 8 steps with consistency distillation, using simple losses without adversarial training. We also show that our method scales to a text-to-image diffusion model, generating samples that are close to the quality of the original model.

Motivation & Objective

• To close the performance gap between standard diffusion models and low-step consistency models.
• To enable a smooth trade-off between sampling speed and generation quality via multi-step inference.
• To improve training stability and sample quality in consistency modeling by generalizing TRACT and consistency distillation.
• To develop a deterministic sampler (aDDIM) that corrects integration errors and reduces blur in low-step generation.
• To demonstrate strong performance on complex benchmarks like ImageNet128 and text-to-image generation.

Proposed method

• The method splits the diffusion process into predefined segments, training a separate consistency model for each segment while sharing parameters across all segments.
• It generalizes both consistency training (CT) and consistency distillation (CD) to multiple steps, enabling smoother and more learnable denoising trajectories.
• The approach incorporates step schedule annealing and synchronized dropout from consistency modeling to improve training stability.
• A new deterministic sampler, Adjusted DDIM (aDDIM), inflates noise predictions to correct for integration errors, reducing blur in low-step sampling.
• The framework allows interpolation between a 1-step consistency model and a full diffusion model (infinite steps), with the 8-step version matching standard diffusion performance.
• For distillation, the method uses aDDIM to train high-quality 16-step consistency models from a pre-trained diffusion model.

Experimental results

Research questions

• RQ1 Can consistency models be extended beyond single-step sampling to improve sample quality while retaining speed advantages?
• RQ2 Does multi-step consistency training with shared parameters lead to better performance than single-step consistency or standard diffusion?
• RQ3 Can a deterministic sampler like aDDIM achieve competitive FID scores on complex datasets such as ImageNet128 with few steps?
• RQ4 How does step schedule annealing affect performance in low-step multistep consistency models?
• RQ5 Can multistep consistency models match standard diffusion quality on text-to-image generation with minimal sampling steps?

Key findings

• The proposed Multistep Consistency Models achieve a state-of-the-art FID of 1.4 on ImageNet64 and 2.1 on ImageNet128 using only 8 sampling steps.
• With 4 sampling steps, the model attains 1.6 FID on ImageNet64 and 2.3 FID on ImageNet128, demonstrating strong performance at low inference cost.
• The method matches standard diffusion model performance (e.g., 100-step DDIM) using just 8 steps, effectively closing the quality gap.
• Annealing the step schedule significantly improves performance in low-step models, with diminishing returns as step count increases.
• The aDDIM sampler enables competitive FID scores on ImageNet128, a benchmark previously lacking a reliable deterministic low-step baseline.
• In text-to-image generation, the 16-step distilled model produces samples nearly indistinguishable from the original 100-step DDIM model when using the same random seed.