One-step Diffusion with Distribution Matching Distillation

TL;DR

This paper presents Distribution Matching Distillation, a method to convert diffusion models into fast, one-step image generators with minimal quality loss, achieving high speed and competitive image quality.

Contribution

The introduction of Distribution Matching Distillation enables one-step diffusion-based image generation with minimal quality degradation, outperforming previous few-step methods.

Findings

Achieves 2.62 FID on ImageNet 64x64

Generates images at 20 FPS with FP16 inference

Comparable quality to Stable Diffusion

Abstract

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.