Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Score-Optimal Diffusion Schedules

Christopher Williams, Andrew Campbell, Arnaud Doucet, Saifuddin Syed

TL;DR

The paper tackles the challenge of choosing discretisation schedules for score-based diffusion models by introducing score-optimal diffusion schedules that minimize a data-driven transport cost along the diffusion path. It frames the problem via a predictor–corrector decomposition and defines an incremental work measure using Stein/Fisher divergences, then shows the optimal schedule corresponds to a geodesic in a metric induced by the local cost. An online algorithm estimates the optimal schedule from estimated scores (and Hessians when using $\mathcal{L}_p$), with a velocity scaling $v(t)=\sigma(t)$ to align with diffusion dynamics. Empirically, the method recovers known performant schedules on image data, improves sampling efficiency, and enables online schedule learning during training across 1D and high-dimensional datasets, using pre-trained diffusion models as a testbed. The approach offers a hyperparameter-free, scalable way to tailor discretisation to data geometry, with practical impact on sample quality and speed for diffusion-based generative models.

Abstract

Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate a path of probability distributions interpolating between a reference Gaussian distribution and a data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial to obtain high quality samples. However, beyond hand crafted heuristics, a general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to a cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to the dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation of the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that our learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores on image datasets.

Score-Optimal Diffusion Schedules

TL;DR

The paper tackles the challenge of choosing discretisation schedules for score-based diffusion models by introducing score-optimal diffusion schedules that minimize a data-driven transport cost along the diffusion path. It frames the problem via a predictor–corrector decomposition and defines an incremental work measure using Stein/Fisher divergences, then shows the optimal schedule corresponds to a geodesic in a metric induced by the local cost. An online algorithm estimates the optimal schedule from estimated scores (and Hessians when using ), with a velocity scaling to align with diffusion dynamics. Empirically, the method recovers known performant schedules on image data, improves sampling efficiency, and enables online schedule learning during training across 1D and high-dimensional datasets, using pre-trained diffusion models as a testbed. The approach offers a hyperparameter-free, scalable way to tailor discretisation to data geometry, with practical impact on sample quality and speed for diffusion-based generative models.

Abstract

Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate a path of probability distributions interpolating between a reference Gaussian distribution and a data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial to obtain high quality samples. However, beyond hand crafted heuristics, a general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to a cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to the dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation of the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that our learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores on image datasets.

Paper Structure

This paper contains 32 sections, 3 theorems, 55 equations, 12 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. The Cost of Traversing the Diffusion Path
  3. Predictor/Corrector Decomposition of the Diffusion Update
  4. The Incremental Cost of Correction
  5. Corrector and Predictor Optimised Cost
  6. Score-Optimal Schedules
  7. Diffusion Schedule Path Length and Energy
  8. Estimation of Score-Optimal Schedules
  9. Choice of Velocity Scaling
  10. Related Work
  11. Computational Experiments
  12. Sampling the Mollified Cantor Distribution
  13. Adaptive Schedule Learning for Bimodal Example
  14. Scalable Schedule Learning Diffusion
  15. Sampling Pre-Trained Models
  16. ...and 17 more sections

Key Result

Theorem 2.1

Suppose $p_t(x),F_{t,t'}(x),v(t)$ and $G_{t,t'}(x)$ are three-times continuously differentiable in $t,t',x$ and let $\dot F_t(x)=\left.\frac{\partial}{\partial t'}F_{t,t'}(x)\right|_{t'=t}$ and $\dot G_t(x)=\left.\frac{\partial}{\partial t'}G_{t,t'}(x)\right|_{t'=t}$. Suppose the following hold: (1)

Figures (12)

  • Figure 1: Density estimates of the mollified Cantor distribution (left) using a DDM with schedule $\mathcal{T}=\{t_i\}_{i=0}^{100}$ generated with $100$ linearly spaces discretisation times $t_i=i/100$ (middle), compared to the optimised schedule $\mathcal{T}^*=\{t_i^*\}_{i=0}^{50}$ with $50$ discretisation times $t_i^*$ generated by \ref{['alg:beta']} (right). The eight modes present in our true mollified distribution are shown in grey on each plot.
  • Figure 2: Comparison of Linear and Learned Schedules over Training Iterations for the bimodal example. Each point corresponds to 500 training iterations.
  • Figure 3: (Left) Incremental costs $\sqrt{\mathcal{L}(t_{j+1}, t_j)}$ for the cosine schedule and our online adaptive algorithm. Higher learning rates enforce equalisation of costs more quickly. (Right) Progression of the learned schedule during 40k training iterations, depicted through the standard-deviation $\sqrt{1 - \bar{\alpha}_t}$.
  • Figure 4: (Left) Costs associated with different schedule choices for the CIFAR10 dataset. Schedules are ordered from lowest FID to highest FID. We compare our Corrector-optimised (CO) cost and Predictor-optimised (PO) cost versus the Kullback-Leibler Upper Bound (KLUB) from sabour2024align. The minimum value for each cost is highlighted in bold. Note low cost is associated with low FID for our cost and not for the KLUB. (Right) Visualisation of schedules during generative sampling with $100$ timesteps. "rho=3" and "rho=7" refer to Eq \ref{['eq:karras_schedule']} with $\rho=3$ and $\rho=7$ respectively. LogLinear from lu_dpm-solver_2022 and a convex schedule are also shown. We show our cost optimised schedules for CIFAR10 both using the corrector optimised cost and the predictor optimised cost.
  • Figure 5: Schedules and density estimates for: linear (blue); Stein score optimised (green); and predictor optimised (red) schedules. The predictor optimised schedule identifies a bump along the diffusion path where the reference Gaussian density splits into two modes. In the regions where the score is evaluated (around $\pm 6$), our trained score is accurate compared to the linear schedule score, which fails to match the slope of the true score, resulting in a wider variance density estimate.
  • ...and 7 more figures

Theorems & Definitions (10)

  • Definition 2.1
  • Definition 2.2
  • Example 2.1: Annealed Langevin
  • Example 2.2: Probability Flow ODE
  • Theorem 2.1
  • Theorem 3.1
  • proof
  • proof
  • Proposition B.1
  • proof