PaperlayerView exact revision on arXiv

Source-linked AI summary

Riemannian Diffusion Models

Chin-Wei Huang, Milad Aghajohari, Avishek Joey Bose, Prakash Panangaden, Aaron Courville

arXiv:2208.07949v1cs.LG

TL;DR

Diffusion frameworks and their variational likelihood methods are largely designed for Euclidean spaces, although many scientific data types lie on curved manifolds. This paper generalizes continuous-time diffusion models to arbitrary Riemannian manifolds with a variational likelihood framework and new divergence methods. It reports state-of-the-art likelihoods over prior manifold generative models.

  • Problem

    Many scientific data types lie on curved Riemannian manifolds, but existing diffusion-model variational frameworks are largely purpose-built for Euclidean spaces and do not readily translate to general manifolds.

    Source: §1 · p. 1

  • Method

    The paper introduces Riemannian Diffusion Models and a Riemannian continuous-time ELBO, with divergence-estimation methods and an importance-sampling technique for the time integral.

    Source: §1 · pp. 1–2, §6 · p. 9

  • Results

    RDM achieves new state-of-the-art likelihoods over prior manifold generative models across the reported benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RDM extends continuous-time diffusion modeling to arbitrary Riemannian manifolds, including challenging non-compact hyperbolic spaces.

    Source: §6 · p. 9

    The closest-point projection may not always be easy to obtain, and the cited |G| computation may cause numerical instability and high variance.

    Source: §3.1 · pp. 4–5

Abstract

from arXiv · show

Diffusion models are recent state-of-the-art methods for image generation and likelihood estimation. In this work, we generalize continuous-time diffusion models to arbitrary Riemannian manifolds and derive a variational framework for likelihood estimation. Computationally, we propose new methods for computing the Riemannian divergence which is needed in the likelihood estimation. Moreover, in generalizing the Euclidean case, we prove that maximizing this variational lower-bound is equivalent to Riemannian score matching. Empirically, we demonstrate the expressive power of Riemannian diffusion models on a wide spectrum of smooth manifolds, such as spheres, tori, hyperboloids, and orthogonal groups. Our proposed method achieves new state-of-the-art likelihoods on all benchmarks.

1 Introduction

Many scientific data types are naturally represented on curved Riemannian manifolds, while diffusion models and their variational frameworks do not readily extend beyond Euclidean spaces. This paper generalizes continuous-time diffusion models to arbitrary Riemannian manifolds and reports state-of-the-art likelihoods over prior manifold generative models.

1 Introduction

  • Because scientific data can lie on curved manifolds and diffusion frameworks are largely built for Euclidean spaces, the paper develops a manifold-aware alternative.Its Riemannian Diffusion Models use Stratonovich SDEs and an extrinsic ambient-space formulation that simplifies implementation.Source: §1 · pp. 1–2
  • RDM achieves new state-of-the-art likelihoods over prior manifold generative models across the reported manifold benchmarks.The experiments cover spherical manifolds, tori, hyperbolic spaces, and orthogonal groups.Source: §1 · p. 2
  • The paper derives a Riemannian continuous-time ELBO and proves its maximization is equivalent to Riemannian score matching for marginally equivalent SDEs.The framework is built on the Riemannian Feynman-Kac representation and Girsanov’s theorem.Source: §1 · p. 2
  • For the Riemannian divergence needed by the ELBO, the paper proposes QR-based computation for low-dimensional manifolds and projected Hutchinson estimation for scalable unbiased estimation.These methods do not depend on a closest-point projection, which may not be readily available for some manifolds.Source: §1 · p. 2
  • An importance-sampling technique for the time integral reduces objective-estimation variance without requiring careful design of the inference noise schedule.Source: §1 · p. 2

2 Background

This background introduces diffusion processes and likelihood estimation, then reviews the geometric and stochastic-calculus concepts used to formulate diffusion on manifolds.

2 Background

  • The section frames the geometric and diffusion concepts needed to construct Riemannian Diffusion Models.Source: §2 · p. 2

2.1 Euclidean diffusion models

  • A diffusion model is described by an SDE initialized from a prior, and its likelihood can be estimated using a variational lower bound called the continuous-time ELBO.The lower bound follows from a stochastic change-of-variable formula, Feynman-Kac, Girsanov’s theorem, and Jensen’s inequality.Source: §2.1 · p. 2

2.2 Riemannian manifolds

  • The manifold is treated as embedded in Euclidean space, where its metric is induced by the embedding and its tangent space is represented through coordinate-chart Jacobians.The tangential projection maps onto the subspace spanned by those Jacobian columns.Source: §2.2 · p. 3

2.3 SDE on manifolds

  • Manifold SDEs are defined through their action on smooth functions, and the Stratonovich integral is used because it follows the usual chain rule.In embedded coordinates, the SDE coefficients must lie in the tangent space, keeping the dynamics on the manifold.Source: §2.3 · p. 3

3 Riemannian diffusion models

Section 3 develops a variational likelihood framework for diffusion processes on Riemannian manifolds, then addresses computation, training, variance reduction, and the connection to score matching.

3 Riemannian diffusion models

  • A Riemannian stochastic change-of-variable formula yields marginal densities, and change of measure plus Jensen’s inequality gives the CT-ELBO training objective.Source: §3 · pp. 3–4

3.1 Computing Riemannian divergence

  • QR-based tangent bases compute divergence without closest-point projection, while projected Hutchinson estimates offer an alternative when QR is costly.Source: §3.1 · p. 5

3.2 Fixed-inference parameterization

  • Fixing inference avoids optimizing its generating vector fields; with tangential projection, (V · ∇g)V = 0 and the generative drift can use variational parameter a.Source: §3.2 · p. 5

3.3 Variance reduction

  • Importance sampling with proposal q(s) gives an unbiased time-integral estimate, but its zero expected gradient precludes training q by maximizing the estimate, so variance is minimized instead.Source: §3.3 · p. 6

3.4 Connection to score matching

  • For marginally equivalent SDEs, maximizing the plug-in reverse-process Riemannian CT-ELBO equals minimizing Riemannian score matching; λ = 1 yields an equivalent ODE.Source: §3.4 · pp. 6–7

4 Related work

Prior work frames diffusion models through score-based generation, inversion of a fixed noise process, and maximum-likelihood training. Riemannian generative models extend flows and score-based methods beyond Euclidean space; RDMs distinguish themselves through a maximum-likelihood framework and broader manifold scope.

Diffusion models

  • Continuous-time diffusion models in a maximum-likelihood framework are the Euclidean special case of the theory developed here.The related work also describes diffusion models as score-based generators or as learned inversions of a fixed noise-injecting process.Source: §4 · p. 7

Riemannian Generative Models

  • Earlier Riemannian generative work includes constant-curvature models, general-manifold continuous-time flows, Moser Flow, and score-based models for compact manifolds.RDMs are distinguished from concurrent score-based work by their maximum-likelihood framing.Source: §4 · p. 7
  • RDMs support importance-sampling variance reduction and likelihood estimation, and apply to non-compact manifolds including hyperbolic spaces.Experiments cover a wider variety of manifolds, including the orthogonal group and toroids.Source: §4 · p. 7

5 Experiments

Experiments evaluate Riemannian diffusion models on spherical earth-event datasets, molecular tori, synthetic hyperbolic manifolds, and synthetic SO(3). On tori, RDM outperforms the reported baseline across datasets, especially on higher-dimensional RNA data.

5 Experiments

  • The experiments instantiate RDMs by parameterizing a with an MLP and maximizing the CT-ELBO.The reported training procedure and selected hyperparameters are provided in the appendix.Source: §5 · p. 7

5.1 Sphere

  • On spherical earth-event datasets, Table 1 reports test NLL over five runs, while learned importance sampling lowers CT-ELBO variance and speeds training.The datasets cover events including volcanoes, earthquakes, floods, and fires.Source: §5.1 · pp. 7–8, Table 1 · p. 8

5.2 Tori

  • RDM beats the 4,096-power-spherical-distribution baseline across protein and RNA datasets, with the largest difference on higher-dimensional RNA; Table 2 reports test NLL.Figure 4 plots protein model-density contours with test samples. In the integration ablation, fewer than 100 time steps tend to underestimate the loss, while 100-step training is indistinguishable from direct sampling; wall-clock times are similar because the inference path does not involve the neural module.Source: §5.2 · pp. 8–9, Figure 4 · p. 8, Table 2 · p. 9

5.3 Hyperbolic Manifolds

  • For noncompact hyperbolic manifolds, the inference process needs nonzero drift to avoid being dissipative; the experiments model synthetic datasets.Figure 5 places data above the learned density.Source: §5.3 · p. 9, Figure 5 · p. 9

5.4 Special Orthogonal Group

  • On SO(3), the framework models a synthetic multimodal density and visualizes data and learned density using Euler angles, one per rotation axis.The orthogonal-group experiment uses projected Hutchinson to estimate Riemannian divergence; closest-point projection requires singular value decomposition.Source: §5.4 · p. 9

6 Conclusion

The paper extends continuous-time diffusion models to Riemannian manifolds, develops a variational training framework with a connection to Riemannian score matching, and reports state-of-the-art density estimation across several geometric data settings.

6 Conclusion

  • Figure 6 presents a synthetic multimodal density on SO(3) beside the learned density.Source: Figure 6 · p. 9
  • RDMs achieve state-of-the-art density-estimation performance on geoscience, protein/RNA, hyperbolic synthetic, and orthogonal-group synthetic data.Source: §6 · p. 9
  • The Riemannian Continuous-Time ELBO provides a variational training objective whose maximization is connected to Riemannian score matching through marginally equivalent SDEs.Source: §6 · p. 9
  • The framework includes ambient-space SDE parameterization, methods for computing Riemannian divergence, and time-integral importance sampling to reduce loss variance.Source: §6 · p. 9

Checklist

The checklist reports limitations and reproducibility information, while the appendix outlines notation and foundational manifold concepts, including charts, tangent spaces, and embeddings.

Checklist

  • The checklist identifies numerical integration and divergence computation or estimation as framework- and dynamical-model-constrained choices, and cautions that the approach may not work out of the box in higher-dimensional settings.Source: Checklist · p. 13
  • The checklist indicates that theoretical assumptions and proofs, training details, reproducibility materials, and error bars are addressed.Source: Checklist · p. 13

A Riemannian manifolds

  • The appendix distinguishes index-free, intrinsic-coordinate, and extrinsic-coordinate notation, using extrinsic coordinates for implementation.Source: A Riemannian manifolds · p. 15

A.1 Smooth manifolds and tangent vectors

  • A smooth manifold is described through compatible local charts, and its tangent space can be viewed as a subspace of the ambient space when embedded.Source: §A.1 · pp. 15–16

A.2 Riemannian metric

This section introduces the Riemannian metric as an inner product on tangent spaces, explains its ambient-space representation, and uses it to define volume and probability measure.

A.2 Riemannian metric

  • A Riemannian metric assigns an inner product to each tangent space, and a metric tensor field records these inner products across the manifold.Source: §A.2 · p. 16
  • For an embedded submanifold, the metric is induced by pulling back the ambient Euclidean metric, with its matrix representation determined by coordinate basis vectors.Source: §A.2 · pp. 16–17
  • The metric induces a measure on the manifold, with volume density given by the absolute determinant of the metric matrix and probability densities normalized with respect to that measure.Source: §A.2 · p. 17

A.3 Riemannian gradient and divergence

The section defines the Riemannian gradient through directional derivatives and develops divergence properties that support computation in intrinsic or embedded settings.

A.3 Riemannian gradient and divergence

  • The Riemannian gradient is defined by the metric-based relation between a function and its directional derivatives along tangent vector fields.Source: §A.3 · p. 18
  • The Riemannian divergence has a coordinate-free form and, for an embedded manifold, can be computed in the ambient space.Source: §A.3 · p. 18
  • The divergence theorem extends Stokes' theorem to Riemannian manifolds for compactly supported fields.Source: §A.3 · p. 18

A.4 Covariant derivative

An affine connection defines how to differentiate a vector field along another, with the covariant derivative satisfying linearity and a product rule. The Levi-Civita connection is the unique metric-compatible symmetric connection; for embedded submanifolds, the tangential connection provides an example.

A.4 Covariant derivative

  • An affine connection compares vector-field values at nearby points, and ∇_U V is the covariant derivative of V in the U-direction.The connection is written as ∇: X(M) × X(M) → X(M).Source: §A.4 · p. 19
  • The covariant derivative is linear in each vector-field argument and obeys a product rule when its second argument is multiplied by a smooth function.The product rule adds U(f)V to f∇_U V.Source: §A.4 · p. 19
  • The Levi-Civita connection is the unique connection that is symmetric and compatible with the metric.Metric compatibility means the covariant derivative of the metric is zero in every direction.Source: §A.4 · p. 20
  • For an embedded submanifold, projecting the Euclidean connection tangentially gives its Levi-Civita connection, independent of the smooth extensions chosen.The Euclidean connection itself is an affine connection.Source: §A.4 · pp. 19–20

B Proofs

This section proves the density, likelihood, divergence, and shared-marginal results, and specifies geometric projections and metric choices for the manifolds considered. It also clarifies that the hyperbolic projection is not generally closest in Euclidean distance and derives tangent and closest-point projections for orthogonal groups.

B Proofs

  • The marginal density of the SDE is derived by expressing its time evolution as a manifold PDE and applying the Feynman–Kac formula.The proof identifies the governing equation as the Fokker–Planck equation.Source: B Proofs · pp. 21, 23
  • The Riemannian CT-ELBO bound follows by changing measure with Girsanov’s theorem and applying Jensen’s inequality.The bound uses the stated inference SDE, with the generative degrees of freedom evaluated in reversed time.Source: B Proofs · p. 24
  • The appendix proves a Riemannian divergence identity, gives its embedded-manifold form, and shows that the tangential projection matrix satisfies (V · ∇g)V = 0.The embedded form uses ambient vector-field coefficients and projection onto the tangent space.Source: B Proofs · pp. 25–26
  • For λ ≤ 1, the two defined processes have the same marginal density q(·, s), and λ = 1 yields an equivalent ODE.The proof constructs a reverse process by mirroring the diffusion term around zero and changing the time variable.Source: B Proofs · p. 28

C.1 Spheres and tori

  • Spheres are unit-norm submanifolds of R^(d+1), while d-tori are products of circles embeddable in R^(2d); their tangent projection removes the normal component.For spheres, the closest-point projection maps x to x/∥x∥2.Source: §C.1 · pp. 29–30

C.2 Hyperbolic spaces

  • For hyperbolic spaces, the appendix uses a Lorentzian model but computes density under the induced Euclidean metric, whose likelihood can be maximized instead of the Lorentzian density.It also distinguishes the Euclidean-orthogonal tangent projection from the usual Lorentz projection.Source: §C.2 · p. 30
  • The hyperbolic projection is valid but not generally the Euclidean closest-point projection; replacing the Lorentz norm with Euclidean distance leads to a Lagrange multiplier requiring root finding.The paper notes that this contradicts the claim by Skopek et al. (2019), and says the valid projection is used to simulate dynamics in numerical integration.Source: §C.2 · p. 31

C.3 Orthogonal groups

  • The tangent-space projection onto TXO(n) is formulated as a Frobenius-norm closest-point projection, with dXX⊤ skew-symmetric.The skew-symmetric matrices are defined by X⊤ = −X.Source: §C.3 · p. 31
  • For orthogonal groups, the closest-point projection of X is π(X) = UV⊤, where U and V are the left and right singular matrices of X.The derivation formulates the projection under the constraint M⊤M = I using the Frobenius norm.Source: §C.3 · p. 32

D.1 Architecture

The experiments use manifold-specific neural architectures and optimization settings, evaluate with a tighter multi-sample lower bound, and numerically integrate the SDEs with manifold-appropriate projections. The KELBO approaches the true likelihood as the sample count grows, with results changing little beyond K = 100.

D.1 Architecture

  • The variational network is an MLP with sinusoidal or swish activation; hyperbolic models add ActNorm, while orthogonal-group inputs are flattened.The importance sampler uses a deep sigmoidal flow and fixed scaling; disconnecting gradients from the solver saves compute but slightly biases its variance-minimization updates, while still reducing variance substantially.Source: §D.1 · p. 33

D.2 Optimization

  • Training uses Adam, with manifold-specific learning rates and momenta; sphere models use cosine scheduling, and the importance sampler is updated every 500 network-training steps.The torus mixture of power spherical distributions uses learning rate 0.03 with β1 = 0.9 and β2 = 0.999.Source: §D.2 · p. 33

D.3 KELBO

  • KELBO increases monotonically toward the true likelihood as K grows; evaluation uses K = 100, and results change little for K > 100.It is a tighter lower bound constructed from i.i.d. trajectories sampled from Q.Source: §D.3 · p. 34

D.4 Numerical integration of the SDEs

  • SDEs are integrated with Stratonovich-Heun and projected after each step; evaluation approximates the KELBO Itô integral with Euler-Maruyama and uses adaptive step sizes.Density plots use the equivalent ODE's exact log likelihood, integrated with dopri5.Source: §D.4 · p. 34
Loading 2208.07949v1…