Training-Free Stein Diffusion Guidance: Posterior Correction for Sampling Beyond High-Density Regions
This work addresses a bottleneck in diffusion-based sampling for applications like molecular design, offering a principled correction mechanism, though it is incremental as it builds on existing paradigms.
The paper tackles the problem of unreliable guidance in low-density regions for training-free diffusion models by introducing Stein Diffusion Guidance (SDG), which corrects approximate posteriors using Stein variational inference, resulting in improved performance on molecular sampling tasks compared to standard methods.
Training free diffusion guidance provides a flexible way to leverage 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, provides principled posterior simulation but is prohibitively expensive for fast sampling. 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 theoretical bound on the value function, demonstrating the necessity of correcting approximate posteriors to faithfully reflect true diffusion dynamics. Leveraging Stein variational inference, SDG identifies the steepest descent direction that minimizes the Kullback-Leibler divergence between approximate and true posteriors. By incorporating a principled Stein correction mechanism and a novel running cost functional, SDG enables effective guidance in low-density regions. Experiments on molecular low-density sampling tasks suggest that SDG consistently surpasses standard training-free guidance methods, highlighting its potential for broader diffusion-based sampling beyond high-density regions.