Efficient Diffusion Training via Min-SNR Weighting Strategy

By Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, Baining Guo. IEEE International Conference on Computer Vision, 2023

302 citations recorded by the source corpus.

Denoising diffusion models have been a mainstream approach for image generation, however, training these models often suffers from slow convergence. In this paper, we discovered that the slow convergence is partly due to conflicting optimization directions between timesteps. To address this issue, we treat the diffusion training as a multi-task learning problem, and introduce a simple yet effective approach referred to as Min-SNR-γ. This method adapts loss weights of timesteps based on clamped signal-to-noise ratios, which effectively balances the conflicts among timesteps. Our results demonstrate a significant improvement in converging speed, 3.4× faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet 256 × 256 benchmark using smaller architectures than that employed in previous state-of-the-art. The code is available at https://github.com/TiankaiHang/Min-SNR-Diffusion-Training.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Radiology, Nuclear Medicine and Imaging

View Paper: https://www.semanticscholar.org/paper/d37ee0c440e36df133722093f6a824f4964e8ae0

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Radiology, Nuclear Medicine and Imaging

Citations in the corpus

Sora: A Review on Background, Technology, Limitations, and Opportunities of Large Vision Models (2024, methodology): “Furthermore, they also leverage a bag of training strategies like the Adan optimizer [57], Min-SNR weighting [58], dynamic masking ratio, and an improved power-cosine weighting for classifier-free guidance [59].”
Zero123++: a Single Image to Consistent Multi-view Diffusion Base Model (2023, methodology): “We employ the Min-SNR weighting strategy [5] to make the training process more efficient.”
Text-to-Audio Generation using Instruction-Tuned LLM and Latent Diffusion Model (2023, methodology): “(3) using standard normal noise ϵ n , τ is the text encoding (see Section 2.1) for guidance, and γ n is the weight of reverse step n [6], taken to be a measure of signal-to-noise ratio (SNR) in terms of α 1: N .”
Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion (2024, unknown): “For example, if one uses min snr strategy [24], its reweighting factor will always fall between [0 , C ] which we divide by C to get S t ∈ [0 , 1] .”
Masked Diffusion Transformer is a Strong Image Synthesizer (2023, methodology): “…introducing an effective mask latent modeling scheme, we proposed a masked diffusion transformer method, which, for the ﬁrst time, explicitly enhances the contextual learning ability of DPMs. • Experiments show that our method better synthesis images while using much less training time than SOTA.”
Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation (2023, unknown): “5 or 1 Min-SNR-γ [Hang et al., 2023] Table 1: Diffusion model objectives in the literature are special cases of the weighted loss with a weighting function w ( λ ) given in this table.”