Distilling diffusion — the setup & progressive distillation
Lessons 01–02 distilled a model whose output is a distribution over tokens. A diffusion teacher is something else entirely: a multi-step solver. The three knobs survive the move — but every one of them changes meaning.
The first two lessons built the knob framework on autoregressive teachers. The next four carry it to diffusion models, where the goal is narrow and famous: a teacher needs ~1000 (or ~25) network evaluations to make one image; we want a student that needs 1–4. This requires the diffusion machinery from the generative series — the score / PF-ODE picture (lesson 16) and the sampler ladder (lesson 17). We pick up exactly where lesson 17’s “distillation” row left off.
Why a diffusion teacher breaks the lesson-02 picture
In lesson 02 the teacher was a function p(y|x) you could query for a full probability vector per token, and the student matched that distribution. A diffusion teacher gives you none of that. What it gives you is an iterative map: start from noise xT∼N(0,I), call the network εθ(xt,t) dozens of times, and a sample falls out the end. There is no tractable density q(x) over images to KL against.
So the cost we are attacking is not “params” — it is sampling latency = (network calls) × (cost per call). Lesson 17 showed the cost per call drops in latent space (generative lesson 19) and the call count drops with better ODE solvers — until the smoothness wall, below ~10 steps, where no fixed-order solver helps because the trajectory is too curved. Distillation is how you cross that wall: amortize the integral of a curved ODE into the network’s weights so it can jump.
The three knobs, re-read for diffusion
| knob | lesson 02 (LLM) | diffusion |
|---|---|---|
| WHAT the student matches | the teacher’s token distribution / its sampled sequences | the teacher’s denoising trajectory — a multi-step output, a clean-sample estimate x̂₀, or the whole output distribution |
| HOW the gap is measured | forward / reverse KL, temperature | usually L2 in sample space (no tractable density to KL!), later a score-difference or an adversarial loss |
| WHERE the data comes from | fixed set / teacher text / student rollouts | (noise, teacher-trajectory) pairs — off-policy on the teacher’s ODE; later, student-generated samples (on-policy) |
The biggest shock is HOW: with no density over images, the natural lesson-02 move (match distributions by KL) is off the table for the simplest methods. They fall back to the bluntest possible target — make the student’s output numerically close to the teacher’s, in pixel/latent space, under L2. Lessons 05–06 will earn back a distributional objective; this lesson and the next use L2.
Progressive distillation: halve the steps, repeat
The first and simplest method (Salimans & Ho 2022). The teacher is a trained diffusion model run with a deterministic DDIM sampler (lesson 17). Observe: two DDIM steps from xt on the teacher’s sampling grid land you two grid-points further along the trajectory. Train a student to reach that same point in one step.
Concretely, run the two teacher steps to get the target endpoint, convert it to the equivalent clean-sample target x̃₀ the student should predict from xt, and regress:
Now the trained student is a new teacher that needs half as many steps. Freeze it, repeat: 1024 → 512 → 256 → … → 8 → 4. Each round is one finetuning run; after k rounds the step count has halved k times (a 1024·2−k-step sampler).
Imagine the teacher walks from noise to image in 1024 tiny, safe footsteps along a winding path. Progressive distillation teaches an apprentice to cover two footsteps with one stride that lands on the same spot — a slightly bolder shortcut across the curve. Promote the apprentice to teacher and train a new apprentice to take strides twice as long again. Each generation takes bolder shortcuts across the bend; after ten generations the walk is four leaps. You never asked anyone to learn the whole map — only to double the last guy’s stride, which is a gentle, learnable target.
Why halving (and not “just train 1-step directly”)
You could try to regress a 1-step student against the teacher’s full 1024-step output from the start. It fails: from a pure-noise xT the mapping to a specific image is wildly multi-valued and high-curvature — the L2-optimal answer is the blurry average of all images that noise could become (the conditional mean) — the generic failure of L2 regression against a one-to-many target, which is exactly why the generative series preferred ε/v targets and small steps. Halving keeps each distillation target close to what the student can already do, so the regression stays in the well-behaved regime. The price is the logarithmic schedule: ~10 separate training runs to reach 1–2 steps, each depending on the last. That serial cost, and the residual blur at 1 step, are precisely what consistency models (lesson 04) attack.
v-prediction: the parameterization that makes it stable
Salimans & Ho also introduced v-prediction here for a reason worth stating. At the high noise levels a few-step student must handle, ε-prediction is ill-conditioned (a tiny ε error scales up enormously when you solve for x̂₀) and x₀-prediction is ill-conditioned at low noise. The velocity target v = √ᾱt·ε − √(1−ᾱt)·x₀ stays well-scaled at both ends (generative lesson 09’s table), so the L2 loss has stable gradients across the wide step sizes distillation forces. Few-step distillation and v-prediction were born in the same paper because few-step distillation needs a stable target.
Interactive · two teacher steps → one student step
The teacher (oracle DDIM on a two-moons target) takes the full step ladder; the “student” here takes exactly half as many steps but jumps along the same deterministic trajectory. Watch the half-step path track the full path with one big stride per two small ones. Push the halving further (fewer student steps) and watch the strides start to cut the corner of the curve — the error progressive distillation must keep small by only ever halving once per round.