PaperlayerView exact revision on arXiv

Source-linked AI summary

Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation

Mingyuan Zhou, Huangjie Zheng, Zhendong Wang, Mingzhang Yin, Hai Huang

arXiv:2404.04057v3cs.LGcs.AIcs.CVstat.ML

TL;DR

Diffusion models offer high-quality generation but are expensive because reverse diffusion requires multiple network evaluations. SiD uses semi-implicit forward-diffusion distributions and three score identities to distill a pretrained model into a single-step generator, achieving strong FID performance across four benchmarks while reducing FID rapidly. The method’s main practical limitation is its higher memory and iteration-time cost from maintaining and differentiating through additional networks.

  • Problem

    Diffusion models require expensive multi-step reverse diffusion, motivating single-step distillation of knowledge embedded in pretrained score-estimation networks.

  • Method

    SiD reformulates forward diffusion as semi-implicit distributions and combines three score identities with model-based score matching and Monte Carlo estimation to train a generator from synthesized images.

  • Results

    Across four benchmark datasets, SiD achieves state-of-the-art FID, reduces FID exponentially during distillation, and outperforms established approaches across differing step counts and data-access settings.

  • Takeaways & Limitations

    SiD demonstrates that pretrained diffusion models can be distilled into high-quality single-step generators with rapid FID improvement.

  • Takeaways & Limitations

    SiD requires three equal-sized networks, increasing memory use and iteration time relative to methods such as Diff-Instruct.

Abstract

from arXiv · show

We introduce Score identity Distillation (SiD), an innovative data-free method that distills the generative capabilities of pretrained diffusion models into a single-step generator. SiD not only facilitates an exponentially fast reduction in Fréchet inception distance (FID) during distillation but also approaches or even exceeds the FID performance of the original teacher diffusion models. By reformulating forward diffusion processes as semi-implicit distributions, we leverage three score-related identities to create an innovative loss mechanism. This mechanism achieves rapid FID reduction by training the generator using its own synthesized images, eliminating the need for real data or reverse-diffusion-based generation, all accomplished within significantly shortened generation time. Upon evaluation across four benchmark datasets, the SiD algorithm demonstrates high iteration efficiency during distillation and surpasses competing distillation approaches, whether they are one-step or few-step, data-free, or dependent on training data, in terms of generation quality. This achievement not only redefines the benchmarks for efficiency and effectiveness in diffusion distillation but also in the broader field of diffusion-based generation. The PyTorch implementation is available at https://github.com/mingyuanzhou/SiD

1. Introduction

Diffusion models produce high-quality images but require expensive multi-step reverse diffusion. SiD distills a pretrained diffusion model into a single-step generator using score identities and achieves rapid, high-quality generation across four benchmarks.

  • Diffusion models generate diverse, photorealistic images but require repeated reverse-diffusion network evaluations, making inference expensive.
  • SiD trains a single-step generator by matching scores of diffused real and synthesized distributions across noise levels.It reformulates forward diffusion as semi-implicit distributions and combines three score-related identities with score and Monte Carlo estimation.
  • On ImageNet 64x64, FID falls from 153.52 initially to 3.07 after 50 million synthesized images.The progression corresponds to roughly 0 through 49K training iterations under α = 1.0.
  • SiD is named for the three score identities underlying its loss mechanism.
  • Across CIFAR-10, ImageNet 64x64, FFHQ 64x64, and AFHQv2 64x64, SiD achieves state-of-the-art FID and an exponentially fast reduction during distillation.The evaluation uses VP-EDM checkpoints as teacher diffusion models.

2. Related Work

Prior work reduces diffusion inference cost through faster numerical solvers, reverse-chain distillation, adversarial training, and distribution matching. SiD follows the diffused real/fake alignment framework but replaces JSD or KL objectives with an approximated model-based score-matching loss.

  • Numerical SDE/ODE solvers reduce reverse-diffusion steps but face a trade-off between fewer steps and visual quality.
  • Reverse-chain distillation includes progressive distillation and consistency models that seek efficient generation while preserving output consistency.
  • Adversarial methods distill diffusion models into one- or two-step generators using time-dependent discriminators or divergence objectives.Diffusion-GAN minimizes Jensen–Shannon divergence at each diffusion time step, while UFOGen uses a time-step-dependent discriminator.
  • Diff-Instruct and DMD apply diffused real/fake distribution alignment to distill pretrained diffusion models into single-step generators.DMD adds a regression loss to improve distillation quality.
  • SiD differs by using a model-based score-matching loss, a Fisher-divergence variant, and a strategy for approximating its analytically intractable objective.

3. Forward Diffusion as Semi-Implicit Distribution: Exploring Score Identities

The paper casts forward diffusion of real and generated data as semi-implicit distributions whose scores admit expectation-based identities. Three identities, including a score projection identity, are combined to make score-based single-step distillation computationally tractable.

  • A semi-implicit distribution has an analytically defined conditional p(x | z) and a sampleable mixing distribution p(z), while its marginal is difficult to compute.
  • Both real and generated diffused data have scores expressible as expectations, enabling approximation with neural networks or Monte Carlo estimation.
  • Forward diffusion is semi-implicit because q(x_t | x_0) is analytic while the data distribution p_data(x_0) is unknown and represented by empirical samples.
  • If the generator distribution matches the data distribution, their diffused distributions match at every noise level; matching a diffused distribution at any level implies generator-data matching.
  • Score Identities: The first two identities are Tweedie formulas for recovering expected real or generated clean samples from noisy observations through their diffused-distribution scores.The real-data identity underlies denoising score matching, and an analogous identity applies to diffused fake data.
  • Score Identities: The third, score projection identity addresses the intractable generated-data score by estimating an expected product involving that score.
  • Score Identities: Combining all three identities with model-based score matching yields the loss mechanism for single-step pretrained-diffusion distillation.This integration differs from prior uses of the identities in score matching, variational inference, and multi-step reverse diffusion.

4. SiD: Score identity Distillation

SiD formulates diffusion distillation as model-based score matching and develops approximations that provide more stable generator guidance. The resulting alternating optimization uses fused losses and achieves strong single-step distillation performance, while exposing limitations of simpler loss approximations.

  • Model-based Explicit Score Matching: SiD constructs a model-based explicit score-matching loss from the pretrained score Sϕ and evaluates it on diffused generator samples rather than diffused real data.This loss is also called a Fisher divergence and differs from the usual score-matching formulation through its expectation over pθ(xt).
  • Empirical Results: SiD achieves state-of-the-art FID across CIFAR-10, ImageNet, FFHQ, and AFHQv2, with FID decreasing exponentially fast during distillation.The single-step student can exceed the iterative EDM teacher, suggesting that reverse diffusion may accumulate errors even with fine steps and advanced solvers.
  • Loss Approximation: The theoretical loss depends on the optimal score estimator ψ*(θ), making direct minimization potentially a bilevel optimization problem.The approximation replaces ψ*(θ) with ψ, which ignores its dependence on the generator parameters.
  • Loss Approximation: The first approximated loss can be unreliable because score approximation error and score difference may dominate when the generator distribution poorly matches the data distribution.CIFAR-10 experiments and a toy failure case show that this loss can fail to yield meaningful results despite optimizing toward the correct generator parameter in the proposition.
  • Projected Score Matching: The alternative projected-score formulation uses the pretrained model’s denoising of corrupted fake data to provide a more stable gradient signal for the generator.Its gradient is not directly governed by the score-approximation norm or score difference, unlike the first approximation.
  • Fused Loss and Optimization: SiD combines the approximated losses through a weighted fused objective and alternates score-estimator updates with generator updates.The training procedure updates ψ given θ and then θ given ψ, while noise weighting and timestep sampling follow related distillation methodologies.

5. Experimental Results

SiD is evaluated across four benchmark datasets using FID and related metrics, with experiments examining convergence, benchmark performance, and resource costs. The method rapidly improves FID, often surpasses teacher and competing models, but requires substantially more memory and computation than Diff-Instruct.

  • Evaluation Protocol: SiD is evaluated on CIFAR-10, ImageNet 64x64, FFHQ 64x64, and AFHQ-v2 64x64 using FID and Inception Score, with Precision and Recall additionally measured on ImageNet.FID uses 50k generated samples and the EDM teacher’s training set as reference.
  • Benchmark Performance: SiD surpasses the EDM teacher across nearly all tested environments and exceeds a broad range of multi-step, single-step, and GAN-based competitors.The sole reported exception is ImageNet 64x64, where SiD at α = 1.2 attains FID 1.524, above RIN’s 1.23.
  • Benchmark Performance: SiD outperforms CTM while remaining independent of training data, unlike CTM, whose performance uses training data and an auxiliary GAN loss.The approach also avoids the multi-stage, pairwise-regression, or additional-GAN requirements associated with several competing methods.
  • Convergence Speed: SiD rapidly reduces FID during distillation, with FID evolution tracked against thousands of synthesized images across datasets and α = 1 or α = 1.2.The experiments use the generator’s synthesized images to assess convergence speed.
  • Convergence Speed: On ImageNet 64x64 with batch size 1024 and α = 1.0, SiD reaches FID below 5 after 7.5M images, below 4 after 13M, and below 3 after 31M images.It exceeds Progressive Distillation’s FID 15.39 after about 500k synthesized images, and outperforms Diff-Instruct and DMD at fewer than 7M and 40M images, respectively.
  • Limitations: SiD maintains three equal-sized networks during distillation, increasing memory allocation by about 50–70% and iteration time by approximately 28–36% relative to Diff-Instruct.The additional networks are the pretrained score network, generator score network, and generator.

6. Conclusion

SiD distills pretrained diffusion models into single-step generators by aligning diffused real and generative distributions through model-based score matching. Experiments show efficient FID reduction and superiority across varied generation settings.

  • SiD uses semi-implicit distributions to minimize a model-based score-matching loss between diffused real and generative distributions.
  • SiD significantly reduces FID efficiently and outperforms established generative approaches across single-step, few-step, data-free, and training-data-dependent settings.

Impact Statement

Distilled diffusion models can reduce deployment resource requirements, but distilling models trained on harmful content raises ethical and safety concerns. Responsible deployment therefore requires balancing efficiency gains with oversight.

  • Distilled diffusion models may require fewer computational resources, potentially improving energy efficiency and reducing costs in resource-intensive applications.
  • Distilling models trained on violent or pornographic data may facilitate harmful-content generation and perpetuate inappropriate patterns.
  • Responsible deployment requires ethical guidelines and oversight to mitigate safety and societal risks from distilled models.

A. Ablation Study and Parameter Settings

Ablations examine α, optimizer settings, batch size, and training trajectories. Performance is generally stable for α from 0 to 1.2, while negative α and some ImageNet settings degrade or destabilize training.

  • Impact of α: The α ablation evaluates values from −0.25 to 1.5 on CIFAR-10 unconditional distillation while tracking generations as synthetic-image processing progresses.
  • Impact of α: α values from 0 to 1.2 yield stable SiD performance, while negative α produces large FIDs.
  • Impact of α: IS and FID improve gradually as α increases, peaking at α = 1 or α = 1.2, which are selected for subsequent experiments.
  • Setting of β1: Setting β1 = 0.9 for the generator score network often prevents convergence, so β1 = 0 is retained for that network.
  • Batch Size for ImageNet 64x64: ImageNet training with batch size 1024 shows exponential FID decline before sudden divergence near FID 2.62, with FP16 precision suspected as a contributor.
  • Batch Size for ImageNet 64x64: Increasing the overall ImageNet batch size to 8192 and reducing the learning rate produces stable performance for α = 1 but occasional FID spikes for α = 1.2.
  • Batch Size for ImageNet 64x64: Larger batch sizes slow the point at which SiD surpasses Diff-Instruct and DMD.

B. Algorithm Box

Algorithm 1 alternates updates to a generator score network and generator using synthesized noisy images. Training continues until FID plateaus or the allotted budget is exhausted, returning the generator.

  • SiD initializes the generator and generator score network from the pretrained score network before alternating optimization.
  • The algorithm samples noise and timesteps, constructs synthesized noisy images, and updates the generator score network with Equation (10).
  • It then constructs another noisy synthetic sample and updates the generator with Equation (23).
  • Training stops when FID plateaus or the training budget is exhausted, and the output is the trained generator Gθ.

C. Proofs

The proofs establish score identities for Gaussian diffusion and use them to characterize the semi-implicit score construction underlying the method. They also prove Identity 3 and Theorem 5.

  • Score identities: Tweedie’s formula expresses the score of the semi-implicitly constructed diffusion distribution, with a_t = 1 recovering the main-text identity.The proof begins from the Gaussian diffusion identity and exploits the semi-implicit construction of pθ(x_t).
  • Score identities: The appendix includes a proof of Identity 3.
  • Theorem 5: Theorem 5 is proved by expanding the L2 norm and isolating a term involving conditional denoising differences and the score of pθ(x_t).

D. Analytic study of the toy example

The toy example analyzes Gaussian data and generator distributions, deriving their scores and the resulting score difference. It also shows that a naive gradient approximation can fail to provide a meaningful learning signal.

  • Gaussian toy example: For Gaussian data pdata(x0) = N(0, I) and generator pθ(xg) = N(θ, I), the optimal parameter is θ* = 0.
  • Gaussian toy example: The data score is S(x_t) = −x_t/(1 + σ_t^2).
  • Gaussian toy example: The score difference is δϕ,ψ*(θ)(x_t) = −θ/(1 + σ_t^2).
  • Gaussian toy example: The conditional expectation fϕ(x_t, t) equals E[x0 | x_t] = x_t/(1 + σ_t^2), and the corresponding scaled difference is −ψ/(1 + σ_t^2).
  • Gradient behavior: A naive approximation involving (fϕ − fψ)^2 could fail to provide a meaningful gradient for learning θ.

E. Training and Evaluation Details and Additional Results.

The appendix reports implementation settings, computational comparisons, evaluation procedures, and additional SiD results across datasets. It also documents a divergence issue mitigated by increasing batch size and lowering the learning rate.

  • Training setup: SiD development used 8 Nvidia RTX A5000 GPUs, while evaluations used cloud nodes with 16 Nvidia A100-40GB GPUs; gradient accumulation supported ImageNet batches up to 8192.
  • Computational comparison: Table 6 compares hyperparameters, distillation time, memory usage, and per-iteration computation costs for SiD and Diff-Instruct on 16 NVIDIA A100 GPUs.
  • Evaluation protocol: FID evaluation costs are excluded from Table 6, and SiD evaluates the EMA generator after every 500k synthesized images while retaining and re-evaluating the lowest-FID generator across 10 independent runs.
  • Training stability: Increasing the batch size to 8192 and lowering the learning rate to 4e-6 alleviated sudden FID divergence observed for ImageNet 64x64.
  • Additional figures: SiD examples are reported with FID 1.923 for unconditional CIFAR-10, 1.710 for label-conditioned CIFAR-10, 1.524 for label-conditioned ImageNet, 1.550 for FFHQ, and 1.628 for AFHQ-V2.
Loading 2404.04057v3…