CREPE: Controlling Diffusion with Replica Exchange

TL;DR

This paper proposes a flexible alternative based on replica exchange, an algorithm designed initially for sampling problems that demonstrates its versatility across various tasks, including temperature annealing, reward-tilting, model composition and classifier-free guidance debiasing, with competitive performance compared to prior SMC methods.

Abstract

Inference-time control of diffusion models aims to steer model outputs to satisfy new constraints without retraining. Previous approaches have mostly relied on heuristic guidance or have been coupled with Sequential Monte Carlo (SMC) for bias correction. In this paper, we propose a flexible alternative based on replica exchange, an algorithm designed initially for sampling problems. We refer to this method as CREPE (Controlling with REPlica Exchange). Unlike SMC, CREPE: (1) generates particles sequentially, (2) maintains high diversity in the generated samples after a burn-in period, and (3) enables online refinement or early termination. We demonstrate its versatility across various tasks, including temperature annealing, reward-tilting, model composition and classifier-free guidance debiasing, with competitive performance compared to prior SMC methods.

Figures (20)

• Figure 1: Trajectory of images generated using CREPE for prompted reward-tilting on ImageNet-512, thinned every 8 iterations. After burn-in, the samples align closely with the prompt.
• Figure 1: Inference-time tempering performance for Alanine Dipeptide, Tetrapeptide and Hexapeptide.
• Figure 2: Comparison between diffusion inference-time control with SMC and CREPE. We visualise the diffusion process using colour shading: darker colours correspond to higher noise/mask levels (large $t$), while brighter colours indicate states closer to the data distribution (small $t$). SMC (Left): particles are initialised at $t=1$ and progressively denoised towards lower noise levels. During denoising, importance resampling is applied to select particles that better satisfy the imposed constraints. CREPE (Right): particles are initialised at several different diffusion steps, and they undergo local exploration and communication in parallel, evolving them towards desired constraints. The example shows using SMC and CREPE to debias classifier-free guidance.
• Figure 2: Debias ImageNet-64 CFG. We do not discard burn-in samples in CREPE to ensure a fair comparison.
• Figure 3: Histogram of Alanine Hexapeptide annealed to 600K by SMC and CREPE, projected to two TICA axes.