Uncertainty-Aware Scheduling: State-Dependent Training-Free Diffusion Alignment

Yu Cao, Ziquan Liu, Da Li, Shaogang Gong

Keywords: diffusion models, alignment, generative models

Abstract: Training-free guidance has emerged as a promising approach to aligning diffusion models with downstream objectives by steering the denoising process with reward-based signals. Typically, reward functions are trained on clean images then applied to noisy intermediate predictions ($\hat{x}_{0|t}$), suffering from a domain gap that compromises effective guidance. Existing methods improve guidance through handcrafted schedules, such as fixed or time-dependent tempering. However, such state-agnostic scheduling is prone to suboptimal alignment as we discovered that denoising progress varies widely across samples. We propose State-Dependent Adaptive Guidance (SDAG) to schedule guidance through an uncertainty-aware confidence-calibrated assessment. SDAG introduces a lightweight quality predictor that estimates denoising progress from intermediate states, i.e. the closeness between their approximated and the final clean targets. Through a last-layer Laplace approximator, this predictor provides uncertainty estimates, which are used together with the closeness scores to scale guidance reliably. Our SDAG applies to both standard denoising and population-based sampling, such as Sequential Monte Carlo, where coordination by effective sample size ensures robust collective guidance. Experiments demonstrate that SDAG achieves superior alignment while maintaining computational efficiency, establishing a promising paradigm for adaptive guidance in training-free diffusion alignment.

Primary Area: generative models

Submission Number: 9707