Riemannian Score-Based Generative Modelling
Authors: Valentin De Bortoli
Emile Mathieu
Michael John Hutchinson
James Thornton
Yee Whye Teh
Arnaud Doucet
Venue: NeurIPS 2022 Accept (2022)
Keywords: Diffusion models
generative modeling
Riemannian manifolds
score matching

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 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 \emph{Riemannian Score-based Generative Models} (RSGMs), a class of generative models extending SGMs to Riemannian manifolds.  We demonstrate our approach on a variety of compact manifolds, and in particular with earth and climate science spherical data.

Riemannian Score-Based Generative Modelling
Valentin De Bortoli", Émile Mathieu Michael Hutchinson"
James Thornton! Yee Whye Teh: Arnaud Doucet
Abstract
Score-based generative models (SGMs) are a powerful class generative models that exhibit remarkable empirical performance. Score-based generative modelling (SGM) consists of a "noising" stage, whereby diffusion used to gradually add Gaussian noise to data, and generative model, which entails "denoising" process defined by approximating the time-reversal of the diffusion. Existing SGMs assume that data is supported on Euclidean space, i.c. 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 bcsre-based Generative Models (RSGMs), class of generative models extending SGMs 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 (SGMs) also called diffusion models (Song and Ermon, 2019; Song al., 2021; Ho al., 2020; Dhariwal and Nichol, 2021) formulate generative modelling as denoising process. Noise is incrementally added to data using diffusion process until it becomes approximately Gaussian. The generative model 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 also a diffusion whose drift depends on the logarithmic gradients of the noised data densities, i.e. the Stein scores, estimated using neural network via score matching (Hyvärinen, 2005; Vincent, 2011).
SGMs have been primarily applied to data living on Euclidean spaces, i.c. manifolds with flat geometry. However, in large number of scientific domains the distributions of interest are supported on Riemannian manifolds. These include, to name a few, protein modelling (Shapovalov and Dunbrack Jr, 201 1), cell development (Klimovskaia 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).
al., 2020), image recognition (Lui, 2012),
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 defining the forward diffusion process directly on the Riemannian manifold, inducing a manifold-valued reverse process. This requires constructing noising process on the manifold that converges to an casy-to-sample reference distribution. We establish that, as in the Fuclidean case, the corresponding rime-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
equal contribution.
I Dept. of Computer Science ENS, CNRS, PSL University Paris, France. *Dept. of Statistics, University of Oxford, Oxford, UK.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

typically only available as an infinite sum through the Sturm- -iouvile decomposition (Chavel, 1984). To this end, we develop non 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 number of natural disaster occurrence datasets collected by Mathieu and Nickel (2020). We show that RGSMs achieve better performance than recent baselines (Mathicu and Nickel, 2020; Rozen et al., 2021) and scale better to high-dimensional manifolds.
2 Euclidean Score-based Generative Modelling
We recall here briefly the key concepts behind SGMs on the Euclidean space Rd and refer the readers to Song et al. (2021) for more detailed introduction. We consider forward noising process (X4 (20 defined by the following Stochastic Differential Equation (SDE)
dX;- dxtt + 2dd,, Xo Po:
where (B:)c20 is a ddddimensional Brownian motion and po is the data distribution. The available data gives an empirical approximation of po- The process (X,) is simply an Omstein-Ullhebeekk (OU) process which converges with geometric rate to N(0, Id). Under mild conditions on po. the time- reversed process (Y)):20~ (X)_)) also satisfies an SDE (Cattiaux al., 2021; Haussmann and Pardoux, 1986) given by
dY. (Y, + 2V p + 22d,, Yo Pr.
(2)
where Pr denotes the density of X, By construction, the law of Y is equal to the law of X, for t E [0, 7J and in particular Yr ~ Po- Hence, if one could sample from (Yo)re(o.7) then its final distribution would be the data distribution pn- Unfortunately we cannot sample exactly from (2) as pr and the scores (V log pt (x)))=(0,2)) are intractable. Hence SGMs rely on a few approximations. First, Pr is replaced by the reference distribution N(0,Id) Id) as we know that pr converges geometrically towards it. Second, the following denoising score matching identity is exploited to estimate the scores
Vz, (xrs = Jae V xr log Pr (zo) Poje (zo (2, ( )dzo.
where (z (z0) is the transition density of the OU process (1) which is available in closed-form. It follows directly that V log Pt is the minimizer of (s) (X), - -z, PogPpoo (X, [Io over functions s where the expectation is over the joint distribution of Xo, X,. This result can be leveraged by considering neural network sa [0, 7 Rd Rd trained by minimizing the loss function
(()) = 1-' Xe (t, d)t for some weighting function X > 0. Finally, an Euler- Maruyama discretization of (2) is performed using discretization step y such that T >N for NEN
Y=+1=X0 + +(Yn + 2T ny, Yn)) V2 Yo ~ 1d), Z,, itd I(D,
The above showcases the basics of SGMs but we highlight that many improvements have been proposed; see (e.g. Song and Ermon, 2020; Jolicoeur-Martinan et al., 2021: Dhariwal and Nichol, 2021). In particular, selecting an adaptive stepsize (Tm)NEN (Bao et al., 2022: Watson et al., 2021) and using a predictor-corredd scheme (Song al., 2021) instead of a simple Buler-Maruyame discretization drastically improves performance.
3 Riemannian Score-based Generative Modelling
We now move to the Riemannian manifold setting, and more specifically assume that M is a complete, orientable connected and boundaryless Riemannian manifold, endowed with 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 casy-to-sample reference distribution, ii) rime-reversal formula on M which defines a backward generative process, iii) a method for approximating samples of SDEs on manifolds, iv) method to efficiently approximate the drift of the time-reversal process. Notation gathered in App. B.
'Metrics g are sections of I'MM @ M' the rank 2 tensor bundle of the dual tangent space, i.e. smooth varying bilinear maps on T.M, verifying symmetry and positive cemi-definiteness.

---

X,- 3.1 Noising processes on manifolds
The first necessary component a suitable generic noising process on manifolds that will converge to convenient stationary distribution. A simple choice is to use Langevin dynamics described by
dX, = - Vx, dt + dB,
(3)
which admits the invariant density (w.r.t. the volume form) given by dpter (() x (-U(x) (Dur- mus, 2016, Section 2.4), where V is the Riemannian gradient'.
Two simple choices for (z)) present themselves. Firstly, setting =(z)= (272) where d. is the geodesic distance and μ E M is an arbitrary mean location, induces the drift Vx, (X,) = -expx (p)/722 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) (M /2)
(2>2) + log 0 exp,," (a)]]. In contrast to the Riemannian normal, sampling and evaluating the density of this distribution is easy.
One recovers the standard target distributions when .M- = I& and μ 0 since then the drift b(4, X.) 2 expx, (0) On compact manifolds, the invariant measure Vol has finite volume, thus natural choice is to target the uniform distribution which is given by Volm/IM In this case. Vx, UX+, 0 and the noising process is simply Brownian motion on M.
Uhheebbbc noising process (Song et al., 2021) for both of these
3.2 Time-reversal on Riemannian manifolds
In order to use these noising processes we prove the time-reversal formula for manifolds. a general- isation of the results in the Euclidean case, e.g. see Cattiaux al. (2021, Theorem 4.9). Consider an SDE of the form dX, b(X, )dl + dBB where B:M 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))1=(0.z] diffusion process then (X) is also a diffusion process w.r.t. the backward filtration whose coefficients can be computed, and are shown in Eq. (4). The proof relies on an extension of Cattiaux et al. (2021, Theorem 4.9) to the Riemannian manifold case and is postponed App. H.
Theorem 1 (Time-reversed diffusion). Let T 0 and (B,)) )120 be a Brownian motion M such that B&M has distribution the volume form Prer". Let (X1)uejo,1) be associated with the SDE dx, b(X,)dt + dBB.M Let (Yt) where Q is the distribution of (B)ar and P the distribution of (X1)eE(o.T) In addition, assume that P.- L(X.). the distribution of X,, admits a smooth positive density Pt w.r.t. Peet for any t E [0, T). Then, (Yo)cEjo,r) is associated with the SDE
(Xr-t)uE(o.T) and assume that KLP(P)Q < +0o,
dY. ---(())) + V Poggr-((Y.))dl ddB.
(4)
3.3 Approximate sampling of diffusions
Obtaining samples from SDEs on a manifold is non-trivial in general. If M isometrically embedded into I² (with 2 d) one can define (B:M)120 as a R -vllued process, see App. C.3. However, this approach 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 review of their properties. GRWs can approximate any well-behaved diffusion on M. Hence, we introduce GRWs in general framework and consider discrete-time process (X2) new which approximates the diffusion (X,) (2)) defined by
dx,- = btt, dt) + at, X.).BMM......
(5)
This generalisation is key to sampling the backward diffusion process defined in Theorem 1.
M *The (Riemannian) gradient V is defined s.t. for any ʃ M - R, æ E M, u E T.M. (VJ, v). = df()).
'exp, : .T.M + M denotes the exponential mapping on the manifold, see e.g. Lee (2013, Chapter 20). *1 denotes the absolute value the determinant, and 0 the Jacobian f. 'Note that in the case of non-compact manifold Pee is only a measure and not probability measure.
3

---

(a) single step of Geodesic Random Walk.
(b) Many steps yicld an ap- proximate Brownian mo- tion trajectory.
(c) The density of a single step of Gaussian Random Walk [Left] and the Brownian motion density [Right] agree well for small time steps.
Figure 1: Geodesic Random Walks can be used to approximate Brownian motion and more generally SDEs on manifolds. (a) At each step, tangential noise is sampled (red), which is added the drift term (not pictured). This langent vector is then pushed through the exponential map produce geodesics step the manifold (blue). (b) Iterating this procedure yield approximate sample paths from the process.
Algorithm 1 GRW (Geodesic Random Walk)
Require: N2 b,
y T/N
D Step-size
2: for k E {O..... . 1 do
3:
Zk+1 ~ N(0, Id)
> Sample a Gaussian in the tangent space of X (ho 2z) Compute the Faler-Maruyama step tangent space Move along the geodesic defined by WK+1 and X M
4:
WA+1 yb(ky, X2)
5: X expx; [W++1]
6: end for
7: return [X]]
Definition 2 (Geodesic Random Walk), Let X3) be M-valued random variable. For any O, define (XN)EN such that for any n E N, X = expxx (T{b(())) + VrVu+11], where nen sequence of T-M-alued random variables such that for any n E N, [F.]] = and E[Vn+1V, 1 oo" (X7), where Fn is the filtration generated by {X] 1x Ik=o We say that the M-valued process (X7)) is a Geodesic Random Walk.
Algorithm approximately simulates the diffusion (X:) defined in Eq. (5) using GRWs; see Kuwada (2012); Cheng et al. (2022) for quantitative error bounds in the time-homogeneous case and App. 1.2 for novel extentsion for the time-inhomogeneous case. Fig. 1 provides 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 v log Pt which unfortunately intractable. derive an approximation, we first remark that for any 8, E (0, 7) with / > and M, pt (xxx) = SM Ptpa (1p (dP, (ar,), where IP, L(X.), the distribution of X,. Thus, we have that for any s, é E [0, T] with > and E M
Vz, log JM V., s, log Pt a (ral (r.ear.))) ( Pa dr..).
Hence, for any s, E [0, with > s we have that Voog Pt = arg min (ainffelsoo (s)). : st E (F()))),
where ( (s:) Jyz | 2 log (2 (r.)) - s, (2:)) (x., xr), which is referred as the Denois- ing Score Matching (DSM) loss. It can also be written in an implicit fashion.
Proposition Let 1, 8 E (0, with > 8. Then, for any 8, C (M), 6 (s:) = 20( +
Jux I/, Log xr. )1 [l, (x., v.)., where (s.) ) = [ div(ss)
The proof is postponed to App. J. For any E (T, the minimizers of the loss fim on X(M) (where X(M) is the set of vector fields M)) are the same the for Ejja- The loss gim 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 VoogPel= T S <> 2 0}. the forward noising process transition kernels, or approximation of this family, then can the DSM loss to learn

---

Algorithm 2 RSGM (Riemannian Score-Based Generative Model)
M Require: e, C,T, N, (Xo0) / m==1 loss, 8, 0o, Nuer, Peets P
// TRAINING IIf
2: for n E {O....;
.....;,.u Nuer 1} do
3: Xo ~ (1/M) M 5x
Random mini-batch from dataset > Uniform sampling between e and T Approximate forward diffusion with Algorithm 1 b Compute score matching loss from Table b ADAM optimizer step
4:
t~ (())
5:
X, GR, N. Xo. P)
6:
((Om) ((T,, N, Xo; X, loss 50,)
7: 0n+1 = optimizer -pptimizer..uuaateeee updatee(n, ((,))
8: end for
9: 8" BNieer
10: I/ SAMPLING III
Yo ~ Phef
Sample from uniform distribution Reverse process drift
12: (b(1,)) = So- (T L, x) for any E [0. T). x M
13: (Yx) GRW Yo, by. Idl, P)
Approximate reverse difflusion with Algorithm I
14: return 0" 1Y
Table Differences between SGM Euclidean spaces and RSGM Riemannian manifolds.
Ingredient \\ Space
Euclidean
'Generic' Manifold
Compact Manifold
Forward process dX,
d,t + dBM
Vx, (X)))t dB
dB:
Easy-to-sample distribution
Gaussian
Wrapped Gaussian
Uniform
Time reversal
Cattiaux al. (2021)
Theorem 1
Sampling forward process Sampling backward process
Direct
Geodesic Random Walk (Algorithm 1)
Euler-Marupy Maruyama
Geodesic Random Walk (Algorithm 1)
(st E X(M) E [0, f]}. this is not the case then we turn to fimm Note that f. requires the computation of divergence term which requires d facobian-vector calls. In high dimension, stochastic estimator is necessary (Hutchinson, 1989). Following Song and Ermon (2020); Nichol and Dhariwal (2021) the loss be weighted with term A, 0.
Parametric family of vector fields. We approximate (V log by family of functions (so) Joe where θ is a set of parameters and Se [0, T] + (M)) In Euclidean space, vector fields are simply functions So IRd Rd, In manifolds, although for any x E M, T.MM Rd there does not necessarily exist set of d smooth vector fields {E, such that span (=)) T.M (Chapter 8, page 179, Lee, 2006)* 9. Fortunately, one can rely larger set of smooth vector fields {E) |n with >> d that does span the tangent bundle. Then it suffices to construct a neural network Sq [0, T M + g parametrize the score network as (t, x) Di 8 (z). See App. E for a discussion on the different choices of generating sets {E, (z)
Combining this parameterization with the score matching losses, the time reversal formula of The- orem and the sampling of forward and backward processes described in Sec. 3.3, we define our RGSM algorithm in Algorithm 2. This algorithm can also benefit from predictor- cortctor-coorecttor scheme as in (Song et al., 2021), ) see App. G.
4 RSGMs on compact manifolds
Assuming compactness of the manifold M, can leverage number of special properties to implement a specific case of our algorithm. In particular we benefit from the fact that on compact manifolds have proper uniform distribution over the manifold, and have access 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 the stationary distribution. Table I highlights the main differences between RSGMs on compact manifolds, generic manifolds and Euclidean score-based models.
variety of
Manifolds for which there exists such a global frame {E, (x) are referred as parallelizable, $ is well- known example of non-parallelizable manifold as per the Hairy ball theorem.
5

---

the sequence Heat kernel on compact Riemannian manifolds. For any E M and t 2 s 2 0, the heat kernel Pris(-1az, is defined as the density of B,M 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 infinite series. In order to circumvent this issue we consider two techniques: i) a truncation approach, ii) a Taylor expansion around é = 0 called Varadhan asymptoties. First, we recall that in the case of compact manifolds the heat kernel given by the Sturm-Liouville decomposition (Chavel, 1984) given for any > O and To. 24 E M by
P)o(=z|<z)== = DJEN B Ast H5 (20))
(6)
where the convergence occurs in (piet @ Pret). (Ay)JEN and (øy)JEN are the eigenvalues, respec- tively the eigenvectors, of M. the haplace-Beltanne operator in the manifold, in (pe()))) (see Saloff-Coste, 1994, Section 2). When the eigenvalues and eigenvectors are known, we rely on an approximation of the logarithmic gradient of Psjo by truncating the sum Eq. (S8) with E N terms to obtain for any & > 0 and zo, 2t E .M
Vx, log (1))(=))))) == SJ,e(Zo, x1) A Z2, log Eyzo Ast
(7)
Under regularity conditions on .M can be shown that for any Z, E M and t 2 0. lim (20, 224) V xr log (2 (mp) (see Jones 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 approximate Piju for any / 0, see App. F
When the eigenvalues and eigenvectors are unknown or not tractable, we can still derive an approx- imation of the heat kernel for small times t. Using Varadhan's 's asymptotics see Bismut (1984, Theorem 3.8) or Chen et al. (2021, Theorem 2.1)-for any x, E M with y & Cut(z) (where Cut(z) is the cut-locus of in .M (see Lee, 2018, Chapter 10)) we have that
lim VV,, v., log (=) exp=,! (zo).
Using the previously defined score- more-matching losses and the approximations to the heat kernel above, we highlight three methods compute V logg Pt in Table 2.
Table 2: Computational complexity of score matching losses w.r.t score network forward and backward passes. random variable T such that [B[c] = 0 and Elee" " ]= ld.
Requirements
Loss
Approximation
Loss function
Complexity
Prjo
expx.
None
1 (X.) log Po
((1)
Exjo (DSM)
Truncation (7)
1B [Xss Sx,, x.)
asymptotic expansion
O())
Varhadan (8)
E [[|s(X,) - expx (Xo)/
((1)
Esj= (DSM)
Varhadan (8)
E [X)s - exxxx (X. )/(() =]
X
0(1)
Deterministic
[[[]((()))) + (X.)
X
X
((d)
At (ISM)
Stochastic
E
[
) + {x())
0(1)
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 the data distribution Pu- Let denote (Yx) generated by Algorithm 2. This result relies on the following assumption, which is satisfied for a large class of manifolds M such the d-dimensional sphere and torus, compact matrix groups and products of these manifolds.
A1. There exist C, a > O such that for any 1 E (0, and x M. Pt (x]x)) S Ct-a/2, where ( (- is the density of the heat kernel, i.e. the density of B.M with initial condition Vn 10
The diagonal upper- bound is implied by Sobolev inequalities which control of the growth of some functions by the growth of their gradient. Al is satisfied in experiments, Saloff- (1994); Gross (1992).

---

Theorem 4. Assume AI, that Po is smooth and positive and that there exists M 2 0 such that for any E [o and M. ||@- (t,z) V loggp (x(2)] S M. with Sg- E C((0, 27, (M)). Then if
T > 1/2, there exists C 2 0 independent on T such that
W (C(C(Y)). = C(e AIT + +
where W1 is the Wasserstein distance of order one on the probability measures on M.
The proof is postponed to App. I. In particular, for any s > O, choosing T > 0 large enough, M small enough (which can be achieved using the universal property of neural networks) and 7 small enough, we get that (,(C(YN) Pu) < E. This result might seem weaker than the result obtained for Moser flows in (Rozen et al., 2021, Theorem 3), but emphasize that our bound takes into account the dineeeiscttization contrary to Rozen et al. (2021) which considers the continuous time flow. If we consider the time-reversed continuous-timg SDE then recover 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 approximation of the score.
5 Related work
In this section we discuss previous work on parametrizing family of distributions for manifold-valued data. Here, the manifold structure is considered 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 distribu- tions that have been proposed can be categorized into three main approaches (Navarro et al., 2017): wrapping, projecting and conditioning. Wrapped distributions consider parametric distribution R" that is pushed-forwad by marginalizing out some distribution along the normal bundle of M. Conditioning distributions encompass the von Mises-Fisher and Kent distributions (Fisher, 1953; Kent, 1982). Considering mixtures of these distributions is key to increase flexibility (Peel et al., 2001; Mardia et al., 2008).
along a surjective map & R" - >MM Projected distributions are defined
Push-forward of Euclidean normalizing flows. More recently, approaches leveraging the flexibility of normalizing flows (Papamakarios al., 2019) have been proposed. Following the wrapping method described above, these methods parametrize a normalizing flow in IR" before being pushed along an invertible map V R" - M. However, to globally represent the manifold, the map e needs to be homeomorphism. which can only happen if M is topologically equivalent to IR,, hence limiting the scope of that approach. One natural choice for this map is the exponential map exp,, T S Rd. This approach has been taken, for instance, by Falorsi et al. (2019) and Bose et al. (2020), respectively parametrizing distributions Lie groups and hyperbolic space.
Neural ODE on manifolds. To avoid artifacts or numerical instabilities due to the manifold em- bedding, another line of work uses tools from Riemannian geometry to define flows directly on the manifold of interest (Falorsi and Forré, 2020; Mathieu and Nickel, 2020; Falorsi, 2021). Since these methods do not require specific embedding mapping, they are referred Riemannian. They extend continuous normalizing flows (CNFs) (Grathwohl et al., 2019) to the manifold setting, by implicity parametrizing flows solutions of Ordinary Differential Equations (ODEs). As such, the parametric flow is a continuous function of time. This approach has recently been extended by Rozen et al. (2021) introducing Moser flows, whose main appeal being that it eircumvents the need 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 flows.
Optimal transport on manifolds. Another line of work has developed flows on manifolds using tools from optimal transport. Sei (2013) introduced flow that is given by fo x + exp. exp..Vvv)) with v/g c-convex function and d3 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 benefits of this approach depend on the specific choice of parametric family of c-convex functions (Rezende and Racaniere, 2021; Cohen et al., 2021), trading-off expressivity with scalability.
7

---

Table 3: Summary of computational complexity (w.rt. neural network forward and backward passes) for different methods. is the manifold dimension, & the number of Monte Carlo batches in Moser flow's regularizer, N is the number of steps in the (adaptive) ODE solver, whereas N is the number of steps the SDE Marur-Maauyama solver-which can usually be lower than N. Moser flow and RSGM training complexity varies the Hutchinson stochastic estimator is used. See Table 2 for score matching losses complexity.
Method
Training
Likelihood evaluation
Sampling
RCNF
Solving ODE ((dN)
Solving augmented ODE ((dN)
Solving ODE O(N)
Moser flow
Computing div O(dk) (kk)
Solving augmented ODE O(dN)
Solving ODE ((N)
RSGM
Score matching ((d) or (1))
Solving augmented ODE (ddN)
Solving SDE ((N*)
Experiments
In this section we benchmark the empirical performance of RSGMs along with other manifold-valuee methods introduced in Sec. 5. We also compare to 'Stercographic" score-based model, introduced in App. N. First, we assess their modelling capacity on earth and climate science spherical data. Then, we test the methods scalability with respect manifold dimensions with a synthetic experiment on the torus TV. Eventually, we evaluate the models' regularity and time complexity with a synthetic SO₂(R) target. Experimental details are provided in App. O.
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), floods (Brakenridge, 2017) and wild fires (EOSDIS, 2020). We compare 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 density on the sphere (Gemici al., 2016). We evaluate the lig-liieli 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 flows are slow to train and additionally that both Moser flows and stereographic SGMs are computationally expensive to evaluate.
6.2 Synthetic data on tori
We now move to another manifold, that the torus T S x ... X S', so as to assess the scalability of the different methods with respect to the dimension d. We consider wrapped Gaussian target distribution on T with a random mean and unit variance. Moser flows' (Rozen et al., 2021) loss involves regularization term which involves an integral over the manifold, approximated by Monte Carlo (MC) estimator with uniform proposal. This term regularizes Moser flows towards probability measures, i.e. with unit volume. We thus expect Moser flows to fail in high-dimension as the number of samples K required for the MC estimator to be accurate will grows as O(ea), and the
Table 4: Negative log- likeeiilldood scores for each method on the earth and climate science datasets. Bold indicates hest results (up to statistical significance). Means and confidence intervals are computed over 5 different runs. Novel methods are shown with blue shading.
Method
Volcano
Earthquake
Flood
Fire
Mixture of Kent
00.00
0.33
0.73+0.07
-1.1840.06
Riemannian CNF
60.5
0.14
1.11
-0.8040.54
Moser Flow
-4.21
-0.16
0.57
-1.2840.05
Stereographic Score- Based
-3.80
0.19
0.59
-1.2840.12
Riemannian Score-Based
-4.92+0.25
0.19
0.45
-1.3340.06
Dataset size
827
6120
4875
12809
8

---

(a) Volcano
(b) Farthquake
(c) Flood
(d) Fire
Figure Trained score-based generative models on earth scienees data. The learned density colored green-hlue. Blue and red dots represent training and testing datapoints, respectively.
memory required to compute this estimator grows either in (KK) for exact divergences O(K) for approximated divergences (see Table 3).
In Fig. 3, we observe that RSGMs are able to fit well the target distribution even in high dimension, with a linear or constant computational cost-depending on the divergence estimator. In contrast, Moser flows scale poorly with the dimension, the extent that are unable to train them for d 2 10. This is due to the combination of the complexity which grows linearly with 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,(RR {Q E MR(R). QQ Id, dct(Q) = 1}. We consider the synthetic dataset consisting of samples in SO₂ from mixture of wrapped normal distributions with M components.
We compare RSGMs against Moser flows and a wrapped-exponentiall baseline inspired by Falorsi al. (2019)-vv)) where we parametrize a standard Euclidean SGM on s3 that then pushed-fore forward on (O)( RSGMs are trained using the ε (DSM) loss with the Varadhan approximation (see Table 2). From Table 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 flows which only perform well in some range of M. Similarly to Sec. 6.2, we find Moser flows be much slower to train due to the large number of Monte Carlo samples needed in the reguralizer (K = 104). We also note from Table that the number of score network evaluations (NFE) significantly lower for RSGMs, and is particularly detrimental for Moser flows (* 10y).
Exact divergence
Approximate divergence
Stepe por socond
2 16
101
100
dimension
109
10!
102
100
10a
102
Torus dimension
Torus dimension
Moser Flow (K=1000, Moser likelihood)
Moser Flow (K=1000, ODE likelihood)
Moser Flow 5005,00, Moser likelihond)
Maser Flow (K TKK5500 5000, ODE likelihond)
Moser Flow (K=20000, Moser likelihood)
Moser Flow (K=20000, ODE likelihood)
RSGM
Figure 3: Comparison Moser flows and RSGMs training speed and performance on the synthetic high- dimension torus task. Moser flows trained with Amin = 1. We report two likelihoods, the 'Moser' closed form density-non not guaranteed to be normalized-and and the 'ODE' likelihood given by solving an augmented ODE (as in CNFs) with the vector field induced by the Moser flow density-whit is guaranteed to have unit volume.

---

RSGM
Exp-Wrapped SGM
Target
0.18
0.19
0.19
0.18
-0.32
0.06
0.17
0.15
Test log- eig-liklllhood
3.0
0.
Br
0.18
0.20
0.19
0.38
-3.00
0.07
0.16
0.15
15
Model
0.38
0.37
0.20
0.17
0.27
-0.29
0.03
0.15
0.13
10-1
2=100-5 korr 109
2x Le-ts Jo-t
10-
Learning rate
Learning rate
(a) Histograms of (Oa(R) samples from target mix- (b) RSGMs are much more robust to hyperparameters ture distribution with M 4 components, represented than Exp-wrapped SGMs. The diffusion coefficient given by (t. X,) = / 3(1) = Bo + (3, St)t.
via their Euler angles.
Figure 4: Trained score-based generative models on synthetic SOs(R) data.
Table 5: Test log likelihood and associated number of function evaluations (NFE) in 10 the synthetic mixture distribution with M components SO₂ (RR). Bold indicates best results (up statistical significance). Means standard deviations are computed over different runs. Novel methods are shown with blue shading.
M 16
M 32
M 64
Method
log-likelihood
NFE
log-likelihood
NFE
log-likelihood
NFE
Moser Flow exp-wrapped SGM
4..5
2.340.5
0. 0.174
2.3
-0.4940.02
7 7.341.4
10.04
0.540.1
0.16
0.5
-0.5840.04
0.540.0
RSGM
0.89
0.140.0
0.20
0.1
-0.4940.02
0.1
6.4 Synthetic data on hyperbolic space
Finally we demonstrate RSGM on non-compact manifold: the two dimensional hyperbolic space HP² which defined 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 synthetic dataset of samples from mixture of exp- wrapped normal distribution. From Fig. 5, we can qualitatively see that both score- based models are able to fit the target distribution.
(a) Target distribution.
(b) Exp-wrapped SGM.
(c) RSGM.
Figure 5: Samples from different prohahility distributions on 2 coloured w.r.t their density.
7 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. The main benefits of our method stems from its scalability to high dimensions, its applicability to broad class of manifolds due the diversity of available loss functions, its robustness and crucially its capacity to model complex datasets. We also provided theoretical guarantees on the convergence of RSGMs. In future work, we would like explore more generic classes of manifolds, such ones with boundary, along with alternative noising processes. Another promising extension concerns stochastic control on manifolds and more precisely, deriving efficient algorithms solve Schrödinger bridges in the same spirit De Bortoli et al. (2021) on Euclidean state spaces.
10

---

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/R01 16/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/1 CoSines.
References
Allgower, E. L. and Georg, K. (2012). Numerical Continuation Methods: An Introduction, volume 13. Springer Science & Business Media.
Ambrosio, L. (2003). Optimal transport maps in Monge-Kanttorviee problem. arXiv preprint arXiv:0304389vv.
Atkinson, K. and Han, W. (2012). Spherical Harmonics and Approximations the Unit Sphere: An Introduction, volume 2044. Springer Science & Business Media.
Bakry, D., Gentil, I., and Ledoux, M. (2014). Analysis and Geometry of Markov Diffusion Operators, volume 348. Springer.
Bao, F., Li, C., Zhu, J., and Zhang, B. (2022). Analytic-dpm: an analytic estimate of the optimal reverse variance diffusion probabilistic models. arXiv preprint vrXiv2201.06500:
Barfoot, T., Forbes, R., and Furgale, P. Т. (2011). Pose estimation using linearized rotations and quatermion algebra. Acta Astronautica, 68(1):101-11
Batzolis, G., Stanczuk, J., Schönlieb, C.-B., and Etmann, C. (2021). Conditional image generation with score-based diffusion models. arXiv preprint arXiv:2111.13066
Bismut, J.-M. (1984). Large deviations and the Malliavin calculus. Birkhauser Prog. Math., 45.
Bose, J., Smofsky, A., Liao, R., Panangaden, P., and Hamilton, W. (2020). Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke,
A., VanderPlas, J., Wanderman-Milhe S., and Zhang, Q (2018). JAX: composable transformations of Puthhon+NumPy programs.
Brakenridge, G. (2017).
Global active archive
large
flood events.
http://thooooherraaoroocoooorrd wtt.o//www...aa....aeee/eieeeeea html.
Brehmer, J. and Cranmer, K. (2020). Flows for simultaneous manifold learning and density estimation. arXiv preprint urXiv:2003. 13913.
Caterini, A. L., Loaiza-Ganemi G., Pleiss, G., and Cunningham, J. (2021). Rectangular flows for manifold learning, arXiv preprint arXiv:2106.01411
Cattiaux, P.. Conforti, G., Gentil, I., and Léonard, C. (2021). Time reversal of diffusion processes under finite entropy condition. arXiv preprint arXiv:2104.07708.
Chavel, I. (1984). Eigenvalues in Riemannian Geometry. Academic press.
Chen, X., Li, X. M., and Wu, B. (2021). Logarithmic heat kernels: estimates without curvature restrictions. arXiv preprint arXiv:2 106.027466
Chen, Y., Georgiou, T., and Pavon, M. (2016). Entropic and displacement interpolation: a computa- tional approach using the Hilbert metric. SIAM Journal Applied Mathematics, 76(6):2375-2
Cheng, X., Zhang, J., and Sra. S. (2022). Theory and algorithms for diffusion processes Rieman- nian manifolds. arXiv preprint arXiv:2204. 13665.
Choi, K., Meng, C., Song, Y., and Ermon, S. (2021). Density ratio estimation via infinitesimal classification. arXiv preprint arXiv:2111.11010

---

Chung, H., Sim, B., and Ye, J. C. (2021). cefloser-diffue-ffiie Accelerating conditional diffu- sion models for inverse problems through stochastic contraction. arXiv preprint arXiv:2112.05144
Cohen, S., Amos, B., and Lipman, Y. (2021). Riemannian convex potential maps. arXiv preprint arXiv:2106.1022
Croke, C. (1980). Some isoperimetric inequalities and eigenvalue estimates. Annales scien- lifiques de I l'École normale supérieure, volume 13, pages 419-435.
De Bortoli, V., Thortoo,, J., Heng. J., and Doucet, (2021). Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems.
Dhariwal, P. and Nichol, A. (2021). Diffusion models beat GAN on image synthesis. arXiv preprint arXiv:2105.05233.
Dormand, R. J. and Prince. (1980). family of embedded Kunge-Kutta formulae. Journal of Computational and Applied Mathematics, pages 19-26.
Durmus, A. (2016). High Dimensional Markov Chain Monte Carlo Methods: Theory, Methods and Application. Theses, Paris-Sud XI.
EOSDIS (2020). Land, atmosphere real-time capability for (lance) system operated by nasa's 's carth science data and information system (esdis).
hatta//aarreal--immefiimm/aaaat -ire-dddata.
Falorsi, L. (2021). Continuous normalizing flows on manifolds. arXiv:2104.14959
Falorsi, L., de Haan, P., Davidson, T. R., and Forré, P. (2019). Reparameterizing distributions on lie groups. In International Conference on Artificial Intelligence and Statistics, pages 3244-3253.
Falorsi, L. and Forré, (2020). Neural ordinary differential equations on manifolds. arXiv preprint arXiv:2006.0666 066663
Federer, H. (2014). Geometric Measure Theory. Springer.
Feiten, W., Lang, M., and Hirche, S. (2013). Rigid motion estimation using mixtures of projected Gaussians. In International Conference on Information Fusion, pages 1465-1472. IEFF.
Fisher, R. A. (1953). Dispersion sphere. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 217(1130):295-303
Ganea, O.-E., Huang, X., Bunne, C., Bian, Y., Barzilay, R., Jaakkola, T. S., and Krause, A. (2022). Independent Equivaaaan models for end-to-end rigid protein docking. International Conference on Learning Representations.
Garcfa-Zelada, D. and Huguet, B. (2021). Seenier-Scrrüdingge problem compact manifolds with boundary. Stochastic Analysis and Applications, pages 1-29.
Gemici, M. C., Rezende, D., and Mohamed, S. (2016). Normalizing flows on Riemannian manifolds. arXiv preprint arXiv: 16l::02300000
Grathwohl, W., Chen, R. T. Q., Bettencourt, J., and Duvenaud, D. (2019). Scalable reversible generative models with free-form continuous dynamics. In International Conference Learning Representations.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel two-sample test. Journal of Machine Learning Research, 13(mul):723--733
Grigor'yan, A. (1999). Estimates heat kernels on Riemannian manifolds. London Math. Soc. Lecture Note Ser. 273:140-225.
Gross, I.. (1992). Logarithmic sobolev inequalities on lie groups. Illinois journal of mathematics, 36(3):447-49..
12

---

Gunther, M. (1991). Isometric embeddings of Riemannian manifolds, Kyoto, 1990. In Proc. Intern. Congr. Math., pages 1137-1143. Math. Soc. Japan.
Hauberg, S. (2018). Directional Statistics with the Spherical Normal Distribution. In 2018 2 International Conference on Information Fusion (FUSION), pages 704-711, Cambridge, United Kingdom. IEEE.
Haussmann, U. G. and Pardoux, E. (1986). Time reversal of diffusions. The Annals of Probability, 14(4): 188-1205.
He, Y. (2013). A lower bound for the first eigenvalue in the Laplacian operator compact Riemannian manifolds. Journal of Geometry and Physics, 71:73-84 84.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems.
Hou, B., Miolane, N., Khanal, B., Lee, M. C. H., Alansary, A., McDonagh, S., Hajnal, J. V., Rueckert,
D., Glocker, B., and Kainz, B. (2018). Computing CNN Loss and Gradients for Pose Estimation with Riemannian Geometry. In Frangi, A. F., Schnabel, J. A., Davatzikos, C., Alberola-López. C., and Fichtinger, G., editors, Medical Image Computing and Computer Assisted Intervention MICCAI 2018, pages 756-764, Cham. Springer International Publishing.
Hsu, E. (1999). Estimates of derivatives of the heat kernel on compact Riemannian manifold. Proceedings of the american mathematical society, 127(12):3739-377 -3744.
Hsu, E. P. (2002). Stochastic Analysis on Manifolds. Number 38. American Mathematical Society.
Hutchinson, M. F. (1989). A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in CoatitticsSimulatioon and Computation, 18(3):1059-076.
Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4).
Ikeda, N. and Watanabe, S. (1989). Stochastic Differential Equations and Diffiusion Processes, volume 24 of North-Holland Mathematical Library. North-Holland Publishing Co., Amsterdam; Kodansha, Ltd., Tokyo, second edition.
Jolicoeur-Martineau, A., Taillefa,,,,,r, R., Tachet des Combes, R., and Mitliagkas, I. (2021). Adversarial score matching and improved sampling for image generation. International Conference on Learning Representations.
Jones, P. W., Maggioni, M., and Schul, R. (2008). Manifold parametrizations by eigenfunctions of the Laplacian and heat kernels. Proceedings of the National Academy of Sciences of the United States of America, (5566)1803-11
Jørgensen, E. (1975). The central limit problem for geodesic random walks, Zeitschrift für Wahrschein- lichkeitstheorie und verwandte Gebiete, 32(1-2):1-64.
Karpatne, A., Ebert-Uphoff, I., Ravela, S., Babaie, H. A., and Kumar, V. (2018). Machine learning for the geosciences: Challenges and opportunities. IEEE Transactions on Knowledge and Data Engineering, 3 14455154441
Kawar, B., Vaksman, G., and Elad, M. (2021a). Snips: Solving noisy inverse problems stochastically. arXiv preprint 1rX 14951.
Kawar, B., Vaksman, G., and Elad, M. (2021b). Stochastic image denoising by sampling from the posterior distribution. arXiv preprint arXiv:2101.09552.
Kent, J. T. (1982). The Fisher- Binghhr- distribution on the sphere. Journal of the Royal Statistical Society: Series B (Methodological), 44(1):71-80
Kingma, D. P. and Ba, (2015). Adam: A Method for Stochastic Optimization. arXiv: 1412.6980 [cs].
13

---

Klimovskaia, A., Lapez-Paz, D., Bottou, L., and Nickel, M. (2020), Poincaré maps for analyzing complex hierarchies single-cell data. Nature communications, 11(1):1-9.
Kloeden, P. and Platen, E. (2011). Numerical Solution of Stochastic Differential Equations. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg.
Köhler, J., Klein, L., and Noé, F. (2020). Equivariant Flows: Exact likelihood generative learning for symmetric densities. wrXiv:2006.02425 [physics, stat].
Kullback, S. (1997). Information Theory and Statistics. Dover Publications, Inc., Mineola, NY. Reprint of the second (1968) edition.
Kurtz, T. G., Pardoux, É., and Protter, P. (1995). Stratonovich stochastic differential equations driven by general semimartingales. In Annales de I'IHP Probubilités et statistiques, volume 31, pages 351-377.
Kuwada, K. (2012). Convergence of time-inhomogeneouus geodesic random walks and its application to coupling methods. The Annals of Probubility, 1055):194-11979
Lee, J. (2010). Introduction to Topological Manifolds, volume 202. Springer Science Business Media.
Lee, J. M. (2006). Riemannian Manifolds: An Introduction Curvature, volume 176. Springer Science & Business Media.
Lee, J. M. (2013). Smooth manifolds. In Introduction to Smooth Manifolds, pages 1-31. Springer.
Lee, J. M. (2018). Introduction to Riemannian manifolds. Springer.
Lee, S.-g., Kim, H., Shin, C., Tan, X., Liu, C., Meng, Q.. Qin, T., Chen, W., Yoon, S., and Liu, T.-Y. (2021). Priorgrad: Improving conditional denoising diffusion models with data-driven adaptive prior. arXiv preprint arXiv:2106.0644 06406.
Leobacher, G. and Steinicke, A. (2021). Existence, uniqueness and regularity of the projection onto differentiable manifolds. Annals of Global Analysis and Geometry, 60(3):559-587.
Léonard, C. (2012a). From the Schrödinger problem to the Monge-Kantoroviel problem. Journal of Functional Analysis, 187-19988
Léonard, C. (2012b). Girsanov theory under finite entropy condition. In Séminaire de Probabilités XLIV, pages 429 465. Springer.
Léonard, C., Relly, S., Zambrini. J.-C., et al. (2014). Reciprocal processes: a measure-theoretical point of view. Probability Surveys, 11:237-269.
Li, P. (1986). Large time behavior of the heat equation on complete manifolds with non-negative ricci curvature. Annals of Mathematics, 124(1):1-21.
L Liptser, R. S. and Shiryaev, A. N. (2001). Statistics of Random Processes. I, volume of Applications of Mathematics (New York). Springer- Verlag, Berlin, expanded edition. General theory, Translated from the 1974 Russian original by A. B. Aries, Stochastic Modelling and Applied Probability.
Lui, Y. M. (2012). Advances in matrix manifolds for computer vision. Image and Vision Computing, 30(6-7):380-388
Mardia, K. V., Hughes, G., Taylor, C. C., and Singh, H. (2008). A multivariate von Mises distribution with applications to bioinformatics. Canadian Journal of Statistics, 3601)999-109.
Mathicu, E., Lan, L., Maddison, J., Tomioka, R., and Tch, Y. W. (2019). Continuous hicrarchical representations with Poincaré variational auto-encoders. arXiv preprint arXiv:1901.06033
Mathieu, F. and Nickel, M. (2020). Riemannian continuous normalizing flows. arXiv preprint arXiv:2006.10605.
14

---

Miolane, N., Guigui, N., Brigant, A. L., Mathe, J., Hou, B., Thanwerdas, Y., Heyder, S.. Peltre, O., Koep, N., Zaatiti, H., Hajni H., Cabanes, Y., Gerald, T., Chauchat, P., Shewmake, C., Brooks, D.. Kainz, B., Donnat, C., Holmes, S., and Pennec, X. (2020a). Geomstats: A python package for riemannian geometry in machine learning. Journal of Machine Learning Research, (122))11199
Miolane, N., Guigui, N., Brigant, A. L., Mathe, J., Hou, B., Thanwerdas, Y., Heyder, S., Peltre, O.,
Koep, N., Zaatiti, H., Hajri, H., Cabanes, Y., Gerald, T., Chauchat, P., Shewmake, C.,, Brooks, D., Kainz, B., Donnat, C., Holmes, S., and Pennec, X. (2020b). Geomstats: A python package for riemannian geometry in machine learning. Journal Machine Learning Research, (123))1119
Navarro, K., Frellsen, and Turner, R. E. (2017). The multivariate generalised von Mises distribution: inference and upplications. In AAAI Conference on Artificial Intelligence.
(NGDC/WDS), N. D. W. D. S. (2022a). Ncei/wds global signifi- cant earthquake database. pagge/bnfisordddggooooaaaa .gg. htw.aad...g..a.arr.01111...1 h01153
(NGDC/WDS), N. G. D. C.
W. D. S. (2022b). Ncei/wds global significant pagg/biniso"kd=ggo.oooaaaa www.aaagggggg..gga.....oo....
volcanic eruptions database.
Nichol, A. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. arXiv
preprint arXiv:2102.09667.
Papamakarios, G., Nalisnick, E.,, Rezende, D. J., Mohamed, S., and Lakshminaraaanann B. (2019). Normalizing flows for probabilistic modeling and inference. arXiv preprint 1rXiv:19/2.02262.
Peel, D., Whiten, W. J., and McLachlan, G. J. (2001). Fitting mixtures of Kent distributions to aid in joint set identification. Journal of the American Statistical Association, 96(453):56-63,
Pennec, X. (2006). Intrinsic Statisties on Riemannian Manifolds: Basic Tools for Geometric Measurements. Journal of Mathematical Imaging and Vision, 25
Prokudin, S., Gehler, P., and Nowozin, S. (2018). Deep directional statistics: Pose estimation with uncertainty quantification. In European Conference on Computer Vision (ECCV).
Revuz, D. and Yor, M. (1999). Continuous Martingales and Brownian Motion, volume 293 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Mathemattial Sciences]. Springer- Verlag, Berlin, third edition.
Rezende, D. J. and Racanière, S. (2021). arXiv:21
Implicit Riemannian concave potential maps.
Roy, D. M., Kemp, C., Mansinghka, V., and B Tenenbaum, J. (2007). Learning annotated hierarchies from relational data.
Rozen, N., Grover, A., Nickel, M., and Lipman, Y. (2021). Moser flow: Divergence- generative modeling manifolds. Advances in Neural Information Processing Systems.
Saloff-Coste, L. (1994). Precise estimates the rate which certain diffusions tend equilibrium. Mathematische Zeitschrift, 217(1):641-677.
Santambrogio, F. (2017). (Euclidean, metric, and Wasserstein] gradient flows: an overview. Bulletin of Mathematical Sciences, 711):87-154.
Schrödinger, E. (1932). Sur théorie relativiste de l'électron et l'interprétation de la mécanique quantique. Annales de l'Institut Henri Poincaré, 2(4):269-31) 310.
Sei. Т. (2013). A Jacobian inequality for gradient maps on the sphere and its application to directional statistics. Communications in SkaisticcsThhen and Methods, 42(14):2525-2542
Senanayake, R. and Ramos, F. (2018). Directional grid maps: modeling multimodal angular uncer- tainty in dynamic environments. 20/8 IEFE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 3241-3248. IEEE.
15

---

Shapovalov. M. V. and Dunbrack Jr, R. L. (2011). A smoothed backbone- backbone-dependenn rotamer library for proteins derived from adaptive kernel density estimates and regressions. Structure, 19(6):844-3) 858.
Sinha, A., Song, J., Meng, C., and Ermon, (2021). D2c: Diffusion-denoising models for few-shot conditional generation. arXiv preprint arXiv:2106.000 068/9.
Song, Y. and Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution.
In Advances Neural Information Processing Systems.
Song. Y. and Ermon, S. (2020). Improved techniques for training score-based generative models, In Advances in Neural Information Processing Systems.
Song, Y., Bohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations.
Steyvers, M. and Tenenbaum, J. B. (2005). The large-scale structure of semantic networks: Statistical analyses and model of semantic growth. Cognitive science, 29(1):41-78
Sun, Y., Flammarion, N., and Fazel, M. (2019). Escaping from saddle points on Riemannian manifolds. Advances in Neural Information Processing Systems, 32.
Urakawa, H. (2006). Convergence rates to equilibrium of the heat kernels compact Riemannian manifolds. Indiana University Mathematics Journal, pages 259-288.
Vargas, F., Thodoroff, P., Lawrence, N. D., and Lamacraft, (2021). Solving Schrödinger bridges via maximum likelihood. arXiv preprint arXiv:2106.0208//
Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation, 2367).1661-1674
Watson, D., Ho, J., Norouzi, M., and Chan, W. (2021). Learning to efficiently sample from diffusion probabilistic models. arXiv preprint arXiv:21066 100.220000000000 03802.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect 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 Sec. 7.
(c) Did you discuss any potential negative societal impacts of your work? [No] The work presented this paper focuses on the learning of score-based models on manifold. We do not foresee any immediate societal impact of such 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.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions all theoretical results? [Yes] Yes, see AI.
(b) Did you include complete proofs of all theoretical results? [Yes] Yes, proofs are postponed to the supplementary material.
3. 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 as 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 App-
(c) Did you report error bars (e-g., with respect the random seed after running experi-
ments multiple times)? [Yes] Error bars reported for each experiment.
16

---

(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, 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. 6.1.
(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 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 using/curating contains personally identifiable information or offensive content? [No] Not applicable.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [No] Not applicable.
(b) Did you deseribe any potential participant risks, with links to Institutional Review Board (IRB) approvals, 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.


Reviews:
[{"review_id": "r44uVP1qNKx", "reviewer": "Reviewer_X25L", "summary": "This work extends score-based generative modeling with time-reversed diffusion processes from Euclidean geometry to compact Riemannian manifolds. Several modifications have been made to adapt the method to compact Riemannian manifolds, including using Brownian motion in forward pass, uniform base distribution, and geodetic random walks in forward/backward sampling.  Experiments on real and synthetic datasets show the effectiveness of proposed methods in modeling the symmetry in data. ", "questions": "\n#### Strengths\n- This work proposes novel methods for extending score-based generative modeling techniques beyond Euclidean geometry.  \n- This paper is clearly written, covering both the big picture of the considered research topic and the technical details of the proposed method.\n- Both theoretical convergence analysis and empirical experimental results are included in the paper to show the validity of the proposed approach.\n  \n#### Weaknesses\n- It is not quite clear to me whether the assumption of compactness of the Riemannian manifold can be too limited for the consideration of generative modeling for real-world applications. It would be more helpful to provide mode detailed analysis from this perspective.\n- For the consideration of scalability, stochastic approximations are needed to obtain a tractable training objective. \n- Most of the experimental results are on synthetic datasets, and the choices of real-world datasets satisfying the assumptions of this work seem to be limited.\n\n", "limitations": "- Have the authors tried different parametrizations for the score function, i.e. using the clean data or noise as the prediction target similar to DDPM?\n- How general are the proposed approach in terms of the applicability of various samplers or ODE solvers? like DDIM and other advanced solvers.\n- How large the gap is to extend this work to Riemannian manifolds which are non-compact or with boundaries or implicitly defined? ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: the originality and quality of the theoretical part seem good, and the paper is clearly written.  The extension to the Riemannian manifold is a reasonable idea to consider. The experiments show that this algorithm can work in some cases.\n\nWeakness: The main weakness of the paper is its significance. First, very few datasets really lie on some compact Riemannian manifold without boundary (very very restrictive in terms of math, in fact. ). This setting is more like a mathematical toy to me. This is why the experiments are mainly about synthetic datasets. Second, there are some questions about its feasibility. For example, why uniform distribution on a compact manifold is easy to sample? It highly depends on how the manifold is represented in the memory. So this paper is more about math toys not about solving real machine learning problems. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "V5PKEMqsdZt", "reviewer": "Reviewer_Rw15", "summary": "The paper proposes a score-based generative model (SGM) on compact Riemannian manifolds. The paper first points out three major components that contribute to the success of SGMs in Euclidean spaces.\nLinear SDEs in Euclidean spaces (more precisely, OU processes) have stationary distributions and are Normal distributions.\nFor the SDEs, the reverse diffusions exist, whose initial distributions are the stationary distributions above, and the form of the reverse SDEs is known, yielding efficient parameterizations.\nThe solutions of OU processes are Normal distributions; thus, it is easy to sample and compute loss in denoising score matching (DSM).\nThe parametric reverse diffusions can be trained via minimizing various score matching losses, including DSM.\nThus, the paper proposes those counterparts for SGMs on compact Riemannian manifolds, specifically by identifying and solving the non-trivialities of the counterparts.\n\nThen, the paper shows that linear diffusions on compact Riemannian manifolds have reverse diffusions, and the resulting SDEs are the same form as diffusions in Euclidean spaces.\nHowever, the paper identifies that for the linear diffusions on compact Riemannian manifolds, its solution for a given data at time 0 doesn't have analytical forms, unlike OU processes. This makes two problems: the difficulty of simulation x_t given x_0 and the difficulty of computing denoising score matching loss.\n\nTo address this, the paper proposes using Geodesic random walks (GRW) for simulating both forward noising and reverse denoising processes; note that it is approximate sampling, unlike efficient sampling of the forward noising processes in SGMs on Euclidean spaces.\n\nFourth, the paper proposes two approximate methods to approximate the conditional log density gradients (or score) $\\nabla_{x_t} p(x_t|x_0)$ of the forward noising process. When we know the eigenvalues and eigenvectors of the manifold of interest, one can approximate the conditional scores using the first N eigenvectors, i.e., the Truncation method. One can employ Varadhan's asymptotic, trading-off approximation errors when we don't have such information.\n\nNext, the papers introduce DSM and implicit score matching (ISM) losses, whose computation will involve the simulation and approximation techniques above.\n\nFinally, the paper introduces the neural network-based score function parameterizations suitable for compact Riemannian manifolds. Briefly speaking, when d-dimensional Euclidean spaces can describe the manifold, the suitable parameterization requires n > d sets of vector components.\n\nFor the experiments, the paper demonstrates the effectiveness of the proposed method, especially by comparing it with other manifold-valued generative models.", "questions": "For Varadhan's asymptotic, What can we do when $t$ is large? Does Eq. 7 mean that using the approximation results in biased learning signals when we use the asymptotic to estimate score matching losses?\n", "limitations": "Please find the comments on the Weakness section above. In general, the paper's contributions wrt the novelty are clear, and the proposed methods are well-defined. In addition, I found that the paper has a well-organized structure that makes it clear to understand the proposed methods. Thus, I'm inclined to accept the paper.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n\nIn my understanding, the paper's contributions are clear, and I also consider that the results are essential for several reasons:\n  - The paper provides novel manifold-valued generative models based on SGMs.\n - The paper motivates each proposed method well so that readers can understand how each method addresses the challenges of designing manifold-valued SGMs.\n - I found that the paper has a well-organized structure that makes it clear to understand the proposed methods.\n - The paper provides sufficient information about their theoretical claims, helping readers understand the proposed methods' characteristics and limitations.\n\n**Weakness**\nThe followings can be improved.\n - Clarifying the terminologies will help readers to understand the paper more clearly. For instance, one can mention that \"heat kernel\" means \"(fundamental) solutions of (linear) SDEs\". I understand such terms are widely accepted in SDE literature but possibly confuse the ML communities.\n - The papers can introduce notations before introducing the main theoretical claims. As the paper proposes technically demanding methods, the clarifications of the notations seem necessary beforehand—for instance, exponential mapping notation is introduced in line 203 but appears in Eq 7.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Rk8TUcNLHtR", "reviewer": "Reviewer_z9yP", "summary": "In this paper, the authors propose an extension of score based generative modeling to data supported on a compact Riemannian manifold. To that end, concepts the need to be extended from the classical Euclidean setting include an easy to sample reference distribution (uniform on the manifold), that can be further seen as the limit of a noising process (brownian motion on the manifold), as well a an SDE reverting the noising process, in order to be able to generate samples by integrating this SDE thanks to geodesic random walks. The authors also present how to obtain score estimation on the manifold by using the heat kernel, as well as the parametization of vector fields as neural networks, and finally a way to obtain exact likelihood computations, as in the Euclidean case. They obtain a convergence result showing that when the score estimation cost function gets smaller, their model obtains a density that gets arbitrarily close to the true one in terms of Wasserstein-1 distance. They show that their model compares favorably in terms of computional complexity to riemannian continuous normalizing flows or Moser flows. Several experiments real goescience data (on the sphere) and synthetic data on the torus and the special orthogonal group support those claims.", "questions": "- Where does the denoising score matching (equation below equation 2) identity come from ? I couldn’t find the same equation in [Song et al. 2021b], where the score matching objective is obtained from an integration by parts?\n\n- Here it makes sense to consider a uniform distribution as the simple distribution to start from, thanks to the convergence result of Brownian motion on compact manifolds to that distribution. Could other possibilities be envisioned? Has this case already been treated in the Euclidean case, where typically one chooses a normal distribution for the base distribution (and if so, how does it compare to the classical score based with a gaussian base distribution)?", "limitations": "The main limitation of the method is the restriction to compact manifolds, which is clearly stated. Other manifolds of practical interest are not covered. Nevertheless, the construction for compact manifolds is already sound and covers a certain number of practical cases of interest. Especially, this construction translates into what seems to be efficient algorithms, which is often a drawback of geometric methods operating on manifold valued data.\n", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths :\n\n- The authors are able to extend score based generative models that have now become very popular to compact Riemannian manifolds by generalizing or adapting all the necessary ingredients to this setting.\n\n- The theoretical construction and guarantees are extensive and show that the propose method converges and has a better computational complexity than existing alternatives for manifold valued data. In particular one appealing feature is the scalability to higher dimensional manifolds that is verified in the experiments (up to 200 on the torus data and 64 in the special orthogonal data in the synthetic experiments), while competing approaches cannot deal efficiently (or in the extreme case cannot even run) with such high dimensional non euclidean data\n\n\nWeaknesses :\n\n- The approach is only limited to compact manifolds, which encompass a number of interesting cases already but leave out many more (SPD matrices, hyperbolic spaces…) . The main difficulties are the convergence of the Browian motion to the uniform distribution that only holds in the compact case, and the manipulation of the transition kernel of the Brownian motion in that case. \n\n- The only real case study with geoscience data does not show clear improvement over other methods, most likely because the problem is too simple and the dimension of the manifold (2 here) is not high enough for the computational benefits of the proposed approach to be visible. Maybe it would be possible to find a real application in higher dimensions where the method can be shown to clearly outperform others? ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0xajEfylCGW", "reviewer": "Reviewer_uQqV", "summary": "This paper explore continuous-time score-based generative models in non-Euclidean space specifically on the Riemannian manifold. The paper defines the forward process of the SGM model directly on the Riemannian manifold and integrates Geodesic Random Walks (GRW) to project the Brownian motion onto a manifold for sampling on the reverse diffusion process. The authors validate their proposal by testing on multiple datasets (e.g. Earth and climate science dataset, Synthetic data on tori, and orthogonal group).", "questions": "- How are stereographic score-based models projected and could a similar method (albeit a different projection) be possible for comparison on the tori experiment?\n- On the torus experiment, why were the Moser Flows the only benchmark comparison?", "limitations": "The authors discusses and addresses the current limitations on their model (e.g. incompatible manifolds).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The overall paper was clear and organized well. The paper did a good job of translating and introducing the SGM forward noising process for the Euclidean space to the Riemannian forward process. \n- The author provides a nice comprehensive overview of related literature works with useful insight on each past work.\n- The paper is highly motivated and novel as there have not been any past papers that have applied score-based/diffusion models in a non-Euclidean space.\n- Derivations are concise and eloquent, but it would be nice to see the elaborated derivations that are available in the Appendix.\n\nWeaknesses:\n- Many of the experiments do not have many benchmark models to compare with the proposed model. \n- There is not an explanation on the stereographic Score-Based method. A quick explanation on this model would be helpful.\n\nAdditional Remarks:\n- Section 5.2: I believe that \"torii\" should be spelled \"tori\" as it is the plural form of \"torus\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept