This is the official implementation of the paper "PID-controlled Langevin Dynamics for Faster Sampling of Generative Models". Each directory corresponds to a task. Please refer to the README.md file in each directory for more implementation details.
Paper · NeurIPS · Code · WebPage
- Control-theoretic insight: Reinterprets Langevin dynamics as a feedback control system, where energy gradients act as feedback signals.
- PID-enhanced sampling: Integrates Proportional, Integral, and Derivative control terms into Langevin updates:
- P-term: basic gradient guidance;
- I-term: accumulates historical gradients for momentum-like acceleration;
- D-term: anticipates gradient trends for adaptive stabilization.
- Plug-and-play compatibility: Requires no retraining or prior information; integrates with any Langevin-based sampler directly (EBM, SGM, etc.).
- Significant speedup: Achieves up to 10× faster sampling while maintaining or improving generation quality across image and reasoning tasks.
The PID-controlled Langevin dynamics update is given by
where
PIDLD Algorithm Flowchart
-
Require: Score function
$\nabla_x U_\theta(x)=\nabla_x\log p_\theta(x)=\nabla_x(-f_\theta(x))$ ; number of steps$T$ ; step size$\epsilon$ ; control parameters$k_p,k_i,k_d$ ; decay rate$\gamma<1$ ; initial point$x_0$ . - Initialize integral term
$I_0 = 0$ . - Compute initial score
$s_0 = \nabla_x U_\theta(x_0)$ . - For
$t = 0$ to$T-1$ do: $s_t = \nabla_x U_\theta(x_t)$ -
$P_t = s_t$ (Proportional term) -
$I_t = \dfrac{1}{t+1}\big(I_{t-1}\cdot t + s_t\big)$ (Integral term) -
$D_t = s_t - s_{t-1}$ (Derivative term) -
$u_t = k_p P_t + k_i I_t + k_d D_t$ (Control signal) - State update:
$$x_{t+1} = x_t + \epsilon \cdot u_t + \sqrt{2\epsilon}\xi_t,\quad \xi_t \sim \mathcal{N}(0,I)$$ - Decay integral gain:
$k_i = k_i \cdot \gamma$ - End for
-
Return:
$\hat{x} = x_T$
We evaluate PIDLD against standard Langevin-based samplers (vanilla ALD and MILD) across three regimes: toy 2‑D examples, image generation, and reasoning (solution sampling). The focus is on sampling quality versus computational budget (NFE).
-
Toy experiments
- Purpose: validate the roles of P/I/D terms on simple multimodal landscapes.
- Findings: both I and D terms accelerate convergence and reduce KL/divergence; in particular, D term improves stability, and I term reduces steady-state bias.
-
Image generation (CIFAR10, CelebA)
- Setup: apply PIDLD as a plug‑in to pretrained score-based models (NCSNv2) and energy models (IGEBM), vary NFE and tune PID gains (with decaying k_i); compute FID on 10k samples.
- Goal: test whether PIDLD can reach baseline or better image quality at substantially lower NFEs.
- Results: PIDLD consistently matches or outperforms baselines at much lower NFEs, demonstrating clear efficiency gains.
-
Reasoning (Sudoku, Connectivity)
- Setup: use an energy-based solver (IRED) and evaluate solution accuracy under different NFEs.
- Goal: assess whether PIDLD helps navigate complex energy landscapes to find valid solutions faster.
- Results: PIDLD yields higher accuracy with lower NFEs, showing overall computational advantages versus vanilla Langevin dynamics sampling.
For more details, please refer to the paper and the code.
If you find the idea useful for your research, please consider citing:
@inproceedings{chen2025pidcontrolled,
title={{PID}-controlled Langevin Dynamics for Faster Sampling on Generative Models},
author={Hongyi Chen and Jianhai Shu and Jingtao Ding and Yong Li and Xiao-Ping Zhang},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
url={https://openreview.net/forum?id=y9LHDCKeeN},
}