Stein Diffusion Guidance: Training-Free Posterior Correction for Sampling Beyond High-Density Regions
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Molecular Generation | 5ht1b | Novel Hit Ratio22.69 | 21 | |
| Molecular Generation | fa7 | Novel Hit Ratio1.156 | 21 | |
| Molecular Generation | jak2 | Novel Hit Ratio9.167 | 21 | |
| Molecular Sampling | parp1 | Novel Hit Ratio8.78 | 9 | |
| Novel docking score optimization | 5ht1b | Top 5% DS7.37 | 7 | |
| Novel docking score optimization | jak2 | Top 5% Docking Score10.178 | 7 | |
| Novel docking score optimization | parp1 | Top 5% DS9.583 | 7 | |
| Novel docking score optimization | fa7 | Top 5% Docking Score7.794 | 7 | |
| Label Guidance | Label Guidance Evaluation Set | Accuracy54 | 6 | |
| Super-Resolution | Super Resolution Evaluation Set | LPIPS0.228 | 6 |