PaperlayerView exact revision on arXiv

Source-linked AI summary

EM Distillation for One-step Diffusion Models

Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, Ruiqi Gao

arXiv:2405.16852v2cs.LGcs.AIstat.ML

TL;DR

Diffusion sampling requires many score evaluations, while existing distillation approaches can struggle with one-step generation or full distribution coverage. EMD uses an EM-inspired maximum-likelihood approach with joint sampling to distill a diffusion model into a one-step generator. It reports strong ImageNet results and favorable performance on one-step text-to-image generation.

  • Problem

    Diffusion sampling is computationally costly, and existing distillation methods face challenges in the one-step regime or capturing the full distribution.

    Source: §1 · p. 1

  • Method

    EMD uses an EM-inspired maximum-likelihood objective and joint MCMC sampling to train a one-step latent-variable student from a pretrained diffusion teacher.

    Source: §1 · pp. 1–2, §3.1 · p. 3

  • Results

    FID 2.20 on ImageNet-64 and 6.0 on ImageNet-128; EMD outperforms existing one-step generation approaches and performs favorably on one-step text-to-image generation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    EMD demonstrates strong performance in class-conditional ImageNet generation and text-to-image generation.

    Source: §6 · p. 10

    EMD still requires teacher-model initialization to perform competitively and is sensitive to the initialization timestep t*.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Diffusion sampling is costly, and existing distillation approaches struggle either with one-step generation or full distribution coverage. EMD uses an EM-inspired maximum-likelihood objective and joint sampling to distill a diffusion teacher into a one-step generator, with strong reported results across image and text-to-image generation.

1 Introduction

  • Existing trajectory distillation methods remain challenging in the one-step regime, while distribution-matching methods can miss modes because of mode-seeking divergences.Reducing diffusion sampling to a single score-function evaluation could enable real-time generation.Source: §1 · p. 1
  • EMD approximates a mode-covering divergence between a pretrained diffusion teacher and a latent-variable student that maps noise to data in one step.Its maximum-likelihood approach uses EM-style expectation and maximization steps.Source: §1 · p. 1
  • FID 2.20 on ImageNet-64 and 6.0 on ImageNet-128; EMD outperforms existing one-step generation approaches and performs favorably for one-step text-to-image generation.The method uses a multi-step MCMC formulation for the more competitive results.Source: §1 · p. 2
  • Jointly updating data and latent pairs with MCMC supports gradient estimation; reparameterization simplifies tuning, and removing a tractable linear noise term reduces gradient variance.The paper also connects EMD to Variational Score Distillation and Diff-Instruct, and describes how MCMC strength interpolates between mode-seeking and mode-covering divergences.Source: §1 · p. 2

2 Preliminary

The preliminaries introduce diffusion and score-based sampling, Langevin dynamics, maximum-likelihood learning of latent-variable models with EM, and score-distillation objectives that match noisy student and teacher distributions.

2.1 Diffusion models and score matching

  • Diffusion models gradually add noise to data and reverse the process to recover the data distribution; a time-dependent score function supports reverse-time SDE sampling.The score function is estimated by a neural network through denoising score matching.Source: §2.1 · p. 2

2.2 MCMC with Langevin dynamics

  • Langevin dynamics samples a target distribution by iterating gradient-based updates with noise, but convergence can require many steps, especially in high dimensions.Short-run variants with early termination have been used to learn energy-based models.Source: §2.2 · p. 2

2.3 Maximum Likelihood and Expectation-Maximization

  • EM learns latent-variable models by maximizing marginal likelihood, using posterior-based expectations to estimate gradients before optimizing model parameters.The marginal is generally analytically intractable, so the expectation can be approximated with Monte Carlo samples, including MCMC samples.Source: §2.3 · pp. 2–3

2.4 Variational Score Distillation and Diff-Instruct

  • Variational Score Distillation and Diff-Instruct match noisy student and teacher marginals by minimizing reverse KL, alternating generator updates with learning an auxiliary score network.The teacher score comes from the pretrained diffusion model; the auxiliary network estimates the student distribution's score.Source: §2.4 · p. 3

3 Method

EMD frames one-step distillation as EM over a latent-variable student, with joint sampling and noise cancellation supporting the updates. Across noise levels, its sampling choices connect the method to VSD and span mode-seeking to mode-covering behavior [26].

3.1 EM Distillation

  • EMD estimates student M-step gradients from joint MCMC samples targeting ρθ(x, z) = q(x)pθ(z|x), initialized with student samples instead of direct teacher draws.The student can reuse the teacher’s x-prediction function as its generator.Source: §3.1 · pp. 3–4

3.2 Reparametrized sampling and noise cancellation

  • Reparameterizing x = αgθ(z) + σϵ moves joint sampling to (ϵ, z)-space and simplifies step-size tuning; canceling accumulated noise yields cleaner targets and stabilizes generator learning.The paper reports that the reparameterization permits shared step sizes across noise levels and improves performance empirically.Source: §3.2 · pp. 4–5

3.3 Maximum Likelihood across all noise levels

  • A weighted objective matches student marginals to the diffusion process across all noise levels, using a shared generator across those levels.The gradient is expressed as an expectation over noise levels and the corresponding latent-variable models.Source: §3.3 · p. 5

3.4 Connection with VSD and Diff-Instruct

  • With a special sampling scheme, EMD’s gradient is the VSD gradient up to sign; varying sampling from one x-only step to joint (x, z) updates interpolates between mode-seeking and mode-covering divergences [26].The paper describes the one-step x-only sampler as likely biased and many-step joint sampling as closer to a mixing sampler.Source: §3.4 · p. 6

4 Related Work

Diffusion acceleration methods reduce sampling work, but trajectory approaches struggle in the one-step regime without perceptual-quality loss [43] [12]. One-step methods improve generation, though many use mode-seeking objectives that may miss the full distribution [20] [21] [56].

Diffusion acceleration

  • PF-ODE solvers and trajectory distillation accelerate diffusion sampling, but both struggle to reach one-step generation without significant perceptual-quality loss [43, 12, 13–18].Trajectory distillation condenses the PF-ODE sampling trajectory into fewer steps.Source: §4 · p. 6

Single-step diffusion models

  • One-step methods include adversarial fine-tuning, score distillation, and regression, but often minimize mode-seeking divergences that can fail to capture the full distribution [20] [21] [56] [9] [26] [23].Some methods are data-free or allow flexible generator architectures; concurrent Fisher-divergence work reports performance gains [59].Source: §4 · p. 6

5 Experiments

The experiments evaluate EMD on ImageNet and text-to-image generation, examining its sampling components and training overhead. EMD-16 outperforms existing one-step diffusion distillation methods on ImageNet FID; results also report recall, text-to-image quality, and computational cost.

5.1 ImageNet

  • Complete noise cancellation decreases gradient variance and speeds convergence, whereas canceling only the last step's noise is marginally helpful.The authors note that noise cancellation does not theoretically guarantee lower gradient variance.Source: §5.1 · p. 7, Figure 3 · p. 7
  • More Langevin steps monotonically improve FID and Recall; ε-only sampling performs similarly to joint (ε, z) sampling at few steps but falls behind as steps increase.The authors suggest more steps encourage mode coverage by better approximating the mode-covering forward KL.Source: §5.1 · p. 7, Figure 3 · p. 7
  • EMD-16 outperforms existing one-step diffusion distillation methods on FID for ImageNet 64×64 and 128×128; on ImageNet 64×64, its recall is competitive with distribution-matching methods but trails trajectory-distillation approaches.The comparisons use 300k training steps on ImageNet 64×64 and 200k on ImageNet 128×128.Source: Table 2 · p. 8, Table 3 · p. 8, §5.1 · p. 8

5.2 Text-to-image generation

  • For zero-shot generalization on MSCOCO, the one-step text-to-image EMD-8 generator achieved FID-30k 9.66 at tCFG 2.0; higher tCFG produced a model with competitive CLIP Score.It was trained on 256 TPU-v5e for 5 hours (5000 steps).Source: §5.2 · p. 9

5.1 ImageNet

  • The (ε, z)-corrector improves over the baseline, while the (x, z)-corrector struggles to match it; the authors suggest the spaces may differ in MCMC suitability or step-size tuning needs.The comparison uses EMD-8 on ImageNet 64×64 after 100k training steps.Source: Table 1 · p. 7, §5.1 · p. 8

5.3 Computation overhead in training

  • EMD-16 only doubles Diff-Instruct's wall-clock time, and equal or longer Diff-Instruct runs do not realize EMD's performance gain.The reported timing comparison accounts for other overheads; per-step timing was logged for ImageNet 64×64 ablations.Source: Table 6 · p. 10, §5.3 · p. 10

6 Discussion and limitation

EMD distills a pretrained diffusion model into a one-step student through maximum-likelihood optimization, with sampling and noise-cancellation techniques supporting training. The discussion identifies initialization and training-cost constraints, alongside future directions.

6 Discussion and limitation

  • EMD competitively matches pretrained diffusion marginals with a one-step student, but currently depends on teacher initialization and is sensitive to the initial t∗ choice.Training from scratch is theoretically supported but did not achieve competitive results empirically; MCMC sampling adds training cost and requires step-size tuning.Source: §6 · p. 10

A Expectation-Maximization

  • The EM-like gradient transformation uses an expectation over the latent posterior, with the posterior score expectation equal to zero.The passage identifies this equality as the reason for the transformation’s third line.Source: A Expectation-Maximization · p. 15

B Reparametrized sampling and noise cancellation

  • Reparameterizing joint samples as (ϵ, z) addresses slow convergence and complex step-size tuning in joint MCMC; canceling accumulated noise empirically stabilizes generator training.The method tracks sampled Langevin noises and removes them after the MCMC loop; the generator loss is an MSE between the corrected target and generator output.Source: B Reparametrized sampling and noise… · pp. 15–16, C JAX-style Code · pp. 16, 18

C JAX-style Code

  • The JAX-style corrector integrates pre-sampled ϵ-update noises with weights, subtracts them from corrected ˆϵ, and uses the result to parameterize ˆx.The generator loss is the squared difference between the stop-gradient corrected target and the generator output.Source: C JAX-style Code · pp. 16, 18

D.1 ImageNet 64×64

The appendices describe the training setups for ImageNet and text-to-image distillation, then discuss how the initial denoiser setting λ∗ affects samples and performance.

D.1 ImageNet 64×64

  • ImageNet 64×64 distillation ran for 300k steps, roughly eight days, on 64 TPU-v4, using an (ϵ, z)-corrector with 0.1 dropout in both score networks.The teacher followed the best EDM setting with an ADM UNet; the teacher’s noise schedule and score-matching weighting were inherited.Source: §D.1 · p. 18, Table 7 · p. 18

D.2 ImageNet 128×128

  • ImageNet 128×128 distillation ran for 200k steps, roughly ten days, on 128 TPU-v5p, using an (ϵ, z)-corrector with 0.1 dropout in both score networks.The teacher followed VDM++ with U-ViT, L, while the student used a cosine-adjusted schedule and EDM monotonic score-matching weighting.Source: §D.2 · p. 18, Table 8 · p. 18

D.3 Text-to-image generation

  • Text-to-image distillation used the public Stable Diffusion v1.5 checkpoint as teacher and inherited its noise schedule and score-matching weighting.The appendix lists its other hyperparameters in Table 9.Source: §D.3 · p. 18, Table 9 · p. 18

D.4 Choosing t∗and λ∗

  • λ∗ represents the selected denoiser noise level; on ImageNet 128×128, values in [−6, −3] gave similar performance, while the manuscript used a rough grid search for that setting and text-to-image.At initialization, λ∗=0 produced Gaussian-noise-like images, while λ∗=−6 produced more detail than λ∗=−10.Source: §D.4 · p. 19, Figure 6 · p. 19

E.1 Additional ImageNet results

The appendix provides additional qualitative samples from the one-step generator on ImageNet 64×64 and ImageNet 128×128 to support evaluation of generation quality and diversity.

E.1 Additional ImageNet results

  • Additional one-step generator samples are presented for ImageNet 64×64 and ImageNet 128×128.Source: §E.1 · p. 19
  • The samples are intended to help evaluate generation quality and diversity.Source: §E.1 · p. 19
  • The section adds qualitative examples rather than reported quantitative comparisons.The passage characterizes the material as additional qualitative samples.Source: §E.1 · p. 19

E.2 Additional text-to-image results

The appendix presents visual comparisons of EMD with few- or single-step text-to-image methods. It reports stronger quality and diversity than SD-turbo, LCM, and InstaFlow, with visual quality on par with 50-step teacher generation.

E.2 Additional text-to-image results

  • The SD-turbo comparison uses a checkpoint fine-tuned from Stable Diffusion 2.1, unlike EMD’s Stable Diffusion 1.5 teacher.Source: §E.2 · p. 21
  • EMD’s one-step samples significantly outperform LCM and InstaFlow and show better diversity and quality than SD-turbo, with visual quality on par with 50-step teacher generation.Tables 10–13 provide side-by-side comparisons for text prompts.Source: §E.2 · pp. 19, 21, Table 10 · p. 22, Table 11 · p. 23
  • Additional EMD samples cover a more diverse set of prompts beyond the four side-by-side comparison prompts.Tables 14 and 15 contain these additional qualitative results.Source: §E.2 · p. 21, Table 14 · p. 26, Table 15 · p. 27

Related papers

Loading 2405.16852v2…