PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality

Daniel Watson, William Chan, Jonathan Ho, Mohammad Norouzi

arXiv:2202.05830v1cs.LG

TL;DR

Diffusion models offer strong generation quality but are expensive because sampling requires many forward passes. The paper introduces DDSS and GGDM to optimize few-step samplers through differentiable perceptual quality objectives, reporting improved unconditional image quality without fine-tuning the pretrained model.

  • Problem

    DDPM sampling typically requires hundreds of forward passes, while objectives such as ELBO can mismatch sample quality, especially with few inference steps.

  • Method

    DDSS differentiates through sampling to optimize a flexible GGDM family using perceptual sample-quality scores, reparametrization, and gradient rematerialization.

  • Results

    DDSS consistently finds few-step samplers with better image-generation quality than strong literature baselines across qualitative and quantitative unconditional-generation evaluations.

  • Takeaways & Limitations

    Pretrained DDPMs can produce high-quality samples in very few steps, and better sampling paths can be found without fine-tuning their parameters.

  • Takeaways & Limitations

    Finding samplers with inference budgets K < 10 and little apparent quality loss remains challenging, and GGDM does not guarantee matching the original DDPM marginals.

Abstract

from arXiv · show

Diffusion models have emerged as an expressive family of generative models rivaling GANs in sample quality and autoregressive models in likelihood scores. Standard diffusion models typically require hundreds of forward passes through the model to generate a single high-fidelity sample. We introduce Differentiable Diffusion Sampler Search (DDSS): a method that optimizes fast samplers for any pre-trained diffusion model by differentiating through sample quality scores. We also present Generalized Gaussian Diffusion Models (GGDM), a family of flexible non-Markovian samplers for diffusion models. We show that optimizing the degrees of freedom of GGDM samplers by maximizing sample quality scores via gradient descent leads to improved sample quality. Our optimization procedure backpropagates through the sampling process using the reparametrization trick and gradient rematerialization. DDSS achieves strong results on unconditional image generation across various datasets (e.g., FID scores on LSUN church 128x128 of 11.6 with only 10 inference steps, and 4.82 with 20 steps, compared to 51.1 and 14.9 with strongest DDPM/DDIM baselines). Our method is compatible with any pre-trained diffusion model without fine-tuning or re-training required.

1 INTRODUCTION

Diffusion models produce high-quality samples but typically require hundreds of denoising passes. DDSS optimizes flexible few-step samplers directly against perceptual sample-quality objectives.

  • DDPMs can synthesize high-quality images, audio, and 3D shapes, while also supporting likelihood computation and stable training.
  • Hundreds of forward passes through a U-Net make DDPM generation less efficient than typical GAN generation.
  • The minimum inference-step budget for competitive quality depends on the problem, including dataset complexity and conditioning strength.
  • DDSS unrolls diffusion sampling and uses reparametrization with gradient rematerialization to optimize parametric few-step samplers.Its GGDM sampler family includes DDIM as a special case.
  • A mismatch between ELBO-based training and sample quality motivates optimizing a perceptual loss, with KID yielding fast samplers with better image quality than prior work.The method is reported as robust to different KID kernel choices.
  • The paper identifies DDSS, GGDM, and KID-based sampler optimization as its main contributions.

2 BACKGROUND ON DENOISING DIFFUSION IMPLICIT MODELS

DDPMs learn to reverse a Markovian noising process trained with an ELBO, whereas DDIM constructs faster samplers from alternative ELBOs that preserve DDPM marginals.

  • DDPMs gradually add Gaussian noise to data through a pre-specified Markovian forward process.
  • The learned reverse model pθ(xt−1|xt) inverts the forward diffusion process and is trained by maximizing a possibly reweighted ELBO.
  • DDPM parameterization makes ELBO maximization equivalent to minimizing a weighted sum of denoising score-matching objectives.
  • DDIM defines alternative ELBOs and forward processes that share DDPM marginals while permitting arbitrary posterior variances.
  • DDIM can use only a subsequence of DDPM timesteps, producing a faster sampler compatible with the same selected marginals.
  • DDIM(η = 0), combined with suitable timestep strides, was reported as state-of-the-art for few-step sampling at the smallest inference budgets.

3 DIFFERENTIABLE DIFFUSION SAMPLER SEARCH (DDSS)

DDSS learns few-step diffusion samplers by differentiating through stochastic sampling and optimizing a perceptual sample-quality objective. The approach includes flexible GGDM sampler families and uses gradient rematerialization to make optimization feasible under memory constraints.

  • DDSS approach: DDSS backpropagates through a K-step stochastic sampling process to optimize differentiable sampler parameters with stochastic gradient descent.The method uses the reparameterization trick and can optimize any differentiable loss over model samples.
  • Sampler families: GGDM is a non-Markovian sampler family with additional degrees of freedom, combining information from previous noisier images at each denoising step.The paper also considers GGDM +PRED, which learns marginal coefficients in addition to GGDM parameters.
  • Sampler families: Learning sampler timesteps remains helpful even when the pretrained DDPM uses discrete training timesteps with continuous embeddings.The authors report that such models behave well at timesteps absent during training.
  • Differentiable sample quality scores: KID optimization compares model and real samples in pretrained InceptionV3 feature space using a differentiable perceptual objective.The paper reports that perceptual features are crucial, because direct pixel-space optimization consistently worsened apparent sample quality during training.
  • Differentiable sample quality scores: Gradient rematerialization reduces memory demands because storing forward-pass states otherwise scales linearly with the number of inference steps.This addresses the memory cost of differentiating through large DDPM architectures.

4 GENERALIZED GAUSSIAN DIFFUSION MODELS

GGDM is a flexible non-Markovian Gaussian sampler family that generalizes DDIM by allowing broader parameterization of the diffusion process. Unlike DDIM, it does not require matching the original DDPM marginals, a relaxation supported by empirical comparisons.

  • GGDM defines a novel family of Gaussian diffusion processes that includes DDIM as a special case.
  • GGDM introduces free parameters µ_tu and σ_t, allowing each less-noisy prediction to use previous noisier images and independently control sampler variances.
  • The GGDM process is defined by µ_tu and σ_t rather than by the β_t or ᾱ_t coefficients used to define a usual DDPM forward process.
  • 4.1 IGNORING THE MATCHING MARGINALS CONDITION: GGDM does not guarantee matching the original DDPM’s marginals, unlike DDIM.
  • 4.1 IGNORING THE MATCHING MARGINALS CONDITION: VARS achieves similar or better sample-quality scores than DDIM with the same number of parameters and can outperform the DDIM(η = 0) baseline.

5 EXPERIMENTS

Experiments evaluate DDSS on pretrained diffusion models using FID and Inception Score across datasets and inference budgets. Ablations show that learned sampler components modestly improve results, while GGDM-based searches outperform less flexible alternatives more substantially.

  • 5 EXPERIMENTS: DDSS is evaluated on pretrained CIFAR10 and ImageNet 64x64 DDPMs, comparing discovered samplers with DDPM and DDIM baselines.
  • 5 EXPERIMENTS: The evaluation reports FID and Inception Score using 50K model samples and training-data samples.
  • 5 EXPERIMENTS: DDSS experiments use inference budgets of 5, 10, 15, 20, and 25 steps.
  • 5.1 ABLATIONS FOR KID KERNEL AND GGDM VARIANTS: The linear KID kernel performs slightly better than the cubic kernel used by Bińkowski et al. (2018) in the main results.
  • 5.1 ABLATIONS FOR KID KERNEL AND GGDM VARIANTS: Learning the linear kernel, timesteps, and x0-prediction coefficients each slightly improves FID and Inception Score, while removing any still outperforms the strongest baselines.
  • 5.2 SEARCH SPACE ABLATION: DDIM sigma optimization does not outperform DDIM(η = 0) on CIFAR10, whereas VARS can outperform DDSS applied to DDIM and some DDIM(η = 0) baselines.

6 DISCUSSION

DDSS finds better few-step samplers within a flexible GGDM family, but very small budgets and highly diverse ImageNet data remain challenging. The paper also argues that effective sampling paths need not preserve the original DDPM forward-process marginals.

  • DDSS consistently improves image generation quality over strong baselines when applied to sufficiently flexible sampler families such as GGDM.
  • K < 10 inference steps remains challenging because the proposed search family still shows apparent quality loss at such small budgets.
  • ImageNet 64x64 shows less pronounced relative improvement than CIFAR10, which the authors hypothesize reflects ImageNet’s high sample diversity.
  • Pre-trained score-based models can produce high-quality few-step samples without parameter fine-tuning, although default DDPM sampling is usually suboptimal in this regime.
  • Alternative variational lower bounds reveal better sampling paths that need not share the original DDPM forward process’s marginals.

7 OTHER RELATED WORK

The related work covers dynamic solvers, learned step selection, alternative score-based samplers, and qualitative comparisons against DDIM. These approaches target reduced inference cost but differ in sampling procedure and quality trade-offs.

  • A dynamic-step SDE solver reduces score-function calls to approximately 150 on CIFAR10 but falls behind DDIM(η = 0) at 50 steps.
  • Watson et al. propose dynamic programming to choose log-likelihood-optimized inference steps, representing a likelihood-focused alternative to quality-optimized sampler search.
  • Score-based alternatives include reverse-SDE discretizations, Langevin corrector steps, and probability-flow ODE sampling.
  • Figure 3 compares non-cherrypicked CIFAR10 samples from DDIM(η = 0) and the proposed approach using identical random seeds, with full-step DDPM and real images as references.

8 CONCLUSION AND FUTURE WORK

The paper introduces DDSS for optimizing few-step diffusion samplers through perceptual loss and presents GGDM as a broad search family. It reports improved unconditional image quality without retraining, while identifying computationally demanding optimization as future work.

  • DDSS optimizes a perceptual loss over diffusion processes using samples from a pre-trained DDPM, enabling few-step sampler search.
  • The method achieves qualitative and quantitative sample-quality improvements over prior efficient DDPM sampling methods for unconditional image generation.
  • GGDM provides a wider sampler family than prior DDIM-focused families, and this flexibility is described as central to DDSS’s success.
  • DDSS requires no pre-trained DDPM fine-tuning or retraining, is applied once, and uses few hyperparameters.
  • Future work includes reducing DDSS’s extensive memory or computation requirements through alternative perceptual losses or gradient-free optimization.

A.1 ADDITIONAL IMAGENET 64X64 SAMPLES

Additional ImageNet 64x64 samples compare linear-stride DDPM and DDIM baselines with DDSS + GGDM samples using a learned stride. Full-step DDPM and real samples provide references.

  • The additional ImageNet 64x64 figure places DDPM samples on the left, DDIM(η = 0) samples in the middle, and DDSS + GGDM samples on the right.
  • DDPM and DDIM(η = 0) use a linear stride, whereas DDSS + GGDM uses a learned stride.
  • Full-step DDPM samples using all 4,000 steps and real samples are included as reference panels.

A.2 PROOF FOR THEOREM 1

The section completes Theorem 1’s induction proof and reports supplementary LSUN and ImageNet experiments evaluating DDSS across inference-step budgets.

  • A.2 PROOF FOR THEOREM 1: Theorem 1 is established by proving a base case and then an inductive step using prior conjugacy.The proof concludes that the result holds for every i ∈ {1, ..., T −t}.
  • A.2 PROOF FOR THEOREM 1: The proof derives each required conditional distribution from the preceding definition and the rewritten inductive hypothesis.The argument explicitly invokes q(x_t+j+1|x_>t+j+1, x_0) and prior conjugacy.
  • Additional ablations: Supplementary ImageNet 64x64 ablations use a step budget K = 15 and fix non-learned timesteps to a linear stride.The linear stride is selected because Table 2 reports it performs best on ImageNet 64x64.
  • LSUN experiments: LSUN bedroom and church models are trained at 128x128 with matched hyperparameters apart from training duration.The models are trained for 400K and 200K steps, respectively, using Adam, batch size 2048, gradient clipping, dropout, and EMA.
  • LSUN experiments: DDSS is evaluated with GGDM alone at 5, 10, and 20 inference steps against DDPM and DDIM(η = 0) at matching budgets.The supplementary figures use identical random seeds and a linear stride, with 1,000-step DDPM samples and real images included for reference.
Loading 2202.05830v1…