CREPE: Controlling Diffusion with Replica Exchange

2025-09-27

AI Summary

This work addresses the challenge of dynamically steering diffusion models during inference—without retraining—to satisfy diverse constraints, including temperature annealing, reward biasing, model ensembling, and classifier-free guidance debiasing. To this end, we propose CREPE: a replica-exchange-based sequential sampling framework that enables online optimization and early termination. Unlike conventional sequential Monte Carlo (SMC) methods, CREPE requires no architectural modifications or changes to the training objective, yet achieves high sample diversity, strong adaptability to heterogeneous constraints, and flexible computational trade-offs. Empirical evaluations demonstrate that CREPE matches or surpasses SMC in multiple constraint-guided generation tasks, while significantly improving controllability and deployment efficiency of diffusion models in practical applications.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Cognitive Modeling & Cognitive Systems: Computational Creativity

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 the 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.

Problem

• Steering diffusion model outputs to satisfy constraints without retraining
• Providing flexible alternative to heuristic guidance and Sequential Monte Carlo
• Enabling versatile control across annealing, reward-tuning, and debiasing tasks

Innovation

• Uses replica exchange for diffusion model control
• Generates particles sequentially with high diversity
• Enables online refinement and early termination