DiffRGD: An Inference-Time Diffusion Guidance Through Riemannian Gradient Descent

Recently, diffusion models have been widely adopted in gen-erative modeling and have served as foundational models for many imagegeneration tasks. To control the generation without costly re-training orfine-tuning, many works seek inference-time guidance methods to steerthe latent via a differentiable objective at inference time. However, thesemethods cannot effectively preserve the original Gaussian distributionbecause they introduce distributional drift, thereby degrading the sam-ple quality. To address this gap, we propose DiffRGD, a distribution-aware guidance framework that explicitly preserves the latent Gaussianstructure. DiffRGD formulates each sampling step as a constrained op-timization problem on a spherical manifold induced by the latent Gaus-sian distribution, and solves it efficiently via Riemannian Gradient De-scent (RGD). DiffRGD is a plug-and-play method that can be seamlesslyintegrated into any pre-trained diffusion model. Extensive experimentsdemonstrate that DiffRGD outperforms previous methods in most im-age restoration and conditional generation tasks. Our project page isavailable at https://diffrgd.github.io/.