Source-linked AI summary
Learning to Efficiently Sample from Diffusion Probabilistic Models
Daniel Watson, Jonathan Ho, Mohammad Norouzi, William Chan
TL;DR
DDPM sampling is expensive because high-fidelity generation often requires hundreds to thousands of refinement steps, while prior schedule selection relied on handcrafted choices. The paper uses exact dynamic programming over decomposable ELBO terms to optimize inference schedules for pretrained DDPMs, achieving 32-step schedules with at most 0.1 bits per dimension sacrificed relative to long schedules.
Problem
DDPMs often require hundreds-to-thousands of refinement steps, and prior sampling-speed improvements commonly select inference schedules by handcrafted experimentation.
Method
The paper applies dynamic programming to the ELBO's decomposed KL terms to find optimal inference schedules for a pretrained DDPM without retraining.
Results
32-step schedules sacrifice ≤0.1 bits per dimension compared with state-of-the-art DDPMs using hundreds-to-thousands of refinement steps on CIFAR10 and ImageNet 64x64.
Takeaways & Limitations
Inference schedules can be optimized once after training, providing efficient few-step sampling for both time-discrete and time-continuous DDPMs.
Takeaways & Limitations
Improved log likelihoods do not necessarily translate into improved FID scores.
Abstract
from arXiv · showhide
Denoising Diffusion Probabilistic Models (DDPMs) have emerged as a powerful family of generative models that can yield high-fidelity samples and competitive log-likelihoods across a range of domains, including image and speech synthesis. Key advantages of DDPMs include ease of training, in contrast to generative adversarial networks, and speed of generation, in contrast to autoregressive models. However, DDPMs typically require hundreds-to-thousands of steps to generate a high fidelity sample, making them prohibitively expensive for high dimensional problems. Fortunately, DDPMs allow trading generation speed for sample quality through adjusting the number of refinement steps as a post process. Prior work has been successful in improving generation speed through handcrafting the time schedule by trial and error. We instead view the selection of the inference time schedules as an optimization problem, and introduce an exact dynamic programming algorithm that finds the optimal discrete time schedules for any pre-trained DDPM. Our method exploits the fact that ELBO can be decomposed into separate KL terms, and given any computation budget, discovers the time schedule that maximizes the training ELBO exactly. Our method is efficient, has no hyper-parameters of its own, and can be applied to any pre-trained DDPM with no retraining. We discover inference time schedules requiring as few as 32 refinement steps, while sacrificing less than 0.1 bits per dimension compared to the default 4,000 steps used on ImageNet 64x64 [Ho et al., 2020; Nichol and Dhariwal, 2021].
1 Introduction
DDPMs are easy to train and faster to sample than autoregressive models, but their hundreds-to-thousands of denoising steps remain costly. The paper treats inference-schedule selection as an optimization problem and uses dynamic programming to find ELBO-optimal schedules for pretrained models.
- DDPMs use a simple denoising objective and non-autoregressive inference, contrasting with GAN instability and autoregressive computational cost.
- Hundreds-to-thousands of refinement steps can remain computationally prohibitive, especially for high-dimensional data.
- Prior acceleration methods include manually crafted schedules, non-Markovian DDIM processes, hand-crafted subsequences, and dynamically adjusted schedules.
- DDPMs decouple training and inference schedules, allowing a model trained with many steps to use fewer inference steps.
- The proposed dynamic programming algorithm finds inference paths maximizing the ELBO by exploiting its decomposition into individual KL terms.
- With 32 refinement steps, discovered schedules sacrifice only 0.1 bits per dimension relative to original 1,000-step CIFAR10 and 4,000-step ImageNet 64x64 models.
2 Background on Denoising Diffusion Probabilistic Models
DDPMs define generation through a forward Gaussian diffusion process and a learned reverse denoising process. Their tractable Gaussian marginals and posteriors support efficient ELBO-based training and closed-form KL evaluation.
- A DDPM combines a forward Markovian process q that adds Gaussian noise with a learned reverse process pθ that removes it.
- The forward transition q(x_t | x_t−1) is Gaussian, with α_t controlling the variance of noise added at each step.
- The reverse process starts from pure Gaussian noise x_T and progressively reconstructs the data distribution.
- Reverse-process parameters are optimized by maximizing a variational lower bound on the training data.
- Gaussian posteriors make each KL term in the ELBO analytically computable, avoiding high-variance Monte Carlo estimation.
3 Linking DDPMs to Continuous Time Affine Diffusion Processes
The paper connects DDPMs to continuous-time affine diffusion processes to make its sampling approach applicable across discrete and continuous formulations. The construction supports exact marginals and ancestral paths with continuous timesteps.
- An affine diffusion SDE defines q(x_t|x_0) for t ∈ [0, 1], linking DDPMs to continuous-time processes.
- The SDE coefficients f_sde and g_sde are integrable and satisfy f_sde(0)=1 and g_sde(0)=0.
- Exact marginals q(x_t|x_s) can be computed for arbitrary 0 ≤ s < t ≤ 1, unlike marginals restricted to a discretized diffusion grid.
- The formulation generalizes discrete-time DDPM marginals and supports training by sampling continuous t uniformly, as well as arbitrary SDE choices.
- Inference can use ancestral sampling paths with continuous-valued timesteps, and the reverse process is mathematically equivalent to subsequence-based discrete DDPM sampling.
- For s=0, the Ho et al. parametrization is used to obtain discretized log likelihoods for fair comparison with prior work.
4 Learning to Efficiently Sample from DDPMs
The method treats inference-schedule selection as an ELBO optimization problem and uses dynamic programming to find optimal paths under fixed refinement-step budgets. It efficiently evaluates transition costs and can recover optimal schedules across budgets without retraining.
- ELBO optimization: The ELBO decomposes into terms associated with contiguous timestep transitions, enabling schedule optimization after DDPM training.The inference path must start at t = 0, end at t = 1, and remain contiguous.
- ELBO optimization: For a fixed budget K, the optimal schedule is the path with the best ELBO among valid paths containing K refinement steps.The schedule is selected from timestep paths satisfying the endpoint and connectivity constraints.
- Dynamic programming algorithm: The algorithm solves schedule selection as a least-cost path problem on a timestep graph whose edge costs are transition-specific ELBO terms.Shared transitions share costs, allowing memoization of L(t, s) before path optimization.
- Dynamic programming algorithm: Dynamic programming restricts paths to exactly K + 1 nodes with fixed start and end nodes, storing optimal costs and predecessor timesteps.Matrices C and D represent least-path costs and the preceding timestep for reconstruction.
- Dynamic programming algorithm: Running the algorithm through all budgets extracts optimal paths for every possible K rather than requiring a separate optimization for each budget.The vectorized procedure computes paths for all budgets and Algorithm 2 reconstructs a selected path.
- Dynamic programming algorithm: Although transition costs appear to require O(T^2) computation, one neural-network forward pass computes all costs from a given timestep, reducing the table to O(T) passes.This reduction also applies when the pretrained DDPM learns its variances.
5 Experiments
The experiments evaluate dynamic-programming schedules across pretrained DDPMs and datasets, comparing likelihood, FID, samples, and Monte Carlo estimation. The method improves few-step likelihoods substantially, including a 32-step ImageNet schedule within 0.1 bits/dim of the 4,000-step baseline, although likelihood gains do not necessarily improve FID.
- Experimental setup: The study evaluates pretrained DDPMs on CIFAR10 and ImageNet 64x64, including Lsimple and Lhybrid configurations.The experiments use prior-work checkpoints and models trained with the authors’ configurations.
- Likelihood results: 32 steps achieve a score within ≤0.1 bits/dim of the original 4,000-step budget for Lhybrid ImageNet 64x64.The reported comparison uses the same model and negative log likelihood in bits per dimension.
- Likelihood results: Dynamic-programming strides produce much better log likelihoods than even and quadratic hand-crafted strides, particularly in the few-step regime.The comparison is visualized for Lvlb CIFAR10 and Lhybrid ImageNet 64x64.
- FID comparison: Optimizing strides for log likelihood does not necessarily improve FID scores.The paper reports FID on 50,000 model samples and notes a weakened correlation between log likelihood and FID.
- Monte Carlo ablation: Using 128 Monte Carlo datapoints yields negligible likelihood differences on CIFAR10 but a visible yet slight ImageNet improvement with more samples.The authors attribute the ImageNet sensitivity to its higher diversity.
- Qualitative comparison: The sample comparisons use non-cherrypicked outputs from even, quadratic, and dynamic-programming strides with the same eight random seeds.This comparison is shown for both Lsimple CIFAR10 and Lhybrid ImageNet 64x64 across computation budgets.
6 Related Work
The related work situates DDPMs among generative models and prior few-step sampling methods. It highlights dynamic-programming schedules as part of a broader effort to accelerate diffusion sampling while preserving model quality.
- Diffusion models: DDPMs have been applied to image generation, shape generation, text-to-speech, and single-image super-resolution.The passage presents DDPMs as a broad class of iterative denoising generative models.
- Few-step sampling: Prior few-step DDPM research includes non-Markovian acceleration and manually crafted schedules for faster sampling.The cited work includes DDIMs, WaveGrad, and other few-step approaches.
- Schedule structure: The paper’s dynamic-programming schedules are illustrated as allocating steps toward both the end and, potentially, the beginning of diffusion.The authors associate late steps with finer image details and speculate that early steps may help break modes.
7 Conclusion and Discussion
The paper presents dynamic programming as an efficient way to optimize inference schedules for pretrained DDPMs, achieving strong likelihoods with few refinement steps. It also notes that likelihood improvements do not necessarily improve FID.
- Conclusion and Discussion: The method optimizes inference schedules for pretrained DDPMs using dynamic programming under a fixed computation budget.The schedule is discovered post-training and requires only forward passes, without retraining.
- Conclusion and Discussion: 32 refinement steps suffice on CIFAR10 and ImageNet 64x64 while sacrificing ≤0.1 bits per dimension relative to DDPMs using hundreds-to-thousands of steps.
- Conclusion and Discussion: The approach achieves strong likelihoods with few refinement steps and outperforms prior methods using hand-crafted strides.
- Conclusion and Discussion: Improved log-likelihoods do not necessarily translate into improved FID scores, reflecting a mismatch between the two metrics.
A.1 Proof for Equation 12
The proof derives identities by implicit differentiation and by using the relation ψ(t, s) = ψ(t,0).
- A.1 Proof for Equation 12: Implicit differentiation is used to derive the result associated with Equation 12.
- A.1 Proof for Equation 12: The two stated steps provide the proof's algebraic basis for Equation 12.
- A.1 Proof for Equation 12: The derivation also uses the identity ψ(t, s) = ψ(t,0).
A.2 Proof for Equations 13 and 14
The proof identifies Gaussian conditional-distribution parameters through conjugacy, showing that f_ts is a conditional mean and g^2_ts is a conditional variance.
- A.2 Proof for Equations 13 and 14: f_ts is the mean of q(x_t|x_s).
- A.2 Proof for Equations 13 and 14: The variance of q(x_t|x_s) is derived from Equation 10.
- A.2 Proof for Equations 13 and 14: The conditional mean of q(x_s|x_t, x_0) is obtained using the Gaussian conjugate prior formula.
- A.2 Proof for Equations 13 and 14: The proof represents the prior over μ and a likelihood with mean μ before deriving the conditional variance.