Learn to Guide Your Diffusion Model

TL;DR

This work tackles the mismatch between classifier-free guidance (CFG) and the true conditional distribution in diffusion models by learning time- and conditioning-dependent guidance weights ω_{c,(s,t)}. By enforcing consistency-based objectives and training a guidance network, the approach yields better distributional alignment (lower FID) across image datasets and enables reward-guided sampling (e.g., CLIP-based rewards) to improve text-to-image alignment. The results demonstrate consistent gains over unguided and static CFG baselines in image generation, with notable improvements in FID, and variable but promising gains in CLIP-aligned text-to-image tasks. The method offers a scalable, flexible framework to refine conditional diffusion sampling and opens directions for theoretical analysis and more sophisticated reward functions.

Abstract

Classifier-free guidance (CFG) is a widely used technique for improving the perceptual quality of samples from conditional diffusion models. It operates by linearly combining conditional and unconditional score estimates using a guidance weight $ω$. While a large, static weight can markedly improve visual results, this often comes at the cost of poorer distributional alignment. In order to better approximate the target conditional distribution, we instead learn guidance weights $ω_{c,(s,t)}$, which are continuous functions of the conditioning $c$, the time $t$ from which we denoise, and the time $s$ towards which we denoise. We achieve this by minimizing the distributional mismatch between noised samples from the true conditional distribution and samples from the guided diffusion process. We extend our framework to reward guided sampling, enabling the model to target distributions tilted by a reward function $R(x_0,c)$, defined on clean data and a conditioning $c$. We demonstrate the effectiveness of our methodology on low-dimensional toy examples and high-dimensional image settings, where we observe improvements in Fréchet inception distance (FID) for image generation. In text-to-image applications, we observe that employing a reward function given by the CLIP score leads to guidance weights that improve image-prompt alignment.

Figures (10)

• Figure 1: Learned guidance weights on ImageNet 64x64. Left, guidance weights $\omega^\phi_{(t-dt,t)}$ (conditioning-agnostic) for baselines as well as for self-consistency (\ref{['eq:distributional_objective']}) and $\ell_2$ (\ref{['eq:l2_objective']}) objectives, where $dt=1/100$. X-axis is time. Right, guidance weights $\omega^\phi_{c,(t-dt,t)}$ for specific ImageNet classes.
• Figure 2: Ablation over $\delta$ and $S_{\min}$ on ImageNet $64 \times 64$. On the X-axis we report values of $S_{\min}$ and on Y-axis we show FID. Each column denotes a method while a color corresponds to a value of $\delta$.
• Figure 3: Learned guidance weights on MS COCO $512\times512$ trained with self-consistency (\ref{['eq:distributional_objective']}) and CLIP reward loss. Please refer to \ref{['fig:ms-coco-visualization', 'fig:ms-coco-visualization-additional']} for the corresponding images.
• Figure 4: T2I Results on MS-COCO. Part 1. (left-to-right) We provide results of images generated from the given text prompt without CFG, with CFG $\omega = 7.5$, our method with self-consistency loss and our method with self-consistency loss and CLIP score reward.
• Figure 5: T2I Results on MS-COCO. Part 2 (Randomly selected). (left-to-right) We provide results of images generated from the given text prompt without CFG, with CFG $\omega = 7.5$, our method with self-consistency loss and our method with self-consistency loss and CLIP score reward.