Research map / Image generation / Fast diffusion sampling
Adaptive Non-Uniform Timestep Sampling for Accelerating Diffusion Model Training
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 an online learned, adaptive non-uniform timestep sampler (Beta-distributed, parameterized by a network) that is trained to maximize the per-step reduction in the diffusion variational bound, accelerating diffusion training versus heuristic timestep weighting/sampling methods.
Problem
Diffusion model training converges slowly because uniform timestep sampling wastes iterations; existing non-uniform weighting/sampling heuristics lack a clear link to the underlying objective and are not robust across settings.
Approach
Motivated by gradient variance imbalance across timesteps and strong interdependence between timestep losses, the method learns a timestep sampler online via a likelihood-ratio gradient that maximizes the estimated reduction in the VLB, with the per-step reduction approximated using a small feature-selected subset of timesteps (|S|=3).
Key results
- On CIFAR-10 with linear schedule at 1.0M iterations, the method reaches FID 2.94 versus 3.19 for DDPM and 3.16 for Min-SNR.
- On CelebA-HQ 256x256, the method reaches FID 6.62 at 300 epochs versus 6.83 for Min-SNR and 8.72 for LDM.
- Combining the method with Min-SNR weighting on ImageNet 256x256 lowers FID to 4.33 at 1.0M iterations versus 4.54 for Min-SNR alone.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use adaptive non-uniform timestep sampling instead of uniform sampling | DDPM on CIFAR-10 · CIFAR-10, Linear schedule | FID: 8.16 -> 3.99 | Helped |
| Use adaptive non-uniform timestep sampling instead of uniform sampling | DDPM on CIFAR-10 · CIFAR-10, Cosine schedule | FID: 3.74 -> 3.87 | Helped |
| Use adaptive non-uniform timestep sampling instead of uniform sampling | DDPM on CIFAR-10 · CIFAR-10, Quadratic schedule | FID: 4.32 -> 3.96 | Helped |
| Use adaptive timestep sampling instead of LDM baseline training | LDM · CelebA-HQ 256×256, 300 epochs | FID: 8.72 -> 6.62 | Helped |
| Add adaptive timestep sampling to Min-SNR weighting | U-ViT · ImageNet 256×256, 1.0M | FID: 4.54 -> 4.33 | Helped |
| Sample timesteps proportional to gradient variance | DDPM · CIFAR-10 | FID: 3.99 -> 21.3 | Hurt |
| Vary the number of selected timesteps |S| used to approximate the objective change (ablation) | Adaptive timestep sampler · CIFAR-10 | FID | Mixed |
Surprising findings
- Sampling timesteps in proportion to gradient variance performs poorly (FID 11.3 at 1.0M vs 2.94 for the learned sampler), attributed to interdependence among timestep losses.
- Heuristic methods converge faster initially, while the learned sampler is slower early but reaches better FID later.
Limitations
- Higher computational cost than heuristic approaches, growing with problem size.
- Not explored within score-based diffusion frameworks.
- Slower initial convergence than heuristic methods.
Experimental setup
Datasets: CIFAR-10 CelebA-HQ 256x256 ImageNet 256x256
Benchmarks: CIFAR-10 CelebA-HQ 256x256 ImageNet 256x256
Metrics: FID wall-clock time
Baselines: DDPM P2 Min-SNR Log-Normal SpeeD
Builds on: DDPM Min-SNR P2 weighting SpeeD Log-Normal timestep sampling
Closest papers by meaning
- Improving Training Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architecture (CVPR 2024 · similarity 0.91)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.91)
- Early-Bird Diffusion: Investigating and Leveraging Timestep-Aware Early-Bird Tickets in Diffusion Models for Efficient Training (CVPR 2025 · similarity 0.90)
- Towards More Accurate Diffusion Model Acceleration with A Timestep Tuner (CVPR 2024 · similarity 0.90)
- Entropic Time Schedulers for Generative Diffusion Models (NeurIPS 2025 · similarity 0.90)
- TimeStep Master: Asymmetrical Mixture of Timestep LoRA Experts for Versatile and Efficient Diffusion Models in Vision (ICML 2025 · similarity 0.90)
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.