Skip to yearly menu bar Skip to main content


Poster
in
Workshop: Deep Generative Model in Machine Learning: Theory, Principle and Efficacy

Revisiting Noise Schedule Design for Diffusion Training

Tiankai Hang · Shuyang Gu · Xin Geng · Baining Guo

Keywords: [ Diffusion Model ] [ Noise Schedule ] [ Efficient Training ]


Abstract: Diffusion models have emerged as the de facto choice for generating high-quality visual signals across various domains.However, training a single model to predict noise across various levels poses significant challenges, necessitating numerous iterations and incurring significant computational costs. Various approaches, such as loss weighting strategy design and architectural refinements, have been introduced to expedite convergence and improve model performance. In this study, we propose a novel approach to design the noise schedule for enhancing the training of diffusion models. Our key insight is that importance sampling of $\log \text{SNR}$, equivalent to a modified noise schedule, improves training efficiency by focusing around $\log \text{SNR}=0$. This sampling helps the model focus on the critical transition between signal and noise dominance, leading to more robust predictions. We empirically demonstrate the superiority of our noise schedule over the standard cosine schedule. Furthermore, our noise schedule shows consistent improvements on ImageNet across different prediction targets.

Chat is not available.