Adaptive Non-uniform Timestep Sampling for Accelerating Diffusion Model Training

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kim, Myunsoo, Ki, Donghyeon, Shim, Seong-Woong, Lee, Byung-Jun
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909866463854592
author Kim, Myunsoo
Ki, Donghyeon
Shim, Seong-Woong
Lee, Byung-Jun
author_facet Kim, Myunsoo
Ki, Donghyeon
Shim, Seong-Woong
Lee, Byung-Jun
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2411_09998
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Adaptive Non-uniform Timestep Sampling for Accelerating Diffusion Model Training
Kim, Myunsoo
Ki, Donghyeon
Shim, Seong-Woong
Lee, Byung-Jun
Machine Learning
Computer Vision and Pattern Recognition
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.
title Adaptive Non-uniform Timestep Sampling for Accelerating Diffusion Model Training
topic Machine Learning
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2411.09998