Score-Optimal Diffusion Schedules

TL;DR

This paper introduces an adaptive algorithm for selecting optimal diffusion schedules in denoising diffusion models, improving sample quality without hyperparameter tuning by evaluating the Stein score.

Contribution

It proposes a novel, scalable method for automatically choosing discretisation schedules based on a derived cost function, enhancing diffusion model performance.

Findings

Recovers high-quality schedules comparable to manual search

Achieves competitive FID scores on image datasets

Does not require hyperparameter tuning

Abstract

Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate a path of probability distributions interpolating between a reference Gaussian distribution and a data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial to obtain high quality samples. However, beyond hand crafted heuristics, a general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to a cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to the dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation of the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that our learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores on image datasets.