Source-linked AI summary
Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed
Eric Luhman, Troy Luhman
TL;DR
Iterative score-based generative models offer high-quality samples but require many denoising evaluations, making them substantially slower than single-step models. The paper uses knowledge distillation to train a Denoising Student that maps Gaussian noise directly to samples in one step. It reports GAN-comparable quality on CIFAR-10 and CelebA, 256 × 256 LSUN generation, and large sampling-speed gains, while noting blur at higher resolutions.
Problem
Iterative score-based models use expensive stochastic or multi-step sampling, while deterministic distillation requires a suitable deterministic generative process.
Method
Knowledge distillation trains a single-step Denoising Student to predict the output of a deterministic DDIM teacher from the same initial noise vector.
Results
Denoising Student achieves GAN-comparable quality on CIFAR-10 and CelebA, generates 256 × 256 LSUN samples, and reports an FID of 9.36 on CIFAR-10.
Takeaways & Limitations
The method reduces sampling to one evaluation while retaining competitive image quality and meaningful latent-space interpolation without adversarial training.
Takeaways & Limitations
At 256 × 256 resolution, the model produces blurry samples with poorly defined textures, which the authors associate with the pixel-level replication objective.
Abstract
from arXiv · showhide
Iterative generative models, such as noise conditional score networks and denoising diffusion probabilistic models, produce high quality samples by gradually denoising an initial noise vector. However, their denoising process has many steps, making them 2-3 orders of magnitude slower than other generative models such as GANs and VAEs. In this paper, we establish a novel connection between knowledge distillation and image generation with a technique that distills a multi-step denoising process into a single step, resulting in a sampling speed similar to other single-step generative models. Our Denoising Student generates high quality samples comparable to GANs on the CIFAR-10 and CelebA datasets, without adversarial training. We demonstrate that our method scales to higher resolutions through experiments on 256 x 256 LSUN. Code and checkpoints are available at https://github.com/tcl9876/Denoising_Student
1 Introduction
Score-based generative models produce high-quality images without adversarial training, but iterative sampling is much slower than single-step alternatives. The paper proposes distilling this process into a single evaluation while retaining image quality and latent-space properties.
- Score-based generative models gradually reverse noise and can rival leading GANs without adversarial training.
- Sampling can require 20-100 evaluations for DDIMs and up to three orders of magnitude more time than single-evaluation GANs.
- Knowledge distillation trains a Denoising Student to predict a DDIM’s output from the same initial Gaussian noise in one function evaluation.
- The approach is reported as 20× to 1000× more efficient than existing score-based methods, without adversarial training or surrogate losses.
- Denoising Student produces samples comparable to GANs on CIFAR-10 and CelebA 64 × 64 and generates 256 × 256 LSUN images in one step.
2 Knowledge Distillation in Deterministic Generative Models
Knowledge distillation is applicable to deterministic iterative generative processes, motivating DDIMs as teacher models whose multi-step mappings can be compressed into faster students.
- 2.1 Knowledge Distillation: Knowledge distillation compresses an expensive, high-performing teacher into a smaller student that approximates the teacher’s output distribution.
- 2.1 Knowledge Distillation: Deterministic teacher functions are required for knowledge distillation, whereas stochastic MCMC procedures make score-based and energy-based models infeasible teachers.
- 2.2 Denoising Diffusion Implicit Models: DDIMs provide suitable teachers because their generative processes are deterministic while retaining iterative-model sample quality and offering somewhat faster sampling.
- 2.2 Denoising Diffusion Implicit Models: DDIM inference starts from a noisy latent and reverses a parameterized noise-adding process through latent variables with decreasing α1:T.
- 2.2 Denoising Diffusion Implicit Models: Each DDIM sampling step is deterministic given the current latent, so the entire process is determined by the initial latent xT and can be modeled as pθ(x0|xT).
- 2.2 Denoising Diffusion Implicit Models: The DDIM teacher remains computationally expensive because evaluating its deterministic mapping requires T neural-network forward passes.
3 Generative Modeling with Denoising Student
Denoising Student training matches a student conditional distribution to a deterministic DDIM teacher using a Gaussian formulation and synthetic teacher-generated examples.
- 3 Generative Modeling with Denoising Student: The student is trained by minimizing the discrepancy between its conditional distribution pstudent(x0|xT) and the teacher distribution pteacher(x0|xT).
- 3 Generative Modeling with Denoising Student: The student uses a trainable mean Fstudent(xT) and unit variance, while Gaussian noise is added to the deterministic teacher output.
- 3 Generative Modeling with Denoising Student: Training samples random xT from the prior, computes the corresponding DDIM teacher output, and minimizes the resulting objective.
- 3 Generative Modeling with Denoising Student: The method avoids jointly training two networks, unlike approaches requiring a GAN discriminator or VAE inference model.
- 3 Generative Modeling with Denoising Student: Initializing the student with the teacher network’s architecture and weights allows it to inherit teacher knowledge and speeds training.
- 3 Generative Modeling with Denoising Student: The student predicts noise and is conditioned on timestep T to address the mismatch between predicting data and inheriting a time-conditioned teacher network.
4 Experiments
Experiments evaluate Denoising Student on four datasets, covering low- and high-resolution image generation, sample quality, sampling speed, and latent-space behavior. The method produces competitive low-resolution results, scales to 256 × 256 LSUN with some blurriness, and substantially accelerates sampling.
- Experimental Setup: Experiments cover CIFAR-10 32 × 32, CelebA 64 × 64, LSUN Bedroom 256 × 256, and LSUN Church 256 × 256.The experiments use pretrained teacher models, with 100-step DDIM teachers for CIFAR-10 and CelebA and a 50-step DDIM teacher for LSUN.
- Image Generation: CIFAR-10 and CelebA samples are reported as high quality and diverse, with a CIFAR-10 FID of 9.36 and a CelebA FID of 10.68.The CIFAR-10 FID is lower than several GANs and far lower than NVAE; the CelebA result is described as competitive.
- Image Generation: At 256 × 256, LSUN Bedroom and Church models learn structure, color, and larger details but produce blurry samples without defined textures.The paper attributes this bluriness to the pixel-level objective used to replicate the teacher’s output on higher-dimensional data.
- Sampling Speed: Denoising Student is about 100× faster than its teacher and 1000× faster than a DDPM when generating 50k CIFAR-10 images.The timing comparison used a batch size of 250 on a single Nvidia V100 GPU.
- Latent Space Manipulation: On CelebA, spherical interpolation between random xT values produces meaningful image interpolations, and changing latent size yields coherent images at different resolutions.These observations are presented as evidence that the student learns the teacher’s latent mapping and retains advanced generalization capabilities.
5 Related Work
The paper situates its approach within knowledge distillation and score-based generative modeling, emphasizing a shift from reducing network size toward reducing sampling evaluations.
- Existing knowledge-distillation research largely targets supervised tasks including image classification, language modeling, and speech recognition.
- Most knowledge-distillation work focuses on transferring knowledge effectively, while typically reducing network parameters rather than evaluation count.
- NCSNs and DDPMs estimate data-density scores and use denoising auto-encoder objectives within score-based generative modeling.
- NCSN and DDPM generation can be viewed as reversing a noise-adding stochastic differential equation, whereas DDIM sampling more closely resembles ODE integration.
6 Conclusion
The conclusion presents knowledge distillation as a simple, stable way to reduce sampling time in deterministic iterative generative models while preserving useful generation and latent-space capabilities.
- The method vastly reduces sampling time for certain iterative generative models with only a minor degradation in performance.
- It applies easily to trained models with deterministic generative processes because it uses a stable objective and needs no extra architectural considerations.
- The model produces significantly better samples than other non-adversarial, single-step methods such as VAEs.
- Learned latent representations support data manipulation through the latent space.
- Future work targets narrowing the teacher–student gap through advanced distillation and producing sharper high-resolution images.
A Samples
The appendix presents uncurated samples, nearest-neighbor visualizations, interpolations, and cross-resolution comparisons for CIFAR-10, CelebA, LSUN Bedroom, and LSUN Church.
- Samples: Additional uncurated samples are shown for CIFAR-10, CelebA, LSUN Bedroom, and LSUN Church.
- Nearest neighbors: Nearest-neighbor visualizations are included because the student trains on teacher outputs rather than real data, so memorization of training examples is not expected.
- Nearest neighbors: CIFAR-10 generated images are compared with their nearest neighbors in the training set in Inception feature space.
- Interpolation: Extended interpolation results cover CelebA, LSUN Bedroom, and LSUN Church across separate row groups.
- Cross-resolution generalization: A 50-step DDIM teacher and the student produce decent 256×384 samples despite neither being trained on that image size.
B Experimental Details
The experimental details specify optimization choices and training observations, including dataset-specific loss, gradient, and learning-rate adjustments.
- Optimization: Experiments use Adam, learning-rate warmup without decay, no dropout or other regularization, and test loss on unseen images for early stopping.
- Training observations: For CIFAR-10, L1 distance performed better than L2 during training.
- Training observations: For LSUN Bedroom, removing value-based gradient clipping improved stability after gradients consistently exceeded norm 10^5.
- Training observations: For LSUN Bedroom, increasing β1 from 0.9 to 0.98 and lowering the learning rate from 2 × 10^-5 to 5 × 10^-6 stabilized training; these settings were applied to Church without sweeping.
- Hyperparameters: Table 3 lists the hyperparameters used in the experiments.