Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent
It addresses a key issue in diffusion model training for generative AI, offering a novel solution that improves sampling accuracy, though it appears incremental as it builds on existing methods.
The paper tackles the problem of sampling drift in diffusion models caused by imperfect score-matching, proposing a consistency property to train on drifted data, and achieves state-of-the-art results in conditional and unconditional generation on CIFAR-10 with baseline improvements in AFHQ and FFHQ.
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. Yet, 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 which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that if the score is learned perfectly on some non-drifted points (via DSM) and if the consistency property is enforced everywhere, then the score is learned accurately everywhere. Empirically we show that our novel training objective yields state-of-the-art results for conditional and unconditional generation in CIFAR-10 and baseline improvements in AFHQ and FFHQ. We open-source our code and models: https://github.com/giannisdaras/cdm