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

2025-09-25

AI Summary

This work addresses the distribution misalignment problem during diffusion model inference. Existing guidance methods introduce bias, while particle correction approaches suffer from weight degeneracy and high computational overhead. To overcome these limitations, we propose DriftLite—a lightweight, retraining-free method that leverages the degrees of freedom in the drift term and particle potential within the Fokker–Planck equation. We instantiate DriftLite via two strategies—Variational Control Guidance (VCG) and Energy-Controlled Guidance (ECG)—both theoretically guaranteed to achieve optimal stability. By online modulating the drift term during diffusion, jointly guided by variance and energy terms, DriftLite enables low-cost, dynamic particle system correction. Experiments on Gaussian mixture modeling, synthetic particle systems, and protein–ligand co-folding demonstrate significantly reduced sampling variance and superior generation quality compared to conventional guidance and sequential Monte Carlo methods.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Adapting pre-trained diffusion models to new distributions without retraining
• Addressing bias in guidance methods and computational costs in particle corrections
• Developing lightweight training-free control for diffusion model inference dynamics

Innovation

• Lightweight particle-based drift control method
• Optimal stability control via Fokker-Planck equation
• Variance- and Energy-Controlling Guidance implementations