Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yinuo Ren, Wenhao Gao, Lexing Ying, Grant M. Rotskoff, Jiequn Han

TL;DR

DriftLite is introduced, a lightweight, training-free particle-based approach that steers the inference dynamics on the fly with provably optimal stability control, and yields two practical instantiations: Variance- and Energy-Controlling Guidance (VCG/ECG) for approximating the optimal drift with minimal overhead.

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. Our source code is publicly available at https://github.com/yinuoren/DriftLite.

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

TL;DR

DriftLite is introduced, a lightweight, training-free particle-based approach that steers the inference dynamics on the fly with provably optimal stability control, and yields two practical instantiations: Variance- and Energy-Controlling Guidance (VCG/ECG) for approximating the optimal drift with minimal overhead.

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. Our source code is publicly available at https://github.com/yinuoren/DriftLite.

Paper Structure

This paper contains 72 sections, 7 theorems, 91 equations, 18 figures, 16 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Our Contributions.
  3. Preliminaries
  4. Diffusion Models
  5. Inference-Time Scaling
  6. Distribution Path Selection.
  7. Guidance-Based Dynamics.
  8. Weighted Particle Method.
  9. Method: Lightweight Drift Control
  10. Degree of Freedom in the Fokker-Planck Equation
  11. In Search of Optimal Control
  12. Variance-Controlling Guidance (VCG).
  13. Energy-Controlling Guidance (ECG).
  14. Practical Implementation
  15. Choice of Bases.
  16. ...and 57 more sections

Key Result

Proposition 2.1

[proposition]prop:afdps_pde The exact time evolution of the density $(q_t)_{t\in[0, T]}$ follows the following Feynman-Kac-type Fokker-Planck equation: where $\widetilde{{\bm{v}}}_t$ is the same drift as in pure guidance eq:modified_drift_afdps, and the reweighting potential $g_t({\bm{x}}) = G_t({\bm{x}}) - \mathbb{E}_{q_t} [G_t(\cdot)]$ is given by:

Figures (18)

  • Figure 1: Qualitative comparison of sampling methods on the GMM annealing task ($\gamma=2.5$).
  • Figure 2: Evolution of ESS and potential variance during inference on the GMM annealing task ($\gamma=2.2$). Our methods (VCG/ECG) substantially reduce variance and stabilize ESS.
  • Figure 3: Performance metrics versus number of particles for the GMM annealing task ($\gamma=2.0$). Our methods consistently outperform baselines and show strong scaling.
  • Figure 4: Qualitative comparison of sampling methods on the GMM reward-tilting task ($\sigma=200.0$).
  • Figure 5: Comparison of generated distributions for the LJ-13 annealing task ($\gamma=2.5$). VCG-SMC is the only method that successfully recovers all three peaks in the (a) RDF and closely matches the (b) Energy Distribution. Insets provide a zoomed-in view.
  • ...and 13 more figures

Theorems & Definitions (13)

  • Proposition 2.1: Guidance-Based Dynamics
  • Proposition 3.1: Degree of Freedom
  • proof : Proof Sketch
  • Proposition 3.2: Optimal Control, Informal Version
  • Proposition A.1: Guidance-Based Dynamics
  • proof : Proof
  • Remark A.2
  • Proposition A.3: Weighted Particle Simulation
  • proof : Proof Sketch
  • Proposition A.4: Degree of Freedom
  • ...and 3 more