Research map / Image generation / Fast diffusion sampling
Align Your Steps: Optimizing Sampling Schedules in Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Align Your Steps (AYS), a principled method that optimizes diffusion model sampling schedules by minimizing a KL-divergence upper bound (KLUB) derived via Girsanov's theorem, improving few-step sample quality across image, video and 2D toy tasks.
Problem
Diffusion model sampling is slow, and the discretized noise-level schedule used by solvers is almost always a hand-crafted heuristic (EDM, log-SNR, time-uniform), leaving the schedule's effect on output quality, especially in few-step regimes, unoptimized.
Approach
Reinterpret stochastic solvers as exactly solving a linearized SDE on each sub-interval, bound the KL divergence between the true learnt reverse SDE and the discretized solver output using Girsanov's theorem (KLUB), and optimize the schedule points iteratively by zeroth-order candidate search with importance-sampled Monte Carlo estimates, using a hierarchical 10→20→40-step scheme.
Key results
- On SDE-DPM-Solver++(2M), an optimized 20-step AYS schedule on CIFAR10/FFHQ reaches FIDs comparable to a 30-step default schedule, a 1.5x speedup.
- In a 42-participant user study on Stable Diffusion 1.5 with 10 steps, participants preferred AYS over EDM and time-uniform schedules for image fidelity and image-text alignment.
- In a Stable Video Diffusion user study, AYS was preferred 58% vs 42% over EDM.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the default EDM schedule with an AYS-optimized schedule | Stable Video Diffusion with DDIM at 10 steps · Stable Video Diffusion user study | preference share: 42\% -> 58\% | Helped |
Surprising findings
- Minimizing the KLUB objective does not translate directly into minimizing final output KL; the KLUB-optimal schedule differs from the KL-optimal schedule for forward Euler, so early stopping is necessary.
- Optimized schedules for the Gaussian toy case differ substantially from popular hand-crafted schedules even for unit variance data.
Limitations
- The KLUB objective is an upper bound on discretization error rather than the error itself, requiring early stopping based on perceptual metrics or manual inspection.
- Schedules are optimized for specific model, dataset, solver and guidance value, and higher-order ODE solvers such as Heun or Runge-Kutta are not covered.
- Evaluation on text-to-image and video relies mainly on small user studies rather than standard automatic metrics.
Experimental setup
Datasets: CIFAR10 FFHQ ImageNet
Benchmarks: CIFAR10 FFHQ ImageNet
Metrics: FID CLIP score user preference
Baselines: EDM schedule time-uniform schedule LogSNR schedule quadratic time schedule cosine LogSNR schedule DDIM DPM-Solver++(2M)
Builds on: EDM DDIM DPM-Solver++ ER-SDE-Solver Girsanov's theorem
Closest papers by meaning
- Score-Optimal Diffusion Schedules (NeurIPS 2024 · similarity 0.92)
- A Unified Sampling Framework for Solver Searching of Diffusion Probabilistic Models (ICLR 2024 · similarity 0.92)
- Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models (ICLR 2025 · similarity 0.91)
- Accelerating Diffusion Sampling with Optimized Time Steps (CVPR 2024 · similarity 0.91)
- CCS: Controllable and Constrained Sampling with Diffusion Models via Initial Noise Perturbation (NeurIPS 2025 · similarity 0.91)
- Learning to Discretize Denoising Diffusion ODEs (ICLR 2025 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.