Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models
This work addresses a significant bottleneck in diffusion models for machine learning practitioners and researchers, providing an incremental yet impactful improvement.
The authors tackled the problem of long sampling times in diffusion models, achieving improved generation performance by disentangling total-variance and signal-to-noise-ratio, with significant gains across various applications. Their approach reveals that constant total-variance schedules can often outperform exploding ones.
The long sampling time of diffusion models remains a significant bottleneck, which can be mitigated by reducing the number of diffusion time steps. However, the quality of samples with fewer steps is highly dependent on the noise schedule, i.e., the specific manner in which noise is introduced and the signal is reduced at each step. Although prior work has improved upon the original variance-preserving and variance-exploding schedules, these approaches $\textit{passively}$ adjust the total variance, without direct control over it. In this work, we propose a novel total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework, where TV and SNR can be controlled independently. Our approach reveals that schedules where the TV explodes exponentially can often be improved by adopting a constant TV schedule while preserving the same SNR schedule. Furthermore, generalizing the SNR schedule of the optimal transport flow matching significantly improves the generation performance. Our findings hold across various reverse diffusion solvers and diverse applications, including molecular structure and image generation.