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. Our source code is publicly available at https://github.com/yinuoren/DriftLite.
翻译:本研究探讨扩散模型的推理时缩放问题,其目标是在无需重新训练的情况下,使预训练模型适应新的目标分布。现有的基于引导的方法虽然简单但会引入偏差,而基于粒子的校正方法则存在权重退化与计算成本高昂的问题。我们提出DriftLite,一种轻量级、无需训练的基于粒子的方法,它通过可证明最优的稳定性控制实时引导推理动态。DriftLite利用了福克-普朗克方程中漂移项与粒子势之间一个此前未被探索的自由度,并衍生出两种实用实现:方差控制引导与能量控制引导,以最小开销近似最优漂移。在高斯混合模型、粒子系统以及大规模蛋白质-配体共折叠问题上的实验表明,相较于纯引导方法与序列蒙特卡洛基线,DriftLite能持续降低方差并提升样本质量。这些结果揭示了一条实现扩散模型可扩展推理时自适应的原理清晰且高效的路径。我们的源代码已公开于https://github.com/yinuoren/DriftLite。