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

2026-06-25

AI Summary

Existing inference-time guidance methods often disrupt the Gaussian structure of the latent space in diffusion models, leading to degraded generation quality. This work proposes a distribution-aware guidance framework that, for the first time, integrates Riemannian gradient descent into the diffusion model inference process. By formulating guidance as a constrained optimization problem on a spherical manifold, the method explicitly preserves the underlying latent Gaussian structure and prevents distributional drift. The approach is plug-and-play and demonstrates significant performance gains over current state-of-the-art methods across diverse image restoration and conditional generation tasks, achieving higher control accuracy while effectively maintaining sample fidelity.

Abstract

Recently, diffusion models have been widely adopted in generative modeling and have served as foundational models for many image generation tasks. To control the generation without costly re-training or fine-tuning, many works seek inference-time guidance methods to steer the latent via a differentiable objective at inference time. However, these methods cannot effectively preserve the original Gaussian distribution because they introduce distributional drift, thereby degrading the sample quality. To address this gap, we propose DiffRGD, a distribution-aware guidance framework that explicitly preserves the latent Gaussian structure. DiffRGD formulates each sampling step as a constrained optimization problem on a spherical manifold induced by the latent Gaussian distribution, and solves it efficiently via Riemannian Gradient Descent (RGD). DiffRGD is a plug-and-play method that can be seamlessly integrated into any pre-trained diffusion model. Extensive experiments demonstrate that DiffRGD outperforms previous methods in most image restoration and conditional generation tasks. Our codebase is available at https://github.com/jwliao1209/DiffRGD.

Problem

• diffusion models
• inference-time guidance
• distributional drift
• Gaussian distribution
• sample quality

Innovation

• Diffusion models
• Inference-time guidance
• Riemannian Gradient Descent
• Distribution preservation
• Manifold optimization