PaperlayerView exact revision on arXiv

Source-linked AI summary

Progressive Distillation for Fast Sampling of Diffusion Models

Tim Salimans, Jonathan Ho

arXiv:2202.00512v2cs.LGcs.AIstat.ML

TL;DR

Diffusion models produce high-quality samples but often require hundreds or thousands of evaluations, limiting sampling speed. The paper stabilizes few-step diffusion models and progressively distills deterministic samplers by halving their step count, reaching competitive quality in as few as 4 steps. The method reduces test-time computation while retaining strong image-generation performance, though the demonstrated student retains the teacher’s architecture and parameter count.

  • Problem

    Unconditional and class-conditional diffusion sampling remains slow because high-quality generation requires hundreds or thousands of network evaluations.

  • Method

    Progressive distillation repeatedly trains a student to reproduce a deterministic teacher sampler with half as many steps, alongside stable few-step parameterizations and loss choices.

  • Results

    Across four image datasets, progressive distillation produces near-optimal results through 4 or 8 sampling steps and compares favorably with alternative fast-sampling methods.

  • Takeaways & Limitations

    The method makes deterministic diffusion sampling substantially cheaper at test time while preserving competitive sample quality in low-step regimes.

  • Takeaways & Limitations

    The student currently uses the same architecture and number of parameters as the teacher, and broader data modalities remain future work.

Abstract

from arXiv · show

Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into a new diffusion model that takes half as many sampling steps. We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, ImageNet, and LSUN, we start out with state-of-the-art samplers taking as many as 8192 steps, and are able to distill down to models taking as few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time than it takes to train the original model, thus representing an efficient solution for generative modeling using diffusion at both train and test time.

1 INTRODUCTION

Diffusion models achieve strong generative-modeling results, but unconditional and class-conditional sampling remains slow because it requires hundreds or thousands of network evaluations. Progressive distillation addresses this bottleneck by compressing deterministic diffusion sampling into fewer steps while retaining sample quality.

  • 1 INTRODUCTION: Diffusion models already achieve strong results on image-generation benchmarks and likelihood-based density estimation.
  • 1 INTRODUCTION: Hundreds or thousands of evaluations make unconditional and standard class-conditional diffusion sampling impractical for many applications.
  • 1 INTRODUCTION: Progressive distillation converts an N-step DDIM sampler into an N/2-step model and repeats this process to reach as few as 4 steps.The procedure is described as learning to integrate the probability flow ODE in fewer steps, or amortizing that integration.

2 BACKGROUND ON DIFFUSION MODELS

Diffusion models learn to denoise progressively corrupted data and can generate samples by numerically integrating a probability flow ODE. Numerical accuracy improves with more integration steps, creating a quality–sampling-time trade-off that motivates distillation.

  • 2 BACKGROUND ON DIFFUSION MODELS: The diffusion model trains a denoising function to estimate clean data from noisy latent variables using a weighted reconstruction loss.The noisy latent is formed from data and Gaussian noise, and the weighting may be justified through likelihood or denoising score matching.
  • 2 BACKGROUND ON DIFFUSION MODELS: The probability flow ODE deterministically maps Gaussian noise to data samples using the learned denoising model.DDIM can be interpreted as an integration rule for this ODE.
  • 2 BACKGROUND ON DIFFUSION MODELS: DDIM provides a deterministic sampling update that performs better than standard Euler or Runge–Kutta integration rules in the authors’ setting.
  • 2 BACKGROUND ON DIFFUSION MODELS: More integration steps reduce numerical error and can improve sample quality, but hundreds or thousands of steps make high-quality sampling costly.

3 PROGRESSIVE DISTILLATION

Progressive distillation trains a student to reproduce two teacher DDIM steps with one step, then repeatedly makes the student the next teacher while halving the sampling budget. The procedure remains close to standard diffusion training and uses sharper student targets.

  • 3 PROGRESSIVE DISTILLATION: Each progressive-distillation iteration halves the required sampling steps by distilling a slow teacher into a faster student.After training, the student becomes the next teacher and the process repeats with N/2 steps.
  • 3 PROGRESSIVE DISTILLATION: The student is initialized as a copy of the teacher and is trained on noisy data with targets derived from two teacher DDIM updates.
  • 3 PROGRESSIVE DISTILLATION: Sharper student predictions enable faster sampling than the original denoising target, which averages multiple possible clean data points consistent with a noisy input.
  • 3 PROGRESSIVE DISTILLATION: Progressive distillation uses discrete time with the highest time index set to zero signal-to-noise ratio, matching the pure-noise test-time input distribution.

4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS

Few-step distillation makes unstable low-signal-to-noise predictions especially important, so the paper evaluates parameterizations and loss weightings designed to stabilize the implied clean-data prediction. Direct x, v, and combined (x, ϵ) specifications work well, while one ϵ/truncated-SNR combination is unstable.

  • 4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS: As signal-to-noise ratio approaches zero, predicting ϵ becomes unstable because converting it to x divides by α_t, whose value approaches zero.The instability matters increasingly as sampling steps decrease because later updates have less opportunity to correct early errors.
  • 4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS: The authors therefore seek parameterizations whose implied x prediction remains stable across changing signal-to-noise ratios.
  • 4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS: Direct x prediction, joint x-and-ϵ prediction, and v prediction are proposed as stable parameterizations for progressive distillation.The v parameterization predicts v = α_tϵ − σ_tx and reconstructs x from v.
  • 4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS: The standard signal-to-noise-ratio weighting is unsuitable for distillation because it assigns zero weight at zero signal-to-noise ratio.The study instead examines alternative loss weightings, including truncated-SNR weighting and a v-space loss.

5 EXPERIMENTS

Experiments show that stable parameterizations support progressive distillation across image-generation benchmarks, preserving sample quality down to 4–8 sampling steps while standard samplers degrade sharply below 128 steps.

  • 5.1 MODEL PARAMETERIZATION AND TRAINING LOSS: Predicting x directly, v, or the combination (ϵ, x) gives strong specifications, while ϵ prediction with truncated-SNR weighting is unstable.Predicting v is the most stable because DDIM step sizes become independent of SNR, although direct x prediction performs slightly better in the ablation.
  • 5.2 PROGRESSIVE DISTILLATION: FID remains near optimal through 4 or 8 sampling steps across CIFAR-10, ImageNet, and LSUN, whereas DDIM and stochastic baselines degrade sharply below 128 steps.At 1 or 2 steps, distilled sample quality declines more quickly; the comparison covers four datasets and uses FID.
  • 5.2 PROGRESSIVE DISTILLATION: Distilled ImageNet samples preserve the mapping from fixed input noise to output images as sampling steps decrease.The figure conditions 64 × 64 ImageNet samples on the malamute class while varying the number of steps.

6 RELATED WORK ON FAST SAMPLING

The paper situates progressive distillation among sampler, diffusion-process, and model-specification approaches to fast generation, emphasizing logarithmic distillation cost relative to prior one-step distillation.

  • 6 RELATED WORK ON FAST SAMPLING: Figure 4 evaluates FID on CIFAR-10, ImageNet, and LSUN by comparing distilled models with DDIM and an optimized stochastic sampler.CIFAR-10 uses four random seeds; the other datasets use single runs, and the stochastic baseline is separately tuned for each step count.
  • 6 RELATED WORK ON FAST SAMPLING: Unlike prior DDIM-to-one-step distillation, progressive distillation avoids running the original model at its full step count and scales logarithmically in total distillation time.Its per-iteration model-evaluation cost is independent of the teacher’s number of sampling steps.
  • 6 RELATED WORK ON FAST SAMPLING: Related approaches adjust samplers or diffusion processes, reduce timesteps by optimization, or train diffusion models in latent space.Examples include stochastic and SDE integrators, timestep dynamic programming, modified forward/reverse processes, and latent-space diffusion.

7 DISCUSSION

Progressive distillation drastically reduces the sampling steps required for high-quality diffusion-model generation, while parameterization and loss choices also contribute to the reported progress.

  • Progressive distillation drastically reduces the sampling steps required for high-quality image generation with deterministic diffusion samplers like DDIM.The authors frame cheaper test-time execution as important for practical applications constrained by running time and computation.
  • The current method keeps the student architecture and parameter count identical to the teacher, leaving smaller-student models as future work.The authors also propose extending progressive distillation beyond images to modalities such as audio.
  • The reported progress reflects both the distillation procedure and alternative diffusion-model parameterizations and training-loss weightings.

ETHICS STATEMENT

The supplied passages combine an ethics statement with technical material on diffusion sampling, DDIM, angular parameterization, architectures, and stochastic evaluation. They state general risks of generative models while describing methods intended to reduce sampling cost.

  • ETHICS STATEMENT: Generative models can be used to create fake content and may exhibit bias when trained on data sets that are not carefully curated.The paper says its specific contribution focuses on reducing test-time computational demands and reports no specific concerns about that contribution.
  • DDIM can be interpreted as a first-order integration rule for the probability flow ODE, linking deterministic sampling to numerical integration.
  • The angular parameterization uses φ and predicted velocity v̂φ to express DDIM updates as circular motion in the (zφt, v̂φt) basis.The relationship among zφt, velocity, αt, σt, x, and noise is visualized in Figure 5.
  • The distilled model can also be evaluated with stochastic sampling, whose performance lies between the undistilled stochastic model and the distilled DDIM model.The stochastic results use the best FID selected by a grid search over 11 noise levels.

G DERIVATION OF THE DISTILLATION TARGET

Progressive distillation derives a student denoising target so one deterministic DDIM step matches two teacher steps, then trains the student to predict that target. This exact matching relies on deterministic sampling, while repeated distillation narrows the student's time coverage and generality.

  • The distillation target differs from standard diffusion training: it is the ˜x required to match the teacher, rather than the clean data x.
  • One student DDIM step can match two teacher DDIM steps when the student predicts the derived target ˜x and is trained with squared-error denoising loss.The student’s one-step sample becomes identical to the teacher’s two-step sample when ˜x is predicted exactly.
  • Progressive distillation can match two teacher steps with one student step because DDIM teacher and student updates are deterministic neural-network mappings.Composing two stochastic DDPM steps produces a non-Gaussian distribution that a single Gaussian DDPM step cannot represent without fidelity loss.
  • As distillation progresses, student models lose generality because they are evaluated only at a small discrete set of times rather than arbitrary continuous times.This reduced generality frees modeling capacity to match the teacher without increasing model size.

H ADDITIONAL RANDOM SAMPLES

Additional fixed-seed samples compare distilled diffusion models across sampling-step counts on CIFAR-10, ImageNet, and LSUN. The deterministic input-to-output mapping is mostly preserved as steps decrease, with a slight quality loss at one step.

  • The samples cover distilled models using 256, 4, and 1 sampling steps, and all samples are uncurated.
  • The input-noise-to-output mapping is mostly preserved when reducing distilled samplers from many steps to one, with only a slight image-quality loss.Using identical random input noise enables direct side-by-side comparison across step counts.
  • Fixed-seed sample comparisons are provided for CIFAR-10, 64 × 64 ImageNet conditioned on ‘coral reef’ and ‘sports car’, and LSUN bedrooms and church-outdoor scenes.

I ABLATION WITH FASTER DISTILLATION SCHEDULES

Faster progressive-distillation schedules reduce computational requirements while retaining excellent results on ImageNet and LSUN. Under limited computation, reducing updates per halving performs better than skipping distillation iterations by dividing steps by four.

  • CIFAR-10 ablations compare 25k, 10k, and 5k updates per student against the proposed 50k-update schedule.The study evaluates how reducing optimization steps affects progressive distillation efficiency.
  • With limited computation, fewer parameter updates per halving outperform skipping distillation iterations by reducing the sampling-step count fourfold.
  • Excellent results are achieved on ImageNet and LSUN with a faster distillation schedule using fewer parameter updates.Figure 13 compares the proposed schedule with a 10k-update fast schedule across reported sampling-step counts.
Loading 2202.00512v2…