Riemannian Score-Based Generative Modelling
Abstract
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 deﬁned by approximating the time-reversal of the diffusion.
Existing SGMs assume that data is supported on a Euclidean space, i.e. a manifold with ﬂat 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.
Introduction
Score-based Generative Models (SGMs) also called diffusion models (Song and Ermon, 2019; Song et al., 2021; Ho et al., 2020; Dhariwal and Nichol, 2021) formulate generative modelling as a denoising process.
Noise is incrementally added to data using a diffusion process until it becomes approximately Gaussian.
The generative model is then obtained by simulating an approximation of the corresponding time-reversal process which progressively denoises a Gaussian sample to obtain a data sample.
This process is also a diffusion whose drift depends on the logarithmic gradients of the noised data densities, i e. the Stein scores, estimated using a neural network via score matching (Hyvärinen, 2005; Vincent, 2011).
SGMs have been primarily applied to data living on Euclidean spaces, i e. manifolds with ﬂat geometry.
However, in a large number of scientiﬁc domains the distributions of interest are supported on Riemannian manifolds.
These include, to name a few, protein modelling (Shapovalov and
Dunbrack Jr, 2011), cell development (Klimovskaia et al., 2020), image recognition (Lui, 2012), geological sciences (Karpatne et al., 2018; Peel et al., 2001), graph-structured and hierarchical data (Roy et al., 2007; Steyvers and Tenenbaum, 2005), robotics (Feiten et al., 2013; Senanayake and
Ramos, 2018) and high-energy physics (Brehmer and Cranmer, 2020).
We introduce in this work Riemannian Score-based Generative Models (RSGMs), an extension of
SGMs to Riemannian manifolds which incorporate the geometry of the data by deﬁning the forward diffusion process directly on the Riemannian manifold, inducing a manifold-valued reverse process.
This requires constructing a noising process on the manifold that converges to an easy-to-sample reference distribution.
We establish that, as in the Euclidean case, the corresponding time-reversal process is also a diffusion whose drift includes the Stein score which is intractable but can similarly be estimated via score matching.
Methodological extensions are required as in most cases the transition kernel of the noising process cannot be sampled exactly.
For example on compact manifolds it is typically only available as an inﬁnite sum through the Sturm–Liouville decomposition (Chavel, 1984).
To this end, we develop non-standard techniques for score estimation and rely on the use of Geodesic
Random Walks for sampling (Jørgensen, 1975).
We provide theoretical convergence bounds for
RSGMs on compact manifolds and demonstrate our approach on a range of manifolds and tasks, including modelling a number of natural disaster occurrence datasets collected by Mathieu and Nickel (2020).
We show that RGSMs achieve better performance than recent baselines (Mathieu and Nickel, 2020; Rozen et al., 2021) and scale better to high-dimensional manifolds.
Euclidean Score-based Generative Modelling
We recall here brieﬂy the key concepts behind SGMs on the Euclidean space R d and refer the readers to Song et al.
for a more detailed introduction.
We consider a forward noising process ( X t ) t ≥deﬁned by the following Stochastic Differential Equation (SDE)
[Equation]
where ( B t ) t ≥ 0 is a d -dimensional Brownian motion and p 0 is the data distribution.
The available data gives us an empirical approximation of p 0 .
The process ( X t ) t ≥ 0 is simply an Ornstein–Ulhenbeck (OU) process which converges with geometric rate to N(0 , Id) .
Under mild conditions on p 0 , the timereversed process ( Y t ) t ≥
= ( X T − t ) t ∈ [0 ,T ] also satisﬁes an SDE (Cattiaux et al., 2021; Haussmann and Pardoux, 1986) given by
[Equation]
where p t denotes the density of X t .
By construction, the law of Y T − t is equal to the law of X t for t ∈ [0 , T ] and in particular Y T ∼ p 0 .
Hence, if one could sample from ( Y t ) t ∈ [0 ,T ] then its ﬁnal distribution would be the data distribution p 0 .
Unfortunately we cannot sample exactly from (2) as p T and the scores ( ∇ log p t ( x ))
t ∈ [0 ,T ] are intractable.
Hence SGMs rely on a few approximations.
First, p T is replaced by the reference distribution N(0 , Id) as we know that p T converges geometrically towards it.
Second, the following denoising score matching identity is exploited to estimate the scores
[Equation]
where p t | 0 ( x t | x 0 ) is the transition density of the OU process (1) which is available in closed-form.
It follows directly that ∇ log p t is the minimizer of (cid:96) t ( s ) =
E [ (cid:107) s ( X t ) − ∇ x t log p t |
( X t | X 0 ) (cid:107) 2 ] over functions s where the expectation is over the joint distribution of X 0 , X t .
This result can be leveraged by considering a neural network s θ : [0 , T ] × R d → R d trained by minimizing the loss function (cid:96) ( θ ) = (cid:82) T {+
+} λ t (cid:96) t ( s θ ( t, · ))
d t for some weighting function λ t > 0 .
Finally, an Euler–Maruyama discretization of (2) is performed using a discretization step γ such that T =
γN for N ∈ N
[Equation]
The above showcases the basics of SGMs but we highlight that many improvements have been proposed; see (e.g. Song and Ermon, 2020;
Jolicoeur-Martineau et al., 2021; Dhariwal and Nichol, 2021).
In particular, selecting an adaptive stepsize ( γ n ) n ∈ N (Bao et al., 2022; Watson et al., 2021) and using a predictor-corrector scheme (Song et al., 2021) instead of a simple Euler–Maruyama discretization drastically improves performance.
Riemannian Score-based Generative Modelling
We now move to the Riemannian manifold setting, and more speciﬁcally assume that M is a complete, orientable connected and boundaryless Riemannian manifold, endowed with a Riemannian metric g 4 .
Four components are required to extend SGMs to this setting: i) a forward noising process on
M which converges to an easy-to-sample reference distribution, ii) a time-reversal formula on M which deﬁnes a backward generative process, iii) a method for approximating samples of SDEs on manifolds, iv) a method to efﬁciently approximate the drift of the time-reversal process.
Notation are gathered in App.
B.
3.1 Noising processes on manifolds
The ﬁrst necessary component is a suitable generic noising process on manifolds that will converge to a convenient stationary distribution.
A simple choice is to use Langevin dynamics described by
[Equation]
which admits the invariant density (w.r.t.
the volume form) given by d p ref / d Vol M ( x ) ∝ e − U ( x ) (Durmus, 2016, Section 2.4), where ∇ is the Riemannian gradient 5 .
Two simple choices for U ( x ) present themselves.
Firstly, setting U ( x ) =
d M ( x, µ ) 2 / (2 γ 2 ) , where d M is the geodesic distance and µ ∈ M is an arbitrary mean location, induces the drift ∇ X t U ( X t ) = − exp − 1 X t ( µ ) /γ 2 6 .
This is the potential of the ‘Riemannian normal’ (Pennec, 2006) distribution, from which it is in general neither trivial to sample nor to compute the normalisation constant (Hauberg, 2018; Mathieu et al., 2019).
An alternative is to target the ’exponential wrapped’ Gaussian.
This is the pushforward of a Gaussian distribution in the tangent space at the mean location along the exponential map.
The potential is given by U ( x ) =
d M ( x, µ ) 2 / (2 γ 2 ) + log | ∂ exp − 1 µ ( x ) | 7 .
In contrast to the
Riemannian normal, sampling and evaluating the density of this distribution is easy.
One recovers the standard Ornstein–Uhlenbeck noising process (Song et al., 2021) for both of these target distributions when M = R d and µ = 0 since then the drift b ( t, X t ) = 1 2 exp − 1 X t (0) =
− 1 2 X t .
On compact manifolds, the invariant measure Vol M has ﬁnite volume, thus a natural choice is to target the uniform distribution which is given by Vol M / |M| .
In this case, ∇ X t U ( X t ) =
and the noising process is simply a Brownian motion on
M .
3.2 Time-reversal on Riemannian manifolds
In order to use these noising processes we prove the time-reversal formula for manifolds, a generalisation of the results in the Euclidean case, e.g. see Cattiaux et al.
(2021, Theorem 4.9).
Consider an SDE of the form d X t = b ( X t )d t + d B M t where B M t is a Brownian motion on M .
We refer to
App.
C.3 for an introduction to Brownian motions on manifolds.
This result shows that if ( X t ) t ∈ [0 ,T ] is a diffusion process then ( X T − t ) t ∈ [0 ,T ] is also a diffusion process w.r.t.
the backward ﬁltration whose coefﬁcients can be computed, and are shown in Eq.
The proof relies on an extension of
Cattiaux et al.
(2021, Theorem 4.9) to the Riemannian manifold case and is postponed to App. H.
Theorem 1 (Time-reversed diffusion) .
Let T ≥ 0 and ( B M t ) t ≥ 0 be a Brownian motion on M such that B M 0 has distribution the volume form p ref8 .
Let ( X t ) t ∈
[0 ,T ] be associated with the SDE d X t
=
b
( X t )d t + d B M t .
Let ( Y t ) t ∈ [0 ,T ] =
( X T − t ) t ∈ [0 ,T ] and assume that KL( P | Q ) < + ∞ , where Q is the distribution of ( B M t ) t ∈ [0 ,T ] and P the distribution of ( X t ) t ∈ [0 ,T ] .
In addition, assume that P t =
L ( X t ) , the distribution of X t , admits a smooth positive density p t w.r.t.
p ref for any t ∈ [0 , T ] .
Then, ( Y t ) t ∈ [0 ,T ] is associated with the SDE
[Equation]
3.3 Approximate sampling of diffusions
Obtaining samples from SDEs on a manifold is non-trivial in general.
If M is isometrically embedded into R p (with p ≥ d ) one can deﬁne ( B M t ) t ≥ 0 as a R p -valued process, see App. C.3.
However, this approach is extrinsic , as it requires the knowledge of the projection operator to place points back on the manifold at each step which can accumulate errors.
Here we consider an intrisic approach based on Geodesic Random Walks (GRWs), see Jørgensen (1975) for a review of their properties.
GRWs can approximate any well-behaved diffusion on M .
Hence, we introduce GRWs in a general framework and consider a discrete-time process ( X γn ) n ∈ N which approximates the diffusion ( X t ) t ≥ 0 deﬁned by
[Equation]
This generalisation is key to sampling the backward diffusion process deﬁned in Theorem 1.
[Figure]
[Figure] [Figure] [Figure] [Figure]
Require: T, N, X γ 0 , b, σ, P 1: γ = T/N
2: for k ∈ { 0 , . . .
, N − 1 } do 3:
Z k +1 ∼ N( Id)
Deﬁnition 2 (Geodesic Random Walk) .
Let X γ 0 be a M -valued random variable.
For any γ > 0 , we deﬁne ( X γn ) n ∈ N such that for any n ∈ N , X γn +1 =
exp X γn [ γ { b ( X γn ) + √ γV n +1 } ] , where ( V n ) n ∈ N is a sequence of T M -valued random variables such that for any n ∈ N , E [ V n +1 |F n ] = 0 and E [ V n +1 V (cid:62) n +1 |F n ] = σσ (cid:62) ( X γn ) , where F n is the ﬁltration generated by { X γk } nk =0 .
We say that
the
M -valued process ( X γn ) n ∈ N is a Geodesic Random Walk.
Algorithm 1 approximately simulates the diffusion ( X t ) t ∈ [0 ,T ] deﬁned in Eq.
(5) using GRWs; see Kuwada (2012); Cheng et al.
(2022) for quantitative error bounds in the time-homogeneous case and App.
I.2 for a novel extentsion for the time-inhomogeneous case.
Fig.
1 provides a graphical illustration of this procedure.
3.4 Score approximation on Riemannian manifolds
Score matching and loss functions.
The reverse process from Eq.
(4) involves the Stein score ∇ log p t which is unfortunately intractable.
To derive an approximation, we ﬁrst remark that for any s, t ∈ (0 , T ] with t > s and x t ∈ M , p t ( x t ) =
(cid:82) M p t | s ( x t | x s )
d P s ( x s ) , where P s = L ( X s ) , the distribution of X s .
Thus, we have that for any s, t ∈ [0 , T ] with t > s and x t ∈ M
[Equation]
Hence, for any s, t ∈ [0 , T ] with t > s we have that ∇ log p t = arg min { (cid:96) t | s ( s t ) :
s t ∈ L 2 ( P t ) } , where (cid:96) t | s ( s t )
= (cid:82) M 2 (cid:107)∇ x log p t | s ( x t | x s ) − s t ( x t ) {+(cid:107) 
+}
d P s,t ( x s , x t ) , which is referred as the Denois- ing Score Matching (DSM) loss.
It can also be written in an implicit fashion.
Proposition 3.
Let t
, s ∈ (0 , T ] with t > s .
Then, for any s t ∈ C ∞ ( M ) , (cid:96) t | s ( s t ) = 2 {+(cid:96) im t ( s t ) + (cid:82) M 2 (cid:107)∇ x t log p t | s ( x t | x s ) (cid:107) 
+}
d P s,t ( x s , x t ) , where (cid:96) im t ( s t ) =
(cid:82) M { 12 (cid:107) s t ( x t ) (cid:107) 2 +div( s t )( x t ) }
d P t ( x t ) .
The proof is postponed to App.
J. For any t ∈ (0 , T ] the minimizers of the loss (cid:96)
im t on X ( M ) (where
X ( M ) is the set of vector ﬁelds on M ) are the same as the ones for (cid:96) t | s .
The loss (cid:96) im t is referred to as the implicit score matching (ISM) loss (Hyvärinen, 2005).
These losses are direct analogous to the versions typically used in Euclidean space.
In the case where we have access to {∇ log p t | s : T ≤ t > s ≥ 0 } , the forward noising process transition kernels, or an approximation of this family, then we can use the DSM loss to learn
Require: ε, T, N, { X m 0 } Mm =1 , loss , s , θ 0 , N iter , p ref , P 1: ///
TRAINING /// 2: for n ∈ { 0 , . .
, N iter − 1 } do 3:
X 0 ∼ (1 /M )
(cid:80) Mm =1 δ X m(cid:46)
Random mini-batch from dataset 4: t ∼ U ([ ε, T ]) (cid:46) Uniform sampling between ε and T 5:
X t = GRW ( t, N, X 0 , b, Id , P) (cid:46) Approximate forward diffusion with Algorithm 1 6: (cid:96) ( θ n ) =
(cid:96) t ( T, N, X 0 , X t , loss , s θ n ) (cid:46) Compute score matching loss from Table 2 7: θ n +
=
optimizer_update ( θ n , (cid:96) ( θ n ))
(cid:46) ADAM optimizer step 8: end for 9: θ [-⋆ = θ N epoch 10: /// SAMPLING /// 
-](cid:63)
= θ N epoch 10: /// SAMPLING /// 11: Y 0 ∼ p ref (cid:46) Sample from uniform distribution 12: b (cid:63)θ ( t, x ) =
s θ (cid:63) ( T − t, x ) for any t ∈ [0 , T ] , x ∈ M (cid:46) Reverse process drift 13: { Y k } Nk =
=
GRW ( T, N, Y 0 , b θ (cid:63) , Id , P) (cid:46) Approximate reverse diffusion with Algorithm14: return θ (cid:63) , { Y k } Nk =
[Table]
d X t = − 12 X t d t + d B M t − 12 ∇ X t U ( X t )
d t + d B M t d B M t
Easy-to-sample distribution
[Table]
Euler–Maruyama
Geodesic Random Walk (Algorithm 1) { s t ∈ X
( M ) : t ∈ [0 , t ] } .
If this is not the case then we turn to (cid:96) im t .
Note that (cid:96) im t requires the computation of a divergence term which requires d Jacobian-vector calls.
In high dimension, a stochastic estimator is necessary (Hutchinson, 1989).
Following Song and Ermon (2020);
Nichol and
Dhariwal (2021) the loss can be weighted with a term λ t > 0 .
Parametric family of vector ﬁelds.
We approximate ( ∇ log p t ) t ∈ [0 ,T ] by a family of functions { s θ } θ ∈ Θ where Θ is a set of parameters and s θ : [0 , T ] → X ( M ) .
In a Euclidean space, vector ﬁelds are simply functions s θ : R d → R d .
In manifolds, although for any x ∈ M , T x M ∼ = R d , there does not necessarily exist a set of d smooth vector ﬁelds { E i } di =1 such that span (cid:0) { E i ( x ) } di =1 (cid:1) =
T x M (Chapter 8, page 179, Lee, 2006) 9 .
Fortunately, one can rely on a larger set of smooth vector ﬁelds { E i ( x ) } ni =1 with n > d that does span the tangent bundle.
Then it sufﬁces to construct a neural network s θ : [0 , T ] × M → R n to parametrize the score network as s θ ( t, x ) = (cid:80) ni
=1 s iθ ( t, x ) E i ( x ) .
See App.
E for a discussion on the different choices of generating sets { E i ( x ) } ni
=1 .
Combining this parameterization with the score matching losses, the time-reversal formula of Theorem 1 and the sampling of forward and backward processes described in Sec.
3.3, we deﬁne our
RGSM algorithm in Algorithm 2.
This algorithm can also beneﬁt from a predictor-corrector scheme as in (Song et al., 2021), see App.
G.
RSGMs on compact manifolds
Assuming compactness of the manifold M , we can leverage a number of special properties to implement a speciﬁc case of our algorithm.
In particular we beneﬁt from the fact that on compact manifolds we have a proper uniform distribution over the manifold, and have access to a variety of approximations of the heat kernel.
As highlighted in Sec. 3.1, in the compact setting we use Brownian motion as the noising SDE, which targets the uniform distribution as the stationary distribution.
Table 1 highlights the main differences between RSGMs on compact manifolds, generic manifolds and Euclidean score-based models.
Heat kernel on compact Riemannian manifolds.
For any x 0 ∈ M and t ≥ s ≥ 0 , the heat kernel p t | s ( ·| x s ) is deﬁned as the density of B M t w.r.t.
the uniform measure on the manifold.
Contrary to the Gaussian transition density of the OU process (or the Brownian motion) in the
Euclidean setting, it is typically only available as an inﬁnite series.
In order to circumvent this issue we consider two techniques: i) a truncation approach, ii) a Taylor expansion around t = 0 called a
Varadhan asymptotics.
First, we recall that in the case of compact manifolds the heat kernel is given by the Sturm–Liouville decomposition (Chavel, 1984) given for any t > 0 and x 0 , x t ∈ M by
[Equation]
(6) where the convergence occurs in L 2 ( p ref ⊗ p ref ) , ( λ j ) j ∈ N and ( φ j ) j ∈ N are the eigenvalues, respectively the eigenvectors, of − ∆ M , the Laplace-Beltrami operator in the manifold, in L 2 ( p ref ) (see
Saloff-Coste, 1994, Section 2).
When the eigenvalues and eigenvectors are known, we rely on an approximation of the logarithmic gradient of p t | 0 by truncating the sum in Eq.
(S8) with J ∈ N terms to obtain for any t > 0 and x 0 , x t ∈ M [Equation]
[Equation]
Under regularity conditions on M it can be shown that for any x, y ∈ M and t ≥ 0 , lim
[Equation]
( x t | x 0 ) (see Jones et al., 2008, Lemma 1).
In the case of the d -dimensional torus or sphere the eigenvalues and eigenvectors are computable (see Saloff-Coste, 1994, Section 2) and we can apply this method to approximate p t | 0 for any t > 0 , see App.
F
When the eigenvalues and eigenvectors are unknown or not tractable, we can still derive an approximation of the heat kernel for small times t .
Using Varadhan’s asymptotics—see Bismut (1984,
Theorem 3.8) or Chen et al.
(2021, Theorem 2.1)—for any x, y ∈ M with y / ∈ Cut( x ) (where
Cut( x ) is the cut-locus of x in M (see Lee, 2018, Chapter 10))
we have that
[Equation]
Using the previously deﬁned score-matching losses and the approximations to the heat kernel above, we highlight three methods to compute ∇ log p t in Table 2.
[Table]
Convergence results in the compact setting We now provide a theoretical analysis of RSGM under the assumption that M is compact.
The following result ensures that RSGM generates samples whose distribution is close to the data distribution p 0 .
Let us denote { Y k } n ∈{ 0 ,...,N } the sequence generated by Algorithm 2.
This result relies on the following assumption
[Table]
which is satisﬁed for a large class of manifolds M such as the d -dimensional sphere and torus, compact matrix groups and products of these manifolds.
A1.
There exist C, α > 0 such that for any t ∈ (0 , 1] and x ∈ M , p t | 0 ( x | x ) ≤ Ct − α/ 2 , where p t | 0 ( ·| x 0 ) is the density of the heat kernel, i.e. the density of B M t with initial condition x 0 10 .
Theorem 4.
Assume A 1 , that p 0 is smooth and positive and that there exists M ≥ 0 such that for any t ∈ [0 , T ] and x ∈ M , (cid:107) s θ (cid:63) ( t, x ) − ∇ log p t ( x ) (cid:107) ≤ M , with s θ (cid:63) ∈ C([0 , T ] , X ( M )) .
Then if T > 1 / 2 , there exists C ≥ 0 independent on T such that
[Equation]
where W 1 is the Wasserstein distance of order one on the probability measures on M .
The proof is postponed to App.
I.
In particular, for any ε > 0 , choosing T > 0 large enough, M small enough (which can be achieved using the universal property of neural networks) and γ small enough, we get that W 1 ( L ( Y N ) , p 0 ) ≤ ε .
This result might seem weaker than the result obtained for Moser ﬂows in (Rozen et al., 2021, Theorem 3), but we emphasize that our bound takes into account the time-discretization contrary to Rozen et al.
(2021) which considers the continuous-time ﬂow.
If we consider the time-reversed continuous-time SDE then we recover a bound in total variation distance, see App.
I. Note that the upper bound M encompasses both the bias introduced by the use of a neural network and the bias introduced by the use of an approximation of the score.
Related work
In this section we discuss previous work on parametrizing family of distributions for manifold-valued data.
Here, the manifold structure is considered to be prescribed, in contrast with methods that jointly learn the manifold structure and density (e.g. Brehmer and Cranmer, 2020; Caterini et al., 2021).
Parametric family of distributions.
The various parametric families of manifold-valued distributions that have been proposed can be categorized into three main approaches (Navarro et al., 2017): wrapping, projecting and conditioning.
Wrapped distributions consider a parametric distribution on R n that is pushed-forward along a surjective map ψ : R n → M .
Projected distributions are deﬁned by marginalizing out some distribution along the normal bundle of
M
Conditioning distributions encompass the von Mises-Fisher and Kent distributions (
Considering
Push-forward of Euclidean normalizing ﬂows.
More recently, approaches leveraging the ﬂexibility of normalizing ﬂows (Papamakarios et al., 2019) have been proposed.
Following the wrapping method described above, these methods parametrize a normalizing ﬂow in R n before being pushed along an invertible map ψ : R n → M .
However, to globally represent the manifold, the map ψ needs to be a homeomorphism, which can only happen if M is topologically equivalent to R n , hence limiting the scope of that approach.
One natural choice for this map is the exponential map exp x : T x M ∼ = R d .
This approach has been taken, for instance, by Falorsi et al.
and Bose et al.
(2020), respectively parametrizing distributions on Lie groups and hyperbolic space.
Neural ODE on manifolds.
To avoid artifacts or numerical instabilities due to the manifold embedding, another line of work uses tools from Riemannian geometry to deﬁne ﬂows directly on the manifold of interest (Falorsi and Forré, 2020; Mathieu and Nickel, 2020; Falorsi, 2021).
Since these methods do not require a speciﬁc embedding mapping, they are referred as Riemannian .
They extend continuous normalizing ﬂows (CNFs)
(Grathwohl et al., 2019) to the manifold setting, by implicity parametrizing ﬂows as solutions of Ordinary Differential Equations (ODEs).
As such, the parametric ﬂow is a continuous function of time.
This approach has recently been extended by Rozen et al.
(2021) introducing Moser ﬂows, whose main appeal being that it circumvents the need to solve
an
ODE in the training process.
We refer to App.
K for an in-depth discussion on the links between our work and Moser ﬂows.
Optimal transport on manifolds.
Another line of work has developed ﬂows on manifolds using tools from optimal transport.
Sei (2013) introduced a ﬂow that is given by f θ : x (cid:55)→ exp x ( ∇ ψ c θ ) with ψ cθ a c -convex function and c = d 2 M the squared geodesic distance.
This approach is motivated by the fact that the optimal transport map takes such an expression (Ambrosio, 2003).
These methods operate directly on the manifold, similarly to CNFs, yet in contrast they are discrete in time.
The beneﬁts of this approach depend on the speciﬁc choice of parametric family of c -convex functions (Rezende and Racanière, 2021;
Cohen et al., 2021), trading-off expressivity with scalability.
[Table]
Experiments
In this section we benchmark the empirical performance of RSGMs along with other manifold-valued methods introduced in Sec.
We also compare to a ‘Stereographic‘ score-based model, introduced in App.
[-AD is also partially supported by the EPSRC grant EP/R034710/
-]
N. First, we assess their modelling capacity on earth and climate science spherical data.
Then, we test the methods scalability with respect to manifold dimensions with a synthetic experiment on the torus T d .
Eventually, we evaluate the models’ regularity and time complexity with a synthetic
SO 3 ( R ) target.
Experimental details are provided in App.
O
[Table]
6.1 Earth and climate science datasets on the sphere
We start by evaluating RSGMs on a collection of simple datasets, each containing an empirical distribution of occurrences of earth and climate science events on the surface of the earth.
These events are: volcanic eruptions (NGDC/WDS), earthquakes (NGDC/WDS), ﬂoods (Brakenridge, 2017) and wild ﬁres (EOSDIS, 2020).
We compare to previous baseline methods: Riemannian
Continuous Normalizing Flows (Mathieu and Nickel, 2020), Moser Flows (Rozen et al., 2021) and a mixture of Kent distributions (Peel et al., 2001).
Additionally, we consider a standard SGM on the 2D plane followed by the inverse stereographic projection which induces a density on the sphere (Gemici et al., 2016).
We evaluate the log-likelihood of each model, extending to the manifold setting the likelihood computation techniques of SGMs, see App.
D. We observe from Table 4, that all benchmarked methods have comparable performance when evaluated on these simple tasks with RSGM performing marginally better on most datasets.
However, we empirically notice that
Moser ﬂows are slow to train and additionally that both Moser ﬂows and stereographic SGMs are computationally expensive to evaluate.
6.2 Synthetic data on tori
We now move to another manifold, that is the torus T d = S 1 × · · · × S 1 , so as to assess the scalability of the different methods with respect to the dimension d .
We consider a wrapped Gaussian target distribution on T d with a random mean and unit variance.
[-We say that X ∈ X ( γ ) is parallel to γ if for any t ∈ [0 , 1] , ∇ ˙ γ X ( t ) = 0 . In local coordinates, let X ∈ X ( γ ) be given for any t ∈ [0 , 1] by X = (cid:80) di =
-]Moser ﬂows’ (Rozen et al., 2021) loss involves a regularization term which involves an integral over the manifold, approximated by a
Monte Carlo (MC) estimator with uniform proposal.
This term regularizes Moser ﬂows towards probability measures,
i
[Table]
e. with unit volume.
We thus expect Moser ﬂows to fail in high-dimension as the number of samples K required for the MC estimator to be accurate will grows as O (e d ) , and the
[Table]
[Figure]
memory required to compute this estimator grows either in O ( Kd ) for exact divergences or O ( K ) for approximated divergences (see Table 3).
In Fig.
3, we observe that RSGMs are able to ﬁt well the [-Stratanovitch integral as follows for any t ≥ [Equation] In particular, denoting ( Z 1 t ) t ≥ 0 and ( Z 2 t ) t ≥ 
-]target distribution even in high dimension, with a linear or constant computational cost—depending on the [-processes such that for any t ≥ 0 , Z 1 t = (cid:82) t 
-]divergence estimator.
[-d Y s and Z 2 t = (cid:82) t 
-]
In contrast,
Moser ﬂows scale poorly with the dimension, to the extent that we are unable to train them for d ≥ 10 .
This is due to the combination of the complexity which grows linearly with [-X 
 a M -valued random variable if for all f ∈ C ∞ ( M )-]both the dimension d and the number of MC samples K , which itself ought to grow exponentially with d —as discussed in the previous paragraph.
This is illustrated by the gap between the ‘Moser’ and ‘ODE’ likelihoods which increases with the manifold dimension (see left Fig. 3).
6.3 Synthetic data on the Special Orthogonal group
In order to demonstrate the broad range of applicability of our model we now turn to the task of density estimation on the special orthogonal group SO d ( R ) =
{ Q ∈ M d ( R ) : QQ (cid:62) = Id , det(Q) = 1 } .
We consider the synthetic dataset consisting of samples in SO 3 ( R ) from a mixture of wrapped normal distributions with M components.
We compare RSGMs against Moser ﬂows and a wrapped-exponential baseline inspired by Falorsi et al.
(2019)—where we parametrize a standard Euclidean SGM on so (3) that is then pushed-forward on SO 3 ( R ) .
[-Then, we recover the M -valued Brownian motion ( B M t ) t ≥ 0 upon letting ( B M t ) t ≥ 
-]
RSGMs are trained using the (cid:96) t | 0 (DSM) loss with the Varadhan approximation (see
[-In addition for any x ∈ M , we denote { P i ( x ) } pi =
-]
Table 2).
[-We denote { X k , X i,j } 1 ≤ i,j,k ≤ d such that for any u ∈ F M , { X k ( u ) , X i,j ( u ) } 1 ≤ i,j,k ≤ d is a basis of T u F M , Using properties of the horizontal lift, see (Hsu, 2002, Chapter 2), we get that ( U t ) t ≥ 
-]
[-Let ( M t ) t ≥ 
-]
From Table 5 we observe that, RSGMs perform consistently, whether the target distribution has few or many mixture components M , as opposed to Exp-wrapped SGMs and Moser ﬂows which only perform well in some range of M .
Similarly to Sec.
[-For any γ > 0 , we define ( X γt ) t ≥ 
-]
[-such that X γ 
-]
6.2, we ﬁnd Moser ﬂows to be much slower to train due to the large number of Monte Carlo samples needed in the reguralizer ( K = 10 4 ).
[-P ( X γn ∈ A ) , with X γ 
-]
We also note from Table 5 that the number of score network evaluations (NFE) is signiﬁcantly lower for
[-Since we have assumed that M is compact we have that there exist (Φ k ) k ∈ N an orthonormal basis of − ∆ M in L 2 ( p ref ) , ( λ k ) k ∈ N such that for any i, j ∈ N , i ≤ j , λ i ≤ λ j and λ 0 = 0 , Φ 
-]
RSGMs, and is particularly detrimental for Moser ﬂows ( (cid:29) 10 3 )
[Figure]
(a) Histograms of SO 3 ( R ) samples from a target mixture distribution with M = 4 components, represented via their Euler angles.
[Figure]
(b) RSGMs are much more robust to hyperparameters than Exp-wrapped SGMs.
The diffusion coefﬁcient is given by σ ( t,
[-p ref and log q 
-]
[-= 
-]
[Figure]
Method
M
[Figure]
[Table]
=
M =
[-, d } , div( E i ) = 
-]
[Table]
[Figure]
[Table]
[-The case of the torus Let { b i } di =1 be a basis of R d . We consider the associated lattice on R d , i.e. Γ = { (cid:80) d i =
 α i b i : { α i } d i =1 ∈ Z d } .-][Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
6.4 Synthetic data on hyperbolic space
Finally we demonstrate RSGM on a non-compact manifold: the two dimensional hyperbolic space H 2 , which is deﬁned as the simply connected space of constant negative curvature.
We use Langevin dynamics as the noising process (Eq. (3)) and target a wrapped Gaussian as the invariant distribution.
We again consider a synthetic dataset of samples from a mixture of exp-wrapped normal distribution.
From Fig.
5, we can qualitatively see that both score-based models are able to ﬁt the target distribution.
[Figure]
Discussion and limitations
In this paper we introduced Riemannian Score-Based Generative Models (RSGMs), a class of deep generative models that represent target densities supported on manifolds, as the time-reversal of
Langevin dynamics.
[-X γk +1 / 
-]
The main beneﬁts of our method stems from its scalability to high dimensions, its applicability to a broad class of manifolds due to the diversity of available loss [-from Table 2 7: θ n +
-]functions, its robustness and crucially its capacity to model complex datasets.
[-SAMPLING /// 
-]
[-s θ ⋆ ( T − t, x ) for any t ∈ [0 , T ] , x ∈ M ▷ Reverse process drift 13: { Y k } Nk =
-]
We also provided theoretical guarantees on the convergence of RSGMs.
In future work, we would like explore more generic classes of manifolds, such a ones with a boundary, along with alternative noising processes.
Another promising extension concerns stochastic control on manifolds and more precisely, deriving efﬁcient algorithms to solve
Schrödinger bridges in the same spirit as De Bortoli et al.
(2021) on Euclidean state spaces.
Acknowledgments
We are grateful to the anonymous reviewers for their insightful comments and the for fruitful discussion more generally.
We thank the geomstat team Miolane et al.
(2020a).
and Engineering and Physical Research Council (EPSRC) under grant EP/R013616/1.
This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research
Initiative.
AD is also partially supported by {+the EPSRC grant EP/R034710/
+}
[-( M ) × M → R such that for any t ∈ [0 , T ] , u ∈ C 
-]
[-Note that since for any x ∈ R p with ∥ x ∥ ≥ R + 1 , φ ( x ) = 
-]
CoSines
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] Our main contribution is the extension of diffusion models on Riemannian manifolds.
(b) Did you describe the limitations of your work?
[-First note that using that P 
-]
[-π 
-]
[Yes] See Sec.
(c) Did you discuss any potential negative societal impacts of your work?
[No] The work presented in this paper focuses on the learning of score-based models on manifold.
[-. , N } , ¯ Y k 
-]
[-= ˆ Y 
-]
[-, N − 1 } , E 
-]
[-Proposition S21, X 
-]
[-Note that we have { ¯ Y Njγ, 0 } Nj =
-]
We do not foresee any immediate societal impact of such a study.
[-M ) 
-]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[-For any x ∈ ¯B( x 0 , κ ( M/ 2)) , we have | ˆ p T ( x ) − p 
-]
[Yes] We have read the ethics review guidelines and our paper conforms to them.
[-. , 2 ℓ } , ℓ ∈ N } such that X 
-]
[-X 01 = exp X 
-]
[-Γ X 01 X 00 E 
-]
[-In addition, we have that for any ℓ ∈ N with ℓ ≥ 1 , X ℓ 
-]
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[-ℓk X ℓ +
-]
[-Therefore, combining this result and that t − (2 k + 1) γ ℓ +1 + γ ℓ +
-]
[-there exists C 1 ≥ 
-]
[-In particular, since E ℓk = Γ X ℓk X ℓ +12 k E ℓ +12 k using (S20) and t − (2 k + 1) γ ℓ +1 + γ ℓ +
-]
[-ℓk X ℓ +
-]
[-there exists C 2 ≥ 
-]
[-= (1 − t )ˆ p 0 + t ˆ p 1 , with ˆ p 
-]
[-Note that u θ also enables density estimation using that ˆ p 
-]
[-Let ˆ p 
-]
[-(1 − t )ˆ p 0 + t ˆ p 1 for any t ∈ [0 , 1] , whereas RSGM performs the interpolation in sample space , i.e. p t = (cid:82) M p 
-]
[-( y ) p t | 
-]
[Yes] Yes, see A 1.
[-d p ref ( x ) ≤ ε and that there exists x 0 ∈ M such that | p T ( x ) − 1 | > κε with κ > 0 and let T ≥ T 0 with T 
-]
[-Let T 0 ≥ 0 such that for any T ≥ T 
-]
[-Assume A 
-]
(b) Did you include complete proofs of all theoretical results?
[Yes] Yes, proofs are postponed to the supplementary material.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Experimental details are given in App.
O. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Experimental details are given in App.
O. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[-Regarding the weighting function, for DSM loss ℓ t | 
-]
[Yes] Error bars are reported for each experiment.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Experimental details are given in App.
O. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Sec.
(b) Did you mention the license of the assets?
[Yes] See App.
O. (c) Did you include any new assets either in the supplemental material or as a URL?
[No] Not applicable.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[No] Not applicable.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[No] Not applicable.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[No] Not applicable.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[No] Not applicable.
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[No] Not applicable.