PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved Distribution Matching Distillation for Fast Image Synthesis

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

arXiv:2405.14867v2cs.CV

TL;DR

DMD’s regression loss requires expensive teacher-generated data and ties distillation to teacher sampling paths, limiting scalable distribution matching. DMD2 removes that requirement, stabilizes training, adds GAN supervision from real images, and simulates inference-time inputs for multi-step generation, achieving state-of-the-art one-step results and megapixel SDXL distillation.

  • Problem

    DMD requires costly teacher-generated noise-image pairs for stable training, tying the student to teacher sampling paths and limiting scalable distribution matching.

  • Method

    DMD2 removes regression loss, uses a two time-scale update rule, integrates a GAN objective with real images, and simulates inference-time generator inputs for multi-step training.

  • Results

    DMD2 achieves state-of-the-art one-step image generation, reporting FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014, and produces megapixel images from SDXL.

  • Takeaways & Limitations

    DMD2 enables scalable, few-step distribution-matching distillation without paired teacher trajectories and can outperform the teacher using as few as 4 sampling steps.

  • Takeaways & Limitations

    The distilled generator has slightly lower image diversity than the teacher, still requires four steps to match the largest SDXL model, and uses a fixed guidance scale during training.

Abstract

from arXiv · show

Recent approaches have shown promises distilling diffusion models into efficient one-step generators. Among them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, without enforcing a one-to-one correspondence with the sampling trajectories of their teachers. However, to ensure stable training, DMD requires an additional regression loss computed using a large set of noise-image pairs generated by the teacher with many steps of a deterministic sampler. This is costly for large-scale text-to-image synthesis and limits the student's quality, tying it too closely to the teacher's original sampling paths. We introduce DMD2, a set of techniques that lift this limitation and improve DMD training. First, we eliminate the regression loss and the need for expensive dataset construction. We show that the resulting instability is due to the fake critic not estimating the distribution of generated samples accurately and propose a two time-scale update rule as a remedy. Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images. This lets us train the student model on real data, mitigating the imperfect real score estimation from the teacher model, and enhancing quality. Lastly, we modify the training procedure to enable multi-step sampling. We identify and address the training-inference input mismatch problem in this setting, by simulating inference-time generator samples during training time. Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64x64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost. Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods.

1 Introduction

DMD2 removes DMD’s costly regression-data requirement, stabilizes distribution matching, adds real-data supervision, and supports multi-step generation. These changes produce fast generators that outperform their teachers and reach strong one-step benchmarks.

  • Diffusion sampling commonly requires dozens of neural-network evaluations, making high-resolution text-to-image synthesis slow and expensive.
  • DMD requires millions of teacher-generated noise-image pairs for regression-based stabilization, which is costly and ties the student to teacher sampling paths.
  • DMD2 removes the regression loss while preserving training stability, eliminating costly data collection and enabling more flexible, scalable training.
  • 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014 are DMD2’s reported one-step FID scores, while SDXL distillation produces megapixel images.
  • A two time-scale update rule addresses instability caused by an insufficiently trained fake diffusion critic that poorly estimates generated-sample distributions.
  • Integrating a GAN objective lets DMD2 distinguish student samples from real images, providing distribution-level supervision that mitigates teacher approximation errors and enhances quality.
  • DMD2 simulates inference-time generator inputs during training to avoid domain mismatch and supports multi-step generators.

2 Related Work

Diffusion distillation methods accelerate generation by approximating teacher sampling trajectories, while GAN-based approaches align students with teacher distributions more broadly. DMD instead uses score-based distribution matching, but originally depended on regression data for stability.

  • Diffusion distillation methods train generators to approximate a teacher’s ODE sampling trajectory in fewer steps.
  • GANs: GAN-based distillation aligns students with teacher distributions through adversarial objectives, including image-space and latent-space diffusion discriminators.
  • Score Distillation: Score distillation uses pretrained text-to-image diffusion scores to align rendered views or student outputs with a text-conditioned image distribution.
  • Diffusion Distillation: DMD parameterizes real and fake score functions with diffusion models and achieved strong one-step synthesis performance, but required precomputed noise-image pairs for stability.

3 Background: Diffusion and Distribution Matching Distillation

Diffusion models generate high-quality images through iterative denoising, but their many sampling steps make generation slow. DMD instead distills a many-step teacher into a one-step generator by matching diffused output distributions with score-function gradients, while retaining a costly regression regularizer for stability.

  • Diffusion models: Diffusion models progressively add noise to real samples and learn to reverse this corruption through iterative denoising.The learned denoised estimate is related to the score function of the diffused data distribution.
  • Diffusion models: Dozens to hundreds of denoising steps are typically required to sample an image, increasing generation cost.
  • Distribution Matching Distillation: DMD distills a many-step diffusion model into a one-step generator by minimizing approximate KL divergences between diffused teacher and generator distributions.
  • Distribution Matching Distillation: DMD computes its training gradient from the difference between real and fake score functions, using a frozen teacher score and an online fake score estimator.The fake score is trained with denoising score matching on samples produced by the generator.
  • Distribution Matching Distillation: DMD’s additional regression loss compares generator outputs with teacher-generated images from the same noise inputs, requiring paired noise-image data.For large-scale synthesis, constructing these pairs can become a major computational bottleneck and encourages adherence to teacher sampling paths.

4 Improved Distribution Matching Distillation

DMD2 removes DMD’s regression loss, stabilizes distribution matching with more frequent fake-score updates, adds GAN supervision from real images, and supports multi-step generation through backward simulation. These changes target costly data construction, teacher-score approximation errors, and training–inference mismatch.

  • Training procedure: DMD2’s method alternates generator optimization with fake-score and discriminator training, updating the fake score more frequently than the generator.
  • Removing regression loss: DMD2 removes the regression loss and its costly precomputed noise-image pairs, enabling more flexible and scalable distillation.The method revisits DMD design choices to avoid tying the student to teacher sampling paths.
  • Removing regression loss: Naively removing regression causes unstable training because the fake diffusion critic fails to accurately track generated-sample statistics.Generated brightness and other statistics can fluctuate without converging.
  • Removing regression loss: Five fake-score updates per generator update stabilize regression-free training and match original DMD quality on ImageNet while converging faster.
  • GAN loss and real data: DMD2 adds a GAN objective whose discriminator distinguishes student-generated samples from real images, providing distribution-level supervision independent of teacher sampling trajectories.Training with real data is intended to mitigate approximation errors in the teacher’s real score function and improve image quality.
  • Multi-step generator simulation: Multi-step training replaces noisy real-image inputs with noisy synthetic images generated by the current student, matching inference-time inputs through backward simulation.This addresses the mismatch created when inference inputs after the first step come from previous generator sampling steps.

5 Experiments

DMD2 is evaluated across ImageNet, COCO, SDXL, user studies, and component ablations, showing strong quality, alignment, and fast sampling. The experiments also identify the contributions of TTUR, GAN supervision, distribution matching, and backward simulation.

  • 5 Experiments: DMD2 significantly outperforms existing distillation techniques and the ODE-sampled teacher on ImageNet-64×64 with a single forward pass.The authors attribute this to removing the ODE regression loss and adding a GAN term.
  • 5 Experiments: 4-step SDXL distillation achieves FID 19.32 and CLIP score 0.332 on COCO 2014, rivaling the teacher in image quality and prompt coherence.The model uses 4 sampling steps, compared with 50 for the teacher in the user study.
  • 5 Experiments: The user study finds DMD2 best for image quality and prompt alignment among distilled models using 4 steps, while the teacher uses 50.DMD2 outperforms its teacher in image quality for 24% of samples and requires 25× fewer forward passes.
  • 5 Experiments: Removing DMD’s regression loss degrades ImageNet FID to 3.48, while TTUR restores baseline performance and the GAN loss adds a further 1.1-point FID improvement.The integrated method also outperforms GAN-only training.
  • 5 Experiments: SDXL ablations show that removing distribution matching, GAN supervision, or backward simulation harms stability, realism, text alignment, or image quality.Without GAN loss, outputs become oversaturated and oversmoothed; omitting backward simulation reduces patch FID.
  • 5 Experiments: Visual comparisons report superior realism and text alignment for DMD2 using identical noise and prompts, with 4-step distilled models compared against a 50-step teacher.The evaluation includes ImageNet and COCO benchmarks, FID, CLIP Score, patch FID, and human evaluations.

6 Limitations

The authors report limitations involving diversity, sampling steps, guidance flexibility, computational cost, and the scope of the optimization objective.

  • 6 Limitations: The distilled generator has slightly lower image diversity than the teacher models.This limitation is reported alongside otherwise superior image quality and text alignment.
  • 6 Limitations: The largest SDXL model still requires four sampling steps to match its quality.The authors identify this as an area for further improvement.
  • 6 Limitations: Fixed guidance-scale training limits user flexibility, while human feedback or reward functions are not incorporated.The authors suggest variable guidance scales and reward-based enhancements as future directions.
  • 6 Limitations: Training large-scale generative models remains computationally intensive and inaccessible to many researchers.The authors frame improved accessibility as a future goal.

7 Broader Impact

The work may benefit creative, educational, and entertainment applications through faster high-quality image synthesis, while also posing misuse and bias risks.

  • 7 Broader Impact: Fast image synthesis could benefit graphic design, animation, digital art, education, and entertainment.The authors mention efficient visual creation, personalized learning materials, and immersive experiences.
  • 7 Broader Impact: Misuse could enable misinformation and fake profiles, while biased training data could perpetuate or amplify unfair impacts on specific groups.The authors express interest in monitoring mechanisms and methods for diversity and fairness.

A SD v1.5 Results

Appendix results compare the one-step generator distilled from SD v1.5 with competing approaches.

  • A SD v1.5 Results: Table 5 presents detailed comparisons between the one-step SD v1.5-distilled generator and competing approaches.

B Text-to-Image Synthesis Further Analysis

The SDXL-based four-step generator achieves strong image quality and text alignment while maintaining competitive diversity against other distilled methods.

  • Our model achieves the best image quality among the compared SDXL-based methods, measured by the lowest FID and Patch FID scores.
  • Its text alignment is comparable to SDXL-Turbo while its diversity score remains better balanced than methods with weaker prompt coherence.
  • SDXL-Lightning has higher diversity, but its lower CLIP score and human evaluation indicate considerably worse text alignment.

C Two Time-scale Update Rule Further Analysis

Removing DMD’s regression objective creates training instability, while updating the fake diffusion critic five times per generator update stabilizes training effectively.

  • Five fake diffusion critic updates per generator update stabilize pixel-brightness statistics compared with the unstable one-update baseline.
  • An update frequency of 10 provides excellent stability but significantly slows training, whereas five updates achieve the best balance.
  • For new models and datasets, the iteration number should be the smallest value that stabilizes general image statistics such as pixel brightness.
  • The five-update rule outperforms both larger fake-critic update counts and an asynchronous learning rate five times larger than the generator’s.

D Additional Text-to-Image Synthesis Results

Additional comparisons evaluate four-step distilled models against the SDXL teacher and competing methods, alongside one-step SDXL samples at 1024×1024 resolution.

  • The paper presents additional visual comparisons for four-step distilled models and one-step generator outputs.
  • Four-step comparisons use identical noise and text prompts for our model, the SDXL teacher, and selected competing methods.The teacher uses 50 sampling steps with classifier-free guidance, while distilled models use four.
  • Our model produces images with superior realism and text alignment in the additional visual comparison.
  • Additional 1024×1024 samples are generated by the one-step generator distilled from SDXL.

E ImageNet Visual Results

The paper presents ImageNet visual results and implementation details spanning one-step distillation, GAN classification, SDXL distillation, and evaluation procedures.

  • The paper also describes qualitative visualizations, prompt sets, human preference evaluation, and COCO image-generation evaluation procedures.
  • One-step ImageNet samples from the distilled generator achieve an FID of 1.28.
  • The GAN classifier attaches a prediction head to the fake diffusion model’s middle-block output and reduces feature maps to a single prediction vector.
  • The ImageNet implementation distills a one-step generator from a pretrained EDM model using AdamW training on seven A100 GPUs.
  • The SDXL experiments train both one-step and four-step generators and address one-step block-noise artifacts with timestep shifting and regression pretraining.
Loading 2405.14867v2…