Unified Control for Inference-Time Guidance of Denoising Diffusion Models
This work addresses the challenge of improving task-specific performance in denoising diffusion models, representing an incremental advancement by integrating existing paradigms.
The authors tackled the problem of aligning diffusion model outputs with downstream objectives by proposing UniCoDe, a unified algorithm that combines sampling-based and gradient-guided methods, resulting in more efficient sampling and better trade-offs between reward alignment and divergence from the prior.
Aligning diffusion model outputs with downstream objectives is essential for improving task-specific performance. Broadly, inference-time training-free approaches for aligning diffusion models can be categorized into two main strategies: sampling-based methods, which explore multiple candidate outputs and select those with higher reward signals, and gradient-guided methods, which use differentiable reward approximations to directly steer the generation process. In this work, we propose a universal algorithm, UniCoDe, which brings together the strengths of sampling and gradient-based guidance into a unified framework. UniCoDe integrates local gradient signals during sampling, thereby addressing the sampling inefficiency inherent in complex reward-based sampling approaches. By cohesively combining these two paradigms, UniCoDe enables more efficient sampling while offering better trade-offs between reward alignment and divergence from the diffusion unconditional prior. Empirical results demonstrate that UniCoDe remains competitive with state-of-the-art baselines across a range of tasks. The code is available at https://github.com/maurya-goyal10/UniCoDe