ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation
stat.ML, cs.LG
Submitted: 2026-10-01
Updated: 2026-10-01
Comments: A shorter version of this work was accepted at the NeurIPS 2026 PriGM Workshop
License: http://creativecommons.org/licenses/by/4.0/
The gist: Inference-time control steers a pretrained generative model towards a target distribution without retraining.
Abstract
Inference-time control steers a pretrained generative model towards a target distribution without retraining. We study tilted targets π 0 proportional to G 0,p 0, where p 0 is the sampler output distribution and G 0 is an evaluable reweighting function. Existing approaches rely on sequential annealing with sequential Monte Carlo (SMC) or parallel annealing with replica exchange (RE). Sequential control is exact but needs large particle populations, whereas no exact parallel control method exists: existing RE corrections approximate an intractable time reversal and are biased. We propose Exact Non-equilibrium COntrol with Replica Exchange (ENCORE), the first exact parallel control method. Each replica stores its generation trajectory, so the upward move is a truncation and the intractable time reversal is never simulated. We prove target invariance and show that the resulting dynamics are those of non-equilibrium replica exchange with the exact time reversal as forward proposal. Under regularity conditions, our diffusion analysis shows that both sequential and parallel control become unstable under refinement of the time discretisation without guidance, whereas guided proposals remain stable and yield diagnostics for tuning the schedule and the computational budget. Across synthetic targets, Boltzmann sampling of biomolecules, and image generation, ENCORE achieves competitive accuracy and diversity, remains robust to sampler perturbations, and applies to distilled samplers where existing RE corrections are unavailable.
Sources
- Classifier-Free Guidance is a Predictor-Corrector
- Are we really tilting? The mechanics of reward guidance in flow and diffusion models
- Rare Event Analysis via Stochastic Optimal Control
- Classifier-Free Diffusion Guidance
- METALICA: METAdynamics and repLICA exchange for enhanced diffusion sampling
- Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis
- Enhanced Diffusion Sampling: Efficient Rare Event Sampling and Free Energy Calculation with Diffusion Models
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Efficient quadratic entropy with distance sketches
- Embedding-Bias in Conditional Independence Testing
- From Log-Odds to Shapley Values: An Explanatory Geometry for the Weighted Naive Bayes Classifier