A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training

2024-05-27

arXiv.org

AI Summary

In diffusion model training, non-uniform timestep distribution causes excessive sampling in the convergence regime—where gradient signals are weak—severely degrading training efficiency. To address this, we propose a dynamic timestep analysis framework that, for the first time, identifies three distinct temporal regimes: acceleration, deceleration, and convergence. Based on this characterization, we design an asymmetric importance sampling strategy and a process-incremental weighting mechanism, enabling targeted optimization of training efficiency at the individual-timestep level. Our method is plug-and-play—requiring no modifications to model architecture or loss function. Extensive experiments across diverse diffusion models (DDPM, DDIM, Score SDE), datasets (CIFAR-10, CelebA, LSUN), and tasks (image generation, super-resolution) demonstrate an average 3.0× training speedup, accompanied by proportional reductions in hardware cost and energy consumption, with zero degradation in generation quality or downstream performance.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Neural Architectures and Foundation Models

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Training diffusion models is always a computation-intensive task. In this paper, we introduce a novel speed-up method for diffusion model training, called, which is based on a closer look at time steps. Our key findings are: i) Time steps can be empirically divided into acceleration, deceleration, and convergence areas based on the process increment. ii) These time steps are imbalanced, with many concentrated in the convergence area. iii) The concentrated steps provide limited benefits for diffusion training. To address this, we design an asymmetric sampling strategy that reduces the frequency of steps from the convergence area while increasing the sampling probability for steps from other areas. Additionally, we propose a weighting strategy to emphasize the importance of time steps with rapid-change process increments. As a plug-and-play and architecture-agnostic approach, SpeeD consistently achieves 3-times acceleration across various diffusion architectures, datasets, and tasks. Notably, due to its simple design, our approach significantly reduces the cost of diffusion model training with minimal overhead. Our research enables more researchers to train diffusion models at a lower cost.

Problem

• Optimizing time step distribution for faster diffusion model training
• Reducing redundant steps in convergence area to cut costs
• Achieving 3x speed-up with asymmetric sampling strategy

Innovation

• Asymmetric sampling strategy for time steps
• Weighting strategy for rapid-change increments
• Plug-and-play architecture-agnostic speed-up method