💻 computer science

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

DiffRGD is a plug-and-play, inference-time guidance framework that preserves the latent Gaussian structure of pre-trained diffusion models by formulating sampling steps as constrained optimization problems on a spherical manifold solved via Riemannian Gradient Descent, thereby outperforming existing methods in image restoration and conditional generation tasks.

Original authors: Jia-Wei Liao, Li-Xuan Peng, Mei-Heng Yueh, Min Sun, Cheng-Fu Chou, Jun-Cheng Chen

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jia-Wei Liao, Li-Xuan Peng, Mei-Heng Yueh, Min Sun, Cheng-Fu Chou, Jun-Cheng Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to recreate a masterpiece painting, but you only have a blurry, noisy sketch to start with. You have a magical "AI artist" (a Diffusion Model) that knows how to turn that sketch into a clear image step-by-step. However, you want to guide this artist to make specific changes—like adding a cat to the scene or fixing a blurry face—without having to retrain the artist from scratch (which would take forever and cost a fortune).

This is where the paper DiffRGD comes in. It proposes a new way to "steer" the AI artist during the painting process without breaking the magic.

Here is the breakdown using simple analogies:

1. The Problem: The "Off-Road" Drift

Most current methods try to guide the AI by simply pushing the image in the right direction. Imagine the AI is driving a car on a very specific, smooth, circular track (this track represents the Gaussian distribution, the mathematical "rules" the AI was trained on).

Old Methods (like DPS): They tell the car, "Turn left to avoid the tree!" But they just yank the steering wheel hard. The car might veer off the track, drive onto the grass, and get stuck in the mud. In AI terms, this is called "off-manifold drift." The image starts to look weird, blurry, or distorted because it's no longer following the natural rules the AI learned.
The Dilemma: If you push gently, the car stays on track but doesn't turn enough. If you push hard, it turns well but crashes off the track.

2. The Solution: The "Spherical Track" (DiffRGD)

The authors realized that the AI's "track" isn't just a flat road; it's actually a sphere (like the surface of a balloon). At every step of the painting process, the AI expects the image to stay on the surface of this specific sphere.

DiffRGD changes the rules of the game:

Instead of letting the car drive off-road, DiffRGD says, "We will only move the car along the surface of the sphere."
They use a mathematical technique called Riemannian Gradient Descent. Think of this as a GPS that knows the terrain is curved. Instead of drawing a straight line through the air (which would take you off the sphere), the GPS calculates the best path along the curve of the sphere to get to your destination.

3. How It Works: The "Polar Decomposition"

To build this spherical track, the authors used a trick called Polar Decomposition.

Imagine the AI's noise as a dart thrown at a bullseye.
The "distance" from the center is the radius (how much noise is left).
The "direction" is where the dart is pointing.
DiffRGD says: "Let's lock the radius (keep the noise level exactly where it should be) and only adjust the direction to fit your request."

By locking the radius and only moving along the surface, the image never loses its "natural" quality. It stays on the track, but it still gets to the right destination.

4. The Results: Better Paintings, No Retraining

The paper tested this on two main types of tasks:

Image Restoration: Fixing damaged photos (like removing scratches, making blurry photos sharp, or filling in missing parts).
Conditional Generation: Creating new images based on specific instructions (like turning a sketch into a photo or changing a face to look like a specific person).

The Outcome:

DiffRGD consistently produced clearer, sharper, and more accurate images than previous methods.
It avoided the "weird artifacts" (glitches) that other methods caused when they pushed the image off its natural track.
It works as a "plug-and-play" tool. You don't need to retrain the AI model; you just plug DiffRGD in, and it guides the existing model better.

Summary Analogy

If previous methods were like trying to steer a boat by throwing a rope from the shore (which often pulls the boat into the rocks), DiffRGD is like having a skilled captain who knows the currents perfectly. The captain steers the boat along the natural flow of the water (the spherical manifold) to reach the destination, ensuring the boat never hits the rocks and the passengers (the image pixels) stay comfortable and safe.

In short: DiffRGD is a smarter way to guide AI image generators that respects the math behind how they "think," resulting in higher-quality images without the need for expensive retraining.

Technical Summary: DiffRGD

Problem Statement

Diffusion models have become foundational for image generation, yet controlling them for specific tasks (e.g., image restoration, conditional generation) without costly re-training or fine-tuning remains a challenge. Existing inference-time guidance methods (e.g., DPS, FreeDoM) steer the sampling process by injecting gradients from a differentiable objective directly into the latent space. However, these methods operate in the ambient space without respecting the underlying Gaussian structure of the diffusion process.

This lack of distributional awareness leads to distributional drift, where the latent variables deviate from the original Gaussian manifold induced by the model. As sampling steps accumulate, this "off-manifold" phenomenon degrades the quality of the final clean sample. Previous attempts to constrain updates, such as MPGD (assuming local linearity) or DSG (projecting to a fixed-radius hypersphere), either rely on infeasible perfect encoders or fail to fully adhere to the specific Gaussian properties of the latent distribution, sometimes degrading to a uniform distribution.

Methodology: DiffRGD

The authors propose DiffRGD, a plug-and-play inference-time guidance framework that explicitly preserves the latent Gaussian structure. The core insight is that at each DDIM sampling timestep $t$ , the latent distribution is an isotropic Gaussian centered at a predicted mean $\mu_t$ .

1. Polar Decomposition and Manifold Construction

The method leverages a polar decomposition property of isotropic Gaussian distributions. At any timestep $t$ , a latent sample $x_t$ can be decomposed into a statistically independent radial component $r$ and a directional component $u$ :
$x_t = \mu_t + \sigma_t r u$
where $r$ follows a scaled chi-distribution and $u$ is uniformly distributed on a unit sphere.

Instead of optimizing in the ambient space or on a fixed-radius sphere, DiffRGD constructs a spherical manifold $\mathcal{S}_{t,r}$ induced by the specific latent distribution:
$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$
Here, the radius $\sigma_t r$ is sampled from the distribution's natural scale, ensuring the feasible set respects the Gaussian properties of the diffusion process.

2. Riemannian Gradient Descent (RGD)

The guidance problem at each timestep is formulated as a constrained optimization problem:
$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$
where $\mathcal{L}$ is the guidance loss (e.g., reconstruction loss for restoration or feature alignment for conditional generation).

To solve this, DiffRGD employs Riemannian Gradient Descent on the spherical manifold $\mathcal{S}_{t,r}$ :

Tangent Space Projection: The Euclidean gradient of the loss is projected onto the tangent space of the sphere at the current point, removing the radial component to ensure the update direction is tangent to the manifold.
Retraction: After taking a step along the tangent space, a retraction operation maps the updated point back onto the spherical manifold $\mathcal{S}_{t,r}$ , preserving the radius $\sigma_t r$ .

This process is repeated for $K$ inner iterations (typically $K=3$ ) per diffusion timestep. The authors prove that under $L$ -smoothness assumptions, this optimization converges to a stationary point.

Key Contributions

Distribution-Aware Geometry: The authors identify a polar-decomposition property that reveals a distribution-induced geometry for diffusion guidance. They propose a Riemannian framework that respects the latent distribution geometry, mitigating the sample degradation caused by off-manifold drift.
Theoretical Analysis: The paper provides an analysis of distribution preservation and proves the convergence of the Riemannian optimization algorithm to a stationary point. In practice, the proposed geometry-aware optimization converges faster than ADMM-based methods.
Plug-and-Play Performance: DiffRGD is demonstrated to be a plug-and-play method that requires no re-training or fine-tuning, achieving state-of-the-art results across diverse tasks.

Experimental Results

The authors evaluated DiffRGD against baselines (DPS, FreeDoM, MPGD, DSG, ADMMDiff) on a single NVIDIA GeForce RTX 4090 GPU using the HuggingFace Diffusers framework.

Image Restoration

Evaluated on FFHQ 256×256 and ImageNet 256×256 for four tasks: Inpainting, Super-Resolution (4×), Gaussian Deblurring, and Motion Deblurring.

Quantitative: DiffRGD achieved statistically significant improvements in PSNR, SSIM, LPIPS, and FID over all baselines in most tasks. For example, in Inpainting, DiffRGD achieved a PSNR of 34.04 compared to 32.38 for the second-best method (ADMMDiff).
Robustness: DiffRGD maintained strong performance even with fewer sampling steps (100 steps), whereas methods like DPS and ADMMDiff showed noticeable degradation.
Qualitative: Visual results showed fewer artifacts and better background consistency compared to DPS and MPGD, with comparable or superior quality to ADMMDiff but with less computational time.

Conditional Generation

Evaluated on CelebA-HQ 256×256 for Segmentation Map, Sketch, and FaceID guidance.

DiffRGD consistently outperformed existing inference-time methods (FreeDoM, DSG, ADMMDiff) in terms of alignment metrics (mIoU, Sketch- $\ell_2$ , FaceID- $\ell_2$ ) and generation quality (FID, KID).
Qualitative results indicated that DiffRGD produced more natural and visually appealing results with better alignment to conditions compared to DSG, which often suffered from unstable control.

Guidance Strength

Experiments on guidance strength ( $\eta_t$ ) showed that DiffRGD achieves a better trade-off between alignment (mIoU) and generation quality (FID) across a wider range of strengths compared to DSG, which performed poorly across all settings.

Significance and Claims

The paper claims that DiffRGD addresses a critical gap in inference-time guidance: the failure of existing methods to preserve the stepwise Gaussian latent structure. By formulating guidance as a constrained optimization on a distribution-induced spherical manifold, DiffRGD prevents the "off-manifold" phenomena that degrade sample quality.

The authors assert that geometry-consistent guidance is a viable and superior alternative to directly applying gradients in the ambient space or using restrictive, non-distribution-aware manifolds. The method offers a simple, plug-and-play solution that improves task metrics while preserving the high-fidelity sample quality inherent to pre-trained diffusion models, without the need for expensive re-training.