Revisiting Noise Schedule Design for Diffusion Training

Tiankai Hang, Shuyang Gu, Xin Geng, Baining Guo

Track: long paper (up to 8 pages)

Keywords: Diffusion Model, Noise Schedule, Efficient Training

TL;DR: This research proposes a new noise schedule for training diffusion models based on importance sampling of log SNR (Signal-to-Noise Ratio) around zero. The approach improves training efficiency compared to standard methods on ImageNet benchmark.

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.

Submission Number: 3