Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

TL;DR

This work addresses the slow sampling problem in diffusion models by optimizing the sampling schedule itself. It introduces Align Your Steps (AYS), a principled, dataset/model/solver-specific framework that minimizes a KL-divergence-based upper bound (KLUB) between true and discretized SDEs to derive better schedules. The approach yields consistent quality improvements across multiple modalities (2D, image, video, text-to-image) and solver types, with optimized schedules released for popular models. Overall, AYS demonstrates that carefully shaped sampling schedules can substantially boost few-step diffusion performance, expanding the practical utility of diffusion-based generative systems.

Abstract

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called $\textit{Align Your Steps}$. We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

TL;DR

This work addresses the slow sampling problem in diffusion models by optimizing the sampling schedule itself. It introduces Align Your Steps (AYS), a principled, dataset/model/solver-specific framework that minimizes a KL-divergence-based upper bound (KLUB) between true and discretized SDEs to derive better schedules. The approach yields consistent quality improvements across multiple modalities (2D, image, video, text-to-image) and solver types, with optimized schedules released for popular models. Overall, AYS demonstrates that carefully shaped sampling schedules can substantially boost few-step diffusion performance, expanding the practical utility of diffusion-based generative systems.

Abstract

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called . We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.
Paper Structure (29 sections, 3 theorems, 48 equations, 26 figures, 7 tables, 2 algorithms)

This paper contains 29 sections, 3 theorems, 48 equations, 26 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Optimizing Sampling Schedules
  4. The Need for Optimized Schedules
  5. Analyzing the Discretization Errors
  6. Practical Considerations of KLUB Estimation
  7. Related Work
  8. Experiments
  9. Toy Experiments
  10. CIFAR10, FFHQ, ImageNet
  11. Text-to-Image
  12. Video Generation Models
  13. Conclusions and Future Work
  14. Theoretical Details
  15. Optimal Schedule for Isotropic Gaussian
  16. ...and 14 more sections

Key Result

Theorem 3.1

Let $p_{data}({\mathbf{x}}) = {\mathcal{N}}(\mathbf{0}, c^2 \mathbf{I})$. Sample ${\mathbf{x}}_{t_{max}}{\sim} p({\mathbf{x}}, t_{max})$ and solve the probability flow ODE using $n$ forward euler steps along the schedule $t_{max} = t_n > t_{n-1} > \dots > t_{1} > t_{0} = t_{min}$ to obtain $\Bar{{\m

Figures (26)

  • Figure 1: Sample fidelity (FID $\downarrow$, sFID $\downarrow$, Inception Score $\uparrow$) on the ImageNet $256 \times 256$ dataset.
  • Figure 2: Align Your Steps. We minimize an upper bound on the Kullback-Leibler divergence ($\mathrm{KLUB}$) between the true and linearized generative SDEs to find optimal DM sampling schedules.
  • Figure 3: Comparing popular sampling schedules against the optimal schedules for Gaussian data.
  • Figure 4: Modeling a 2D toy distribution: (a) Ground truth samples; (b), (c), and (d) are samples generated using 8 steps of SDE-DPM-Solver++(2M) with EDM, LogSNR, and AYS schedules, respectively. Each image consists of 100,000 sampled points. The colors denote the local density of the samples where warmer colors correspond to higher density regions. See \ref{['appendix:extra_2d_experiments']} for details.
  • Figure 5: Side-by-side comparison of selected images generated with Stable Diffusion 1.5 with SDE-DPM-Solver++(2M) over 10 steps with different sampling schedules.
  • ...and 21 more figures

Theorems & Definitions (3)

  • Theorem 3.1: Proof in \ref{['appendix:optimal_gaussian']}
  • Theorem 3.2: KL-divergence Upper bound ($\mathrm{KLUB}$), proof in \ref{['appendix:proving_klub']}
  • Lemma 3.3: Proof in \ref{['appendix:early_stopping_needed']}