Controlling Diffusion Models Through Blockwise Gradient Guidance and Sampling

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, CoDeX, which brings together the strengths of blockwise sampling and gradient-based guidance into a unified framework. Building on the blockwise sampling paradigm of CoDe, CoDeX integrates local gradient signals during sampling, thereby addressing the sampling inefficiency inherent in complex reward-based sampling approaches like CoDe. At the same time, it overcomes the limited applicability of traditional gradient-guided methods, which often struggle with non-differentiable rewards. By cohesively combining these two paradigms, CoDeX enables more efficient sampling while offering better trade-offs between reward alignment and divergence from the diffusion unconditional prior. Empirical results demonstrate that CoDeX consistently outperforms CoDe and remains competitive with state-of-the-art baselines across a range of tasks.