Gradient-Guided Diffusion Improves Constrained Optimization.

Diffusion models are increasingly used not only for generating data but also for optimizing task-specific objectives. A common technique involves guiding the reverse diffusion process using gradients from an external objective. However, when the underlying data distribution resides on a structured feasible set, such as a manifold or a constraint set, this gradient guidance can inadvertently push generated samples away from the learned data geometry. This research proposes a novel "projected-gradient-guided diffusion update" to address this challenge. The core idea is to incorporate the objective gradient directly into the denoising step, leveraging the observation that the Stein denoising operator can act as an approximate projection onto the data geometry. This approach allows for inference-time optimization using only a pre-trained denoiser and gradient evaluations. The paper provides theoretical analysis of this update as an inexact projected-gradient method for constrained optimization. It offers descent and finite-time convergence guarantees across various settings, including linear manifolds, compact convex feasible sets, and compact Riemannian submanifolds. Numerical experiments support these theoretical interpretations, demonstrating how the proposed update effectively balances the goal of objective descent with the crucial need to preserve the learned data geometry.