PaperlayerView exact revision on arXiv

Source-linked AI summary

Multistep Distillation of Diffusion Models via Moment Matching

Tim Salimans, Thomas Mensink, Jonathan Heek, Emiel Hoogeboom

arXiv:2406.04103v1cs.LGcs.AIcs.CVcs.NE

TL;DR

Diffusion sampling is computationally expensive because it typically requires many neural network evaluations. This paper distills many-step models into few-step models by matching conditional expectations, and reports that multi-step distilled models outperform one-step versions and often their teachers, with state-of-the-art ImageNet results.

  • Problem

    Diffusion models can require many sampling steps and neural network evaluations, making inference computationally expensive.

    Source: §1 · p. 2, §2.1 · p. 2

  • Method

    The method distills many-step diffusion models into few-step models by matching conditional expectations of clean data given noisy data along the sampling trajectory.

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

  • Results

    With multiple sampling steps, distilled models consistently outperform their one-step versions and often exceed many-step teachers, setting state-of-the-art ImageNet results.

    Source: §5.1 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    The moment-matching framework extends related distillation methods to multiple sampling steps, with reported gains over one-step versions and often over many-step teachers.

    Source: §6 · p. 9

    The alternating-optimization algorithm is not generally guaranteed to converge, though it is stable with suitable hyperparameters, especially at k ≥ 8 steps.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

We present a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models by matching conditional expectations of the clean data given noisy data along the sampling trajectory. Our approach extends recently proposed one-step methods to the multi-step case, and provides a new perspective by interpreting these approaches in terms of moment matching. By using up to 8 sampling steps, we obtain distilled models that outperform not only their one-step versions but also their original many-step teacher models, obtaining new state-of-the-art results on the Imagenet dataset. We also show promising results on a large text-to-image model where we achieve fast generation of high resolution images directly in image space, without needing autoencoders or upsamplers.

1 Introduction

Diffusion models generate complex outputs through iterative denoising, but sampling commonly requires hundreds of network evaluations. This paper proposes a distributional distillation method that matches conditional expectations to extend one-step approaches to few-step generation.

1 Introduction

  • Sampling diffusion models often requires hundreds of neural network evaluations, shifting their computational burden to inference.The models generate outputs through iterative denoising.Source: §1 · p. 2
  • The proposed distributional method matches conditional expectations of clean data given noisy data along the diffusion sampling trajectory.The authors contrast distributional methods with deterministic methods that approximate the iterative denoising output in fewer steps.Source: §1 · p. 2
  • Extending one-step distillation to few steps improves output quality and can outperform the many-step teacher; a two-minibatch variant avoids an auxiliary model.The two-minibatch variant processes two independent minibatches per parameter update.Source: §1 · p. 2

2 Background

Diffusion models learn to invert a gradual noise process and generate samples by iterative denoising, which can make inference computationally expensive. The background also introduces generalized moment matching as a way to fit distributions by matching expectations of chosen functions.

2.1 Diffusion Models

  • Diffusion training corrupts clean data with time-dependent noise, and sampling reverses that process through repeated denoising steps.The variance-preserving specification sets z_0 = x and z_1 to standard Gaussian noise.Source: §2.1 · p. 2
  • Hundreds of sampling steps may be needed to approach the ideal diffusion model, making inference computationally expensive.In the stated ideal, the clean-data prediction equals the conditional expectation and the number of steps tends to infinity.Source: §2.1 · p. 2

2.2 Generalized method of moments

  • The method of moments fits a distribution by matching its moments to those of the data; for a Gaussian, the first two are mean and variance.The passage also notes that matching moments extends to multivariate distributions, including matching a Gaussian’s mean and covariance matrix.Source: §2.2 · pp. 2–3
  • The Generalized Method of Moments matches expectations of high-dimensional functions under the model and data distributions, and the proposed distillation method is a special case.The moment vector is defined as the expectation of a chosen function f.Source: §2.2 · p. 3

3 Moment Matching Distillation

The method distills a teacher into a few-step generator by matching conditional expectations of clean data along its sampling trajectory. It offers an alternating auxiliary-denoiser approach and an instantaneous parameter-space alternative, with hyperparameter choices for stable few-step training.

3 Moment Matching Distillation

  • Distillation trains a few-step generator to match the conditional expectation of clean data under its sampling distribution to that under the data distribution.The student uses the same sampling algorithm as the teacher but with a strongly reduced number of steps.Source: §3 · p. 3

3.1 Alternating optimization of the moment matching objective

  • Alternating optimization uses an auxiliary denoiser to estimate the student’s conditional clean-data expectation, then updates the student using the difference between auxiliary and teacher predictions.The auxiliary denoiser is regularized toward the teacher; this makes it predict their average at convergence. The algorithm is not guaranteed to converge, though it is reported stable with suitable hyperparameters, especially at k ≥8.Source: §3.1 · p. 4

3.2 Parameter-space moment matching

  • Parameter-space moment matching minimizes the difference between expected teacher-loss gradients on generated and real data, whose gradient is zero if the teacher reaches a training-loss minimum.This instantaneous alternative avoids tracking two evolving models.Source: §3.2 · pp. 4–5

3.3 Hyperparameter choices

  • For k=1 or 2, stable training uses less low-signal emphasis than EDM weighting; classifier-free guidance is usually unnecessary for FID but increases Inception and CLIP scores.The setup follows EDM choices for the time schedule and loss weighting, shifts log-SNR with resolution, and keeps the Adam gradient preconditioner fixed during distillation.Source: §3.3 · p. 6

4 Related Work

The paper relates its one-step method to score-matching distillation methods, including Variational Score Distillation and Diff-Instruct. It argues that moment matching provides a rationale for the multistep extension, while also distinguishing the method from adversarial approaches.

4 Related Work

  • In the one-step case, the method is a special case of Variational Score Distillation, Diff-Instruct, and related methods that approximately minimize KL divergence between student and teacher.The connection follows when the starting noise is independent of the generated sample; this correspondence does not extend to the multistep case.Source: §4 · p. 6
  • For multistep sampling, the auxiliary regression objective no longer corresponds to score matching, but moment matching motivates the extension and also works with deterministic DDIM samplers.Source: §4 · p. 6
  • Adversarial multistep methods also use the same conditional sampling distribution, but train a discriminator to distinguish student-generated samples from base-model samples.Their formulations differ in divergence, timestep sampling, and additional losses; examples use unconditional or conditional discriminators.Source: §4 · p. 6

5 Experiments

Experiments benchmark class-conditional ImageNet generation, ablate conditional sampling and guidance, track the distillation loss, and test scaling to a large text-to-image model.

5.1 Class-conditional generation on ImageNet

  • At 8+ steps, moment-matching distillation achieves state-of-the-art FID, beating its teacher and best undistilled models using more than 1,000 steps; optimized guidance beats the 512-step teacher at both ImageNet resolutions.The authors suggest that differing prediction errors across timesteps may allow averaging predictions to improve results.Source: §5.1 · pp. 7–8

5.2 Ablating conditional sampling

  • Conditioning generator sampling on z_t yields much better sample diversity than applying noise independently in the multistep setup.The authors report that both conditional and marginal sampling are valid choices for multistep distillation.Source: §5.2 · p. 8, Figure 3 · p. 8

5.3 Effect of classifier-free guidance

  • Increasing classifier-free guidance typically raises Inception and CLIP scores while reducing FID.The method can use guidance; for alternating optimization it is applied to the teacher, not the generator or auxiliary denoising model.Source: §5.1 · p. 7, §5.3 · p. 8

5.4 Distillation loss is informative for moment matching

  • The instant 2-batch method's simple loss is useful for monitoring distillation progress: it may rise initially, then falls exponentially toward zero with updates.It is minimized without adversarial techniques, bootstrapping, or other tricks.Source: §5.4 · p. 9

5.5 Text to image

  • On 512 × 512 MS-COCO text-to-image, the alternating-optimization 8-step model exceeds its base model on zero-shot FID and CLIP Score.The experiment uses a pixel-space model without an encoder or decoder; the instant 2-batch version performs somewhat less well at eight steps.Source: §5.5 · p. 9

6 Conclusion

Moment Matching Distillation makes diffusion sampling faster by distilling many-step models into few-step models. The authors report that multi-step distilled models outperform their one-step versions and often exceed their many-step teachers, while noting that image-quality metrics need human evaluation.

6 Conclusion

  • Table 3 reports text-to-image results at 512 × 512 resolution.Source: Table 3 · p. 9
  • Multi-step distilled models consistently outperform their one-step versions and often exceed their many-step teachers, setting new state-of-the-art ImageNet results.The conclusion describes the result without supplying specific benchmark values in these passages.Source: §6 · p. 9
  • Moment Matching Distillation makes diffusion models faster to sample by matching conditional expectations of clean data given noisy data along the sampling trajectory.The framework extends related distillation methods to the multi-step setting.Source: §6 · p. 9
  • Automated image-quality metrics are highly imperfect, and the authors plan a full set of human evaluations to complement their reported metrics.Source: §6 · p. 9
  • The authors present alternating updates of a distilled generator and auxiliary denoiser, and a two-minibatch version that updates only the generator.They plan to explore the costs and benefits of these algorithm choices further.Source: §6 · p. 9

A Instant moment matching = matching expected teacher gradients

The instantaneous moment-matching variant avoids alternating optimization of an auxiliary denoising model. It can be interpreted as matching expected teacher gradients in parameter space, under an assumption about the teacher's gradient on training data.

A Instant moment matching = matching expected teacher gradients

  • The instantaneous variant uses a moment-matching loss that does not require alternating optimization of an auxiliary denoising model.The loss is presented as an alternative to the approach that alternates updates of the generator and auxiliary denoiser.Source: A Instant moment matching = matching… · p. 13
  • The instantaneous objective recasts moment matching in teacher-parameter space by matching expected teacher gradients on generated samples against real-data gradients.The expression uses stop-gradients on z_s; the auxiliary gradient is evaluated on an independent minibatch.Source: A Instant moment matching = matching… · p. 13
  • This equivalence assumes the teacher's expected training-loss gradient on training data is zero, as when the teacher has reached a training-loss minimum.Source: A Instant moment matching = matching… · p. 13

B Experimental details

The experiments used TPUv5e hardware and UViT models, with separate batch sizes and distillation budgets for ImageNet and text-to-image. The reported FID was selected from periodic evaluations using fixed seeds and data.

B Experimental details

  • Experiments used 256 TPUv5e chips per run, with global batch sizes of 2048 for ImageNet and 512 for text-to-image.Source: B Experimental details · p. 13
  • Base models used UViT and were trained for 1M steps, taking between 2 days for ImageNet 64 and 2 weeks for text-to-image.The ImageNet experiments used the small model variant.Source: B Experimental details · p. 13
  • Distillation ran for up to 200,000 steps on ImageNet and 50,000 on text-to-image; the best FID was selected from evaluations every 5,000 steps.Evaluation seeds and data were fixed to minimize bias.Source: B Experimental details · p. 13

C More model samples

The appendix presents generated samples from the 8-step distilled models across ImageNet resolutions and for text-to-image. A single-class ImageNet panel is included to help assess sample diversity.

C More model samples

  • Figure 4 shows random-class ImageNet samples at 64 × 64 from the 8-step model using alternating optimization.Source: Figure 4 · p. 14
  • Figure 5 shows 64 × 64 samples from one ImageNet class to help assess sample diversity.Source: Figure 5 · p. 14
  • Figure 6 shows random-class ImageNet samples at 128 × 128 from the 8-step model using alternating optimization.Source: Figure 6 · p. 15
  • Figure 7 presents selected samples from the 8-step distilled text-to-image model.Source: Figure 7 · p. 16

D Relationship to score matching

One-step moment matching is equivalent to score matching, but that equivalence does not generally hold for multistep sampling. Under an approximation that makes the forward processes match, the objectives become equivalent; both methods share a correct-target fixed point, though convergence to it is not guaranteed.

  • In the multistep case, the sampling distribution of z_s depends on z_t, so one-step equivalence between moment matching and score matching no longer holds.In the one-step case, z_1 is pure noise and independent of z_s, making the conditional q(z_s|x̃, z_1) equal to q(z_s|x̃).Source: D Relationship to score matching · p. 17
  • Denoising toward μ_t→s(x̃, z_t), a linear combination of x̃ and z_t, was tried as an alternative score-matching target but did not produce good results.Source: D Relationship to score matching · p. 17
  • Under the approximation p_η(z_t|z_s) ≈ q(z_t|z_s), moment matching and score matching become equivalent again.This follows when the forward processes match and the auxiliary score is fit using the denoising target x̃.Source: D Relationship to score matching · p. 17
  • Both multistep methods have a fixed point at the correct target distribution q under the stated conditions, but whether optimization reaches it is not guaranteed.Before convergence, p_η(z_t|z_s) generally differs from q(z_t|z_s), so the methods optimize different objectives.Source: D Relationship to score matching · p. 17
Loading 2406.04103v1…