Source-linked AI summary
On Distillation of Guided Diffusion Models
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik P. Kingma, Stefano Ermon, Jonathan Ho, Tim Salimans
TL;DR
Classifier-free guidance improves image generation but is costly because it evaluates conditional and unconditional diffusion models many times. The paper distills these guided models into a guidance-conditioned student and then progressively reduces its sampling steps. The resulting models retain reported quality with substantially fewer steps across pixel-space, latent-space, editing, and inpainting tasks.
Problem
Classifier-free guidance is computationally expensive because generating one sample requires evaluating conditional and unconditional diffusion models tens to hundreds of times.
Method
The method first trains a single guidance-conditioned model to match the combined teacher prediction, then progressively distills it by halving the sampling steps.
Results
The approach produces visually comparable results with 4 to 16 pixel-space steps and 1 to 4 latent-space steps, while achieving up to 256× speedup and at least 10× fewer latent-space steps.
Takeaways & Limitations
The distilled models make classifier-free guided diffusion substantially faster while preserving the guidance-strength quality–diversity trade-off and supporting image editing and inpainting.
Takeaways & Limitations
Stochastic sampling requires small modifications to the training algorithm for edge cases.
Abstract
from arXiv · showhide
Classifier-free guided diffusion models have recently been shown to be highly effective at high-resolution image generation, and they have been widely used in large-scale diffusion frameworks including DALLE-2, Stable Diffusion and Imagen. However, a downside of classifier-free guided diffusion models is that they are computationally expensive at inference time since they require evaluating two diffusion models, a class-conditional model and an unconditional model, tens to hundreds of times. To deal with this limitation, we propose an approach to distilling classifier-free guided diffusion models into models that are fast to sample from: Given a pre-trained classifier-free guided model, we first learn a single model to match the output of the combined conditional and unconditional models, and then we progressively distill that model to a diffusion model that requires much fewer sampling steps. For standard diffusion models trained on the pixel-space, our approach is able to generate images visually comparable to that of the original model using as few as 4 sampling steps on ImageNet 64x64 and CIFAR-10, achieving FID/IS scores comparable to that of the original model while being up to 256 times faster to sample from. For diffusion models trained on the latent-space (e.g., Stable Diffusion), our approach is able to generate high-fidelity images using as few as 1 to 4 denoising steps, accelerating inference by at least 10-fold compared to existing methods on ImageNet 256x256 and LAION datasets. We further demonstrate the effectiveness of our approach on text-guided image editing and inpainting, where our distilled model is able to generate high-quality results using as few as 2-4 denoising steps.
1. Introduction
Classifier-free guidance improves diffusion-model sample quality but is inefficient because each update evaluates conditional and unconditional models. The paper distills guided models into faster samplers while retaining quality across pixel-space and latent-space settings.
- Classifier-free guidance requires evaluating two diffusion models tens to hundreds of times per sample, limiting sampling efficiency.The two models are a class-conditional model and an unconditional model.
- The framework applies to both pixel-space and latent-space diffusion models, including Stable Diffusion.
- 4 sampling steps produce visually comparable samples and comparable FID/IS scores for pixel-space models on ImageNet 64x64 and CIFAR-10.The reported pixel-space results use 4 to 16 steps for comparable FID/IS scores across guidance strengths.
- 1 to 4 sampling steps produce comparable visual quality and teacher-matching FID on latent-space models at ImageNet 256x256 and LAION 512x512.The latent-space setting uses at least 10× fewer steps than the base model.
- The distilled model can trade off sample diversity and quality by varying guidance weight w, with good results using one sampling step.
2. Background on diffusion models
Diffusion models learn to denoise progressively corrupted data and can be sampled with DDIM updates. Classifier-free guidance combines conditional and unconditional predictions using a guidance weight, while progressive distillation targets faster sampling.
- A diffusion model is trained from noisy latent states and denoising targets using a weighted squared-error objective.The objective samples times, data points, and noisy states, then compares the model prediction with the data sample.
- DDIM sampling starts from Gaussian noise and applies discrete updates until the final sample is generated from the model prediction at time zero.
- Classifier-free guidance combines conditional and unconditional diffusion predictions using guidance weight w to trade off sample quality and diversity.Both diffusion models are evaluated at each update step.
- Progressive distillation repeatedly distills diffusion models to reduce sampling steps, but had not directly addressed classifier-free guided models or non-deterministic samplers.
3. Distilling a guided diffusion model
The method first learns one guidance-conditioned student to match the combined teacher prediction, then progressively halves the number of sampling steps. It supports deterministic and stochastic sampling across pixel-space and latent-space models.
- A single student model conditioned on guidance strength w captures a range of classifier-free guidance levels while reducing sampling steps.
- Stage-one distillation: Stage one trains a student to match the guided teacher output at each time step, supporting either continuous-time or discrete-time teachers.
- Stage-one distillation: The student is conditioned on guidance strength w so users can retain the quality–diversity trade-off across a specified guidance interval.The training interval is [wmin, wmax].
- Stage-two distillation: Stage two progressively halves sampling steps by training one student step to match two DDIM steps of the current teacher.The distilled student is repeatedly promoted to teacher for further halving.
- N-step deterministic and stochastic sampling: The procedure also supports stochastic sampling, but edge cases require small modifications to the training algorithm.
4. Experiments
Experiments evaluate the distillation approach on pixel-space and latent-space diffusion models, including generation, text-guided editing, and inpainting. Across these settings, the distilled models retain strong quality with substantially fewer sampling steps.
- Experimental scope: The experiments cover pixel-space diffusion, latent-space diffusion, text-guided image generation, image-to-image translation, and inpainting.The study uses ImageNet 64x64, CIFAR-10, ImageNet 256x256, LAION, and Stable Diffusion-based tasks.
- Pixel-space generation: Up to 256× speedup is achieved while matching a teacher guided DDIM model with 1024×2 sampling steps using only 4–16 steps.Performance is evaluated using FID and Inception Score on ImageNet 64x64, with extended results on CIFAR-10.
- Text-guided generation: The distilled Stable Diffusion model produces sharper and more coherent text-to-image samples than DDIM at comparable low step counts.At four steps, DDIM samples contain clear artifacts; at eight steps they improve but remain blurrier and less consistent than distilled samples using fewer steps.
- Quantitative text-to-image evaluation: For 512×512 text-to-image generation, the distilled latent diffusion model achieves similar or better FID and CLIP scores with significantly fewer sampling steps than the original model.The strongest gains are reported in the low-step regime on 5000 COCO2017 validation captions.
- Image editing and inpainting: The distilled model generates high-quality style-transfer results with 3 denoising steps and high-quality inpainting results with 4 steps on unseen data.The image-to-image experiment uses stochastic latent encoding and deterministic distilled decoding.
5. Related Work
Prior work accelerates diffusion sampling through alternative samplers, higher-order solvers, and distillation. However, existing distillation methods were not directly applicable to classifier-free guided models.
- Sampling acceleration: DDIM, probability-flow samplers, fast SDE integrators, higher-order solvers, exponential integrators, and dynamic programming methods target faster diffusion sampling.These approaches improve sampling speed without being presented as direct solutions to classifier-free guidance distillation.
- Distillation methods: Existing distillation methods mainly address non-classifier-free guided diffusion models.Examples include one-step DDIM inversion, higher-order solver distillation, and progressive distillation.
- Open gap: Progressive distillation reduces sampling steps while preserving the model architecture, but had not been directly applied to classifier-free guided models or non-DDIM samplers.The paper presents its approach as resolving both limitations.
6. Conclusion
The paper introduces a two-stage distillation approach for guided diffusion models and evaluates it across generation and image-editing tasks. It reports major inference-cost reductions while retaining strong sample quality at very low step counts.
- Contribution: The two-stage approach reduces classifier-free guided pixel-space and latent-space diffusion inference cost by at least an order of magnitude.The method first distills guidance into a single model and then progressively reduces sampling steps.
- Results: Comparable FID to the teacher is achieved with as few as 4–8 steps, while visually appealing results can be produced with only 2 steps.Applications include text-guided image-to-image translation and inpainting.
- Future work: The authors identify improving performance in the two- and one-step regimes as future work.This specifies the remaining performance focus rather than a demonstrated limitation across all settings.
A. Results overview
The appendix overview describes the experimental and algorithmic setup for the distillation pipeline. It uses a U-Net-based teacher-student procedure with noise, time, guidance, and optimization steps specified for stage-one distillation.
- Results overview: The overview includes speedup results for pixel-space and latent-space models, plus additional comparisons with DDIM, DPM, and DPM++ solvers.Further experimental details are provided separately for pixel-space and latent-space distillation.
- Teacher and architecture: The model architecture is a U-Net similar to prior guided diffusion models, with v-prediction used in the reported training setting.The same number of channels and attention are used for ImageNet 64x64 and CIFAR-10.
- Stage-one distillation: Stage-one distillation requires a trained classifier-free guidance teacher, a dataset, and a loss-weight function.The algorithm samples data, time, guidance strength, and Gaussian noise before constructing a noisy input.
- Stage-one distillation: The stage-one target combines conditional and unconditional predictions as (1 + w)ˆx_c,θ(z_t) − wˆx_θ(z_t).The guidance strength w is sampled during training, and the resulting target is used for optimization.
- Stage-one distillation: The student parameters are updated by gradient descent on the stage-one loss.This completes the optimization loop specified in the algorithm excerpt.
B.3. Stage-two distillation for deterministic sampler
Stage-two distillation progressively reduces a teacher’s sampling steps by training a student to match two teacher updates, then repeatedly halving the step count. The procedure supports deterministic and stochastic samplers, with the stochastic target using a larger effective sampling step.
- Deterministic sampler: The deterministic procedure initializes a student from the Stage-one teacher and progressively distills 1024 DDIM steps to one step.The student is trained for 50,000 updates, or 100,000 updates at one or two steps.
- Deterministic sampler: Each deterministic training example samples data, a timestep, guidance strength, and noise before applying two DDIM teacher steps to form the target.The guidance strength is sampled from a training interval, and the target is used in the student loss.
- Progressive reduction: After convergence, the trained student becomes the next teacher and the number of sampling steps is halved.This teacher-student loop repeats until the desired sampling budget is reached.
- Stochastic sampler: The stochastic variant uses the same progressive framework but distills a target corresponding to a sampling step twice as large as in the deterministic sampler.It handles the one-step boundary case separately and visualizes results across guidance strengths.
- Stochastic sampler: The stochastic teacher target uses two DDIM steps when possible and one step at the boundary before the student is trained and the step count is halved.The extra samples compare the resulting approach with DDIM baselines.
B.6. Extra distillation results
Extra experiments evaluate distilled pixel-space models, style transfer, and latent-space generation. The results show strong low-step performance, while a naive jointly distilled conditional/unconditional model performs poorly at small step counts.
- Pixel-space results: FID and IS evaluations cover ImageNet 64x64 and CIFAR-10 across multiple guidance strengths.The evaluated guidance strengths are {0, 0.3, 1, 2, 4} for ImageNet and {0, 0.1, 0.2, 0.3, 0.5, 0.7, 1, 2, 4} for CIFAR-10.
- Text-guided generation: Two-step distilled text-guided generation on LAION-5B produces more realistic and higher-quality images than DDIM and DPM++ baselines using more steps.The distilled model evaluates one model instead of both conditional and unconditional models at each step.
- Guidance behavior: Varying guidance strength in one-step ImageNet sampling trades sample diversity against quality while retaining visually pleasant results.The comparison includes class-conditional samples from the stochastic approach and DDIM baselines at 8×2 and 16×2 sampling steps.
- Pixel-space results: The distilled pixel-space model significantly outperforms DDIM with fewer steps and matches teacher performance with 8 to 16 steps.Training over a guidance interval performs comparably to training at a fixed guidance strength.
- Naive distillation: The naive distillation approach remains weak because it retains separate conditional and unconditional components and therefore incurs twice the evaluations or memory.The reported CIFAR-10 FID/IS scores do not show strong performance.
- Latent-space results: For latent diffusion, two- and four-step sampling nearly matches DDIM at 32 steps in precision and significantly outperforms it in recall at low step counts.The evaluation uses ImageNet 256x256 class-conditional generation and 5,000 samples.
C.2.1 Training details
The latent-diffusion experiments fine-tune Stable Diffusion-style models for LAION generation and compare distilled sampling with DDIM and DPM-Solver. Distillation improves low-step image quality while retaining guidance control and reducing model evaluations.
- Training setup: The experiments cover LAION-5B at 256x256 and 512x512 resolutions.Both settings start from a DDIM teacher with 512 sampling steps.
- Baselines: DPM-Solver is used as a training-free comparison that performs well at 10 or more steps but still requires classifier-free guidance.This doubles U-Net evaluations relative to the distilled guidance-conditioned approach.
- Class-conditional generation: The distilled latent model nearly matches DDIM precision at 2 and 4 steps and significantly exceeds DDIM recall at low step counts.These precision and recall results are evaluated for ImageNet 256x256 on 5,000 samples.
- Text-guided generation: On LAION text-guided generation, the method approaches DDIM after only a few thousand training steps.FID and Inception Score are evaluated on 5,000 COCO2017 captions at 256x256.
- Text-guided generation: At 512x512, the distilled Stable Diffusion model outperforms DPM-Solver in the 2- and 4-step regime on FID and CLIP ViT-g/14.The evaluation uses 5,000 COCO2017 validation captions; the text notes that longer training may close the CLIP gap above 10 steps.
- Efficiency and quality: Qualitative comparisons report sharper and more coherent low-step results than training-free samplers, with a twofold reduction in maximum memory or denoising-step overhead.The distilled model uses one guidance-conditioned model rather than separate conditional and unconditional evaluations.
- Image editing: Guided image translation exhibits a trade-off: adding more initial noise improves text alignment but reduces faithfulness to the input image.The analysis varies efficiency, realism, controllability, and initial perturbation noise.
C.4.1 Training details
The inpainting and image-editing experiments train distilled models on LAION-5B at 512x512 resolution and evaluate low-step sampling. They combine quantitative inpainting evaluation with analyses of efficiency, realism, controllability, and faithfulness.
- Training setup: The inpainting model is trained on LAION-5B at 512x512 using the LAMA large-mask generation scheme.Training begins from a DDIM teacher with 512 sampling steps and uses its output as the distillation target.
- Evaluation: Inpainting quality is quantitatively compared with DDIM using FID on 2,000 COCO2017 examples.The DDIM baseline uses classifier-free guidance and therefore twice the function evaluations listed under its setting.
- Image editing: The editing analysis examines trade-offs among sample quality, controllability, and efficiency for distilled text-guided models.The experiments use 4-step and 8-step models with different initial perturbation noise levels.
- Additional results: Additional quantitative comparisons with DDIM are reported for low sampling-step counts, alongside further samples.The additional results are presented in Table 8 and Figure 32.
D. Extra samples for pixel-space distillation
The appendix provides extra pixel-space distillation samples using both deterministic and stochastic samplers. Samples are shown on CIFAR-10 and ImageNet 64x64 across distilled sampling budgets from 256 steps down to 1 step.
- Sampling procedures: Extra samples are generated with both deterministic and stochastic samplers.The section explicitly references the deterministic sampler in Algorithm 2 and the stochastic sampler in Algorithm 3.