Optimizing Noise Schedules of Generative Models in High Dimensionss

TL;DR

This paper investigates noise schedules in high-dimensional diffusion models, demonstrating that tailored schedules can recover both high- and low-level features efficiently, reducing the number of steps needed for accurate sampling.

Contribution

It introduces specific noise schedules for VP and VE diffusion models that improve feature recovery and computational efficiency in high dimensions.

Findings

VP recovers low-level features but misses high-level features.

VE captures high-level features but not low-level features.

Tailored schedules enable $ ext{O}(1)$ steps in dimension d for sampling.

Abstract

Recent works have shown that diffusion models can undergo phase transitions, the resolution of which is needed for accurately generating samples. This has motivated the use of different noise schedules, the two most common choices being referred to as variance preserving (VP) and variance exploding (VE). Here we revisit these schedules within the framework of stochastic interpolants. Using the Gaussian Mixture (GM) and Curie-Weiss (CW) data distributions as test case models, we first investigate the effect of the variance of the initial noise distribution and show that VP recovers the low-level feature (the distribution of each mode) but misses the high-level feature (the asymmetry between modes), whereas VE performs oppositely. We also show that this dichotomy, which happens when denoising by a constant amount in each step, can be avoided by using noise schedules specific to VP and VE that allow for the recovery of both high- and low-level features. Finally we show that these schedules yield generative models for the GM and CW model whose probability flow ODE can be discretized using $\Theta_d(1)$ steps in dimension $d$ instead of the $\Theta_d(\sqrt{d})$ steps required by constant denoising.