PID-controlled Langevin Dynamics for Faster Sampling of Generative Models

Conference: NeurIPS 2025 arXiv: 2511.12603 Code: GitHub Area: Diffusion Models / Sampling Acceleration Keywords: Langevin Dynamics, PID Control, Sampling Acceleration, Energy-Based Models, Training-Free

TL;DR

This work introduces PID control theory into Langevin dynamics sampling, leveraging gradient history (integral term) to build momentum for traversing energy barriers and gradient trends (derivative term) to suppress oscillations, achieving fast and stable convergence. The approach requires no additional training and delivers over 10× sampling acceleration on both SGMs and EBMs.

Background & Motivation

Langevin dynamics is the core sampling method in EBMs (energy-based models) and SGMs (score-based generative models), but suffers from a fundamental speed bottleneck—requiring a large number of fine-grained iterations to converge to the target distribution. For instance, NCSNv2 requires 1000+ neural function evaluations (NFEs).

The physical root cause lies in the complex high-dimensional energy landscape: particles frequently encounter near-zero gradient regions (local minima or unstable equilibria) and must rely on stochastic noise to escape, which is extremely inefficient. Directly increasing step size introduces larger noise fluctuations, severely degrading sample quality.

The authors' key insight comes from control theory: standard Langevin dynamics is equivalent to a simple proportional (P) feedback control system that exploits only the current gradient. Control theory has long established that adding integral (I) and derivative (D) terms can substantially improve the response speed and stability of control systems. This analogy opens an entirely new perspective on sampling acceleration.

Method

Overall Architecture

Langevin sampling is recast as a feedback control problem—the energy gradient \(\nabla_x U_\theta(x)\) serves as the feedback signal and the particle as the controlled system. An integral term (gradient history) and a derivative term (gradient trend) are introduced on top of the standard proportional term (current gradient), forming a complete PID controller that drives the sampling process.

Key Designs

1. PID-Controlled Langevin Dynamics: The core update rule is: \(x_{t+1} = x_t + \epsilon\left(k_p \nabla_x U_\theta(x_t) + \frac{k_i}{t}\sum_{s=0}^{t}\nabla_x U_\theta(x_s) + k_d(\nabla_x U_\theta(x_t) - \nabla_x U_\theta(x_{t-1}))\right) + \sqrt{2\epsilon}\xi_t\)

   • P term (\(k_p\)): Standard gradient guidance, responding to the current error.
   • I term (\(k_i/t \cdot \sum\)): Accumulated gradient history creates a momentum effect, helping particles traverse energy barriers and unstable equilibria. The \(1/t\) normalization prevents the integral term from dominating over time.
   • D term (\(k_d\)): Captures the gradient rate of change, accelerating motion in directions of consistently decreasing gradient and suppressing motion where gradient consistently increases, thereby reducing overshoot.

2. Exponentially Decaying Integral Gain Schedule: The integral gain decays over time as \(k_i(t) = \gamma^t \cdot k_i\) (\(\gamma < 1\)), inspired by gain scheduling in control theory. High integral gain is needed early to build momentum for barrier traversal, while low gain is needed later to avoid instability. The decay ensures that sampling eventually reverts to standard Langevin dynamics, preserving theoretical convergence guarantees.

3. Seamless Integration with Annealed Langevin Dynamics: PIDLD directly replaces the standard Langevin steps within each noise scale \(\sigma_i\). A key innovation is the continuous state transfer across noise scales—the final \(I_t\) and gradient (for the D term) from the previous scale serve as the initial state for the next scale \(\sigma_{i-1}\), ensuring that historical information is propagated throughout the entire sampling process.

Convergence Guarantee

The paper provides a theoretical guarantee that the derivative term preserves convergence: under a locally strongly convex energy landscape, when the step size satisfies \(\epsilon < \frac{1}{(1+2k_d)m}\) (where \(m\) is the strong convexity parameter), the system is asymptotically stable and converges to a unique stationary distribution. This establishes the D term as a valid stabilizing component that does not compromise the sampling process.

Key Experimental Results

Main Results: Image Generation FID Comparison

Dataset	Method	NFE=25 (SGM)	NFE=100 (SGM)	NFE=232×5 (SGM)	NFE=20 (EBM)	NFE=40 (EBM)
CIFAR10	Vanilla ALD	46.8	17.2	12.5	58.1	35.3
	MILD	-	-	13.0	49.9	34.4
	PIDLD	18.3	12.1	11.4	46.1	33.2
CelebA	Vanilla ALD	25.0 (50)	13.6 (250)	9.5 (500×5)	63.5 (20)	35.4 (30)
	MILD	-	-	11.0	41.1	32.9
	PIDLD	8.0 (50)	5.7 (250)	5.6 (500×5)	38.9 (20)	30.0 (30)

Reasoning Tasks: Sudoku and Connectivity Accuracy

Task	Method	NFE=5	NFE=10	NFE=30
Sudoku	Vanilla	45.99%	51.00%	50.77%
	MILD	49.75%	54.82%	55.25%
	PIDLD	50.54%	55.48%	55.94%
Connectivity	Vanilla	86.16% (1)	87.22% (2)	87.49% (5)
	MILD	86.16% (1)	88.54% (2)	90.15% (5)
	PIDLD	86.16% (1)	91.32% (2)	92.95% (5)

Ablation Study

Configuration	CIFAR10 FID (NFE=25)	CelebA FID (NFE=50)	Notes
P only (baseline)	46.8	25.0	Standard ALD
P + I	~30	~15	I term aids barrier traversal
P + D	~22	~10	D term is the primary contributor in image generation
P + I + D (PIDLD)	18.3	8.0	Complementary; full model is optimal

Key Findings

• Over 10× acceleration: On SGMs, PIDLD reaches the baseline's best performance (requiring 1000+ NFEs) with only 100 NFEs; the improvement on CelebA is even more pronounced (38.3% FID reduction).
• D term dominates in image generation: The annealed Langevin energy landscape is smoothed by noise in early stages, making barriers shallow and limiting the I term's traversal benefit; the D term facilitates rapid convergence to local well centers at each noise scale.
• I term dominates in reasoning tasks: Reasoning tasks require finding global energy minima, and the momentum effect of the I term yields continuously growing advantages at higher NFEs (on the Connectivity task, PIDLD with 2 NFEs surpasses the vanilla baseline at 10 NFEs).
• Advantage over MILD: PIDLD outperforms the momentum-only method MILD across all configurations with greater stability, benefiting from the comprehensive feedback mechanism of PID control.

Highlights & Insights

• An elegant intersection of control theory and generative modeling: Recasting sampling as a control problem yields a natural and verifiable interpretation for the P, I, and D terms.
• Training-free plug-and-play: No additional data, prior information, or retraining is required; the method directly replaces the sampler.
• Complementarity of I and D terms: The two terms respectively play the dominant role in reasoning and generation tasks, indicating that the full PID controller provides a task-adaptive optimal balance.
• Gain decay ensures theoretical convergence: The \(\gamma^t\) decay causes the system to eventually reduce to standard LD, reconciling early-stage acceleration with late-stage theoretical guarantees.

Limitations & Future Work

• The method primarily targets models that use Langevin sampling (EBMs, SGMs) and does not directly apply to ODE-based samplers such as DDPM/DDIM.
• The \(1/t\) normalization of the I term and the exponential decay parameter \(\gamma\) introduce additional hyperparameters that require tuning.
• Theoretical convergence is only proven under a locally strongly convex setting; global convergence guarantees for multimodal distributions remain absent.
• No direct comparison with ODE-based samplers such as DDIM is provided (the authors note these are not comparable methods).

Related Work & Insights

• CLD (Critically Damped Langevin Diffusion) accelerates sampling via HMC but requires additional learning of diffusion velocity; PIDLD is training-free.
• Matrix preconditioning methods rely on statistics of the target data, limiting generalizability; PIDLD requires no prior information.
• MILD is the most directly comparable baseline, employing only momentum (equivalent to the I term alone); the D term in PIDLD provides additional stabilization.
• PID control has been previously applied to deep learning optimizers (e.g., accelerated gradient descent); this work represents its first application to generative model sampling.

Rating

• Novelty: ⭐⭐⭐⭐⭐ The control-theoretic perspective is entirely novel, with clearly verifiable roles for each P/I/D term.
• Experimental Thoroughness: ⭐⭐⭐⭐ Covers SGM, EBM, and reasoning tasks, though evaluation is limited to low-resolution images.
• Writing Quality: ⭐⭐⭐⭐ Theoretical derivations and experimental analyses are well balanced, with effective use of toy experiments for motivation.
• Value: ⭐⭐⭐⭐ Offers significant improvements for Langevin sampling, though applicability is limited to LD-based samplers.

Related Papers

• [NeurIPS 2025] Preconditioned Langevin Dynamics with Score-Based Generative Models for Infinite-Dimensional Linear Bayesian Inverse Problems
• [NeurIPS 2025] Cross-fluctuation Phase Transitions Reveal Sampling Dynamics in Diffusion Models
• [NeurIPS 2025] Physics-Constrained Flow Matching: Sampling Generative Models with Hard Constraints
• [NeurIPS 2025] Understanding Representation Dynamics of Diffusion Models via Low-Dimensional Models
• [NeurIPS 2025] Elucidated Rolling Diffusion Models for Probabilistic Forecasting of Complex Dynamics