Bridging the Gap: Addressing Discrepancies in Diffusion Model Training for Classifier-Free Guidance

TL;DR

This paper identifies a discrepancy in diffusion model training for classifier-free guidance and proposes a new loss function to improve sample quality, robustness, and efficiency, validated on CIFAR-10 and preliminary results on Stable Diffusion.

Contribution

The paper introduces an improved loss function that aligns training with sampling behavior, enhancing diffusion model performance and robustness.

Findings

Higher quality samples with fewer timesteps

Increased robustness to guidance scale variations

Preliminary benefits observed on large diffusion models

Abstract

Diffusion models have emerged as a pivotal advancement in generative models, setting new standards to the quality of the generated instances. In the current paper we aim to underscore a discrepancy between conventional training methods and the desired conditional sampling behavior of these models. While the prevalent classifier-free guidance technique works well, it's not without flaws. At higher values for the guidance scale parameter $w$, we often get out of distribution samples and mode collapse, whereas at lower values for $w$ we may not get the desired specificity. To address these challenges, we introduce an updated loss function that better aligns training objectives with sampling behaviors. Experimental validation with FID scores on CIFAR-10 elucidates our method's ability to produce higher quality samples with fewer sampling timesteps, and be more robust to the choice of guidance scale $w$. We also experiment with fine-tuning Stable Diffusion on the proposed loss, to provide early evidence that large diffusion models may also benefit from this refined loss function.