PaperlayerView exact revision on arXiv

Source-linked AI summary

One-Step Diffusion Distillation through Score Implicit Matching

Weijian Luo, Zemin Huang, Zhengyang Geng, J. Zico Kolter, Guo-jun Qi

arXiv:2410.16794v1cs.CVcs.AIcs.LG

TL;DR

Diffusion models often require many sampling steps, and reducing those steps can substantially degrade generation quality. SIM distills pretrained diffusion models into one-step generators by implicitly minimizing score-based divergences. It reports strong results on CIFAR10 and text-to-image generation, including a one-step model matching its teacher's performance.

  • Problem

    Diffusion models typically need many network evaluations to generate samples, while reducing sampling steps substantially degrades performance.

    Source: §1 · p. 2

  • Method

    SIM distills pretrained diffusion models into one-step generators by implicitly minimizing score-based divergences using exactly computed gradients under mild regularity conditions.

    Source: §1 · p. 3, §3.2 · p. 5, §5 · p. 10

  • Results

    6.42 aesthetic score for SIM-distilled PixelArt-α with no decline versus its multi-step teacher; SIM also reports strong one-step CIFAR10 results.

    Source: §1 · p. 3, §4.1 · p. 7

  • Takeaways & Limitations

    SIM's data-free one-step distillation is presented as a route to more affordable deployment across domains without compromising the underlying generative model's performance.

    Source: §5 · p. 10

    The 0.6B SIM-DiT sometimes produces poor tiny human faces or arms and fingers, and may generate incorrect object counts or prompt content.

    Source: §4.2 · p. 10

Abstract

from arXiv · show

Despite their strong performances on many generative tasks, diffusion models require a large number of sampling steps in order to generate realistic samples. This has motivated the community to develop effective methods to distill pre-trained diffusion models into more efficient models, but these methods still typically require few-step inference or perform substantially worse than the underlying model. In this paper, we present Score Implicit Matching (SIM) a new approach to distilling pre-trained diffusion models into single-step generator models, while maintaining almost the same sample generation ability as the original model as well as being data-free with no need of training samples for distillation. The method rests upon the fact that, although the traditional score-based loss is intractable to minimize for generator models, under certain conditions we can efficiently compute the gradients for a wide class of score-based divergences between a diffusion model and a generator. SIM shows strong empirical performances for one-step generators: on the CIFAR10 dataset, it achieves an FID of 2.06 for unconditional generation and 1.96 for class-conditional generation. Moreover, by applying SIM to a leading transformer-based diffusion model, we distill a single-step generator for text-to-image (T2I) generation that attains an aesthetic score of 6.42 with no performance decline over the original multi-step counterpart, clearly outperforming the other one-step generators including SDXL-TURBO of 5.33, SDXL-LIGHTNING of 5.34 and HYPER-SDXL of 5.85. We will release this industry-ready one-step transformer-based T2I generator along with this paper.

1 Introduction

Diffusion models offer high-quality generation but often require many sampling steps, motivating distillation into efficient generators. The introduction presents SIM as a flexible score-divergence framework and reports strong one-step results.

1 Introduction

  • Diffusion models corrupt the data distribution and learn scores of noisy distributions across varying noise levels.The introduction describes their applications across data synthesis, density estimation, text-to-image generation, text-to-3D creation, and image editing.Source: §1 · p. 1
  • Diffusion generation often needs many network evaluations, and reducing sampling steps substantially degrades performance, limiting deployment where quick inference matters.The cited settings include mobile and edge devices and applications requiring rapid responses.Source: §1 · p. 2
  • Earlier divergence-based distillation studies focused on KL or Fisher divergence, leaving a more general framework for diffusion distillation lacking.The cited studies include KL-based methods by Luo et al. [42] and Yin et al. [81] and Fisher-divergence work by Zhou et al. [92].Source: §1 · p. 3
  • SIM computes gradients for a broad class of score-based divergences between generator and diffusion-model scores, enabling implicit minimization despite intractable scores.The framework applies to arbitrary distance functions between the two score functions.Source: §1 · p. 3
  • SIM includes KL-based Diff-Instruct and recovers SiD as a squared-L2 special case; its Pseudo-Huber variant converges faster and is more hyperparameter-robust than L2.The comparison concerns different distance functions used to define the divergence.Source: §1 · p. 3
  • 6.42 aesthetic score for SIM-distilled PixelArt-α matched its teacher and exceeded SDXL-TURBO (5.33), SDXL-LIGHTNING (5.34), and HYPER-SDXL (5.85).On CIFAR10, FID was 2.06 for unconditional and 1.96 for class-conditional generation.Source: §1 · p. 3

2 Diffusion Models

The section describes diffusion as progressively corrupting data toward simple noise and learning scores that support reversing this process. Sampling uses the learned score network to simulate an SDE.

2 Diffusion Models

  • The forward process transforms the data distribution toward a simple noise distribution.It begins from q0 = qd.Source: §2 · p. 3
  • A continuous-indexed score network approximates marginal scores of the forward process and is trained with a weighted denoising score-matching objective.The forward SDE is specified by a predefined drift, scalar diffusion coefficient, and independent Wiener process.Source: §2 · p. 3
  • After training, the score network approximates the diffused-data marginal score and generates samples by simulating the SDE, which is slower than one-step generation.The weighting function controls learning importance across time levels, and the conditional transition describes the forward diffusion.Source: §2 · pp. 3–4

3 Score Implicit Matching

SIM formulates one-step distillation as matching student and teacher diffusion-time score distributions, then uses a gradient-equivalent objective to train the generator without directly differentiating its intractable score.

3 Score Implicit Matching

  • The student maps prior noise to a single sample, and the approach assumes the pretrained diffusion model sufficiently approximates the data distribution.The method treats the pretrained model as the only item of consideration.Source: §3 · p. 4

3.1 General Score-based Divergences

  • The objective integrates a proper distance between student and teacher marginal scores across diffusion times, with time weights and a sampling distribution covering both.The divergence is zero if the marginal scores agree at all times, which implies agreement of the initial distributions almost surely under the sampling distribution.Source: §3.1 · p. 4

3.2 Score Implicit Matching

  • Under mild regularity conditions, SIM uses a stop-gradient choice to replace an intractable score derivative with evaluations from a separate approximation network.Training alternates between learning the online score network on generator samples and updating the generator using that network and the pretrained diffusion model.Source: §3.2 · pp. 4–6

3.3 Instances of Score Implicit Matching.

  • Squared L2 recovers SiD’s delta loss [92], while SIM permits alternative distances and defaults to Pseudo-Huber.The authors report that Pseudo-Huber adaptively normalizes the score difference, which can stabilize training and yield robust, fast convergence.Source: §3.3 · p. 6

3.4 Related Works

  • Related work groups distillation into trajectory, distributional-matching, and other methods, and characterizes SIM by its flexible distances and Pseudo-Huber performance.The passage contrasts trajectory and adversarial approaches that require real samples with some divergence-based approaches that do not.Source: §3.4 · p. 6

4 Experiments

Experiments evaluate SIM as a one-step distillation method on CIFAR10 and a transformer-based text-to-image model. They report sample quality, training behavior and cost, alongside generation failure cases.

4.1 One-step CIFAR10 Generation

  • 2.06 FID for CIFAR10 unconditional generation and 1.96 FID for class-conditional generation place one-step SIM among top-performing models; it outperforms SiD and DI in the unconditional setup.The unconditional result is on par with CTM; the class-conditional result is among top-performing models.Source: §4.1 · p. 7, Table 2 · p. 7
  • At 1e−4, SIM converges steadily while DI and SiD are unstable early, and it shows better FID and Inception Scores than SiD at matched iterations across tested configurations.The authors suggest SIM’s loss normalization may keep objective scale from changing abruptly during training.Source: Figure 2 · p. 8, §4.1 · p. 8

4.2 Transformer-based One-step Text-to-Image Generator

  • CIFAR10 experiments distill pretrained EDM under DI- and SiD-matched settings; T2I experiments distill 0.6B PixelArt-α, comparing against its 14-step teacher on SAM-LLaVA-Caption10M and few-step models on COCO-2017.The T2I comparison does not include DI because the authors observed fast convergence but mode collapse.Source: §4.1 · p. 7, §4.2 · p. 8
  • SIM-DiT-600M is trained data-free with 4 A100-80G GPUs for 2 days, versus hundreds of A100 GPU-days for other Table 3 models; qualitative comparisons report higher aesthetics than other leading few-step models.The authors attribute efficiency to data-free distillation and Pseudo-Huber loss normalization.Source: §4.2 · p. 9, Figure 3 · p. 9
  • 99.6% of PixArt-α’s aesthetic score and 100% of its PickScore are recovered on SAM Caption; SIM-DiT-600M also has a dominant aesthetic score over leading few-step models.Its Image Reward is slightly smaller, which the authors say could potentially be optimized with more training compute.Source: §4.2 · pp. 8–9, Table 3 · p. 10
  • The 0.6B one-step model sometimes produces poor tiny human faces, arms, and fingers, incorrect object counts, or content that does not strictly follow prompts.The authors suggest scaling model size and teacher diffusion models may help address these issues.Source: Figure 4 · p. 9, §4.2 · p. 10

5 Conclusion and Future Works

SIM distills pretrained multistep diffusion models into data-free one-step generators, with the stated aim of preserving their performance. The conclusion identifies broader model families and data incorporation as directions for further research.

5 Conclusion and Future Works

  • SIM transforms pretrained multistep diffusion models into data-free one-step generators.The conclusion says this could support more affordable deployment without compromising the underlying models' performance.Source: §5 · p. 10
  • The authors identify generalizing SIM to other pretrained generative models, such as flow-matching models, as an open direction.They also say adding data to SIM may improve images that the teacher model fails to generate well; this remains unexplored.Source: §5 · p. 10

A.1 Proof of Theorem 3.1

The appendix proof relies on the score-projection identity, which connects denoising score matching and denoising auto-encoders. It restates the identity without proof and points readers to prior work for its proof.

A.1 Proof of Theorem 3.1

  • The score-projection identity was first found by Vincent to bridge denoising score matching and denoising auto-encoders.Zhou et al. [92] later applied it to derive distillation methods based on Fisher divergences.Source: §A.1 · p. 17
  • The appendix restates the score-projection identity under mild conditions as Theorem A.1.The authors refer readers to Zhou et al. [92] for a complete proof.Source: §A.1 · p. 17
  • The proof of Theorem 3.1 proceeds by invoking the score-projection identity for vector-valued functions.The appendix states the identity holds for arbitrary u(·, θ) and parameter θ.Source: §A.1 · p. 17

A.2 Pytorch style pseudo-code of Score Implicit Matching

The appendix provides PyTorch-style pseudocode for Algorithm 1 using the Pseudo-Huber distance. Its displayed steps initialize an online diffusion model from the teacher and update it alongside the generator.

A.2 Pytorch style pseudo-code of Score Implicit Matching

  • Algorithm 1 pseudocode uses the Pseudo-Huber distance, while a separate algorithm covers CIFAR10 with EDM.The listing identifies the code as SIM pseudocode.Source: §A.2 · pp. 18–19
  • The pseudocode initializes an online diffusion model by copying the teacher and alternates which model is trainable.It shows two online-model update iterations and random diffusion-time sampling.Source: §A.2 · pp. 18–19
  • The generator update uses a normalized offset coefficient in its loss before backpropagation and an optimizer step.The listing also shows the online diffusion model's loss being backpropagated and optimized.Source: §A.2 · pp. 18–19

A.3 Instances of SIM with different distance functions

Appendix A.3 presents SIM losses for different distance functions, and Appendix B.1 identifies which images in the human-preference figure came from SIM-DiT-600M.

A.3 Instances of SIM with different distance functions

  • Table 4 presents SIM loss instances for different distance functions, including Huber, L1, and powered exponential choices.The paper uses Pseudo-Huber as its default distance unless otherwise specified.Source: Table 4 · p. 19, §A.3 · p. 19

B.1 Answer for the human preference study

  • The human-preference figure identifies SIM-DiT-600M outputs as the middle image in row one and the leftmost images in rows two and three.These are the specific generated images to locate in the figure.Source: §B.1 · pp. 19–20

B.2 Experiment details on CIFAR10 dataset

The CIFAR10 experiments use an EDM formulation for the teacher and generator, then compare time-sampling and weighting choices to select the final SIM configuration.

B.2 Experiment details on CIFAR10 dataset

  • The generator uses the pretrained EDM denoiser architecture with a pre-selected time index and is initialized from the teacher's parameters.The EDM denoiser predicts clean samples from noisy inputs, and its objective is written as a denoising regression loss.Source: §B.2 · p. 20
  • The standard Karr time distribution performs poorly for SIM, while a modified log-normal distribution works better.The paper names the modified schedule as the SIM time distribution.Source: §B.2 · p. 21
  • SIM uses constant-one weighting for generator updates, whereas SiD uses a specially designed weighting function.The authors report that SIM's loss normalization makes the ad-hoc weighting unnecessary.Source: §B.2 · p. 21
  • SIM with no generator-loss weighting and a log-normal time distribution gives the best performance among the compared configurations.The selected configuration is used for the final CIFAR10 experiments.Source: §B.2 · p. 21

B.3 Experiment details on Text-to-Image Distillation

The text-to-image experiments adapt PixArt-α to an EDM formulation and test noise distributions; a separate user study compares distilled and teacher outputs against prompts.

B.3 Experiment details on Text-to-Image Distillation

  • The one-step generator uses the PixArt-α model with iDDPM+DDIM preconditioning and a noise range of [0.01, 156.6155].The generator uses t*=2.5; larger t* values sped convergence in practice, with negligible difference over complete training and minimal impact on final results.Source: §B.3 · p. 21
  • Training used SAM-LLaVA-Caption10M prompts, and the best model was trained for about 16k iterations, taking about two days on four A100-40G GPUs.The reported run was less than two epochs.Source: §B.3 · p. 22
  • Noise distributions concentrated at smaller values produced excessively dark samples, while slightly larger distributions made sample structure unstable.Source: §B.3 · p. 22

B.4 Instruction for Human Preference Study

  • The user study showed participants two randomly ordered images from the student and teacher and asked which better matched the prompt and had higher image quality.The study used 17 prompts and collected approximately 30 responses.Source: §B.4 · p. 22

B.5 Generated Samples on CIFAR10

The supplied figures show one-step SIM samples for CIFAR10 conditional and unconditional generation, alongside a convergence comparison with SiD.

B.5 Generated Samples on CIFAR10

Loading 2410.16794v1…