Overview: Progressive distillation trains a sequence of student models that each learn to emulate multiple teacher sampling steps in one step, halving the sampling steps repeatedly until you reach a target fast sampler. The goal is to preserve sample fidelity (e.g., FID/LPIPS) while reducing sampling latency.1) Clarify setup & notation- Teacher T_k performs S_k sampling steps (e.g., S_0 = 1000).- Student S_{k+1} should perform S_{k+1} = S_k / 2 (or another reduction factor r) steps by learning to map the same noisy input-to-output in fewer steps.- Use either noise-prediction parameterization (εθ) or denoised x0 prediction (x̂0θ) consistently.2) Teacher-student training steps (per distillation stage)- Freeze teacher T_k (pretrained or previous-stage student EMA).- For each training batch: - Sample a clean x0 and a time pair (t1, t2) that correspond to two consecutive timesteps in the student’s target coarse schedule (e.g., t2 < t1, with teacher using intermediate steps between them). - Simulate teacher rollout: apply teacher’s sampler starting from x_{t1} and run the teacher’s internal sub-steps to produce the teacher’s next-state x_{t2}^{teacher} (or the teacher-predicted denoised x̂0^{teacher}). - Student S_{k+1} takes x_{t1} and directly predicts x_{t2} (or x̂0) in one step.- Loss: match teacher outputs. Typical losses: - L_simple = E[||ε̂_student(x_{t1}, t1) - ε_teacher_rollout||^2] when using noise-prediction. - Or L_x0 = E[||x̂0_student(x_{t1}, t1) - x̂0_teacher||^2] for x0 parameterization. - Optionally combine with perceptual loss: L = α L_mse + β L_perceptual (LPIPS) + γ L_vgg, and/or adversarial loss if higher perceptual realism is needed. - Add consistency / regularization: L_noise_perturb = E[||S(x_{t1}+δ) - S(x_{t1})||^2] to stabilize.3) Schedule for reducing sampling steps (progressive plan)- Typical choice: reduce by factor 2 each distillation stage: S_{k+1} = S_k / 2. Start S_0=1000 → 500 → 250 → 125 → 50 → 25 → 8 → 4 → 2 → 1 depending on target.- Alternative: non-uniform reduction (e.g., larger reductions early, smaller near the end) if fidelity degrades.- For each stage: - Train the student for N epochs or until validation metrics plateau. - Use teacher EMA weights for stability and to generate soft targets. - Optionally warm-start student from teacher weights.4) Training details & best practices- Use teacher’s sampler to produce high-quality targets (DDIM/PLMS deterministic samplers often used).- Use mixed-precision and gradient accumulation for large batch effective sizes.- Use EMA of student weights during training for evaluation and to serve as next-stage teacher.- Learning rate: smaller than initial diffusion training (e.g., 1e-5–5e-5) with cosine or linear warmup.- Data augmentation and class-conditioning (if applicable) preserved.5) Validation to ensure fidelity vs speed tradeoff- Quantitative: - FID and KID between generated and real data at each stage. - Precision/Recall and Density/Coverage for distributional fidelity. - LPIPS/SSIM for perceptual quality and diversity metrics (sample diversity vs mode collapse). - Likelihood proxies (ELBO variants) if available. - Inference latency (ms/sample) and throughput (samples/sec) measured on target hardware.- Qualitative: - Visual inspection on held-out prompts/conditions. - Human A/B tests for realism preference vs teacher outputs.- Robustness: - Test across seeds, conditioning prompts, out-of-distribution inputs. - Check for artifacts introduced by larger step sizes.- Stopping criterion: choose distilled student stage where FID degradation is within an acceptable epsilon (e.g., ΔFID ≤ small threshold) while meeting latency targets.6) Diagnostics & ablations- Track L_distill, FID, sample diversity each epoch; abrupt divergence signals optimization or loss mismatch.- Ablate loss choices (noise vs x0 vs hybrid), schedule factor, teacher sampler (stochastic vs deterministic).- If fidelity drops at aggressive reductions, insert intermediate smaller reductions or use a multi-step student that performs m teacher-like micro-steps internally (learned sub-steps).7) Example pseudocode sketch (conceptual)python
# pseudocode: one distillation stage (teacher T, student S)
for batch in dataloader:
x0 = sample_clean_images(batch_size)
t1 = sample_time_indices(batch_size) # coarse times
x_t1 = q_sample(x0, t1) # add noise
# teacher simulates fine-grained steps to t2
x_t2_teacher = teacher_simulate_to(x_t1, t2)
# student predicts x_t2 in one step
pred = student.predict(x_t1, t1)
loss = mse(pred, x_t2_teacher) # or match teacher's x̂0 / ε
loss.backward(); optim.step()
