EM Distillation for One-step Diffusion Models

TL;DR

EM Distillation (EMD) is a novel maximum likelihood approach that efficiently converts diffusion models into one-step generators, maintaining high perceptual quality and outperforming existing methods on standard benchmarks.

Contribution

The paper introduces EM Distillation, a new EM-based method for distilling diffusion models into one-step generators with minimal quality loss and improved stability.

Findings

EMD achieves better FID scores on ImageNet-64 and ImageNet-128.

It outperforms existing one-step generative methods.

It effectively distills text-to-image diffusion models.

Abstract

While diffusion models can learn complex distributions, sampling requires a computationally expensive iterative process. Existing distillation methods enable efficient sampling, but have notable limitations, such as performance degradation with very few sampling steps, reliance on training data access, or mode-seeking optimization that may fail to capture the full distribution. We propose EM Distillation (EMD), a maximum likelihood-based approach that distills a diffusion model to a one-step generator model with minimal loss of perceptual quality. Our approach is derived through the lens of Expectation-Maximization (EM), where the generator parameters are updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents. We develop a reparametrized sampling scheme and a noise cancellation technique that together stabilizes the distillation process. We further reveal an interesting connection of our method with existing methods that minimize mode-seeking KL. EMD outperforms existing one-step generative methods in terms of FID scores on ImageNet-64 and ImageNet-128, and compares favorably with prior work on distilling text-to-image diffusion models.