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

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.

Overview

Most existing inference-time guidance methods inject external gradients directly into the sampling process, formulating the per-step update as:

Because the neighborhood $\mathcal{N}(\hat{\bm{x}}_{t-1})$ is defined in the ambient Euclidean space, these methods cannot preserve the stepwise Gaussian latent distribution, causing distributional drift that ultimately degrades sample quality. DiffRGD addresses this by performing gradient updates directly on a distribution-induced spherical manifold.

Method

Starting from a noisy latent $\bm{x}_t$ at timestep $t$, DDIM defines the next-step latent $\bm{x}_{t-1}$ as sampling from a Gaussian distribution:

where the mean is defined by

and the predicted clean sample is estimated by Tweedie's formula:

Polar Decomposition of Isotropic Gaussian

By fixing the sampled radius $r$, we construct a spherical manifold that captures the Gaussian structure as the feasible set for RGD:

Problem Formulation

We formulate the per-timestep guidance as a constrained optimization problem on $\mathcal{S}_{t,r}$:

Riemannian Gradient Descent for Diffusion Guidance

To solve the constrained optimization on $\mathcal{S}_{t,r}$, we employ RGD with three geometric operations:

Tangent Space

$$\operatorname{T}_{\bm{x}}\mathcal{S}_{t,r} = \left\{ \bm{v} \in \mathbb{R}^n \;\middle|\; (\bm{x} - \bm{\mu}_t)^\top \bm{v} = 0 \right\}.$$

Projection Operator

$$\Pi_{\operatorname{T}_{\bm{x}}\mathcal{S}_{t,r}}(\bm{v}) = \left( \bm{I}_n - \frac{(\bm{x} - \bm{\mu}_t)(\bm{x} - \bm{\mu}_t)^\top}{\|\bm{x} - \bm{\mu}_t\|_2^2} \right) \bm{v}.$$

Retraction

$$\operatorname{R}_{\bm{x}}(\bm{v}) = \bm{\mu}_t + \sigma_t r \cdot \frac{\bm{x} - \bm{\mu}_t + \bm{v}}{\|\bm{x} - \bm{\mu}_t + \bm{v}\|_2}.$$

The Riemannian gradient and retraction-based update at each inner iteration $k$ are:

Algorithm

Experiments

Image Restoration

We evaluate on four inverse problem tasks using 150 images from FFHQ 256×256. Each task is formulated as a noisy linear inverse problem $\bm{y} = \mathcal{A}(\bm{x}) + \bm{n}$, $\bm{n} \sim \mathcal{N}(\bm{0}, 0.05^2\bm{I})$, with guidance objective $\mathcal{L}_\text{inv}(\hat{\bm{x}}_0, \bm{y}) = \|\mathcal{A}(\hat{\bm{x}}_0) - \bm{y}\|_2$. Tasks: (i) Inpainting (70% random mask), (ii) Super-Resolution ×4 (bicubic), (iii) Gaussian Deblurring (31×31, $\sigma=3.0$), (iv) Motion Deblurring (61×61, intensity 0.5).

Table 1. Quantitative results using DDIM 1,000 sampling steps on 150 samples from the FFHQ 256×256 validation set. Metrics averaged over 100 bootstrap runs. Bold = best, underline = second-best, * = statistically significant improvement ($p < 0.05$) over second-best.

Conditional Generation

We evaluate on three conditional generation tasks from CelebA-HQ 256×256: segmentation map, sketch, and FaceID guidance. Given a conditional image $\bm{y}$, we extract task-specific features via a pre-trained model $\psi_\theta$, and define the guidance loss as:

Table 2. Quantitative results using DDIM 100 sampling steps on 150 samples from CelebA-HQ 256×256. Metrics averaged over 100 bootstrap runs. * = statistically significant improvement ($p < 0.05$).

Citation

@inproceedings{liao2026diffrgd,
  title     = {DiffRGD: An Inference-Time Diffusion Guidance
               Through Riemannian Gradient Descent},
  author    = {Liao, Jia-Wei and Peng, Li-Xuan and Yueh, Mei-Heng
               and Sun, Min and Chou, Cheng-Fu and Chen, Jun-Cheng},
  booktitle = {European Conference on Computer Vision (ECCV)},
  year      = {2026},
}