ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation
Jiahao Yu ⋅ Saifuddin Syed ⋅ José Miguel Hernández-Lobato ⋅ Jiajun He
Abstract
Inference-time control steers a pretrained generative model towards a new target without retraining. We study tilted targets $q_0\propto g\,p_0$, where $p_0$ is the sampler output distribution and $g$ is an evaluable reweighting function. Existing approaches typically rely on sequential Monte Carlo (SMC) or replica exchange (RE; also known as parallel tempering). However, SMC can suffer from severe weight degeneracy when the particle population is insufficient. While existing RE constructions can yield higher sample diversity, they rely on approximate intermediate densities, potentially introducing approximation error. To address this error, we propose \textit{\textbf{E}xact \textbf{N}on-equilibrium \textbf{CO}ntrol with \textbf{R}eplica \textbf{E}xchange} (ENCORE). Instead of applying RE to marginals, ENCORE formulates RE directly in trajectory space, eliminating the need to approximate intermediate marginals. We evaluate ENCORE on synthetic targets, Boltzmann distributions of small biomolecules, and image-generation tasks. Across these settings, ENCORE achieves low distributional error while preserving the strong diversity properties of replica exchange.
Chat is not available.
Successful Page Load