CREPE: Controlling Diffusion with Replica Exchange

Jiajun He1*, Paul Jeha2†, Peter Potaptchik3†, Leo Zhang3†, José Miguel Hernández-Lobato1, Yuanqi Du4, Saifuddin Syed5‡ Francisco Vargas6‡
1University of Cambridge, 2Technical University of Denmark, 3University of Oxford, 4Cornell University, 5University of British Columbia, 6Xaira Therapeutics
*Corresponding author †Equal contribution ‡Last authors
ICLR arXiv Code

Overview

Inference-time control of diffusion models aims to steer model outputs to satisfy new constraints without retraining. CREPE is a flexible approach based on replica exchange (parallel tempering): it enables principled control of diffusion models, maintains high diversity of generated samples, and supports online refinement or early termination.

What you can do with CREPE

Reward tilting / prompt alignment

Steer the sample towards a desired reward or prompt.

Inverse problem solving

Solve inverse problems with the diffusion model prior.

Temperature annealing (powering)

Sharpen distributions to a lower temperature level.

Model composition

Compose objectives via product of diffusion experts.

CFG debiasing

Debias the classifier-free guidance (CFG) to achieve better conditioning on the class labels.

Modality-agnostic control

Control models for both continuous and discrete diffusion models within a unified framework.

Key idea: exchange across diffusion times

CREPE runs multiple chains at different diffusion steps (noise levels). Each chain performs local updates, and we periodically propose swaps between neighbouring noise levels with an acceptance rule. This enables principled communication between easy-to-mix (high-noise) and high-fidelity (low-noise) states.

CREPE overview figure
An illustration of the key idea of CREPE and compared it to SMC-based control methods.

Experiments

Inference Time Tempering to sample from Boltzmann distribution of small molecules

We train a diffusion model on samples from a Boltzmann distribution of small molecules at higher temperature. We then use CREPE to temper the model to sample from a different Boltzmann distribution at lower temperature.

Inference time scaling of Boltzmann sampling

Debiasing Classifier-Free Guidance

The standard classifier-free guidance (CFG) in diffusion modelsis biased towards the class labels. We use CREPE to debias the CFG to achieve better conditioning. Compared to the baseline (FKC), CREPE is more stable and can achieve better diversity and performances.

Debiasing Classifier Free Guidance

Prompt Alignment with Reward Tilt

We use CREPE to align the model output with a desired prompt.

IR CLIP scores

Model Product (stitching) for Maze Navigation

We train a diffusion model on short trajectory in a maze. We then use CREPE to control stitching the model to a longer trajectory for more complex navigation.

Task 0

Task 0

Task 1

Task 1

Task 2

Task 2

Task 3

Task 3

Task 4

Task 4

CREPE for Discrete Diffusion Models

CREPE can be applied to discrete diffusion models seamlessly. In this example, we use CREPE to debias the classifier-free guidance (CFG) on a discrete diffusion on MNIST.

CREPE for Discrete Diffusion Models

BibTeX

@article{he2025crepe,
      title={CREPE: Controlling Diffusion with Replica Exchange},
      author={He, Jiajun and Jeha, Paul and Potaptchik, Peter and Zhang, Leo and Hern{\'a}ndez-Lobato, Jos{\'e} Miguel and Du, Yuanqi and Syed, Saifuddin and Vargas, Francisco},
      journal={ICLR},
      year={2026}
    }