Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent

Giannis Daras, Yuval Dagan, Alex Dimakis, Constantinos Costis Daskalakis

Keywords: diffusion models, sampling drift, Fokker-Planck, invariances, Stochastic Differential Equations, Martingales

TL;DR: We propose a novel training objective that enforces the network to have consistent predictions.

Abstract:

Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. However, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a \emph{Consistency} property (CP) which states that predictions of the model on its own
generated data are consistent across time. Theoretically, we show that the differential equation that describes CP together with the one that describes a conservative vector field, have a unique solution given some initial condition. Consequently, if the score is learned well on non-drifted points via DSM (enforcing the true initial condition) then enforcing CP on drifted points propagates true score values. Empirically, we show that enforcing CP improves the generation quality for conditional and unconditional generation on CIFAR-10, and in AFHQ and FFHQ.
We open-source our code and models: https://github.com/giannisdaras/cdm.

Submission Number: 10552