Multistep Consistency Models

1. Introduction

• Diffusion models require many sampling function evaluations despite using one function evaluation per datapoint during training.The passage notes that generation can require hundreds of evaluations.Source: §1 · p. 1
• Consistency models reduce sampling time but, alongside TRACT, had not reached standard diffusion performance in the cited prior work.Consistency Training and Consistency Distillation had improved over earlier works, but the reported gap remained.Source: §1 · p. 1
• Eight-step Multistep Consistency Models match baseline diffusion performance on ImageNet64 and ImageNet128.The method unifies Consistency Models and TRACT, relaxing the single-step constraint and splitting diffusion time into predefined segments modeled by separate consistency models with shared parameters.Source: §1 · p. 1
• Multistep Consistency Models interpolate between single-step consistency models and standard diffusion, with the latter corresponding to the ∞-step trajectory.The illustrated four-step trajectory is smoother than the one-step path and is described as likely easier to learn.Source: Figure 1 · p. 2
• aDDIM inflates the noise prediction to correct integration-error blur, providing a deterministic sampler for datasets such as ImageNet128.The introduction identifies this sampler as a contribution alongside the multistep method.Source: §1 · pp. 1–2
• 1.6 FID on ImageNet64 and 2.3 FID on ImageNet128 are achieved with four sampling steps.The paper also reports matching baseline diffusion performance with eight steps on both datasets.Source: §1 · pp. 1–2

2. Background: Diffusion Models

• Diffusion adds noise to data so that the state approaches a standard normal at t→1 and the data at t→0.The passage expresses the noisy state as a weighted data term plus noise.Source: §2 · p. 2
• Sampling iteratively moves from noise toward data, and image-generation processes empirically use hundreds of iterations.The learned denoising function estimates x from the noisy state and time.Source: §2 · pp. 2–3
• Consistency models learn a direct noise-to-data mapping and aim to predict the same x along the probability ODE trajectory.The text notes that training consistency models is more difficult than training diffusion models.Source: §2 · p. 3
• Consistency Training uses data-derived states, whereas Consistency Distillation can use a pretrained teacher to take a probability-flow ODE step.For distillation, DDIM may define the next state from the current noisy state and an estimate of x.Source: §2 · p. 3
• Larger time gaps propagate solutions faster but introduce bias in Consistency Training; smaller gaps reduce bias but slow propagation.A step schedule anneals the gap over training.Source: §2 · p. 3
• DDIM is a linearization of the probability-flow ODE, and its inverse recovers x from states z_t and z_s.The DDIM equation also serves as a tool for constructing a multistep diffusion algorithm.Source: §2 · p. 4

3.1. General description

• Multistep consistency divides diffusion time into equal segments, replacing one difficult full-path mapping with several shorter integrations.The model learns intermediate targets along the DDIM path rather than predicting the final data point at every step.Source: §3.1 · p. 4
• With as many student steps as teacher steps, multistep consistency training reduces to the diffusion objective at every step.The objectives also coincide at the beginning of a step.Source: §3.1 · p. 5
• The objective can be expressed as weighted prediction error in x-space, where Euclidean distance worked better than squared distance, especially with fewer student steps.The x-space reparameterization also permits use of diffusion losses such as v-loss.Source: §3.1 · p. 5
• A separate Multistep CM must be trained for each desired student-step count, although fine-tuning from a pretrained model is expected to use a fraction of the original training budget.Source: §3.1 · p. 5

What about training in continuous time?

• Continuous-time training is possible if the model receives a step-conditioning signal, since its prediction must change discontinuously between adjacent steps.The authors use a discrete grid for compatibility with pretrained diffusion models and report comparable early experiments with continuous time.Source: What about training in continuous time? · p. 6

3.2. The Adjusted DDIM (aDDIM) sampler.

• DDIM integration error can blur samples because its iterates have too-small norms; stochastic samplers correct this with noise, while second-order deterministic samplers require an extra distillation pass.aDDIM instead increases the deterministic noise contribution in one evaluation; the authors describe the modification as heuristic.Source: §3.2 · p. 6
• aDDIM matches the second-order Heun sampler on ImageNet64 and outperforms it on ImageNet128.Figure 3 compares sampling methods on small models without distillation; Heun adds a second-order correction to DDIM.Source: §3.2 · p. 7, Figure 3 · p. 7
• aDDIM estimates the missing variance and adjusts the deterministic noise prediction; the correction vanishes as adjacent times converge.The variance can be estimated empirically with η=0.75 or approximated analytically.Source: §3.2 · p. 7

4. Related Work

• Multistep Consistency Models combine Consistency Models and TRACT, targeting more sampling steps than TRACT’s one- or two-step endpoint to close the quality gap with standard diffusion.The paper attributes gains to a more conservative step target, an improved sampler, and an annealed schedule.Source: §4 · p. 7
• Progressive Distillation reduces required sampling steps through staged distillation, while DDIM and higher-order or SDE integrators seek fewer function evaluations.The cited approaches include DDIM (Song et al., 2021a), Heun sampling (Karras et al., 2022), SDE integrators (Jolicoeur-Martineau et al., 2021), and Progressive Distillation (Meng et al., 2022; Salimans and Ho, 2022).Source: §4 · p. 8
• Flow Matching and Rectified Flows generally retain many evaluation steps unless paired with distillation, which can reduce sample quality.Rectified Flows and Flow Matching are described as mapping to a standard normal distribution; the passage specifically says Rectified Flow distillation reduces steps at a quality cost.Source: §4 · p. 8
• Unlike Multistep Consistency Models, CTMs integrate to arbitrary endpoints and rely on adversarial training for high-quality few-step samples, whereas Multistep CMs use distance metrics.The authors suggest that learning a handful of fixed trajectories may be easier than learning arbitrary integrations; they also note adversarial training’s mode-collapse and instability risks.Source: §4 · p. 8

5. Experiments

• ImageNet experiments use pretrained diffusion weights, both consistency training and distillation, and classifier-free guidance only for base ImageNet128.Models train for 200,000 steps with batch size 2048 and a teacher-step schedule annealed from 64 to 1280 over 100,000 steps.Source: §5.1 · p. 9

5.1. Evaluation on ImageNet

• At 4 steps, consistency distillation achieves 1.7 FID versus 2.4 for reimplemented Progressive Distillation on ImageNet64.The comparison uses the same base architecture and an aDDIM sampler for the Progressive Distillation reimplementation.Source: §5.1 · p. 10
• The aDDIM teacher is more useful at higher student-step counts, while vanilla DDIM works better at few steps when the student remains farther from the teacher in FID.For teacher-step annealing, gains are strongest for low-step models and diminish as student steps increase.Source: Table 4 · p. 9, §5.1 · p. 10
• State-of-the-art FID scores on ImageNet64 and ImageNet128 are reported with 4-step and 8-step generation; the ImageNet128 base model reaches 2.1 FID in 8 steps.The literature comparison is in Table 5; the 2.1 FID result is reported for the ImageNet128 base model.Source: §5.1 · pp. 9–10, Table 5 · p. 10

5.2. Evaluation on Text to Image modelling

• Text-to-image 16-step distillation nearly matches its teacher in FID; the 8-step model attains 8.1 FID, and shared-seed samples show generally minor differences from 100-step DDIM.The reported CLIP-score loss is attributed to guidance distillation; against the guidance-distilled baseline, the 16-step model has no low-guidance CLIP or FID loss and only minor high-guidance FID degradation.Source: §5.2 · p. 11, Table 6 · p. 11

6. Conclusions

• Multistep Consistency Models unify Consistency Models and TRACT and close the performance gap between standard diffusion and few-step sampling.The method trades sampling speed for quality, reaching performance comparable to standard diffusion in as little as eight steps.Source: §6 · p. 11

A.1. Setup

• The experiments follow a standard UViT setup, using UNets with MLP blocks in self-attention blocks to form transformer blocks, and initialize runs from pretrained diffusion-model parameters.The paper follows the setup from simple diffusion (Hoogeboom et al., 2023).Source: §A.1 · p. 14
• ImageNet runs anneal teacher steps log-linearly from 64 to 1280 over 100,000 training iterations.The schedule is defined by an exponential interpolation of the logarithms of the endpoints, clipped over the 100,000-iteration interval.Source: §A.1 · p. 14
• ImageNet training uses batch size 2048, 200,000 steps, xvar_frac 0.75 for aDDIM, and Huber epsilon 1e-4.The authors note that fixed settings enabled clean comparisons, though they were somewhat excessive for larger student-step models.Source: §A.1 · p. 14
• The ImageNet64 UViT has 394M parameters in the small variant, while ImageNet128 adds 128-resolution ResNet blocks and has 397M parameters in its small variant.The ImageNet64 large model has 1.23B parameters; the ImageNet128 large model has 1.25B parameters.Source: §A.1 · p. 14
• Text-to-image training uses 512 × 512 data, a multiscale loss, and a UViT with 16 middle transformer blocks; distillation fixes teacher steps at 256 for 30,000 training steps.The text-to-image UViT uses an interpolated cosine schedule from noise resolution 32 to 512; training uses batch size 2048 and Huber epsilon 1.Source: §A.1 · pp. 14–15

• Small variants use 64 TPUv5e chips; CT runs at 2.7 steps per second on ImageNet64 and 2.2 on ImageNet128, while CD runs at 2.5 and 1.7, respectively.Source: §A.2 · p. 15
• Large variants use 256 TPUv5e chips; CT runs at 2.9 steps per second on ImageNet64 and 2.2 on ImageNet128, while CD runs at 2.5 and 1.8, respectively.Source: §A.2 · p. 15
• The text-to-image model uses 256 TPUv5e chips, trains at 0.71 steps per second for 30,000 iterations, and uses a batch size of 2048.Source: §A.2 · p. 15

• Image models are trained on ImageNet, while the text-to-image model uses a privately licensed dataset comparable with public text-to-image datasets but filtered for content.Source: §A.3 · p. 15

A.4. Teacher sampling

• With cosine schedules on ImageNet64, aDDIM is the best-performing method overall except at 64 function evaluations; the figure also compares the Karras et al. (2022) sigma schedule.Source: Figure 5 · p. 15
• Noise DDIM breaks DDIM's determinism and is unsuitable for consistent distillation, yet behaves similarly to aDDIM, suggesting the heuristic correction simulates noise injection's positive effects.Source: §A.4 · pp. 15–16
• Sampler quality depends on evaluation schedule: Heun favors the Karras et al. (2022) schedule, while noisy methods work better with cosine; the authors identify this interaction for future investigation.Source: §A.4 · p. 16

A.5. Additional results

• For the same prompt and initial noise, samples from the 16-step distilled model are very similar to those from the 100-step DDIM teacher, with minor variations.Source: §A.5 · p. 16, Figure 6 · p. 16