DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models

ICLR 2026 Conference Submission23218 Authors

Keywords: Diffusion Model, Inference-Time Scaling, Variance Reduction, Sequential Monte Carlo, Guidance

TL;DR: We introduce DriftLite, a lightweight method that reduces variance in particle dynamics, enabling scalable inference-time adaptation of diffusion models.

Abstract: We study inference-time scaling for diffusion models, where the goal is to adapt a pre-trained model to new target distributions without retraining. Existing guidance-based methods are simple but introduce bias, while particle-based corrections suffer from weight degeneracy and high computational cost. We introduce *DriftLite*, a lightweight, training-free particle-based approach that steers the inference dynamics on-the-fly with provably optimal stability control. DriftLite exploits a previously unexplored degree of freedom in the Fokker-Planck equation between the drift and particle potential, and yields two practical instantiations: *Variance- and Energy-Controlling Guidance (VCG/ECG)* for approximating the optimal drift with minimal overhead. Across Gaussian mixture models, particle systems, and large-scale protein-ligand co-folding problems, DriftLite consistently reduces variance and improves sample quality over pure guidance and sequential Monte Carlo baselines. These results highlight a principled, efficient route toward scalable inference-time adaptation of diffusion models.

Primary Area: generative models

Submission Number: 23218