Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality

By Daniel Watson, William Chan, Jonathan Ho, Mohammad Norouzi. 2022

241 citations recorded by the source corpus.

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.

Subfields: Computer Vision and Pattern Recognition, Radiology, Nuclear Medicine and Imaging, Generative Models, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Radiology, Nuclear Medicine and Imaging

In Generative Models

In Optimization & Training

Citations in the corpus

Elucidating the Design Space of Diffusion-Based Generative Models (2022, background): “In other words, the training procedure of D θ should not dictate σ ( t ) , s ( t ) , and { t i } , nor vice versa; from the viewpoint of the sampler, D θ is simply a black box [52, 53].”
DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps (2022, background): “Samplers that need extra training or optimization include knowledge distillation [13, 14], learning the noise level or variance [15, 16, 33], and learning the noise schedule or sample trajectory [17, 18].”
UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models (2023, methodology): “Fast Sampling of DPMs Fast samplers of DPMs can be either training-based (Salimans & Ho, 2022; Bao et al., 2022a; Watson et al., 2021) or training-free (Lu et al.”
Directly Fine-Tuning Diffusion Models on Differentiable Rewards (2023, unknown): “Unfortunately, we found this method to be unstable and ultimately not learn effectively, possibly because it relies on the assumption that ∇ x r ( x 0 , c ) ≈ ∇ x r (ˆ x 0 , c ) , which may not be true in practice.”
ModelScope Text-to-Video Technical Report (2023, background): “To solve the first obstacle, research efforts focused on improving the sampling efficiency by learning-free sampling [54, 55, 35, 69, 37, 69] and learning based sampling [60, 70, 49].”
Align Your Latents: High-Resolution Video Synthesis with Latent Diffusion Models (2023, background): “However, despite advances in model distillation [49,52,73] and accelerated sampling [4,11,12,36,45,46,80,96,100,108], DMs generally require repeated evaluations of a computationally demanding large U-Net.”
A Comprehensive Survey of AI-Generated Content (AIGC): A History of Generative AI from GAN to ChatGPT (2023, background): “There are also works [125, 126] which directly figure out optimal sampling trajectories via dynamic programming.”
Bespoke Solvers for Generative Flow Models (2023, methodology): “…aimed at learning the solver: Dockhorn et al. (2022) (GENIE) introduced a higher-order solver, and distilled the necessary JVP for their method; Watson et al. (2021) (DDSS) optimized a perceptual loss considering a family of generalized Gaussian diffusion models; Lam et al. (2021) improved the…”
DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models (2022, unknown): “Samplers based on extra training include model distillation (Salimans & Ho, 2022; Luhman & Luhman, 2021), learning reverse process variances (San-Roman et al., 2021; Nichol & Dhariwal, 2021; Bao et al., 2022a), and learning sampling steps (Lam et al., 2021; Watson et al., 2022).”
LION: Latent Point Diffusion Models for 3D Shape Generation (2022, background): “[118] Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi.”
A Survey on Generative Diffusion Models (2022, methodology): “Current DP-based methods [85], [113] take O ( T 2 ) of computational cost by optimizing the sum of ELBO losses.”
Diffusion Models: A Comprehensive Survey of Methods and Applications (2022, methodology): “It also utilizes DDSS [213] to optimize fast samplers within a few time steps by differentiating through sample quality scores.”
RealFusion 360° Reconstruction of Any Object from a Single Image (2023, background): “Most of these approaches train deep networks specifically for the goal of inferring NeRFs from a small number of views; examples include IBRNet [44], NeRF-WCE [9], Pixel-2”
I2VGen-XL: High-Quality Image-to-Video Synthesis via Cascaded Diffusion Models (2023, background): “For this purpose, some studies aim to reduce the number of denoising steps in the model by improving sampling efficiency, such as learning-free sampling [24,42,43,56,56] and learning based sampling [37, 48, 60].”
The Role of ImageNet Classes in Fréchet Inception Distance (2022, methodology): “We also suspect that similar interference with FID can happen when the ImageNet pre-trained classifiers are used to curate the training data [9] or as a part of the sampling process [57].”