Distributional Diffusion Models with Scoring Rules

Valentin De Bortoli, Alexandre Galashov, J Swaroop Guntupalli, Guangyao Zhou, Kevin Patrick Murphy, Arthur Gretton, Arnaud Doucet

TL;DR: New losses for diffusion based on scoring rules; Better quality sampling in few steps regime

Abstract: Diffusion models generate high-quality synthetic data. They operate by defining a continuous-time forward process which gradually adds Gaussian noise to data until fully corrupted. The corresponding reverse process progressively ``denoises" a Gaussian sample into a sample from the data distribution. However, generating high-quality outputs requires many discretization steps to obtain a faithful approximation of the reverse process. This is expensive and has motivated the development of many acceleration methods. We propose to speed up sample generation by learning the posterior distribution of clean data samples given their noisy versions, instead of only the mean of this distribution. This allows us to sample from the probability transitions of the reverse process on a coarse time scale, significantly accelerating inference with minimal degradation of the quality of the output. This is accomplished by replacing the standard regression loss used to estimate conditional means with a scoring rule. We validate our method on image and robot trajectory generation, where we consistently outperform standard diffusion models at few discretization steps.

Lay Summary: We propose a new loss for diffusion models based on scoring rules. This new loss allows us to learn a full posterior distribution of clean data given noisy data. Using this learned posterior distribution in DDIM sampler allows us to achieve better quality of samples in few steps regime.

Primary Area: General Machine Learning

Keywords: Diffusion models, energy distance, maximum mean discrepancy, scoring rules, accelerated sampling of diffusion models

Submission Number: 2558