Rectified Diffusion Guidance for Conditional Generation

TL;DR

Classifier-Free Guidance (CFG) boosts conditional sampling in diffusion-based generation but is theoretically incompatible with diffusion reciprocity, leading to an expectation shift in the guided distribution. The authors propose Rectified Classifier-Free Guidance (ReCFG), which relaxes the two-score combination into a two-parameter form with a closed-form per-pixel solution derived from a data-driven lookup, thereby annihilating the expectation shift while preserving inference speed and DDIM compatibility. Theoretical analysis clarifies the source of the shift and shows how a zero-mean condition can be enforced via $\gamma_0=(1-\gamma_1)\mathbb E_{\mathbf x_t}[\boldsymbol\epsilon_\theta(\mathbf x_t,c,t)]/\mathbb E_{\mathbf x_t}[\boldsymbol\epsilon_\theta(\mathbf x_t,t)]$, with pixel-wise coefficients obtained from precomputed lookups. Empirically, ReCFG improves conditional fidelity and sample quality across class- and text-conditioned diffusion models (e.g., on ImageNet and CC12M) without retraining, demonstrating practical utility for open-vocabulary conditional generation.

Abstract

Classifier-Free Guidance (CFG), which combines the conditional and unconditional score functions with two coefficients summing to one, serves as a practical technique for diffusion model sampling. Theoretically, however, denoising with CFG \textit{cannot} be expressed as a reciprocal diffusion process, which may consequently leave some hidden risks during use. In this work, we revisit the theory behind CFG and rigorously confirm that the improper configuration of the combination coefficients (\textit{i.e.}, the widely used summing-to-one version) brings about expectation shift of the generative distribution. To rectify this issue, we propose ReCFG with a relaxation on the guidance coefficients such that denoising with \method strictly aligns with the diffusion theory. We further show that our approach enjoys a \textbf{\textit{closed-form}} solution given the guidance strength. That way, the rectified coefficients can be readily pre-computed via traversing the observed data, leaving the sampling speed barely affected. Empirical evidence on real-world data demonstrate the compatibility of our post-hoc design with existing state-of-the-art diffusion models, including both class-conditioned ones (\textit{e.g.}, EDM2 on ImageNet) and text-conditioned ones (\textit{e.g.}, SD3 on CC12M), without any retraining. Code is available at \href{https://github.com/thuxmf/recfg}{https://github.com/thuxmf/recfg}.

Figures (3)

• Figure 1: Visualization of expectation shift. The demonstrated toy data is simulated by $q_0(\mathbf x_0|c)\sim\mathcal{N}(c,1)$, $q(c)\sim\mathcal{N}(0,1)$, $q_0(\mathbf x_0)\sim\mathcal{N}(0,2)$. Gamma-powered distribution $q_{0,\gamma}(\mathbf x_0|c)$ from CFG ho2022classifierfree fails to recover the same conditional expectation as ground-truth due to expectation shift (i.e., probability density function and histogram by DDIM song2020denoising sampler in red). To make a further step, larger $\gamma$ suggests more severe expectation shift, i.e., the peak of $q_{0,\gamma}(\mathbf x_0|c)$ tends further away from $q_0(\mathbf x_0|c)$ (i.e., probability density function in blue) as $\gamma$ goes from $1.5$ to $2.5$. As a comparison, our ReCFG successfully recovers the ground-truth expectation and smaller variance (i.e., probability density function and histogram by DDIM song2020denoising sampler in green), consistent with the motivation of guided sampling.
• Figure 2: Visualization of the lookup table on LDM rombach2022high, EDM2 Karras2024edm2, and SD3 esser2024sd3, each of which consists of the expectation ratio $\mathbb E_{\mathbf x_t}[\boldsymbol\epsilon_\theta(\mathbf x_t,c,t)]/\mathbb E_{\mathbf x_t}[\boldsymbol\epsilon_\theta(\mathbf x_t,t)]$. Each pixel represents the scale of the pixel-wise ratio, i.e., color red implies that ratio is greater than one, while color blue stands for ratio smaller than one. The darker the color is, the farther the ratio appears away from one. We report in each row the expectation ratios on five timesteps uniformly sampled from the whole trajectory, under different DPMs and NFEs. It is noteworthy that expectation ratios at the same timestep vary largely by different pixels, and there is no general pattern along with timesteps or pixels.
• Figure 3: Qualitative comparison on EDM2 and SD3. Left and right in each cell suggest samples via CFG and ReCFG, respectively.

Theorems & Definitions (15)

• Theorem 1
• Remark 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Lemma 1
• proof : Proof of \ref{['lem:1']}
• Lemma 2
• proof : Proof of \ref{['lem:2']}