One-step Diffusion with Distribution Matching Distillation

In plain words

Image generators such as Stable Diffusion make a picture by removing noise in many small steps, often dozens of them. Each step is a full pass through a large network, so generation is slow.

This paper trains a second network that makes the picture in one step. It does not copy the slow process step by step. Instead it is trained so that its pictures look statistically like the slow model's pictures. To judge this, the training uses two helper models. One knows what the slow model's pictures look like. The other is kept up to date with what the new model is producing. Comparing the two tells the new model which way to nudge each picture: toward the slow model's kind of picture and away from its own current output.

A second, cheaper signal compares the new model's output with the slow model's output for the same starting noise. The authors add it to keep the new model from producing only a few kinds of picture. The result is reported to match the slow model closely in quality while running much faster.

TL;DR

The paper turns a pretrained diffusion model into a one-step generator. Training follows the difference of two diffusion score estimates, one from the frozen teacher and one from a fake-score model that is updated during training §3.2, eq. 7. An LPIPS regression to teacher outputs on paired noise is added with weight 0.25 by default §3.3, eq. 9. Author-reported: 2.62 FID on class-conditional ImageNet 64x64 with one forward pass, against 2.32 for the EDM teacher with 512 tab. 1. Text-to-image distillation of Stable Diffusion v1.5 reaches 11.49 FID on zero-shot COCO-30k at 0.09 s per image, against 8.78 at 2.59 s for the teacher tab. 3.

Problem

Diffusion models need tens to hundreds of network evaluations per sample, which limits interactive use §1. Earlier distillation methods train a one-step student to reproduce the teacher's noise-to-image mapping. This is costly because each loss needs a full teacher trajectory, and the students still lag behind the teacher §1 §2. The paper instead asks only that student samples look like teacher samples in distribution, without matching specific noise-image pairs §1. A pure distribution-matching objective can drop or collapse modes, so a regression term is also used §3.3.

Method

The teacher is a pretrained diffusion denoiser μ_base. The student G_θ has the same architecture but takes no timestep input §3.1. Before training, G_θ(z) = μ_base(z, T−1) for every z, with T = 1000 §3.1. The loss is the sum D_KL + λ_reg · L_reg, with λ_reg = 0.25 unless stated otherwise §3.3.

1. Generator and teacher. The real distribution is the teacher's training distribution. Its score is estimated from a frozen copy of μ_base §3.2, eq. 4.
2. Fake score. A second denoiser μ_fake^φ is initialised from the teacher. It is trained with the ordinary denoising loss on generator samples, so that it tracks the generator's current output distribution §3.2, eq. 5, eq. 6.
3. Distribution matching gradient. Generator samples are diffused with noise at a random time t drawn from [T_min, T_max] = [0.02T, 0.98T]. The gradient for the generator is the difference between the fake and real scores, times a weight w_t and α_t §3.2, eq. 7, eq. 8. The weight w_t normalises the gradient size across noise levels, and the paper reports that this beats the weightings used by DreamFusion and ProlificDreamer on CIFAR-10 §3.2 tab. 2.
4. Regression loss. The paper builds a dataset of pairs (z, y), where y is the teacher's output for noise z from a deterministic ODE sampler §3.3. The generator is penalised with LPIPS between G_θ(z) and y §3.3, eq. 9. The paper says this keeps the modes covered and stabilises training §3.3 fig. 3.
5. Two data streams. The distribution matching term uses unpaired generator samples. The regression term uses the paired dataset §3.3.
6. Training loop. Each iteration updates the generator with the distribution matching loss plus λ_reg times the regression loss. It then updates the fake-score denoiser with its denoising loss on the generator's outputs alg. 1.
7. Classifier-free guidance. For guided teachers, the regression pairs are sampled from the guided model. The real score uses the mean prediction of the guided model. The fake-score formula is unchanged, and the guidance scale is fixed during training §3.4.

The gradient in step 3 is the paper's approximation of the KL gradient. The exact KL gradient is written as a score difference in §3.2, eq. 2. The approximation replaces the exact scores with diffusion-model scores on noised samples and averages over t §3.2. The derivation is in app. F.

# DMD training loop  [alg. 1]
# teacher: frozen mu_base; fake_denoiser: trainable mu_fake (init: teacher)
# generator: one-step G_theta (init: teacher, [§3.1])
# pair_noise, pair_targets: regression dataset D = {(z, y)}  [§3.3]

def train_step(batch_size):
    # Fresh noise and generator samples, unpaired  [alg. 1]
    noise = sample_normal(batch_size)
    images = generator(noise)
    # Paired noise and teacher outputs for the regression loss  [alg. 1]
    pair_noise, pair_targets = sample_pairs(batch_size)
    pair_images = generator(pair_noise)
    # Generator update: distribution matching plus regression  [alg. 1]
    # lambda_reg = 0.25 by default; CIFAR-10 unconditional uses 0.5 [app. B.1]
    dm_loss = distribution_matching_loss(images, teacher, fake_denoiser)
    reg_loss = lpips(pair_images, pair_targets)  # [eq. 9]
    generator_loss = dm_loss + 0.25 * reg_loss  # [§3.3]
    update(generator, generator_loss)
    # Fake-score update: denoising loss on generator outputs  [alg. 1] [eq. 6]
    times = sample_uniform(0.0, 1.0, shape=(batch_size,))  # see Pitfalls
    noisy_images = forward_diffuse(stop_gradient(images), times)
    denoise_loss = timestep_weighted_mse(
        fake_denoiser(noisy_images, times),
        stop_gradient(images),
    )
    update(fake_denoiser, denoise_loss)

# Distribution matching loss for the generator  [eq. 3-5, eq. 7-8]
# images: (B, C, H, W) generator outputs; C channels, S spatial positions

def distribution_matching_loss(images, teacher, fake_denoiser):
    # One diffusion time per image, drawn from [T_min, T_max]  [§3.2]
    times = sample_uniform(T_MIN, T_MAX, shape=(B,))
    # Noise schedule values at those times  [eq. 3]
    alpha, sigma = schedule(times)
    # Diffuse the generator outputs: x_t ~ N(alpha_t x, sigma_t^2 I)  [eq. 3]
    noisy_images = alpha * images + sigma * sample_normal_like(images)
    # Mean predictions of both denoisers; no gradient through them  [§3.2]
    with no_gradient():
        teacher_mean = teacher(noisy_images, times)
        fake_mean = fake_denoiser(noisy_images, times)
    # Real and fake scores from the mean predictions  [eq. 4] [eq. 5]
    real_score = -(noisy_images - alpha * teacher_mean) / sigma**2
    fake_score = -(noisy_images - alpha * fake_mean) / sigma**2
    # Per-sample weight that normalises the gradient size  [eq. 8]
    weight = (sigma**2 / alpha) * (C * S) / l1_norm(teacher_mean - images)
    # Gradient with respect to images: w_t * alpha_t * (fake - real)  [eq. 7]
    grad = weight * alpha * (fake_score - real_score)
    # Surrogate loss whose gradient w.r.t. images is grad  [app. B]
    return 0.5 * mse(images, stop_gradient(images - grad))

Key equations, numbered as in the paper:

\nabla_\theta D_{KL} \simeq \mathbb{E}_{z, t, x, x_t}\left[ w_t\, \alpha_t \big(s_\text{fake}(x_t, t) - s_\text{real}(x_t, t)\big) \frac{dG}{d\theta} \right] §3.2, eq. 7

w_t = \frac{\sigma_t^2}{\alpha_t} \cdot \frac{C S}{\lVert \mu_\text{base}(x_t, t) - x \rVert_1} §3.2, eq. 8

s_\text{real}(x_t, t) = -\frac{x_t - \alpha_t \mu_\text{base}(x_t, t)}{\sigma_t^2} \qquad s_\text{fake}(x_t, t) = -\frac{x_t - \alpha_t \mu_\text{fake}^\phi(x_t, t)}{\sigma_t^2} §3.2, eq. 4-5

\mathcal{L}_\text{reg} = \mathbb{E}_{(z, y) \sim \mathcal{D}}\, \ell\big(G_\theta(z), y\big), \qquad \ell = \text{LPIPS} §3.3, eq. 9

Final objective for the generator: D_\text{KL} + \lambda_\text{reg}\, \mathcal{L}_\text{reg} §3.3. The fake-score denoising loss is \lVert \mu_\text{fake}^\phi(x_t, t) - x_0 \rVert_2^2, weighted by timestep as in the base model's training §3.2, eq. 6.

Hyperparameters

Values are as stated in the paper. "Not stated" means the paper does not say.

Results

Author-reported. Setup: class-conditional ImageNet 64x64 and CIFAR-10 for class-conditional generation §4.1; zero-shot MS-COCO 2014 30k for text-to-image, with the distilled model trained on LAION-Aesthetic §4.3. FID is computed as described in app. D.

Class-conditional ImageNet 64x64, FID (lower is better). Baseline numbers are taken from Song et al. tab. 1.

Ablation, FID tab. 2:

CIFAR-10, FID tab. 6. Baseline numbers are from Song et al.; the EDM teacher row is conditional.

Text-to-image, zero-shot MS-COCO 30k, guidance scale 3 tab. 3. Latency is batch size 1, hardware not stated.

Text-to-image, zero-shot MS-COCO 30k, guidance scale 8, with FID and CLIP score tab. 4. Baselines marked † were evaluated by the authors.

Headline claims as stated: about 30x faster than SD v1.5 for text-to-image §4.3. About 20 frames per second at 512x512 with FP16 inference; the hardware for this figure is not stated §1 §4.3. The ImageNet "512-fold" speed claim §4.1 is a ratio of forward passes, not a measured time. The introduction gives a 100x reduction in network evaluations §1. The teaser caption gives 90 ms per image for DMD against 2590 ms for SD fig. 1.

Paper vs code

No official code.

Searched: the arXiv metadata has no code link. The project page the paper links (tianweiy.github.io/dmd/) names no code repository. That page links the DMD2 project page, which links github.com/tianweiy/DMD2. That repository implements DMD2 (arXiv 2405.14867), a different method. Its README says DMD2 was "developed based on our reimplementation of the original DMD paper", and it is licensed CC BY-NC-SA 4.0 (non-commercial), so it was not used. A GitHub repository search for "distribution matching distillation" returned no repository for this paper.

Pitfalls

• Time range for the fake-score update is inconsistent. Section 3.2 samples t from [0.02T, 0.98T] §3.2. Algorithm 1 samples t from U(0, 1) for the fake-score denoising step alg. 1. The paper does not say which applies to the fake-score update, and there is no code to check.
• Appendix listing and equation (8) differ by a factor α_t. Working through equations (4), (5), (7) and (8) gives a generator gradient of α_t · C S · (μ_fake − μ_real) / ‖μ_base − x‖₁. The appendix listing computes C S · (μ_fake − μ_real) / ‖μ_base − x‖₁, with no α_t app. B. Its comment says the weight "diverges slightly" from the paper's equation to fit the mean-prediction scheme. This is our derivation from the equations; check it before implementing.
• Surrogate loss reduction is not stated. The appendix listing uses a mean squared error and a factor 0.5 app. B. The main text does not give the reduction, so this changes the effective gradient scale.
• Denoising weight for the fake score is described but not written out. The paper says to use the base model's timestep weighting §3.2. The appendix comment gives SNR + 1/0.5² for EDM and SNR for Stable Diffusion v1.5 app. B.
• LAION batches concatenate paired and unpaired samples. Algorithm 1 concatenates the generator samples with the regression samples when the dataset is LAION alg. 1. The paper does not explain why, or whether the distribution matching term sees both sets.
• Batch-size units. For CIFAR-10, ImageNet and LAION 6.25+, the text says total batch size app. B. For the LAION 6+ log, "DM BS" and "Reg. BS" are not given units tab. 5.
• Regression weight is sensitive. In the LAION 6+ log, V1 (weight 0.1) reaches FID 23.88 and V3 (weight 1, otherwise the same settings) reaches 16.10 tab. 5. The final run uses 0.25 tab. 5.
• The LAION 6+ schedule was changed during training. The authors adjusted dataset size, regression weight, VAE type and maximum DM step, and say the schedule "may not be the most efficient or optimal" app. B.4. Do not treat the final row as a tuned recipe.
• VAE and latent space are not fully specified. For LAION 6.25+, a tiny VAE decodes generator outputs for the regression loss app. B.3. For LAION 6+, the small VAE is used in V1–V8 and the standard SD VAE in V9–V12 tab. 5. The paper does not say whether the distribution matching term operates on latents or decoded images.
• Initialisation uses a fixed timestep. G_θ(z) = μ_base(z, T−1) for all z §3.1. The paper does not explain the choice further.
• Guidance formula is not written out. The real score uses the guided model's mean prediction, and the fake score is unchanged §3.4. The exact guided mean is not given; a fixed guidance scale is used during training §3.4.
• The ablation without regression loss uses a smaller learning rate. It uses 1e-5 to prevent divergence app. C.2.
• Regression distance. On CIFAR-10, LPIPS gives FID 2.66 and L2 gives 2.78 app. I.
• FIDs are not directly comparable across rows. ImageNet and CIFAR-10 use 50,000 samples with the EDM evaluation code. COCO uses 30,000 samples, 512 to 256 downsampling with Lanczos, and clean-fid app. D. Most ImageNet and CIFAR-10 baselines are copied from Song et al. tab. 1 tab. 6.
• Costs. LAION 6.25+ took about 36 hours on 72 A100 GPUs §4.3. LAION 6+ took about two weeks on about 80 A100 GPUs app. B.4. The method fine-tunes both the generator and the fake-score network, so training memory is high §5.
• Limitations stated by the authors. There is a small quality gap to finer teacher discretisations (100 or 1000 evaluations) §5. The authors suggest LoRA as a way to reduce training memory §5.
• Figure and table numbers. Figure numbers count the teaser as Figure 1 and the Related Work overview as Figure 2. Table 5 and Table 6 are appendix tables, numbered from the LaTeX source app. B.4 app. E.

Lineage

• Builds on: score-based diffusion models, including the score view of denoisers §3.1 §3.2; ProlificDreamer's Variational Score Distillation, which the paper adapts from test-time 3D optimisation to training a whole generator §1; EDM as the teacher for CIFAR-10 and ImageNet app. B; Stable Diffusion v1.5 as the teacher for text-to-image §4.3. The regression distance follows InstaFlow and Consistency Models §3.3.
• Compared against: Progressive Distillation, Rectified Flow and InstaFlow, Consistency Models, TRACT, BOOT, DFNO, UFOGen, LCM-LoRA, DPM-Solver++, UniPC, GigaGAN and StyleGAN-T tab. 1 tab. 3 tab. 4.
• Superseded by: not stated in this paper.