Adaptive Diffusion Guidance via Stochastic Optimal Control

TL;DR

This work addresses the theoretical gaps in diffusion-model guidance by formalizing how guidance strength $w$ affects classifier confidence and by proving that positive guidance keeps samples within the conditional data manifold. Building on these insights, it introduces a stochastic optimal control (SOC) framework that dynamically selects $w_t(x,c)$ during sampling, supported by a gradient-based learning algorithm that uses Girsanov reweighting to handle trajectory-distribution changes. The SOC objective balances improving class alignment with staying close to unguided trajectories, yielding a learnable policy $w_\theta$ that adapts across time, state, and conditioning class. Empirical results on toy settings demonstrate improved alignment and sample quality for learned schedules, while the authors acknowledge gaps in applying the method to real-world images and outline directions for robustification and scalable implementation.

Abstract

Guidance is a cornerstone of modern diffusion models, playing a pivotal role in conditional generation and enhancing the quality of unconditional samples. However, current approaches to guidance scheduling--determining the appropriate guidance weight--are largely heuristic and lack a solid theoretical foundation. This work addresses these limitations on two fronts. First, we provide a theoretical formalization that precisely characterizes the relationship between guidance strength and classifier confidence. Second, building on this insight, we introduce a stochastic optimal control framework that casts guidance scheduling as an adaptive optimization problem. In this formulation, guidance strength is not fixed but dynamically selected based on time, the current sample, and the conditioning class, either independently or in combination. By solving the resulting control problem, we establish a principled foundation for more effective guidance in diffusion models.

Figures (2)

• Figure 1: Optimal guidance $w^*$ obtained by solving the HJB equation.
• Figure :

Theorems & Definitions (12)

• lemma 1
• corollary 1
• theorem 1
• remark 1
• lemma 1
• proof
• corollary 1
• proof
• lemma 2
• proof