Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models

TL;DR

This paper introduces a TV/SNR disentangled framework for diffusion models, enabling independent control of total variance and SNR, leading to improved sampling efficiency and quality across applications.

Contribution

It proposes a novel TV/SNR framework that allows independent control of variance and SNR, and demonstrates improved diffusion schedules for better performance.

Findings

Constant TV schedules can outperform exponentially exploding TV schedules.

Generalized SNR schedules enhance generation performance.

Framework is effective across various diffusion applications.

Abstract

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.