# Adaptive Non-uniform Timestep Sampling for Accelerating Diffusion Model Training

arXiv: https://arxiv.org/abs/2411.09998 · HTML: https://mlsanity.com/paper/2411.09998/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 29 of its 36 references that are in the map.
- First posted on arXiv 15 Nov 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 29 of 29)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Text2Video-Zero: Text-to-Image Diffusion Models are Zero-Shot Video Generators](https://mlsanity.com/paper/2303.13439/) (2023)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Perception Prioritized Training of Diffusion Models](https://mlsanity.com/paper/2204.00227/) (2022)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [DIFUSCO: Graph-based Diffusion Solvers for Combinatorial Optimization](https://mlsanity.com/paper/2302.08224/) (2023)
- [DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models](https://mlsanity.com/paper/2211.15029/) (2022)
- [Learning to Optimize](https://mlsanity.com/paper/1606.01885/) (2016)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [Learned Optimizers that Scale and Generalize](https://mlsanity.com/paper/1703.04813/) (2017)
- [Neural Optimizer Search with Reinforcement Learning](https://mlsanity.com/paper/1709.07417/) (2017)
- [A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization](https://mlsanity.com/paper/2406.01661/) (2024)
- [3DDesigner: Towards Photorealistic 3D Object Generation and Editing with Text-guided Diffusion Models](https://mlsanity.com/paper/2211.14108/) (2022)
- [A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training](https://mlsanity.com/paper/2405.17403/) (2024)
- [Fast Diffusion Model](https://mlsanity.com/paper/2306.06991/) (2023)
- [ArchiSound: Audio Generation with Diffusion](https://mlsanity.com/paper/2301.13267/) (2023)

## Built on by (most cited first, 9 of 9)

- [TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization](https://mlsanity.com/paper/2412.21037/) (2024)
- [T-LoRA: Single Image Diffusion Model Customization Without Overfitting](https://mlsanity.com/paper/2507.05964/) (2025)
- [Fast3Dcache: Training-free 3D Geometry Synthesis Acceleration](https://mlsanity.com/paper/2511.22533/) (2025)
- [MotionFlux: Efficient Text-Guided Motion Generation through Rectified Flow Matching and Preference Alignment](https://mlsanity.com/paper/2508.19527/) (2025)
- [VADF: Vision-Adaptive Diffusion Policy Framework for Efficient Robotic Manipulation](https://mlsanity.com/paper/2604.15938/) (2026)
- [Stage-adaptive audio diffusion modeling](https://mlsanity.com/paper/2605.04547/) (2026)
- [GeoSPRINT: Geometric Redundancy-Aware Step Pruning for Inference in Diffusion Trajectories](https://mlsanity.com/paper/2609.02160/) (2026)
- [Timestep Weighting: A Hidden Key to Effective ELBO-Based Flow-Matching RL](https://mlsanity.com/paper/2609.32665/) (2026)
- [Scaling Online Optimal Transport Minibatches for Flow Matching](https://mlsanity.com/paper/2610.18838/) (2026)

## Abstract (arXiv)

As a highly expressive generative model, diffusion models have demonstrated exceptional success across various domains, including image generation, natural language processing, and combinatorial optimization. However, as data distributions grow more complex, training these models to convergence becomes increasingly computationally intensive. While diffusion models are typically trained using uniform timestep sampling, our research shows that the variance in stochastic gradients varies significantly across timesteps, with high-variance timesteps becoming bottlenecks that hinder faster convergence. To address this issue, we introduce a non-uniform timestep sampling method that prioritizes these more critical timesteps. Our method tracks the impact of gradient updates on the objective for each timestep, adaptively selecting those most likely to minimize the objective effectively. Experimental results demonstrate that this approach not only accelerates the training process, but also leads to improved performance at convergence. Furthermore, our method shows robust performance across various datasets, scheduling strategies, and diffusion architectures, outperforming previously proposed timestep sampling and weighting heuristics that lack this degree of robustness.
