Unified Control for Inference-Time Guidance of Denoising Diffusion Models

TL;DR

UniCoDe is a universal, inference-time guidance algorithm for diffusion models that combines sampling and gradient-based methods to improve task-specific output alignment efficiently.

Contribution

It introduces a unified framework that integrates local gradient signals with sampling, enhancing efficiency and performance over existing methods.

Findings

Competitive performance across various tasks

Improved sampling efficiency

Better trade-offs between reward alignment and prior divergence

Abstract

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