ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation

Jiahao Yu; Saifuddin Syed; José Miguel Hernández-Lobato; Jiajun He

stat.MLOct 1, 2026

ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation

Authors: Jiahao Yu, Saifuddin Syed, José Miguel Hernández-Lobato, Jiajun He

Organizations: University of Oxford · University of Cambridge · University of British Columbia

Abstract

Inference-time control steers a pretrained generative model towards a target distribution without retraining. We study tilted targets $π_0\propto 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.

Figures & tables

Figure 1: One-dimensional output distributions at nominal transition variance with 10,000 samples.

Figure 2: Replica movement (top) and output chains (bottom) for ENCORE and CREPE.

Method	Annealing	Exact	Parallel over	Moves/iter.	# Iter.	Memory
FKD ( Singhal et al., 2025 )	sequential	yes	$M$ particles	$M$	$N$	$O(Md)$
TDS ( Wu et al., 2023a )	sequential	yes	$M$ particles	$M$	$N$	$O(Md)$
RNE ( He et al., 2026a )	sequential	no / yes ∗	$M$ particles	$M$	$N$	$O(Md)$
CREPE ( He et al., 2026b )	parallel	no	$N$ replicas	$N$	$M$	$O(Nd)$
ENCORE (ours)	parallel	yes	$N$ replicas	$N/2$	$M$	$O(N^{2}d)$

Table 1: Monte Carlo corrections for inference-time control with

M

samples and

N

levels. ∗ RNE is exact when taking the same construction as TDS ( He et al., 2026a , App. C.5) .

Figure 4: Alanine transition-region sampling at $800\,\mathrm{K}$ ; each method panel shows 5,000 samples.

Figure 5: 1D densities and sampler-specific targets.

Panel of Table 2 — Table 2: A6 Boltzmann-tempering error using 5,000 samples per evaluation (mean $\pm$ standard deviation over three subsampling evaluations from one sampler run).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: 2D densities for smooth (top) and discontinuous (bottom) rewards.

Figure 10: Conditional $\phi/\psi$ densities (top) and joint $\psi$ /energy densities (bottom) within $|\phi|\leq 0.4$ .

Panel of Table 10 — Table 10: ALDP and A4 Boltzmann-tempering accuracy with 5,000 samples per evaluation (mean $\pm$ standard deviation over three subsampling evaluations from one sampler run).

Explore similar work

Mar 18, 2026

Generative Replica-Exchange: A Flow-based Framework for Accelerating Replica Exchange Simulations

Replica exchange (REX) is one of the most widely used enhanced sampling methodologies, yet its efficiency is limited by the requirement for a large number of intermediate temperature replicas. Here we present Generative Replica Exchange (GREX), which integrates deep generative models into the REX framework to eliminate this temperature ladder. Drawing inspiration from reservoir replica exchange (res-REX), GREX utilizes trained normalizing flows to generate high-temperature configurations on demand and map them directly to the target distribution using the potential energy as a constraint, without requiring target-temperature training data. This approach reduces production simulations to a single replica at the target temperature while maintaining thermodynamic rigor through Metropolis exchange acceptance. We validate GREX on three benchmark systems of increasing complexity, highlighting its superior efficiency and practical applicability for molecular simulations.

Shengjie Huang, Sijie Yang, Jianqiao Yi +6

State Key Laboratory of Bioactive Molecules and Druggability Assessment, International Cooperative Laboratory of Traditional Chinese Medicine Modernization and Innovative Drug Discovery of Chinese Ministry of Education (MOE), School of Pharmacy, Jinan University, #855 Xingye Avenue, Guangzhou, 510632, China · Department of Biochemistry and Molecular Pharmacology, NYU Grossman School of Medicine, New York, NY 10016, USA · Howard Hughes Medical Institute, NYU Grossman School of Medicine, New York, NY 10016, USA +1

Molecular Dynamics Simulation Parallel Tempering Markov Chain Monte Carlo

Ask PDF

Sep 15, 2026

METALICA: METAdynamics and repLICA exchange for enhanced diffusion sampling

Many proteins function through transitions between conformational states, yet rare states are rarely sampled by diffusion models trained on an equilibrium ensemble, demanding better sampling methods. We introduce METALICA, which implements Metadynamics on a pretrained diffusion model via Replica Exchange. It accumulates a bias potential along a Collective Variable, repels new samples from previous ones through biased sampling, and reweights samples onto the unbiased distribution. METALICA holds one replica per diffusion level, forming a Markov Chain that evolves through inter-replica communication and is refined in place as the bias grows. METALICA is the dual of sequential control, in which Sequential Monte Carlo parallelizes the sampler over a batch of particles. Parallelism over the levels of the diffusion-time schedule instead allows METALICA to generate samples from long chains, essential for the discovery of rare events, with accuracy set by run length rather than by the memory available. We validate on a bimodal target with known free energies, then apply METALICA to the unfolding of a protein. At a budget for which sequential control yields no unfolded structure, METALICA populates the basin and resolves a second free energy minimum.

Alireza Omidi, Jiajun He, Jörg Gsponer +1

University of British Columbia · University of Cambridge

Diffusion Sampling

Ask PDF

Aug 20, 2026

Discrete Diffusion Inference-Time Control with Nested Sequential Monte Carlo

We study inference-time control for text generation in discrete diffusion language models, where the goal is to steer sampling toward sequence-level rewards without retraining. Prior work in this domain has focused on particle-based methods such as best-of-

n

sampling and bootstrap sequential Monte Carlo, which may suffer from overoptimism and weight degeneracy, respectively. We address these limitations using nested sequential Monte Carlo methods. We formulate nested SMC (NSMC) and fully-adapted nested SMC (FA-NSMC) for Feynman--Kac steering, identifying and correcting errors in prior formulations that lead to biased final estimates. We evaluate these methods on toxicity and fluency steering tasks, showing that NSMC and FA-NSMC consistently outperform best-of-

n

and bootstrap SMC.

Lohithsai Yadala Chanchu, Hany Abdulsamad, Christian A. Naesseth

University of Amsterdam

Inference-Time Guidance Diffusion Language Model Inference Diffusion Language Models

Ask PDF