A Diffusion Model Induced by MSE Training

Piotr Pokarowski

Keywords: diffusion models, score-based generative modeling, generative models, deep learning

TL;DR: MSE-Induced Diffusion

Abstract:

In the DDPM paper, Ho, Jain, and Abbeel introduced two reversible diffusion processes parameterized by a noise schedule—a generator and an oracle process that the generator learns from—and derived a formula for the Kullback-Leibler divergence (KL) in the form of a time-weighted Mean Squared Error (MSE). However, they empirically found that omitting the weights improved performance on image-synthesis benchmarks, a result later corroborated by many studies. More recently, removing the stochastic component at generation time has proved effective. (1) In this work, we provide a theoretical justification for these practices. We consider a broader class of diffusion processes (not necessarily reversible) parameterized by a noise schedule and a diffusion size b that share the same marginals. Since the weight associated with the MSE depends on b, omitting the weight is equivalent to solving the equation weight(b)=1, which yields a unique "MSE-diffusion". For SOTA models, we checked that b is close to zero; that is, the learned MSE-diffusion is nearly a flow, and we confirm this observation by comparing generators on ImageNet 512×512.
Therefore, flows beat reversible diffusions because training of SOTA models is an implementation of KL minimization for MSE-diffusions, which are nearly flows. The models that succeed are the ones that are really trained. (2) Moreover, by generalizing the diffusion process to both discrete and continuous time, we obtained a novel representation of the diffusion state as the sum of an explicit linear component, an unweighted pathwise integral of the denoiser, and a noise term. This representation offers the advantages of DPM-solvers while enabling the use of classical numerical methods for ODEs.

Primary Area: generative models

Submission Number: 20837