Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality

TL;DR

This paper introduces a differentiable approach to optimize fast samplers for diffusion models, significantly reducing inference steps while maintaining high sample quality, applicable to any pre-trained model without re-training.

Contribution

It proposes Differentiable Diffusion Sampler Search (DDSS) and Generalized Gaussian Diffusion Models (GGDM), enabling gradient-based optimization of samplers for improved efficiency and quality.

Findings

Achieves high-quality image generation with fewer inference steps.

Outperforms existing methods on FID scores across datasets.

Compatible with pre-trained models without re-training.

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.