PaperlayerView exact revision on arXiv

Source-linked AI summary

Riemannian Score-Based Generative Modelling

Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, Arnaud Doucet

arXiv:2202.02763v3cs.LGmath.PRstat.ML

TL;DR

Existing score-based generative models mainly assume Euclidean data, although many scientific distributions lie on Riemannian manifolds. The paper introduces RSGMs, which define geometry-aware forward and reverse diffusions, and reports better baseline performance and improved scaling to high-dimensional manifolds. The approach is supported by convergence guarantees and demonstrations across manifold types and datasets.

  • Problem

    Existing SGMs primarily target Euclidean spaces, despite scientific distributions supported on Riemannian manifolds.

  • Method

    RSGMs define forward diffusion directly on Riemannian manifolds, use the corresponding manifold-valued reverse process, and estimate its Stein score through score matching.

  • Results

    RSGMs achieve better performance than recent baselines and scale better to high-dimensional manifolds across reported manifold tasks.

  • Takeaways & Limitations

    The framework represents target densities on manifolds as time reversals of Langevin dynamics and applies across a range of manifolds and complex datasets.

  • Takeaways & Limitations

    The approach includes an extrinsic sampling alternative requiring a projection operator, which can accumulate errors, and its theoretical compact-manifold result assumes a heat-kernel bound.

Abstract

from arXiv · show

Score-based generative models (SGMs) are a powerful class of generative models that exhibit remarkable empirical performance. Score-based generative modelling (SGM) consists of a ``noising'' stage, whereby a diffusion is used to gradually add Gaussian noise to data, and a generative model, which entails a ``denoising'' process defined by approximating the time-reversal of the diffusion. Existing SGMs assume that data is supported on a Euclidean space, i.e. a manifold with flat geometry. In many domains such as robotics, geoscience or protein modelling, data is often naturally described by distributions living on Riemannian manifolds and current SGM techniques are not appropriate. We introduce here Riemannian Score-based Generative Models (RSGMs), a class of generative models extending SGMs to Riemannian manifolds. We demonstrate our approach on a variety of manifolds, and in particular with earth and climate science spherical data.

1 Introduction

Score-based generative models progressively noise and denoise data, but existing methods primarily target flat Euclidean spaces. RSGMs extend this framework to Riemannian manifolds by incorporating data geometry into both processes.

  • SGMs add noise through diffusion until data becomes approximately Gaussian, then simulate the time-reversed diffusion to denoise Gaussian samples.
  • Existing SGMs primarily model distributions supported on Euclidean spaces, while many scientific datasets naturally live on Riemannian manifolds.
  • RSGMs define the forward diffusion directly on a Riemannian manifold, inducing a manifold-valued reverse process whose drift includes an estimated Stein score.
  • The method uses score-estimation techniques and Geodesic Random Walks, provides convergence bounds on compact manifolds, and is demonstrated across manifolds and tasks.
  • RSGMs outperform recent baselines and scale better to high-dimensional manifolds in the reported experiments.

2 Euclidean Score-based Generative Modelling

Euclidean SGMs combine a forward Ornstein–Uhlenbeck noising process with an approximated time-reversal. They replace intractable terminal distributions and scores using reference-distribution and score-matching approximations, then discretize sampling.

  • The forward process is an Ornstein–Uhlenbeck diffusion that converges geometrically to N(0, Id).
  • The reverse process has a drift depending on logarithmic gradients of the noised data distributions, but exact sampling is infeasible because terminal distributions and scores are intractable.
  • SGMs replace pT with N(0, Id), exploiting geometric convergence of the forward process toward the reference distribution.
  • Denoising score matching estimates the score using the forward transition density, which is available in closed form for the Ornstein–Uhlenbeck process.
  • Sampling uses an Euler–Maruyama discretization, while adaptive stepsizes and predictor-corrector schemes are described as performance improvements.

3 Riemannian Score-based Generative Modelling

RSGMs extend score-based modelling to manifold-valued data through manifold diffusions, time reversal, Geodesic Random Walk sampling, and score matching. The framework combines intrinsic geometry with tractable approximations for training and generation.

  • 3 Riemannian Score-based Generative Modelling: The manifold formulation assumes a complete, orientable, connected, boundaryless Riemannian manifold and requires four extensions of Euclidean SGMs.
  • 3.1 Noising processes on manifolds: Langevin dynamics provides a generic manifold noising process designed to converge to a convenient stationary distribution.
  • 3.1 Noising processes on manifolds: On compact manifolds, targeting the uniform distribution makes the noising process Brownian motion because the potential gradient vanishes.
  • 3.2 Time-reversal on Riemannian manifolds: The manifold time-reversal theorem shows that reversing a diffusion yields another diffusion whose coefficients define the backward generative process.
  • 3.3 Approximate sampling of diffusions: Geodesic Random Walks provide an intrinsic way to approximate well-behaved manifold diffusions and therefore sample the backward process.
  • 3.4 Score approximation on Riemannian manifolds: The approach addresses intractable manifold transition kernels with score approximations, including spectral truncation, Varadhan approximations, and score-matching losses.

4 RSGMs on compact manifolds

On compact manifolds, RSGMs use Brownian motion toward a uniform reference distribution and approximate manifold heat kernels for score matching. Under regularity and score-approximation assumptions, their discretized samples approach the data distribution in Wasserstein distance.

  • Forward noising: Compact-manifold RSGMs use Brownian motion as a noising process targeting the uniform stationary distribution.Compact manifolds provide a proper uniform distribution and heat-kernel approximations.
  • Score approximation: Heat-kernel scores are approximated either by truncating an eigenfunction expansion or using Varadhan asymptotics for small t.The eigenfunction approach applies when eigenvalues and eigenvectors are tractable, including spheres and tori.
  • Score matching: RSGMs offer three score-matching methods whose computational costs are compared by score-network forward and backward passes.The methods use different approximations to the heat kernel.
  • Theory: The convergence theorem assumes a heat-kernel diagonal upper bound, smooth positive data density, and uniformly bounded score-estimation error.The heat-kernel assumption holds for a large class of manifolds, including spheres, tori, and compact matrix groups.
  • Theory: For any ε > 0, sufficiently large T, small score error M, and small discretization error γ yield W1(L(YN), p0) ≤ ε.The bound includes both neural-network and score-approximation bias, as well as time-discretization effects.

5 Related work

Related manifold generative models include pushed-forward Euclidean flows, Riemannian neural ODEs, and optimal-transport flows. These approaches differ in global representational scope, manifold treatment, temporal discretization, and computational complexity.

  • Push-forward of Euclidean normalizing flows: Push-forward Euclidean normalizing flows require a global homeomorphism between R^n and the manifold, limiting them to manifolds topologically equivalent to R^n.The exponential map is a natural choice but does not generally provide the required global representation.
  • Neural ODE on manifolds: Riemannian neural ODEs define continuous normalizing flows directly on the manifold, avoiding artifacts or numerical instabilities from manifold embeddings.They extend continuous normalizing flows to manifold-valued data without requiring a specific embedding map.
  • Optimal transport on manifolds: Optimal-transport manifold flows are discrete in time and trade expressivity against scalability through their chosen family of c-convex functions.Their maps use the exponential map applied to gradients of c-convex potentials.
  • Computational comparison: The computational-complexity comparison accounts for manifold dimension d, Monte Carlo batches k, and solver steps N and N ∗ across methods.RSGM and Moser-flow training complexity varies when the Hutchinson stochastic estimator is used.

6 Experiments

The experiments evaluate RSGMs across spherical earth and climate data, high-dimensional tori, SO3(R), and hyperbolic space, comparing them with manifold-valued and Euclidean baselines. RSGMs fit high-dimensional targets efficiently, require fewer evaluations on SO3(R), and qualitatively fit distributions on H2.

  • Experimental setup: Experiments benchmark RSGMs on earth and climate spherical data, synthetic tori, SO3(R), and H2 against several manifold-valued and Euclidean baselines.The spherical datasets represent volcanic eruptions, earthquakes, floods, and wild fires.
  • Earth and climate spherical data: Figure 2 visualizes learned earth-science densities alongside training and testing datapoints, colored blue and red, respectively.The learned density uses a green-blue color scale.
  • Synthetic data on tori: RSGMs fit the wrapped Gaussian target on Td even in high dimension with linear or constant computational cost, depending on the divergence estimator.Moser flows could not be trained for d ≥10 and their likelihood gap increased with manifold dimension.
  • Synthetic data on the Special Orthogonal group: RSGMs perform consistently across mixture sizes M on SO3(R), whereas Exp-wrapped SGMs and Moser flows perform well only in some ranges.The comparison uses test log-likelihood and associated number of function evaluations.
  • Synthetic data on the Special Orthogonal group: RSGMs use significantly fewer score network evaluations than the compared methods on the SO3(R) mixture task, while Moser flows require ≫10^3 evaluations.The NFE comparison is reported alongside test log-likelihood in Table 5.
  • Synthetic data on hyperbolic space: On non-compact H2, both score-based models qualitatively fit the target distribution in the synthetic mixture experiment.The experiment uses Langevin dynamics for noising and a wrapped Gaussian as the invariant distribution.

7 Discussion and limitations

The paper presents RSGMs as generative models for manifold-supported densities, emphasizing scalability, broad manifold applicability, robustness, complex-data modelling, and convergence guarantees. It identifies more generic manifolds and alternative noising processes as future directions.

  • Discussion: RSGMs represent target densities supported on manifolds as the time-reversal of Langevin dynamics.The paper frames this as its central generative-modelling construction.
  • Discussion: The reported benefits include high-dimensional scalability, broad applicability across manifolds, robustness, complex-dataset modelling, and theoretical convergence guarantees.These benefits are stated as the paper’s main advantages.
  • Limitations and future work: Future work includes manifolds with boundaries, alternative noising processes, and efficient stochastic-control algorithms for Schrödinger bridges on manifolds.These directions extend the method beyond the settings addressed in the paper.

Checklist

The checklist reports affirmative answers for accurately reflecting contributions, stating assumptions, providing proofs and experimental details, reporting error bars, and citing existing assets. It reports no new assets and marks participant-related and societal-impact items as not applicable or absent.

  • Contributions and scope: The authors identify extending diffusion models to Riemannian manifolds as the main contribution.They report that the paper describes its contributions and scope.
  • Societal impact: The checklist reports no immediate societal impact anticipated for the study of score-based models on manifolds.This is the authors’ stated assessment.
  • Theoretical results: The checklist states that assumptions and complete proofs for theoretical results are provided, with proofs postponed to supplementary material.The assumptions are referenced in A1.
  • Experiments: Experimental details, training details, and error bars are reported, while compute resources are also addressed in the checklist.The relevant details are provided in Appendix O.
  • Assets and data: The paper cites existing assets and discusses their licensing, while reporting no new assets in the supplementary material or via a URL.The checklist also marks consent and personally identifiable-information discussions as not applicable.

Riemannian Score-Based Generative Modelling

The supplementary material introduces notation for smooth manifolds, Riemannian geometry, and stochastic processes, while directing detailed background and derivations to later appendices. It assumes a smooth, connected, complete manifold and defines the metric-based geometric objects used in the work.

  • Supplementary overview: The supplementary material gathers notation, proofs, diffusion-model derivations, and Riemannian-manifold background across dedicated appendices.It also covers likelihood computation, parametric vector fields, and stochastic Riemannian geometry.
  • Background: The paper postpones detailed introductions to Riemannian geometry and manifold stochastic processes to supplementary appendices.The main text states that the relevant section primarily establishes notation.
  • Riemannian setting: The setting assumes a smooth, connected, complete manifold M equipped with a Riemannian metric g.The metric defines scalar products, gradients, and distances on the manifold.
  • Notation: The notation includes the tangent bundle TM, cotangent bundle T⋆M, smooth functions C∞(M), and vector fields X(M).These objects are introduced for a d-dimensional smooth manifold M.

C.1 Tensor field, metric, connection and transport

This section develops the geometric and stochastic calculus needed to define diffusions on Riemannian manifolds. It introduces metrics, connections, parallel transport, geodesics, projections, and Stratonovich-to-Itô conversion.

  • Tensor field and Riemannian metric: A Riemannian metric assigns each tangent space a positive-definite inner product, represented locally by the matrix G.The metric is written as g and induces the inner product ⟨X,Y⟩_M.
  • Connection: A connection differentiates vector fields, while the Levi–Civita connection is uniquely determined by being torsion-free and metric-compatible.Christoffel symbols represent the connection in local coordinates.
  • Parallel transport and geodesics: Parallel transport preserves a vector along a curve and provides linear isomorphisms between the corresponding tangent spaces.For nearby points, the paper uses the unique geodesic to define transport between tangent spaces.
  • Parallel transport and geodesics: Geodesics are curves whose velocity is parallel along the curve, satisfying the coordinate equation Γ^k_i,j(x(t)) ẋ^i(t) ẋ^j(t) = 0.The exponential map sends an initial tangent vector to the endpoint of its geodesic at unit time.
  • Orthogonal projection: For embedded closed manifolds, a smooth nearest-point projection extends the orthogonal projection onto each tangent space.The construction uses the unique nearest point in a neighborhood of the manifold.

H.1 Informal derivation

The informal derivation establishes the manifold time-reversal formula by showing that the reversed process has the appropriate martingale structure. It assumes smooth positive time-marginal densities and separately notes the regularity needed for rigor.

  • Informal derivation: The derivation is non-rigorous unless regularity properties of v are established, which the paper addresses separately in a rigorous proof.The authors explicitly identify this regularity step as necessary for completeness.
  • Informal derivation: The reversed process is defined by Y_t = X_{T−t}, and the derivation aims to prove its martingale characterization for smooth test functions.The argument works with the generator of the reversed process.
  • Informal derivation: The derivation implicitly assumes that every X_t has a smooth positive density p_t relative to the reference measure.This assumption permits the use of ∇log p_t in the time-reversal expression.
  • Informal derivation: A backward Kolmogorov equation for v(u,x) = E[g(X_{T−s})|X_u = x] supplies the key identity used in the martingale calculation.The derivation assumes v is regular enough for the equation to hold.
  • Informal derivation: The divergence theorem transforms the generator terms into an expression involving the score ∇log p_u and the reversed generator.A change of variables then yields the stated time-reversal relation.
  • Informal derivation: A second proof approach embeds the manifold processes into Euclidean space and applies an integration-by-parts formula together with Girsanov theory.This provides a separate route to the time-reversal theorem.

I.1 Main results

The section presents the discrete RSGM sequence and establishes convergence guarantees under bounded score-estimation error, compactness, and sufficiently long diffusion time. It also explains how geometric ergodicity, reverse-process approximation, and Geodesic Random Walk discretization contribute to the bound.

  • Algorithmic construction: The RSGM sequence starts from the reference distribution π∞ and applies a sequence of transition kernels to generate Y_N.The construction defines Y_k through kernels Q_k derived from the reverse-process transitions.
  • Error decomposition: The convergence analysis decomposes the error into reference-process mixing, score approximation, and Geodesic Random Walk discretization terms.The proof separately controls these contributions using geometric ergodicity, Girsanov theory, and GRW convergence.
  • Continuous-time comparison: The continuous-time result omits discretization error and controls a density-level L1 bound, while the cited Moser-flow result controls an L∞ density bound.The comparison concerns continuous-time guarantees rather than the fully discretized RSGM bound.
  • Discretization: The GRW analysis constructs progressively finer grids and proves bounds for the approximation using geodesics, parallel transport, and frame-based updates.The discretization results extend prior GRW analysis to time-inhomogeneous drift terms.

J Proof of proposition 3.3

This section states regularity conditions used to prove a proposition and contrasts RSGM's sample-space interpolation with Moser flows' density-space interpolation. The comparison highlights differing density-evolution mechanisms and density-estimation access.

  • Regularity assumptions: The proof assumes regularity of the transition-density and score-related terms, including a C1 vector-field condition.The text states that the score network and noising process can be chosen to control these conditions.
  • Proof mechanism: The proof applies a divergence theorem on non-compact manifolds to manipulate the divergence term in the transition-density calculation.The argument explicitly introduces a smooth vector field and analyzes the final divergence term.
  • Comparison with Moser flows: Moser flows interpolate in density space through ˆp_t = (1 − t)ˆp_0 + tˆp_1, whereas RSGM interpolates in sample space through its diffusion process.The paper characterizes the former as spontaneous creation of density and the latter as displacement of density.
  • Comparison with Moser flows: Moser-flow vector fields can support density estimation directly, while density estimation is not directly accessible using RSGM without an additional Fisher-score-based procedure.The paper proposes an auxiliary estimation approach for RSGM in an appendix.

L Density estimation with Fisher score

This section develops Fisher-score-based density estimation on a Riemannian manifold and relates it to the spatial score used by RSGMs. It also states convergence and approximation results under compactness and smoothness assumptions.

  • Convergence: Compact-manifold mixing yields |pT(x) − 1| ≤ C exp[−λ1T/2] for sufficiently large T.Here λ1 is the first non-negative eigenvalue of −∆M in L2(pref).
  • Fisher score estimation: Fisher score estimation minimizes a functional whose optimizer is the time derivative of the log density.The paper notes that this score is unavailable in practice and therefore estimated through a score-matching objective.
  • Connection to RSGMs: The spatial Stein score and Fisher score can be estimated jointly using the stated decomposition.This connects the density-estimation construction to the score required by RSGMs.
  • Limitations and extensions: RSGM sampling initializes the reverse dynamics from the invariant reference distribution, which must approximate the forward process terminal law.For RSGMs on compact manifolds, the reference distribution is uniform; accurate reverse sampling can require large T and small GRW stepsizes.
  • Limitations and extensions: Classical SGMs target an easy-to-sample reference distribution and therefore do not directly interpolate between two dataset distributions.The section identifies Schrödinger bridges as a route to pinning both initial and terminal distributions.

O Experimental details

The experiments use JAX and modified Geomstats, fixed neural architectures, specified score-matching losses, and standardized optimization and evaluation procedures. Earth datasets are modeled on a spherical approximation with held-out splits.

  • Architecture: Models use multilayer perceptrons with sinusoidal activations and 512 units per hidden layer.The Earth and SO(3) experiments use five hidden layers, while high-dimensional experiments use three.
  • Loss: The default loss is sliced score matching, while denoising score matching produced similar results.For DSM, the weighting uses the Euclidean conditional variance as a proxy for the compact-manifold setting.
  • Optimization: All models use Adam with β1 = 0.9, β2 = 0.999, and batch size 512.The learning rate uses a linear warm-up followed by a cosine schedule.
  • Evaluation: ODE likelihood calculations use Dormand–Prince with absolute and relative tolerance 1e−5, while SDE rollouts use Euler–Maruyama.Unless otherwise stated, SDE rollouts use 100 steps and no corrector.
  • Data: Earth data are split into training, validation, and test sets with proportions 0.8, 0.1, and 0.1.In each experiment, the Earth is approximated as a perfect sphere.
  • Baselines: Kent-mixture component counts are selected by validation likelihood from a grid spanning 5 to 100 components.The selected counts are 25 for Volcano, 50 for Earthquake, and 100 for both Flood and Fire.

Additional experimental results

Additional experiments evaluate forward-sampling approximations, DSM hyperparameters, toroidal synthetic data, and SO(3) distributions. They compare intrinsic RSGMs with several manifold and Euclidean baselines.

  • Forward sampling: Forward samples from geodesic random walks are evaluated against high-step reference samples with N = 1000.The study varies the discretization-step count N to assess approximation discrepancy.
  • DSM loss: On Flood, the Varadhan approximation alone gives decent performance, while combining it with heat-kernel truncation can improve results.Performance is relatively robust when the DSM hyperparameters τ and J are sufficiently large.
  • Torus experiments: Synthetic torus experiments use wrapped Gaussian data with uniformly random means and standard deviation 0.2.RSGM models are trained with both implicit and sliced score-matching losses.
  • SO(3) experiments: The experiments cover applications involving orthogonal constraints, including protein docking, robotics, computer vision, and medical imaging.These applications motivate modeling distributions on rotation and orthogonal-structure manifolds.
  • SO(3) experiments: The SO(3) synthetic data use K = 32 mixture components and include conditional sampling targeting individual components.The Riemannian model uses the DSM loss with the Varadhan asymptotic approximation.
  • SO(3) experiments: SO(3) models are compared with Moser flows and a wrapped-exponential baseline.The Riemannian SGM uses Varadhan-approximated DSM, whereas the wrapped-exponential model uses exact DSM.
Loading 2202.02763v3…