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

arXiv: https://arxiv.org/abs/2509.21655 · HTML: https://mlsanity.com/paper/2509.21655/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 54 of its 117 references that are in the map.
- First posted on arXiv 25 Sep 2025.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 54)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models](https://mlsanity.com/paper/2112.10741/) (2021)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [Directly Fine-Tuning Diffusion Models on Differentiable Rewards](https://mlsanity.com/paper/2309.17400/) (2023)
- [Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem](https://mlsanity.com/paper/2206.04119/) (2022)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [A General Framework for Inference-time Scaling and Steering of Diffusion Models](https://mlsanity.com/paper/2501.06848/) (2025)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Manifold Preserving Guided Diffusion](https://mlsanity.com/paper/2311.16424/) (2023)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control](https://mlsanity.com/paper/2402.15194/) (2024)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts](https://mlsanity.com/paper/2503.02819/) (2025)
- [Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo](https://mlsanity.com/paper/2404.17546/) (2024)
- [Optimizing DDPM Sampling with Shortcut Fine-Tuning](https://mlsanity.com/paper/2301.13362/) (2023)
- [Provably Robust Score-Based Diffusion Posterior Sampling for Plug-and-Play Image Reconstruction](https://mlsanity.com/paper/2403.17042/) (2024)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [InverseBench: Benchmarking Plug-and-Play Diffusion Priors for Inverse Problems in Physical Sciences](https://mlsanity.com/paper/2503.11043/) (2025)
- [Gradient Guidance for Diffusion Models: An Optimization Perspective](https://mlsanity.com/paper/2404.14743/) (2024)
- [What does guidance do? A fine-grained analysis in a simple setting](https://mlsanity.com/paper/2409.13074/) (2024)

## Built on by (most cited first, 23 of 23)

- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [Enhanced Diffusion Sampling: Efficient Rare Event Sampling and Free Energy Calculation with Diffusion Models](https://mlsanity.com/paper/2602.16634/) (2026)
- [Robust Inference-Time Steering of Protein Diffusion Models via Embedding Optimization](https://mlsanity.com/paper/2602.05285/) (2026)
- [Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Control](https://mlsanity.com/paper/2604.08580/) (2026)
- [How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance](https://mlsanity.com/paper/2604.27147/) (2026)
- [Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching](https://mlsanity.com/paper/2601.06514/) (2026)
- [Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion](https://mlsanity.com/paper/2605.23346/) (2026)
- [SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate](https://mlsanity.com/paper/2605.18745/) (2026)
- [Compression as Adaptation: Implicit Visual Representation with Diffusion Foundation Models](https://mlsanity.com/paper/2603.07615/) (2026)
- [Refining Compositional Diffusion for Reliable Long-Horizon Planning](https://mlsanity.com/paper/2605.03075/) (2026)
- [Diffusion-Based Posterior Sampling: A Feynman-Kac Analysis of Bias and Stability](https://mlsanity.com/paper/2605.06538/) (2026)
- [Simple Approximation and Derivative Free Inference-Time Scaling for Diffusion Models via Sequential Monte Carlo on Path Measures](https://mlsanity.com/paper/2605.17850/) (2026)
- [Generative Modeling in Protein Design: Neural Representations, Conditional Generation, and Evaluation Standards](https://mlsanity.com/paper/2603.26378/) (2026)
- [NoiseRater: Meta-Learned Noise Valuation for Diffusion Model Training](https://mlsanity.com/paper/2605.08144/) (2026)
- [Catastrophic Compositional Generation: Why Vanilla Diffusion Models Fail to Extrapolate](https://mlsanity.com/paper/2606.23920/) (2026)
- [METALICA: METAdynamics and repLICA exchange for enhanced diffusion sampling](https://mlsanity.com/paper/2609.17823/) (2026)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models](https://mlsanity.com/paper/2608.08770/) (2026)
- [Interacting particle guidance for sampling reward-tilted generative priors](https://mlsanity.com/paper/2609.37227/) (2026)
- [Specificity-Aware Diffusion Steering via Variance-Reduced Sequential Monte Carlo](https://mlsanity.com/paper/2610.00395/) (2026)
- [Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models](https://mlsanity.com/paper/2610.01933/) (2026)
- [Learn Feasibility Once, Optimize All Objectives: Derivative-Free Diffusion Models for Chance-Constrained Programming](https://mlsanity.com/paper/2610.03071/) (2026)
- [Risk-Calibrated Proposal Transport for Finite-Particle Diffusion Steering](https://mlsanity.com/paper/2610.04171/) (2026)

## Abstract (arXiv)

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.