What's Inside Your Diffusion Model? A Score-Based Riemannian Metric to Explore the Data Manifold

TL;DR

The paper tackles the geometry of data learned by diffusion models by introducing a score-based ambient metric derived from the Stein score $s(\boldsymbol{x}) = \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x})$, forming the metric tensor $g(\boldsymbol{x}) = \mathbf{I} + λ \, s(\boldsymbol{x}) s(\boldsymbol{x})^T$ to stretch normal directions. Geodesics are computed through a three-stage pipeline (noise perturbation, energy minimization under $g$, and denoising) to obtain manifold-consistent interpolations and extrapolations. Empirically, the method improves perceptual and distribution metrics on synthetic manifolds, Rotated MNIST, and Stable Diffusion MorphBench, offering a principled view of diffusion-model geometry and a tool for natural image morphing and editing. The approach provides a framework to navigate the learned data manifold for high-quality image morphing, editing, and analysis of diffusion-model representations.

Abstract

Recent advances in diffusion models have demonstrated their remarkable ability to capture complex image distributions, but the geometric properties of the learned data manifold remain poorly understood. We address this gap by introducing a score-based Riemannian metric that leverages the Stein score function from diffusion models to characterize the intrinsic geometry of the data manifold without requiring explicit parameterization. Our approach defines a metric tensor in the ambient space that stretches distances perpendicular to the manifold while preserving them along tangential directions, effectively creating a geometry where geodesics naturally follow the manifold's contours. We develop efficient algorithms for computing these geodesics and demonstrate their utility for both interpolation between data points and extrapolation beyond the observed data distribution. Through experiments on synthetic data with known geometry, Rotated MNIST, and complex natural images via Stable Diffusion, we show that our score-based geodesics capture meaningful transformations that respect the underlying data distribution. Our method consistently outperforms baseline approaches on perceptual metrics (LPIPS) and distribution-level metrics (FID, KID), producing smoother, more realistic image transitions. These results reveal the implicit geometric structure learned by diffusion models and provide a principled way to navigate the manifold of natural images through the lens of Riemannian geometry.

Figures (8)

• Figure 1: (A) Data points $\mathbf{x}$ sampled from a probability distribution $p(\mathbf{x})$ (left) concentrate on a lower-dimensional manifold $\mathcal{M}$ (right). (B) Linear interpolation (black dashed) versus geodesic interpolation (red curve) between MNIST digits. Geodesics follow the manifold surface, producing valid digit transitions, while linear paths yield superpositions. (C) Geometric deformation thorne2000gravitation: data manifold $\mathcal{M}$ embedded in $\mathbb{R}^N$ (left) transforms flat Euclidean space into curved metric space $\tilde{\mathbb{R}}^N$ (middle, right). In this deformed space, geodesics of $\mathcal{M}$ can be computed directly as geodesics of $\tilde{\mathbb{R}^N}$.
• Figure 2: Embedded sphere experiments: (A) von Mises-Fisher distribution on a 2-sphere. (B) Example 10×10 pixel images from the embedding. (C) Score vectors (purple) point normally to the manifold. (D) Linear interpolation (orange) versus geodesics (blue) between endpoints. (E) Extrapolation comparison with geodesics (purple) following the manifold versus linear paths (green) departing from it. (F) Image space results showing our geodesic paths maintain sample quality while linear paths produce blurring (interpolation) or artifacts (extrapolation).
• Figure 3: Rotated MNIST. A Interpolation Example (Best LERP by PSNR) comparing LERP, SLERP, Noise Diffusion and Geodesic (our method). B Three examples with our extrapolation method.
• Figure 4: MorphBench interpolation example with Stable Diffusion 2.1. Comparing LERP, SLERP, Noise Diffusion and Geodesic (Our method)
• Figure 5: Relative error as a function of penalty parameter $\lambda$. The dashed red line shows the exponential fit $A \exp(-\lambda/\lambda_0) + B$.