[Paper Review] Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality

Daniel Watson, William Chan+2 more|arXiv (Cornell University)|Feb 11, 2022

Generative Adversarial Networks and Image Synthesis23 citations

TL;DR

The paper introduces Differentiable Diffusion Sampler Search (DDSS) to optimize fast, non-Markovian samplers (GGDM family) by differentiating through sample quality, achieving high-quality few-step diffusion sampling without retraining the model.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing inference steps in diffusion models without retraining or fine-tuning the original model.
• Propose a differentiable optimization framework to search for fast samplers.
• Introduce the Generalized Gaussian Diffusion Model (GGDM) as a flexible sampler family.
• Demonstrate that optimizing samplers with a perceptual loss (KID) yields superior sample quality in few-step regimes.

Proposed method

• Unroll the diffusion sampling process and backpropagate using the reparameterization trick and gradient rematerialization to optimize a parametric sampler.
• Define sampler families (DDIM, VARS, GGDM, and GGDM variants) with trainable parameters controlling means and variances.
• Use a perceptual loss based on Kernel Inception Distance (KID) computed in Inception features to align with human-perceived quality (Eq. 7–9).
• Backpropagate through a stochastic sampler by differentiating through the sampling chain with mini-batch SGD (Adam).
• Introduce Generalized Gaussian Diffusion Models (GGDM) that incorporate information from all previous (noisier) images at each denoising step.
• Allow learning of timestep selection (TIME) and predictive coefficients (PRED) to improve few-step performance.

Experimental results

Research questions

• RQ1 Can a differentiable optimization procedure identify fast samplers that outperform existing few-step baselines for pre-trained diffusion models?
• RQ2 Does optimizing a perceptual loss (KID) yield samplers with higher visual fidelity than optimizations based on likelihood or ELBO alone?
• RQ3 How flexible is the GGDM family in discovering high-quality few-step samplers compared to DDIM or VARS?

Key findings

• DDSS discovers fast samplers that achieve substantially better FID/IS scores than strong baselines at the same small step budgets (e.g., GGDM +PRED+TIME with 5–25 steps outperforms DDPM/DDIM baselines).
• Optimizing with KID as the perceptual loss yields high-fidelity samples and robust improvements across datasets (CIFAR-10, ImageNet-64).
• GGDM, a broader sampler family, consistently yields better few-step performance than narrower families like DDIM or VARS, even when marginals differ from the original forward process.
• DDSS does not require fine-tuning or retraining the pre-trained DDPM; it operates as a one-shot post-hoc sampler search.
• Across CIFAR-10 and ImageNet-64, DDSS variants achieve strong reductions in FID and competitive IS with very few steps (e.g., CIFAR-10: 5–25 steps; ImageNet-64: 5–25 steps).
• Qualitative samples show non-cherrypicked improvements over DDIM(η=0) at low step counts.