🤖 AI Summary
This study addresses the degradation in generation quality and diversity caused by high-strength classifier-free guidance (CFG) in diffusion models, which violates sampling dynamics principles. To resolve this, the work reformulates CFG as a continuous-time optimal control problem, establishing for the first time its theoretical connection to the Hamilton-Jacobi-Bellman (HJB) equation. Building upon this formulation, the authors propose curvature-aware CFG, which introduces hyperspherical control constraints based on Gaussian regularization to correct unconstrained biases and optimize generative trajectories. By aligning the guidance mechanism with principled optimal control theory, this approach significantly enhances image generation quality under moderate-to-high guidance strengths. Ultimately, the proposed method achieves a superior trade-off between sample fidelity and diversity compared to standard CFG, offering a theoretically grounded solution for controllable diffusion-based synthesis.
📝 Abstract
Diffusion models generate samples by learning to reverse a fixed corruption process, and classifier-free guidance (CFG) is the standard mechanism for conditioning this process on a desired class or prompt. CFG can be applied at varying guidance strengths, and while higher strengths improve image quality and conditional alignment, too high a guidance strength can degrade image quality and diversity. Furthermore, CFG violates principled diffusion sampling dynamics, and existing explanations for why it works despite the violation disagree on the underlying theory or do not extend to deterministic samplers used in practice. We address both these issues. We first frame CFG sampling as a continuous-time optimal control problem, treating the sampling trajectory as a sequence of controls chosen to maximise the probability of the desired condition. Solving the resulting Hamilton--Jacobi--Bellman equation shows that CFG is recovered under specific path costs when using an unconstrained control set. We argue this lack of constraint is responsible for CFG's failure at high guidance strengths, since it permits the sampling path to move arbitrarily far from the current image estimate. To fix this, we propose curvature-aware CFG (CACFG), which constrains the control set to a hypersphere informed by the Gaussian regularisation used when training variational autoencoders. We show that the control inputs produced by CFG sampling routinely violate this bound, and that across diffusion models, datasets, and guidance schedules, CACFG achieves superior generative quality at mid-to-high guidance strengths with a less severe quality-diversity tradeoff than regular CFG.