Primal-Dual Constrained Diffusion Sampling

Many applications call for generative models to provide samples that satisfy requirements, including target semantic or visual features in image generation or observation consistency for inverse problems. Although diffusion models excel at capturing complex distributions, their stochastic, iterative nature makes it difficult to impose such requirements. Typical approaches, known as guidance, tend to encourage rather than enforce desired properties and are often sensitive to hyperparameters such as guidance weights and step size schedules. Moreover, they typically incur substantial memory and computation costs due to retraining or backpropagating through the diffusion model, which limits the type and number of requirements they can impose. This work addresses these challenges by casting constrained diffusion sampling as an optimization problem in which each requirement is described as generalized moment constraint. Using duality, it shows that the solution to this problem is achieved by a reverse dynamic akin to a parametrized Doob’s h-transform of the original diffusion process, whose parameters are the problem’s Lagrange multipliers. It then introduces Primal-Dual Constrained Diffusion Sampling (PD-CDS), a training-free algorithm that automatically adjusts these parameters throughout the diffusion, and derives convergence guarantees. The generality of this approach is illustrated by inverse problems, semantic control tasks, and visual property constraints, separately and combined. By precluding backpropagation through the diffusion model, PD-CDS is able to match and even outperform existing methods while reducing computation and memory costs.