Riemannian Score-Based Generative Modelling
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Score-based generative models (SGMs) are a powerful class of generative modelsthat 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” processdeﬁned by approximating the time-reversal of the diffusion.
Existing SGMs assumethat 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 oftennaturally 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 tocompact Riemannian manifolds.
We demonstrate our approach on a variety ofmanifolds, and in particular with earth and climate science spherical data.
Introduction
Score-based Generative Models (SGMs) (Song and Ermon, 2019; Song et al., 2021b;
Ho et al.
,2020; Dhariwal and Nichol, 2021) formulate generative modelling as a denoising process.
Noiseis incrementally added to data using a diffusion process until it gets approximately Gaussian.
Thegenerative model is then obtained by simulating an approximation of the corresponding time-reversalprocess, which progressively denoises a Gaussian sample to obtain a data sample.
This process isalso 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 ﬂatgeometry.
However, in a large number of scientiﬁc domains the distributions of interest are supportedon 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 compact Riemannian manifolds which incorporate the geometry of the data by deﬁningthe forward diffusion process directly on the Riemannian manifold, inducing a manifold-valuedreverse process.
This requires constructing a noising process on the manifold that converges to aneasy-to-sample reference distribution.
We focus on the Brownian motion which converges withgeometric rate to the uniform distribution on compact Riemannian manifolds (Saloff-Coste, 1994).
We establish that, as in the Euclidean case, the corresponding time-reversal process is also a diffusionwhose drift includes the Stein score which is intractable but can similarly be estimated via scorematching.
Methodological extensions are required as the transition kernel of the Brownian motioncannot be sampled exactly and is typically only available as an inﬁnite sum through the Sturm– Liouville decomposition (Chavel, 1984) on compact Riemmanian manifolds.
To this end, we developnon-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 and demonstrateour approach on a range of manifolds and tasks, including modelling a number of natural disasteroccurrence datasets collected by Mathieu and Nickel (2020).
We show that RGSMs achieve betterperformance than recent baselines (Mathieu and Nickel, 2020; Rozen et al., 2021) and scale better tohigh-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 readersto Song et al. (2021b)
for a more detailed introduction.
We consider a forward noising process( X t ) t ≥ 0 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 datagives 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 52 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 fort ∈ [0 , T ] and in particular Y T ∼ p 0 .
Hence, if one could sample from ( Y t ) t ∈ [0 ,T ] then its ﬁnal 55 distribution would be the data distribution p 0 .
Unfortunately we cannot sample exactly from (2) as p Tand 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 geometricallytowards 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.
Itfollows 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 62 functions s where the expectation is over the joint distribution of X 0 , X t .
This result can be leveragedby 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–Maruyamadiscretization 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 beenproposed; 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., 2021b) instead of a simple Euler–Maruyama discretization drastically improves performance.
Finally, SGMs can also be derived through variationaland maximum likelihood techniques (Ho et al., 2020;
Huang et al., 2021; Song et al., 2021a).
Riemannian Score-based Generative Modelling
We now move to the manifold setting, and more specially assume that M is a compact and connected
Riemannian manifold.
As discussed in Section 1, three components are required to extend SGM: i) aforward noising process on M which converges to an easy-to-sample reference distribution, ii) atime-reversal formula on M which deﬁnes a backward generative process, iii) a method to efﬁcientlyapproximate the drift of the time-reversal process.
We address all these problems in what follows todesign RGSMs.
The key differences between SGM and RSGM are summarized in Table
3.1 Brownian motion and time-reversal on compact Riemannian manifolds
We use a Brownian motion ( B M t ) t ≥ 0 as the forward noising process on M since it converges geomet- 81 rically to the uniform distribution p ref under mild assumptions on M , see Urakawa (2006, Proposition
[Table]
For most manifolds of interest, one can either sample exactly from p ref or approximately withhigh accuracy.
We refer to Appendix B.5 for an introduction to Brownian motions on manifolds.
We now derive a time-reversal formula for the manifold-valued forward diffusion.
It is the [-tangent space T x M 5: W k +
-]Riemanniancounterpart to the Euclidean time-reversal formula— e.g. see Cattiaux et al.
(2021, Theorem 4.9).
Itshows 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.
The proof relies on an extension of Cattiaux et al.
(2021, Theorem 4.9) to the Riemannian manifold case and is postponed to Appendix G.
This result and the other theoretical results presented in this paper rely on the following assumption,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 , wherep t | 0 ( ·| x 0 ) is the density of the heat kernel, i.e. the density of B M t with initial condition x 0 1 .
Theorem 1 (Time-reversed diffusion) .
Assume A 1 . Let T ≥ 0 and ( B M t ) t ≥ 0 be a Brownianmotion on M such that B M 0 has distribution p ref .
Let ( X t ) t ∈ [0 ,T ] 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 Qis the distribution of ( B M t ) t ∈ [0 ,T ] and P the distribution of ( X t ) t ∈ [0 ,T ] .
In addition, assume that 98 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
[Equation]
3.2 Approximate sampling of diffusions
Obtaining samples from a Brownian motion on a manifold is non-trivial in general.
If M isisometrically embedded into R p (with p ≥ d ) one can deﬁne ( B M t ) t ≥ 0 as a R p -valued process, see
Appendix B.5.
However, this approach is extrinsic and requires the knowledge of the projectionoperator.
Here we consider an intrisic approach based on Geodesic Random Walks (GRWs), see
Jørgensen (1975) for a review of their properties.
GRWs are not restricted to approximating the
Brownian motion on M but in fact 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 whichapproximates the diffusion ( X t ) t ≥ 0 deﬁned by [Equation]
[Equation]
This generalization will be key when sampling the backward diffusion process deﬁned in Theorem 1.
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 )
+ (1 / √ γ )( V n +1 − b ( X γn )) } ] 2 , where ( V n ) n ∈ N is a sequence of T M -valued random variables such that for any n ∈ N ,
113 E
[ V n +1 |F n ] =
b ( X γn ) 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 Equation (4) using GRWs; Kuwada Cheng et al
(2022) for quantitative error bounds in the time-homogeneous case and Appendix H.1 for the time-inhomogeneous case
Require: T, N, X γ 0 , b, σ, P 1: γ = T/N (cid:46) Step-size 2: for k ∈ { 0 , . . .
, N − 1 } do 3: ¯ Z k +1 ∼ N(0 , Id) (cid:46) Standard Gaussian ambient R p 4: Z k +1 =
P( X γk )
¯ Z k + (cid:46) tangent
T x M 5:
{+W k +
 = γb ( kγ, X γk ) + √ γσ ( kγ, X γk )+}
Z k + (cid:46) space 6:
X γk +1 =
exp X γk [ W k +1 ]
(cid:46)
M 7: end for 8: return { X γk
= 3.3 Score approximation on Riemannian manifolds
Heat kernel on compact Riemannian manifolds.
For any x 0 ∈ M and t ≥ s ≥ 0 , the heat 120 kernel p t | s ( ·| x s ) is deﬁned as the density of B M t w.r.t.
p ref when B M s = x s .
Contrary to the Gaussian 121 transition density of the OU process (or the Brownian motion) in the Euclidean setting, it is typicallyonly 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–Liouvilledecomposition (Chavel, 1984) given for any t > 0 and x 0 , x t ∈ M by [Equation] (5) where the convergence occurs in L 2 ( p ref ⊗ p ref ) , ( λ j ) j ∈ N and ( φ j ) j ∈ N are the eigenvalues, respectivelythe eigenvectors, of − ∆ M in L 2 ( p ref ) (see Saloff-Coste, 1994, Section 2). When the eigenvalues andeigenvectors are known, we rely on an approximation of the logarithmic gradient of p t | 0 by truncatingthe sum in Equation (S9) 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 , 131 lim
[Equation]
) (see Jones et al., 2008, Lemma 1).
In the case of 132 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 Appendix E.
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]
Score-matching and loss functions.
The reverse process from Equation (3) involves the Steinscore ∇ log p t which is unfortunately intractable.
To derive an approximation, we ﬁrst remark thatfor 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 ) .
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 ) [-▷ Compute score matching loss from Table 2 7: θ n +
 = optimizer_update ( θ n , ℓ-]− s t ( {+x t ) (cid:107) 
+}
[-SAMPLING /// 
-]
[-s θ ⋆ ( T − t, x ) for any t ∈ [0 , T ] , x ∈ M ▷ Reverse process drift 13: { Y k } Nk =
-]
d P s,t ( x s , x t ) , which is referred as the Denois- 144 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 )
+ 146 (cid:82)
M 2 (cid:107)∇ x t log p t | s ( x t | x s ) (cid:107) 2 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 Appendix I. For any t ∈ (0 , T ] the minimizers of the loss (cid:96) im t on ( M )
( M ) the of vector ﬁelds
M access the computation of a divergence term which requires d Jacobian-vector calls.
In high dimension, astochastic estimator is necessary (Hutchinson, 1989).
Following Song and Ermon (2020); Nicholand Dhariwal (2021) the loss can be weighted with a term λ t > 0 .
Using the results of the previousparagraph, we highlight three methods to compute ∇ log p t which are summarized in Table
[Table]
[Table]
[Table]
t ) ε (cid:3) (cid:55) (cid:55)
O (1) 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 , theredoes 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) 3 .
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 162 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 Appendix D for a discussion on the 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 Section 3.1, we deﬁne our
RGSM algorithm in Algorithm 2.
This algorithm can also beneﬁt from a predictor-corrector schemeas in (Song et al., 2021b), see Appendix F.
The computational complexity of RSGM and some othermethods is presented in Table 3.
[Table]
3.4 Convergence results
We now provide a theoretical analysis of RSGM.
The following result ensures that RSGM generatessamples whose distribution is close to the data distribution p 0 .
Let us denote { Y k } n ∈{ 0 ,...,N } the 172 sequence generated by Algorithm 2.
Theorem 4.
Assume A 1 , that p 0 is smooth and positive and that there exists M ≥
such that forany 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 175 T > 1 / 2 , there exists C ≥ 0 independent on T such that
[Equation]
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 , 0 , 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: θ (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 = where W 1 is the Wasserstein distance of order one on the probability measures on M .
The proof is postponed to Appendix H. In particular, for any ε > 0 , choosing T > 0 large enough, Msmall enough (which can be achieved using the universal property of neural networks) and γ smallenough, 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 accountthe time-discretization contrary to Rozen et al.
(2021) which considers the continuous-time ﬂow.
If weconsider the time-reversed continuous-time SDE then we recover a bound in total variation distance,see Appendix H. One key element of our proof is the use of recent results providing quantitativebounds between GRW and the associated diffusion established by Cheng et al.
Related work
In this section we discuss previous work on parametrizing family of distributions for manifold-valueddata.
Here, the manifold structure is considered to be prescribed, in contrast with methods that jointlylearn 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 192 R n that is pushed-forward along a surjective map ψ : R n → M .
Projected distributions are deﬁned 193 by marginalizing out some distribution along the normal bundle of
M .
Conditioning distributions 194 encompass the von Mises-Fisher and Kent distributions (Fisher, 1953; Kent, 1982).
Consideringmixtures of these distributions is key to increase ﬂexibility (Peel et al., 2001; Mardia et al., 2008).
Push-forward of Euclidean normalizing ﬂows.
More recently, approaches leveraging the ﬂexibilityof normalizing ﬂows (Papamakarios et al., 2019) have been proposed.
Following the wrapping methoddescribed above, these methods parametrize a normalizing ﬂow in R n before being pushed along aninvertible map ψ : R n → M .
However, to globally represent the manifold, the map ψ needs to be ahomeomorphism, which can only happen if M is topologically equivalent to R n , hence limiting thescope of that approach.
One natural choice for this map is the exponential [Equation]
This approach has been taken, for instance, by Falorsi et al.
and Bose et al.
(2020), respectivelyparametrizing 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 themanifold of interest (Falorsi and Forré, 2020; Mathieu and Nickel, 2020; Falorsi, 2021).
Since thesemethods do not require a speciﬁc embedding mapping, they are referred as Riemannian .
They extendcontinuous normalizing ﬂows (CNFs) (Grathwohl et al., 2019) to the manifold setting, by implicityparametrizing ﬂ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
[Table]
ODE in the training process.
We refer to Appendix J for an in-depth discussion on the links betweenour work and Moser ﬂows.
Optimal transport on manifolds.
Another line of work has developed ﬂows on manifolds usingtools 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 bythe fact that the optimal transport map takes such an expression (Ambrosio, 2003).
These methodsoperate directly on the manifold, similarly to CNFs, yet in contrast they are discrete in time.
Thebeneﬁ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.
In thecase of tori and spheres, Rezende et al.
(2020) introduced discrete Riemannian ﬂows based on Möbiustransformations and spherical splines.
Experiments
In this section we benchmark the empirical performance of RSGMs along with other manifold-valuedmethods introduced in Section 4.
First, we assess their modelling capacity on earth and climatescience spherical data.
Then, we test the methods scalability with respect to manifold dimensionswith a synthetic experiment on the torus T d .
Eventually, we evaluate the models’ regularity and timecomplexity with a synthetic SO 3 ( R ) target.
Experimental details are provided in Appendix M.
5.1 Earth and climate science datasets on the sphere
We start by evaluating RSGMs on a collection of simple datasets, each containing an empiricaldistribution of occurrences of earth and climate science events on the surface of the earth.
Theseevents 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 amixture of Kent distributions (Peel et al., 2001).
Additionally, we consider a standard SGM on the2D 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 manifoldsetting the likelihood computation techniques of SGMs, see Appendix C.
We observe from Table 4,that all benchmarked methods have comparable performance when evaluated on these simple taskswith 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 arecomputationally expensive to evaluate.
These observations motivate further empirical investigationsthat are conducted in Sections 5.2 and 5.3.
5.2 Synthetic data on torii
We now move to another manifold, that is the torus T d = S 1 × · · · × S 1 , so as to assess the scalabilityof the different methods with respect to the dimension d .
We consider a wrapped Gaussian targetdistribution on T d with a random mean and unit variance.
Moser ﬂows’ (Rozen et al., 2021) lossinvolves 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 towardsprobability measures, i.e. with unit volume.
We thus expect Moser ﬂows to fail in high-dimension as [Figure] [Figure] S t e p s p e r s ec
[Figure]
nd
[Figure]
i t s p
e r d i m e n s i o n
[Figure]
the number of samples K required for the MC estimator to be accurate will grows as O (e d ) , and thememory required to compute this estimator grows either in O ( Kd ) for exact divergences or O ( K )for approximated divergences (see Table 3).
In Figure 2, we observe that RSGMs are able to ﬁt well the target distribution even in high dimension,with a linear or constant computational cost—depending on the divergence estimator.
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 both the dimension d andthe number of MC samples K , which itself ought to grow exponentially with d —as discussed in theprevious paragraph.
This is illustrated by the gap between the ‘Moser’ and ‘ODE’ likelihoods whichincreases with the manifold dimension (see left Figure 2).
5.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 densityestimation on the special orthogonal group SO d ( R ) =
{ Q ∈ M d ( R ) : QQ (cid:62) = Id , det(Q) = 1 } .
264 Applications with orthogonal constraints span various ﬁelds, such as molecular sciences (Ganeaet al., 2022), computer vision (Prokudin et al., 2018) and medical science (Hou et al., 2018).
Weconsider the synthetic dataset consisting of samples in SO 3 ( R ) from the mixture distribution withdensity p (Q) = 1 M (cid:80) Mm =1 N W (Q | Q m , σ 2 m ) with M ∈ N , where N W (Q m , σ 2 m ) are wrapped Normaldistributions with expectation Q m and variance σ 2 m . (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, [Figure] Method
[Equation]
=
[Equation]
[Equation]
[Table]
[Figure]
[Table]
[Figure]
[Table]
We compare RSGMs against Moser ﬂows and a wrapped-exponential baseline inspired by Falorsiet al.
(2019)—where we parametrize a standard Euclidean SGM on so (3) that is then pushed-forwardon SO 3 ( R ) .
RSGMs are trained using the (cid:96) t | 0 (DSM) loss with the Varadhan approximation (see
Table 2).
From Table 5 we observe that, RSGMs perform consistently, whether the target distributionhas few or many mixture components M , as opposed to Exp-wrapped SGMs and Moser ﬂowswhich only perform well in some range of M .
Similarly to Section 5.2, we ﬁnd Moser ﬂows tobe much slower to train due to the large number of Monte Carlo samples needed in the reguralizer( K = 10 4 ).
We also note from Table 5 that the number of score network evaluations (NFE) issigniﬁcantly lower for RSGMs, and is particularly detrimental for Moser ﬂows ( (cid:29) 10 3 ).
This hasthe direct consequence of dramatically slowing down any sampling and likelihood evaluation of
Moser ﬂows.
Additionally, since the NFE is a proxy of the learnt dynamics smoothness—deﬁned as 280 (cid:82)0 (cid:82) M (cid:107) s θ ( t, x t ) (cid:107) 2 d P t ( x t )d t , we can conclude that RSGMs induce ‘simpler’ models in that sense. We believe that this is due to Moser ﬂows performing the interpolation in density space whereas RSGMs perform the interpolation in sample space (see Appendix J). Additionally, we found RSGMsto be more robust w.r.t.
the choice of hyperparameters than the Exp-wrapped SGMs as shown in
Figure 3b, on top of converging faster.
Discussion and limitations
In this paper we introduced Riemannian Score-Based Generative Models (RSGMs), a class of deepgenerative models that represent target densities supported on compact manifolds, as the time-reversalof a Brownian motion.
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 functions, itsrobustness and crucially its capacity to model complex datasets.
We also provided theoreticalguarantees on the convergence of RSGMs.
We empirically demonstrated that our method outperformsprevious work both in terms of density estimation and scalability on a range of manifolds and datasets. An important future work direction, and a current limitation, is the manifold compactness assumption. Several important manifolds do not ﬁt into this category, such as the special linear group, symmetricpositive deﬁnite matrices or the hyperbolic space (Ungar, 2005).
Another promising extensionconcerns 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.
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?
[Yes] See Section 6.
(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.
We
do not foresee any immediate societal impact of such a study.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes] We have read the ethics review guidelines and our paper conforms to
them.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] Yes, see A 1.
(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 experi-
mental results (either in the supplemental material or as a URL)?
[Yes] Experimental
details are given in Appendix M.
(b) Did you specify all the training details (e g., data splits, hyperparameters,
how they
were chosen)?
[Yes] Experimental details are given in Appendix M.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[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., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Experimental details are given in
Appendix M.
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 Section 5.1.
(b) Did you mention the license of the assets?
[Yes] See Appendix M.
(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.