Stein Diffusion Guidance: Training-Free Posterior Correction for Sampling Beyond High-Density Regions

Training-free diffusion guidance offers a flexible framework for leveraging off-the-shelf classifiers without additional training. Yet, current approaches hinge on posterior approximations via Tweedie’s formula, which often yield unreliable guidance, particularly in low-density regions. Stochastic optimal control (SOC), in contrast, enables principled posterior sampling but remains computationally prohibitive for efficient inference. In this work, we reconcile the strengths of these paradigms by introducing Stein Diffusion Guidance (SDG), a novel training-free framework grounded in a surrogate SOC objective. We establish a new theoretical bound on the SOC value function, revealing the necessity of correcting approximate posteriors to reflect true diffusion dynamics. Building on Stein variational inference, SDG computes the steepest descent direction that minimizes the Kullback-Leibler divergence between approximate and true posteriors. By integrating a principled Stein correction mechanism along with a novel running cost functional, SDG enables effective guidance in low-density regions. Our experiments on diverse image-guidance tasks and on challenging small-ligand sampling for protein docking suggest that SDG consistently outperforms standard training-free guidance methods and highlights its potential for broader posterior sampling problems beyond high-density regimes.

Training-free diffusion guidance is a way to steer AI-generated images or molecules using existing classifiers, without training a new model. This makes it flexible and easy to apply to many tasks. However, current methods often rely on rough mathematical shortcuts, such as Tweedie’s formula, to estimate the endpoint (posterior) samples of the generation process. These shortcuts can become unreliable, especially in low-density regions, where useful but less common samples may exist. A more principled approach, called stochastic optimal control, can provide better guidance, but it is usually too computationally expensive to use efficiently. In this work, we introduce Stein Diffusion Guidance, or SDG, a new training-free method that combines the practicality of existing guidance methods with ideas from stochastic optimal control. SDG improves guidance by correcting the inaccurate estimates produced by these shortcuts so that they are better aligned with true diffusion dynamics. It uses Stein variational inference, a technique for moving a set of samples toward a target distribution, to guide generation more reliably. We test SDG on several image-generation tasks and on small-molecule sampling for protein docking. Across these tasks, SDG performs better than standard training-free guidance methods. These results suggest that SDG is useful for generating high-quality samples in underexplored, low-density regions and may also be valuable for broader posterior sampling problems.