ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation
Abstract
Inference-time control steers a pretrained generative model towards a target distribution without retraining. We study tilted targets π0∝ G0 p0, where p0 is the sampler output distribution and G0 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.
Related
- METALICA: METAdynamics and repLICA exchange for enhanced diffusion sampling
- Safe Distributed Generalized Nash Equilibrium Seeking via Control Barrier Functions
- Equilibrium Singular Dividend Control with Regime Switching and Bankruptcy under Non-Exponential Discounting
- Synthetic Data Augmentation for Functional Data: Wasserstein Risk Bounds and Diffusion-Based Conditional Generation
- MGD: Moment Guided Diffusion for Maximum Entropy Generation
BibTeX
@misc{yu2026encore,
title = {{ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation}},
author = {Jiahao Yu and Saifuddin Syed and José Miguel Hernández-Lobato and Jiajun He},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.02538},
eprint = {2610.02538},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.02538},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation
AU - Jiahao Yu
AU - Saifuddin Syed
AU - José Miguel Hernández-Lobato
AU - Jiajun He
PY - 2026
DA - 2026/10/01
PB - arXiv
DO - 10.48550/arxiv.2610.02538
UR - https://arxiv.org/abs/2610.02538
AB - Inference-time control steers a pretrained generative model towards a target distribution without retraining. We study tilted targets π_0∝ 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.
ER -