One-step Diffusion with Distribution Matching Distillation (DMD)

Created on June 03, 2026

2026   ·   distillation   DMD   distribution_matching_distillation   ·   diffusion_model

Yin et al. 2024


Hozy Summary

  • Goal)
    • Distill a pre-trained diffusion model \(\mu_{\text{base}}(=\mu_{\text{real}})\)
  • Settings)
    • \(G_\theta\) : “one-step” image generator
    • \(\mu_{\text{fake}}^\phi\) : a fake diffusion denoiser that generates fake scores \(s_{\text{fake}}\) for \(\mathcal{L}_{\text{KL}}^\theta\)
  • Idea)
    • Three optimization targets
      1. Distribution Matching Loss : \(\displaystyle\nabla_\theta\mathcal{L}_{\text{KL}}^\theta \simeq \mathbb{E}_{z,t} \left[ w_t \alpha_t \left( s_{\text{fake}}(x_t, t) - s_{\text{real}}(x_t, t) \right) \frac{\text{d}G}{\text{d}\theta} \right]\)
        • Objective)
          • Match the teacher’s distribution \(p_{\text{real}}\) with the model \(G_\phi\)’s distribution
        • Desc.)
          • \(s_{\text{fake}}(x_t, t) = \nabla_{x_t} \log p_{\text{fake}}(x_t) = -\displaystyle\frac{x_t - \alpha_t\mu_{\text{fake}}^\phi(x_t, t)}{\sigma_t^2}\).
      2. (X) Regression Loss : \(\mathcal{L}_{\text{reg}}\)
        • Removed in DMD2.
        • Replaced by two-timescale update rul (TTUR).
          • i.e.) Update \(\mu_{\text{fake}}^\phi\) more frequent than \(G_\theta\).
      3. Denoising Loss : \(\mathcal{L}_{\text{denoise}}^\phi = \left\Vert \mu_{\text{fake}}^\phi(x_t, t) - x_0 \right\Vert_2^2\)
        • Objective)
          • Train \(\mu_{\text{fake}}^\phi\) to model the current generator distribution \(p_{\text{fake}}\), so that \(s_{\text{fake}}\) accurately estimates its score.
        • Desc.)
          • \(x_t = \Psi(x, \epsilon, t)\) : noise injected fake image
            • where
              • \(x = G_\theta(z),\quad z\sim\mathcal{N}(0, \mathbf{I})\) : a fake image
              • \(\epsilon\sim\mathcal{N}(0, \mathbf{I})\) : noise
              • \(\Psi\) : the forward process with a pre-determined noise schedule
    • In DMD2, the GAN loss is introduced.
  • Main Training Algorithm)
    • Pipeline)
      • \(z\sim\mathcal{N} \longrightarrow \underbrace{x = G_\theta(z)}_{\text{fake img}} \stackrel{\Psi(\cdot, \epsilon, t)}{\longrightarrow} \underbrace{x_t}_{\text{noisy}} \longrightarrow \begin{cases} \mu_{\text{real}}(x_t, t) \rightarrow s_{\text{real}}(x_t, t) \\ \mu_{\text{fake}}^\phi(x_t, t) \rightarrow s_{\text{fake}}^\phi(x_t, t) \end{cases}\).
    • Objectives)
      • \(\mathcal{L}_{\text{DMD}} \approx \mathcal{L}_{\text{KL}}\) : Distribution Matching Loss
        • \(\begin{cases} s_{\text{real}}(x_t, t) \\ s_{\text{fake}}^\phi(x_t, t) \end{cases} \longrightarrow \nabla_\theta\;\mathcal{L}_{\text{DMD}} \longrightarrow G_\theta\).
      • \(\mathcal{L}_{\text{denoise}}\) : Denoising Loss
        • \(\begin{cases} \mu_{\text{fake}}^\phi(x_t, t) \\ \text{stopgrad}(x = G_\theta(z)) \end{cases} \longrightarrow \nabla_\phi\;\mathcal{L}_{\text{denoise}} \longrightarrow \mu_\phi\).
      • Omitted \(\mathcal{\text{reg}}\) (not used in DMD2)



Concept) Distribution Based Distillation

  • Settings)
    • \(\mu_{\text{base}}\) : a pre-trained diffusion denoiser
      • Desc.)
        • Mean-prediction.
        • Denoise \(x_T\sim\mathcal{N}(\mathbf{0,I})\) to \(x_0\sim p_{\text{real}}\)
          • where \(T=1000\).
        • Architecture : EDM, Stable Diffusion
    • \(G_\theta\) : “one-step” image generator
      • Desc.)
        • Outputs a fake image
        • Initialized with \(\mu_{\text{base}}\)
          • i.e.) \(G_\theta(z) = \mu_{\text{base}}(z, T-1),\;\forall z\) before training!
        • Two objectives -\(\mathcal{L}_{\text{KL}}\) : Distribution Matching Objective
    • \(\mu_{\text{fake}}^\phi\) : a fake diffusion denoiser to generate fake scores \(s_{\text{fake}}\) for the DMD loss
      • Desc.)
        • Initialized with \(\mu_{\text{base}}\).
        • In the training loop,
          • First used to generate \(s_{\text{fake}}\) to calculate \(\mathcal{L}_{\text{KL}}\), and optimize the generator \(G_\theta\).
          • Then, later (dynamically) updated by optimizing \(\mathcal{L}_{\text{denoise}} = \left\Vert \mu_{\text{fake}}^\phi(x_t, t) - x \right\Vert_2^2\)
            • where \(x = G_\phi(z)\), i.e. the fake image.
  • Loss)
    • \(\mathcal{L} = \mathcal{L}_{\text{KL}} + \lambda_{\text{reg}} \mathcal{L}_{\text{reg}}\).
      • where
        • \(\mathcal{L}_{\text{KL}}\) is optimized by Distribution Matching Objective using \(\nabla_\theta \mathcal{D}_{\text{KL}}\) via \(\mathcal{L}_{\text{denoise}}^{\phi}\) minimization.
        • \(\lambda_{\text{reg}} = 0.25\),
        • \(\mathcal{L}_{\text{reg}}\) is the regression Loss
  • Details)
    • CFG used.


Tech.) Distribution Matching Loss

  • Loss)
    • \(\nabla_\theta \mathcal{D}_{\text{KL}} \simeq \mathbb{E}_{z,t,x,t_t} \left[ w_t \alpha_t \left( s_{\text{fake}}(x_t, t) - s_{\text{real}}(x_t, t) \right) \frac{\text{d}G}{\text{d}\theta} \right]\).
      • where
        • \(z\sim\mathcal{N}(\mathbf{0,I})\),
        • \(x_t\sim q_t(x_t\mid x)\) : a perturbed training sample
          • where
            • \(x=G_\theta(z)\) : one-step sample from the generator
            • \(t\sim\mathcal{U}(T_{\min}, T_{\max})\) : a time step
              • Authors used \(T_{\min} = 0.02 T, T_{\max} = 0.98T\)
            • \(q_t(x_t\mid x) \sim\mathcal{N}(\alpha_t x;\sigma_t^2\mathbf{I})\) for different noise schedules \(\alpha_t\text{ and }\sigma_t\)
        • \(\begin{cases} s_{\text{real}}(x) &= \nabla_x \log p_{\text{real}}(x) &= -\frac{x_t - \alpha_t\mu_{\text{base}}(x_t, t)}{\sigma_t^2} \\ s_{\text{fake}}(x) &= \nabla_x \log p_{\text{fake}}(x) &= -\frac{x_t - \alpha_t\mu_{\text{fake}}^\phi(x_t, t)}{\sigma_t^2} \end{cases}\) : scores derived from Song et al.
        • \(w_t\) : a weight schedule
          • Authors used \(w_t = \displaystyle\frac{\sigma_t^2}{\alpha_t}\frac{CS}{\Vert\mu_{\text{base}}(x_t, t) - x\Vert_1}\)
            • for the number of spatial locations \(S\) and the number of channels \(C\)
  • Derivation)
    • We may start from approximating the fake distribution to the real one.
      \(\begin{aligned} \mathcal{D}_{\text{KL}} = (p_{\text{fake}} \Vert p_{\text{real}}) &= \mathbb{E}_{p\sim p_{\text{fake}}} \left( \log\frac{p_{\text{fake}}}{p_{\text{real}}} \right) \\ &= \mathbb{E}_{z\sim\mathcal{N}(\mathbf{0, I}), x=G_\theta(z)} \left[ - \left( \log p_{\text{real}} - \log p_{\text{fake}} \right) \right] \\ \end{aligned}\).
    • This is intractable, but we only need the gradient w.r.t. \(\theta\) as
      • \(\nabla_\theta \mathcal{D}_{\text{KL}} = \mathbb{E}_{z\sim\mathcal{N}(\mathbf{0, I}), x=G_\theta(z)} \left[ - \left( s_{\text{real}}(x) - s_{\text{fake}}(x) \right) \frac{\text{d}G}{\text{d}\theta} \right]\).
        • for
          • \(s_{\text{real}}(x) = \nabla_x \log p_{\text{real}}(x)\) and \(s_{\text{fake}}(x) = \nabla_x \log p_{\text{fake}}(x)\)
            • cf.)
              • \(s_{\text{real}}\) moves \(x\) towards the modes of \(p_{\text{real}}\)
              • \(-s_{\text{fake}}\) spreads them apart
    • Further perturb the data distribution with random Guassian noise
      • Why?)
        • To alleviate the vanishing gradient problem in most $x\in\mathcal{X}$
      • i.e.)
        • \(x_t\sim q(x_t\mid x)\) where \(q(x_t\mid x)\sim\mathcal{N}(\alpha_t x, \sigma_t^2\mathbf{I})\).
        • Then, the scores can be denoted as
          • \(s_{\text{real}}(x_t, t) = - \displaystyle\frac{x_t-\alpha_t \mu_{\text{base}}(x_t, t)}{\sigma_t^2}\).
          • \(s_{\text{fake}}(x_t, t) = - \displaystyle\frac{x_t-\alpha_t \mu_{\text{fake}}^\phi(x_t, t)}{\sigma_t^2}\).
            • where \(\mu_{\text{fake}}^\phi\)’s parameters are first initialized with the parameters from \(\mu_{\text{base}}\).
            • Then, \(\phi\) are trained minimizing below objective
              • \(\mathcal{L}_{\text{denoise}}^{\phi} = \left\Vert \mu_{\text{fake}}^\phi(x_t, t) - x_0 \right\Vert_2^2\).


Tech.) Regression Loss

  • Goal)
    • Problem)
      • When \(t\approx 0\) (corrected with low level of noise), \(s_{\text{real}}(x_t, t)\) becomes unreliable as \(p_{\text{real}}\approx0\).
      • The optimization is susceptible to mode collapse/dropping, where the fake distribution assigns higher overall density to a subset of the modes.
        • i.e.) the score is invariant to scaling of the original probability density function.
    • Sol.)
      • Minimize the distance between the generator output \(G_\theta(z)\) and the base diffusion model output \(y\).
  • Loss)
    • \(\mathcal{L}_{\text{reg}} = \mathbb{E}_{(z,y)\sim\mathcal{D}} \ell(G_\theta(z), y)\).
      • where
        • \(z\sim\mathcal{N}(\mathbf{0,I})\),
        • \(y\sim\mu_{\text{base}}(z)\) : a sampled image from the base diffusion model
          • Authors used…
            • Heun solver from EDM
              • on CIFAR-10 with 18 steps
              • on Imagenet with 256 steps
            • PNDM solver on LAION with 50 steps
        • \(\ell(\cdot, \cdot)\) : a distance metric
          • Authors used LPIPS following InstaFlow and Consistency Model.


Enjoy Reading This Article?

Here are some more articles you might like to read next:

  • (DM Reconst.) Ch.3 Score-Based Perspective - From EBMs to NCSN
  • Denoising Diffusion Probabilistic Models (DDPM)
  • Guiding a Diffusion Model with a Bad Version of Itself (Autoguidance)
  • Score-Based Generative Modeling through Stochastic Differential Equation
  • Classifier-Free Diffusion Guidance (CFG)