Understanding the Latent Space of Diffusion Models through the Lens of Riemannian Geometry

Yong-Hyun Park; Mingi Kwon; Jaewoong Choi; Junghyo Jo; Youngjung Uh

Source-linked AI summary

Understanding the Latent Space of Diffusion Models through the Lens of Riemannian Geometry

Yong-Hyun Park, Mingi Kwon, Jaewoong Choi, Junghyo Jo, Youngjung Uh

arXiv:2307.12868v2Jul 24, 2023cs.CV

TL;DR

Diffusion-model latent spaces remain insufficiently understood, especially across recursive timesteps and under different conditions. The paper uses a pullback metric and Jacobian-derived local bases to analyze this structure and edit images through a single latent traversal, finding systematic frequency, dataset, and prompt-dependent patterns.

Problem
The latent space of diffusion models remains insufficiently understood, including how its structure evolves across timesteps and responds to conditions.
Method
The paper uses a pullback metric and Jacobian SVD to derive local latent and tangent bases, then traverses a latent basis at one timestep for image editing.
Results
The method enables semantically meaningful single-timestep editing and reveals low-to-high-frequency evolution, increasing cross-sample tangent-space discrepancy, dataset-dependent consistency, and weakening text influence over generation.
Takeaways & Limitations
Latent geometry provides a framework for understanding diffusion generation and for analyzing how timestep, dataset complexity, and text conditions shape local structure.
Takeaways & Limitations
Some latent directions cause entangled edits or abrupt changes, and the method requires further validation beyond its demonstrated Stable Diffusion setting.

Abstract

from arXiv · show

Despite the success of diffusion models (DMs), we still lack a thorough understanding of their latent space. To understand the latent space $\mathbf{x}_t \in \mathcal{X}$, we analyze them from a geometrical perspective. Our approach involves deriving the local latent basis within $\mathcal{X}$ by leveraging the pullback metric associated with their encoding feature maps. Remarkably, our discovered local latent basis enables image editing capabilities by moving $\mathbf{x}_t$, the latent space of DMs, along the basis vector at specific timesteps. We further analyze how the geometric structure of DMs evolves over diffusion timesteps and differs across different text conditions. This confirms the known phenomenon of coarse-to-fine generation, as well as reveals novel insights such as the discrepancy between $\mathbf{x}_t$ across timesteps, the effect of dataset complexity, and the time-varying influence of text prompts. To the best of our knowledge, this paper is the first to present image editing through $\mathbf{x}$-space traversal, editing only once at specific timestep $t$ without any additional training, and providing thorough analyses of the latent structure of DMs. The code to reproduce our experiments can be found at https://github.com/enkeejunior1/Diffusion-Pullback.

1 Introduction

This paper studies the largely unexplored latent space of diffusion models using Riemannian geometry and a pullback metric. It derives local bases for latent and feature spaces, then analyzes their evolution across timesteps, datasets, and text conditions while demonstrating single-timestep image editing.

Motivation: The paper addresses limited understanding of diffusion-model latent spaces, whose structure remains difficult to analyze because training estimates input-independent noise across many recursive timesteps.Prior latent variables lack useful control properties, motivating analysis of the intermediate feature space and latent space together.
Method: A pullback metric and Jacobian SVD yield corresponding local latent and tangent bases, enabling semantically meaningful image editing by traversing a latent basis at one timestep.Parallel transport can transfer basis-based edits to samples with comparable local geometric structures.
Latent-structure evolution: The local latent basis shifts from low- to high-frequency components during generation, while tangent-space differences between samples increase over time.The frequency shift is explicitly confirmed with power spectral density analysis.
Latent-structure evolution: Dataset complexity shapes temporal consistency: tangent spaces remain similar across timesteps for aligned datasets but not for complex datasets such as ImageNet.The paper reports homogeneity for CelebA-HQ and Flowers, but not ImageNet.
Text conditioning: Similar prompts produce similar local tangent spaces, whereas text influence on the tangent space weakens along the generative process.The prompt–tangent-space relationship is reported as a positive correlation between prompt similarity and tangent-space similarity.

2 Related works

Related work positions this paper against prior studies of diffusion-model control, GAN and VAE latent geometry, and feature-space image manipulation. These studies motivate a geometric analysis of diffusion latents, which remain comparatively poorly understood.

Diffusion Models: Diffusion-model research has advanced synthesis, guidance, and score-based interpretations, but the latent space remains largely unexplored.The related work frames latent-space understanding as an open area despite broader progress in diffusion models.
Latent-space geometry: GAN research has used latent-space manipulation and Riemannian geometry to improve understanding and image control, providing precedents for geometric latent analysis.The cited geometric approaches analyze latent spaces in GANs and other deep generative models.
Image manipulation in DMs: Earlier diffusion editing methods manipulate latent variables or U-Net features, but latent variables often lack semantics, creating problems for semantic image editing.Feature-space approaches instead target self-attention, cross-attention, or other U-Net representations.

3 Discovering the latent basis of DMs

The paper defines a local geometry for diffusion latents by pulling back the Euclidean metric of U-Net bottleneck features, then uses the resulting basis for single-timestep, text-aligned image editing.

3.1 Pullback metric: The pullback metric transfers the Euclidean geometry of bottleneck features H to the latent space X through the differentiable map f: X → H.The Jacobian J_x maps latent-space directions into bottleneck tangent directions, providing a measurable structure where X lacks a known metric.
3.2 Finding local latent basis: Local latent directions are the right singular vectors of the Jacobian, while corresponding bottleneck directions are its left singular vectors.The method retains the top n directions, such as n = 50, to form low-rank local latent and tangent spaces.
3.3 Generating edited images with x-space guidance: X-space guidance applies a decoded basis direction through the diffusion model and provides qualitatively similar edits to direct addition with better fidelity.The editing strength is controlled by γ, while ϵ_θ denotes the diffusion model.
3.4 The overall process of image editing: The editing pipeline inverts an image to x_T, denoises to timestep t, identifies the basis, edits x_t, and completes DDIM generation from the modified latent.The basis is computed at the selected timestep and the manipulation follows a single basis vector.
3.4 The overall process of image editing: Prompt-conditioned bases enable semantic edits aligned with text, and parallel transport transfers a direction to another sample’s tangent space when local structures are comparable.The method requires no extra training and adjusts the latent variable only once at a specific timestep.

4 Findings and results

The paper validates a Riemannian latent basis for semantic image editing and uses it to analyze how diffusion-model geometry changes across timesteps, samples, datasets, and text conditions.

4.1 Image editing with the latent basis: Semantic edits include coarse attributes at T and finer attributes at 0.6T, while prompt-conditioned basis vectors align with object, pose, or action semantics.The editing procedure uses DDIM inversion, modifies one latent at a selected timestep, and completes denoising.
4.2 Evolution of latent structures during generative processes: Latent-basis frequency shifts from low to high over generation, supporting the coarse-to-fine behavior of diffusion models.Power spectral density analysis identifies more low-frequency structure at early stages and more high-frequency structure at later stages.
4.2 Evolution of latent structures during generative processes: Tangent spaces across samples diverge as generation progresses, making generally applicable editing directions harder to find at later timesteps.The paper measures subspace discrepancy with a Grassmannian geodesic metric.
4.2 Evolution of latent structures during generative processes: Parallel transport transfers latent-basis edits consistently at T but performs less satisfactorily at 0.6T, where tangent spaces are less homogeneous.The results motivate checking local-subspace similarity before transferring editing directions across samples.
4.2 Evolution of latent structures during generative processes: Simpler datasets produce more consistent tangent spaces across timesteps than complex datasets such as ImageNet.The comparison reports greater timestep-wise similarity for CelebA-HQ than for ImageNet.
4.3 Text-conditioned latent structures: Similar text prompts induce similar tangent spaces, but prompt influence weakens later in generation, especially below 0.7T.Prompt similarity and tangent-space similarity are negatively related through their distance, with the relationship strongest early in generation.

5 Discussion

The method provides semantically meaningful editing but remains limited by attribute entanglement, occasional abrupt changes, and assumptions about the feature-space metric.

The discovered basis can produce disentangled attributes without attribute annotations, but dataset bias can entangle edits such as beard changes with gender.The paper gives female faces with beards becoming male as an example.
Some Stable Diffusion latent vectors occasionally cause abrupt editing changes, indicating unresolved complexity in the latent geometry.
The approach is broadly applicable when the diffusion feature space follows a Euclidean metric, while applicability to transformer-like architectures remains open.
Despite these limitations, the method advances diffusion-model image editing and understanding through geometric analysis and experiments.

6 Conclusion

The paper analyzes diffusion-model latent spaces geometrically, uses pullback metrics for editing and structural analysis, and reports timestep- and prompt-dependent geometric patterns.

Pullback metrics identify latent and tangent bases whose traversal enables image editing in diffusion models.
The latent bases shift from low- to high-frequency components, tangent-space discrepancies grow across samples, and simpler datasets yield more consistent spaces over timesteps.
Similar text prompts produce analogous tangent spaces, but text influence weakens over diffusion timesteps.
The authors argue that understanding diffusion-model geometry may support future applications, while acknowledging potential misuse of image manipulation.

B Implementation details

The experiments use several pretrained diffusion models and datasets with controlled settings, timestep-specific editing, inversion, guidance, low-rank bases, and quality boosting.

The evaluation covers DDPM variants on ImageNet, LSUN, CelebA-HQ, FFHQ, Flowers, and AFHQ, plus Stable Diffusion 2.1.
Comparisons use official implementations with frozen parameters, matching linear schedules and 2562 resolutions except for Stable Diffusion.
Editing is evaluated at selected fractions of T, although the method supports manipulation at any timestep, and real-image editing uses 100-step DDIM inversion.
The experiments use x-space guidance, a low-rank tangent approximation with n = 50, and stochastic quality boosting after t = 0.2T for applicable models.
Power-method approximation with n = 50 takes about 3–4 minutes on an RTX 3090, while n = 3 reduces runtime to about 10 seconds.
Random directions severely distort images, whereas projection onto the discovered latent subspace yields semantically meaningful transformations.

C Ablation study

Ablations show that the discovered latent subspace supports semantic editing and that x-space guidance improves fidelity and smoothness over direct addition, while its mechanism remains incompletely understood.

Projecting random directions onto the discovered latent subspace produces semantic manipulations, unlike direct random-direction editing, which degrades images.
X-space guidance produces higher-quality images with similar meaning than simple addition, especially improving smoothness and artifacts in Stable Diffusion at t = 0.6T.
The underlying principles that make x-space guidance improve manipulation quality remain insufficiently understood.

D Comparative experiment to other state-of-the-art (SoTA) editing methods

The method is qualitatively compared with four text-guided editing baselines and shows similar editing with content preservation, while its computation remains comparable to alternatives.

Qualitative comparisons cover SDEdit, Pix2Pix-zero, PnP, and Instruct Pix2Pix using their official implementations.
x-space guidance achieves qualitatively similar editing while preserving the original image content across AFHQ, Flowers, and Stable Diffusion.
The approach aligns edits with text conditions while preserving object identity, unlike baselines with structure, manipulation, or identity-fidelity deficiencies.
The method’s computation remains comparable to other approaches, requiring about 2.5 seconds for one Jacobian basis vector.The basis vector is identified once at a specific timestep, without some preprocessing or feature-storage steps used by alternatives.

E More Discussions

The analysis links tangent-space similarity to reliable latent-direction transport, while the implementation approximates large Jacobians through subspace iteration and provides reproducible algorithms.

As tangent-space geodesic distance decreases, parallel transport causes less distortion of latent directions.
Similar tangent spaces preserve latent-basis signals during transport, whereas distant timesteps produce transported vectors that deviate from the original.
The appendix provides PyTorch code for the Jacobian subspace iteration and geodesic-metric algorithms.
Jacobian subspace iteration approximates singular values because diffusion-model spaces X and H make direct Jacobian computation infeasible.
The supplied implementation exposes latent variables, timesteps, encoders, low-rank dimension, chunking, and convergence controls for the local pullback computation.

G Additional results

Additional experiments show timestep-dependent editing granularity and prompt-related latent directions, with prompt-conditioned bases producing attributes related to their conditioning text.

Editing at timestep T produces coarse changes, whereas 0.6T produces fine changes across FFHQ, AFHQ, Flowers, and Stable Diffusion.Stable Diffusion follows a similar pattern, with 0.7T producing coarse changes and 0.6T producing fine changes.
Prompt-conditioned latent bases capture attributes related to the condition, with zebra edits more consistent than chimpanzee edits while remaining chimpanzee-related.
Additional examples demonstrate editing with latent basis vectors discovered from varied images and text prompts.

G.3 More discussion on the editing capability of the latent basis discovered with text conditions

The text-conditioned editing analysis identifies instability and preservation limits, while attributing failures to properties of x-space manipulation and timestep-dependent signal content.

At t = T in Stable Diffusion, latent-basis directions can produce rapid, drastic, unexpected changes, although landscapes show desired editing across timesteps.
Pose or action prompts can lose identity, large shape changes can be unsatisfactory, backgrounds may not be preserved, and whole-image changes are difficult.
Identity failures suggest x-space does not disentangle identity effectively, unlike reported identity-preserving results in H.
Because manipulation adds or subtracts signal from the current x_t, edits remain constrained near the original form; broad edits require t = T, where manipulation is unstable.
Despite these limitations, the paper reports direct x_t manipulation at one timestep as, to its knowledge, the first such result.