Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

Distillation and Consistency Models: Getting to Four Steps

6 minute read

Published:

TL;DR: Progressive distillation repeatedly trains a student to reproduce two of its teacher's sampling steps in one, halving the step count per round — eight rounds take 1024 steps down to 4. Consistency models go further by enforcing a property rather than imitating a solver: every point on a probability-flow trajectory must map to that trajectory's origin, so a single evaluation from pure noise is already a sample. Both can be trained with or without a teacher. What you lose is diversity, fine detail, and the freedom to change the guidance scale after the fact.

Why solvers stop helping

An ODE solver of order \(p\) has local error \(O(h^{p+1})\), which is only useful while the trajectory is well approximated by a low-order polynomial over a step. Diffusion trajectories from noise to image curve sharply, and once a step spans a large stretch of the log-SNR range no polynomial fits. Around ten evaluations the solver approach saturates.

The alternative is to stop treating the network as a fixed right-hand side to be integrated and start training a network that is the integral.

Progressive distillation: halve, repeat

Salimans and Ho’s construction is deliberately modest. Start from a trained teacher sampled with a deterministic DDIM solver at \(N\) steps. Train a student, initialised from the teacher’s weights, so that one student step from \(\mathbf{z}_t\) lands where two teacher steps land. The student’s regression target is not the original noise but the teacher’s two-step output — a target the student can hit exactly, since two DDIM steps from \(\mathbf{z}_t\) are a deterministic function of \(\mathbf{z}_t\).

Then make the student the new teacher and repeat.

RoundTeacher stepsStudent steps
11024512
2512256
412864
63216
884

Eight rounds, each a short fine-tune, and each round’s training data is generated by the previous model, so no new images are needed.

One detail is load-bearing. Distilled students are eventually asked to denoise at the very top of the schedule, where \(\bar{\alpha}_t \to 0\). There the \(\boldsymbol{\epsilon}\)-parameterisation is degenerate: recovering the clean image needs \(\hat{\mathbf{x}}_0 = (\mathbf{z}_t - \sqrt{1-\bar{\alpha}_t}\,\hat{\boldsymbol{\epsilon}})/\sqrt{\bar{\alpha}_t}\), and dividing by a vanishing \(\sqrt{\bar{\alpha}_t}\) amplifies any error without bound. Salimans and Ho switch to the \(\mathbf{v}\)-parameterisation,

\[ \mathbf{v}_t \;\equiv\; \sqrt{\bar{\alpha}_t}\,\boldsymbol{\epsilon} \;-\; \sqrt{1-\bar{\alpha}_t}\,\mathbf{x}_0 , \]

which stays a well-conditioned target across the whole schedule because it interpolates between predicting \(\boldsymbol{\epsilon}\) at low noise and predicting \(-\mathbf{x}_0\) at high noise. This is why \(\mathbf{v}\)-prediction shows up in so many later models even when nothing is being distilled.

Consistency models: enforce a property instead

Song et al. asked what function a one-step sampler should be, independent of any solver. Take the probability-flow ODE trajectory \(\{\mathbf{x}_t\}_{t\in[\varepsilon,T]}\) that ends at a data point. Define \(f_\theta(\mathbf{x}_t,t)\) to return that endpoint. Then by construction

\[ f_\theta(\mathbf{x}_t, t) = f_\theta(\mathbf{x}_{t'}, t') \quad \text{for all } t, t' \in [\varepsilon, T] \]

for any two points on the same trajectory. That is the self-consistency property, and it is the entire model. It is anchored by a boundary condition, \(f_\theta(\mathbf{x}_\varepsilon,\varepsilon) = \mathbf{x}_\varepsilon\), enforced architecturally by writing \(f_\theta(\mathbf{x},t) = c_{\text{skip}}(t)\,\mathbf{x} + c_{\text{out}}(t)\,F_\theta(\mathbf{x},t)\) with \(c_{\text{skip}}(\varepsilon)=1\) and \(c_{\text{out}}(\varepsilon)=0\). Without that anchor the trivial constant solution satisfies consistency perfectly.

Self-consistency: every point on one trajectory maps to the same originA curved trajectory runs from a noise point at the right, marked t equals T, leftwards to a data point at the left, marked t equals epsilon. Four intermediate points sit along it. Dashed teal arrows run from each of the four points and from the noise point directly to the same data point, showing that f maps all of them to one output. A second, fainter trajectory below ends at a different data point, showing the mapping is per-trajectory. x_ε x_T (noise) other trajectory f_θ(x_t, t) = x_ε for every t on this trajectory The self-consistency constraint
Notice that the constraint says nothing about intermediate outputs — only that all points sharing a trajectory agree. That is what makes one evaluation from pure noise a complete sampler, and also why two trajectories that pass near each other are the hard case.

Two ways to train it

Consistency distillation uses an existing diffusion model \(\phi\) as a numerical oracle. Take a noisy sample \(\mathbf{x}_{t_{n+1}}\), run one solver step backwards with \(\phi\) to get \(\hat{\mathbf{x}}^\phi_{t_n}\), and penalise disagreement between the two ends:

\[ \mathcal{L}_{\text{CD}} = \mathbb{E}\Bigl[\lambda(t_n)\, d\bigl(f_\theta(\mathbf{x}_{t_{n+1}}, t_{n+1}),\ f_{\theta^-}(\hat{\mathbf{x}}^\phi_{t_n}, t_n)\bigr)\Bigr] \]

with \(d\) a distance (L2 or LPIPS) and \(\theta^-\) a stop-gradient target copy. The boundary condition propagates correctness inward from \(t=\varepsilon\) one adjacent pair at a time.

Consistency training removes the teacher. The only thing the teacher supplied was the score, and denoising score matching gives an unbiased single-sample estimate of it: for \(\mathbf{x}_t = \mathbf{x}_0 + t\mathbf{z}\),

\[ \nabla_{\mathbf{x}}\log p_t(\mathbf{x}_t) = -\,\mathbb{E}\!\left[\frac{\mathbf{x}_t - \mathbf{x}_0}{t^2}\ \middle|\ \mathbf{x}_t\right], \]

so replacing the oracle step with a pair of noisy views of the same \(\mathbf{x}_0\) using the same \(\mathbf{z}\) is unbiased in the limit of small step spacing. Consistency models therefore stand alone as a generative family, not merely as a compression of diffusion.

Key Insight — why multi-step sampling still works: a consistency model can be sampled in more than one step without any extra training. Jump to the origin, then re-noise the estimate to some smaller \(t\) and jump again. Each round replaces the model's guess with a fresh sample that is guaranteed to lie on some valid trajectory, so errors do not compound the way solver errors do. This is why 2 to 4 steps buys a visible quality gain over 1 while 4 to 8 buys much less.

What you give up

One-step generation is a many-to-one map from noise onto the data manifold, trained with a pointwise regression loss. That loss is minimised by conditional averages, so where trajectories crowd together the model hedges — the characteristic result is slightly softer texture and reduced variety, most visible in backgrounds and crowds.

Three practical costs follow. Distilled students inherit the teacher’s errors and cannot exceed it. Guidance is usually baked in at a fixed weight during distillation, so the guidance dial is gone unless the student is explicitly conditioned on \(w\). And distillation is per-checkpoint work: fine-tune the base model and the fast student is stale. LCM-LoRA partially addresses the last point by distilling into a low-rank adapter that transfers across fine-tunes of the same base.

The trap: few-step samples often look better on first inspection — cleaner, more contrasted, more "finished" — because averaging removes noise and idiosyncrasy. Diversity loss does not show up in single images. Compare grids from the same prompt, not individual samples.

References

  1. Salimans, T., & Ho, J. Progressive Distillation for Fast Sampling of Diffusion Models. ICLR 2022.
  2. Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. Consistency Models. ICML 2023.
  3. Song, Y., & Dhariwal, P. Improved Techniques for Training Consistency Models. ICLR 2024.
  4. Luo, S., Tan, Y., Huang, L., Li, J., & Zhao, H. Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference. arXiv 2023.
  5. Sauer, A., Lorenz, D., Blattmann, A., & Rombach, R. Adversarial Diffusion Distillation. ECCV 2024.

You May Also Enjoy

Noether’s Theorem: Symmetry, Conservation, and Equivariant Networks

7 minute read

Published:

Energy is conserved because the laws of physics do not care what time it is. That single sentence is Noether’s theorem, and its machine learning descendant is the reason an equivariant network needs less data than one that must learn the symmetry from examples.

Entropy and Free Energy: From the Second Law to the ELBO

7 minute read

Published:

Thermodynamic entropy and Shannon entropy differ by a constant with units. Once you accept that, the ELBO stops being an inference trick and becomes a free energy — and diffusion models stop being a clever architecture and become a driven non-equilibrium process.

Statistical Mechanics: The Boltzmann Distribution and the Cost of Z

6 minute read

Published:

Counting microstates gives you the Boltzmann distribution, and the Boltzmann distribution gives you softmax, simulated annealing and energy-based models. The partition function is not a bookkeeping constant — it is the object that contains every thermodynamic quantity, and it is intractable for exactly that reason.