PaperlayerView exact revision on arXiv

Source-linked AI summary

One-step Diffusion with Distribution Matching Distillation

Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman, Taesung Park

arXiv:2311.18828v4cs.CV

TL;DR

Diffusion models produce realistic images but require slow iterative sampling, motivating one-step alternatives. DMD matches the teacher and generator distributions with two score models and supplements this objective with regression, achieving strong quality at much higher speed.

  • Problem

    Diffusion models generate realistic images but typically require tens to hundreds of denoising steps, making sampling slow.

  • Method

    DMD trains a one-step generator with a distribution-matching objective based on two score functions, plus regression against precomputed multi-step outputs.

  • Results

    DMD outperforms published few-step diffusion methods, reaching 2.62 FID on ImageNet 64×64 and 11.49 FID on MS-COCO 2014-30k while reducing neural-network evaluations 100×.

  • Takeaways & Limitations

    DMD maintains image fidelity comparable to Stable Diffusion while generating 512×512 images at 20 FPS with FP16 inference.

  • Takeaways & Limitations

    A slight quality gap remains versus finer diffusion discretizations, and training both the fake score function and generator requires significant memory.

Abstract

from arXiv · show

Diffusion models generate high-quality images but require dozens of forward passes. We introduce Distribution Matching Distillation (DMD), a procedure to transform a diffusion model into a one-step image generator with minimal impact on image quality. We enforce the one-step image generator match the diffusion model at distribution level, by minimizing an approximate KL divergence whose gradient can be expressed as the difference between 2 score functions, one of the target distribution and the other of the synthetic distribution being produced by our one-step generator. The score functions are parameterized as two diffusion models trained separately on each distribution. Combined with a simple regression loss matching the large-scale structure of the multi-step diffusion outputs, our method outperforms all published few-step diffusion approaches, reaching 2.62 FID on ImageNet 64x64 and 11.49 FID on zero-shot COCO-30k, comparable to Stable Diffusion but orders of magnitude faster. Utilizing FP16 inference, our model generates images at 20 FPS on modern hardware.

1. Introduction

Diffusion models offer high realism but rely on slow iterative sampling, while existing distillation methods struggle to preserve quality. DMD instead matches the teacher and student distributions, combining distribution matching with regression to achieve fast one-step generation.

  • Diffusion sampling typically requires tens to hundreds of costly neural-network evaluations, limiting interactive image generation.
  • Previous distillation methods compress the diffusion model’s noise→image mapping into a single-pass student, but fitting this complex mapping is demanding.
  • DMD trains the student to make generations indistinguishable from the original diffusion model rather than enforcing individual noise–image correspondences.
  • A precomputed regression loss regularizes distribution matching and aligns one-step generations with the teacher’s large-scale structure.
  • 2.62 FID on ImageNet 64×64 and 11.49 FID on MS-COCO 2014-30k accompany a 100× reduction in neural-network evaluations and 20 FPS 512×512 generation with FP16.

2. Related Work

Diffusion acceleration research includes faster samplers and distillation, while distribution-matching methods provide an alternative to reconstructing individual samples. DMD combines these ideas using score-based distribution matching and teacher-alignment regression.

  • Diffusion acceleration methods either reduce sampling steps with fast samplers or distill multi-step outputs into a single-step student.
  • Fast samplers can reduce diffusion sampling from thousands of steps to 20–50, but further reductions often cause a catastrophic performance decrease.
  • Prior regression-based distillation precomputes denoising trajectories, but running full trajectories for each loss realization is expensive.
  • DMD uses two diffusion denoisers to model real and fake distributions, expressing its distribution-matching update as the difference between their score functions.
  • Distribution-matching generative methods compare synthetic and target distributions rather than relying on sample reconstruction, including GMMD and GANs.
  • VSD connects pretrained text-to-image diffusion models with distribution matching for unpaired, text-conditioned 3D synthesis, motivating DMD’s refinement.

3. Distribution Matching Distillation

DMD distills a pretrained diffusion model into a one-step generator by matching real and generated distributions with approximate score-based KL optimization, while regression preserves the teacher’s large-scale structure and modes.

  • 3.1. Pretrained base model and One-step generator: DMD trains a one-step generator to match the pretrained diffusion model’s output distribution rather than reproduce its exact noise-to-image mapping.The generator is initialized from the pretrained denoiser and removes time conditioning.
  • 3.2. Distribution Matching Loss: Two diffusion denoisers model the perturbed real and fake distributions, with the fake denoiser dynamically updated as generated samples change.Gaussian perturbation makes the distributions overlap, ensuring the score-based objective is well-defined across the ambient space.
  • 3.2. Distribution Matching Loss: The distribution-matching loss minimizes KL divergence through the difference between real and fake distribution scores.Exact scores are replaced with scores estimated by diffusion models on Gaussian-perturbed samples.
  • 3.2. Distribution Matching Loss: A timestep-dependent weighting factor normalizes gradient magnitudes across noise levels and improves the distribution-matching optimization.The weighting uses mean absolute error between the denoised image and input, with Tmin = 0.02T and Tmax = 0.98T.
  • 3.3. Regression loss and final objective: The regression loss regularizes distribution matching and preserves modes that can otherwise be dropped because score optimization is invariant to probability-density scaling.It matches one-step outputs to paired multi-step diffusion outputs generated from the same noise inputs.
  • 3.3. Regression loss and final objective: The final training objective combines distribution matching and regression losses on separate unpaired-fake and paired noise-output data streams.The procedure also supports classifier-free guidance by constructing paired data from the guided model and substituting the corresponding real score.

4. Experiments

Across class-conditional and text-to-image benchmarks, DMD achieves strong image quality while substantially accelerating diffusion generation. Ablations show that both distribution matching and regression losses contribute to realism, diversity, and stable training.

  • Class-conditional Image Generation: DMD nearly matches the original diffusion model on ImageNet-64×64 while increasing generation speed 512-fold.
  • Class-conditional Image Generation: DMD reaches a competitive FID of 2.66 on class-conditional CIFAR-10.
  • Ablation Studies: Removing distribution matching reduces realism and structural integrity, whereas removing regression causes instability, mode collapse, and reduced diversity.
  • Ablation Studies: The proposed sample weighting strategy improves FID by 0.9 by normalizing gradient magnitudes across noise levels and stabilizing optimization.
  • Text-to-Image Generation: With FP16 inference, the text-to-image model generates 20 frames per second, supporting interactive applications.

5. Limitations

The one-step model retains comparable quality to Stable Diffusion while remaining much faster, but a small quality gap remains versus finer diffusion discretizations and training uses substantial memory.

  • Quality gap: A slight quality discrepancy remains between the one-step model and diffusion sampling with 100 or 1000 neural network evaluations.This is identified as an ongoing limitation despite promising results.
  • Training cost: Fine-tuned fake-score and generator weights lead to significant memory usage during training.The paper notes LORA as a potential way to address this issue.
  • Speed-quality trade-off: Around 30× faster inference achieves comparable image quality to the original Stable Diffusion model.The comparison is presented in accompanying video material.

B. Implementation Details

Implementation uses pretrained diffusion models, generated noise-image distillation datasets, specified optimization schedules, and benchmark protocols covering Stable Diffusion, acceleration methods, and COCO, ImageNet, and CIFAR-10 evaluation.

  • Loss construction: The implementation constructs the generator KL loss and trains a fake-score estimator parameterized by µfake.These procedures are detailed in Algorithms 2 and 3.
  • Distillation data: Distillation datasets are generated from pretrained EDM or Stable Diffusion models using deterministic samplers over prescribed trajectories.The configurations include Heun sampling for CIFAR-10 and ImageNet and PNMS sampling for Stable Diffusion.
  • Optimization: Training uses AdamW with explicitly specified learning rates, weight decay, beta parameters, and warmup settings.The reported learning rates range from 2e-6 for ImageNet to 1e-5 for Stable Diffusion configurations.
  • Training logs: Table 5 records training-log settings including maximum distribution-matching timestep, VAE type, and separate batch sizes for distribution matching and regression losses.The table distinguishes small and large VAE decoders and DM versus regression batch sizes.
  • Baselines: Baselines include the base diffusion model, fast diffusion solvers, and few-step diffusion distillation methods.Stable Diffusion uses a 50-step PNMS sampler, while UniPC and DPMSolver++ use default diffusers-library hyperparameters.
  • Evaluation: COCO evaluation generates 30,000 images from MS-COCO2014 validation prompts and computes FID and CLIP score against 40,504 real images.ImageNet and CIFAR-10 evaluation uses 50,000 generated images for each dataset and FID computation.

E. CIFAR-10 Experiments

The CIFAR-10 experiments follow the implementation setup and report sample-quality comparisons against competing methods in Table 6.

  • Experimental setup: CIFAR-10 models are trained following the setup described in Section B.1.
  • Results: Table 6 summarizes the sample-quality comparison on CIFAR-10.
  • Baselines: The CIFAR-10 baseline numbers are derived from Song et al., and some methods use class conditioning.

F. Derivation for Distribution Matching Gradient

The section presents the derivation for Equation 7, while the accompanying Figure 1 materials list the prompts used to produce its image comparisons.

  • Equation 7: The derivation for Equation 7 is presented in this section.
  • Figure 1 prompts: Figure 1 uses prompts including a DSLR photo of a golden retriever in heavy snow and a lightshow at the Dolomites.
  • Figure 1 prompts: The Figure 1 prompt list also includes detailed portrait prompts for an elderly woman and a warrior chief.

H. Equivalence of Noise and Data Prediction

The noise prediction model and data prediction model are interchangeable through a conversion rule.

  • The conversion rule allows the noise prediction model ϵ(x_t, t) and data prediction model µ(x_t, t) to be transformed into each other.

I. Further Analysis of the Regression Loss

The regression component stabilizes DMD training and mitigates mode collapse, while remaining robust to the choice of distance function. On CIFAR-10, LPIPS performs slightly better than standard L2 loss.

  • Regression loss stabilizes training and mitigates mode collapse.The paper mainly uses LPIPS because it is common in prior work.
  • 2.78 FID with L2 versus 2.66 with LPIPS on CIFAR-10 demonstrates robustness to different loss functions.

J. More Qualitative Results

Additional qualitative results cover ImageNet, LAION, and CIFAR-10. The reported samples include strong one-step results, including ImageNet FID=2.62 and LAION quality comparable to Stable Diffusion at substantially higher speed.

  • Qualitative coverage: The additional qualitative results include one-step LAION samples illustrating varied prompts such as a baby parrot, lighthouse, and Olympic swimmer.
  • ImageNet: ImageNet one-step samples achieve FID=2.62 in the class-conditional model.
  • LAION: The Stable Diffusion comparison reports higher image quality than previous few-step generators at the same speed or faster.
  • LAION: LAION one-step samples have comparable image quality to Stable Diffusion at 30× faster speed.
  • CIFAR-10: CIFAR-10 qualitative results report FID=2.66 for the class-conditional model and FID=3.77 for the unconditional model.
Loading 2311.18828v4…