Riemannian Diffusion Models
Authors: Chin-Wei Huang
Milad Aghajohari
Joey Bose
Prakash Panangaden
Aaron Courville
Venue: NeurIPS 2022 Accept (2022)
Keywords: diffusion models
density estimation
generative models
Riemannian manifolds
variational inference

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

Riemannian Diffusion Models
Chin-Wei Huang", Milad Aghajohari', Avishek Joey Bose
Prakash Panangaden, Aaron Courville
University of Montreal & McGill University & Mila
{chin-wei .hangg milad. aagadd.ggaooiii aghajohari, aaron courvillee courvilleeemoontealllcc joey www.emaall......mm bbosemaai.mcgill..... ca, prakash@cs prakash@cs.mcgill.cc mwgiml.cc...
@umontreal.ce
Abstract
Diffusion models recent state-of-the-ar methods for image generation and likelihood estimation. In this work, we generalize continuous-time diffusion mod- els arbitrary Riemannian manifolds and derive variational framework for like- lihood estimation. Computationally. we propose new methods for computing the Riemannian divergence which is needed for likelihood estimation. Moreover, in generalizing the Euclidean case, we prove that maximizing this variational lower- bound is equivalent Riemannian score matching. Empirically, we demonstrate the expressive power of Riemannian diffusion models on a wide spectrum of smooth manifolds, such as spheres, tori, hyperboloids, and orthogonal groups. Our proposed method achieves new state-of-the-ee likelihoods on all benchmarks.
Introduction
By learning to transmute noise, generative models seek uncover the underlying generative factors that give rise to observed data. These factors can often be cast as inherently geometric quantities the data itself need not lie domains such as high-energy physics (Brehmer & Cranmer, 2020), directional statistics (Mardia & Jupp, 2009), geoscience (Mathieu & Nickel, 2020), computer graphics (Kazhdan al., 2006), and linear biopolymer modeling such as protein and RNA (Mardia et al., 2008; Boomsma et al., 2008; Frellsen et al., 2009), data is best represented on Riemannian manifold with non-zero curvature. Naturally, to effectively capture the generative factors of these data, we must take into account the geometry of the space when designing learning framework.
flat Euclidean space. Indeed, in many scientific
Recently, diffusion based generative models have emerged as an attractive model class that not only achieve likelihoods comparable but match the sample quality of GANs without the pains of adversarial optimization (Dhariwal & Nichol, 2021). Succinetly, diffusion model consists of fixed Markov chain that progressively transforms data to a prior defined by the inference path, and generative model which is another Markov chain that is learned invert the inference process (Ho et al., 2020; Song et al., 2021b).
oftte-of-the-art autogressive models (Kingma et al., 2021)
While conceptually simple, the learning framework can have variety of perspectives and goals. For example, Huang et al. (2021) provide variational framework for general continuous-time diffusion processes on Euclidean manifolds as well as a functional Evidence Lower Bound (ELBO) that can be equivalently shown to be minimizing implicit score matching objective. At present, however, much of the success of diffusion based generative models and its accompanying variational framework is purpose built for Euclidean spaces, and more specifically, image data. It does not easily translate to general Riemannian manifolds.
In this paper, we introduce Riemannian Diffusion Models genera-innrllizing conventional diffusion models on Euclidean spaces to arbitrary Riemannian manifolds. Departing from diffusion models on Euclidean spaces, our approach uses the Stratonovich SDE formulation for which the
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

conventional chain rule of calculus holds, which, as we demonstrate in section $3, can be exploited define diffusion on a Riemannian manifold. Furthermore, we take an extrinsic view of geometry by defining the Riemannian manifold of interest as an embedded sub-manifold within a higher dimensional (Fuclidean) ambient space. Such a choice enables us to define both our inference and generative SDEs using the coordinate system of the ambient space, greatly simplifying the implementation of the theory developed using the intrinsic view.
Main Contributions. We summarize our main contributions below:
We introduce variational framework built on the Riemannian Feynman-Kan representation and Giransov's theorem. In Theorem we derive a Riemannian continuous-tinn ELBO, strictly gen- eralizing the CT-ELBO in Huang et al. (2021) and prove in Theorem 4 that its maximization is equivalent to Riemannian score matching for marginally equivalent SDEs (Theorem 3).
To compute the Riemannian CT-ELBO it is necessary compute the Riemannian divergence of our parametrized vector field, for which we introduce a QR-decomposition--oe based method that is computationally efficient for low dimensional manifolds as well projected Hutchinson method for scalable unbiased estimation. Notably, our approach does not depend on the closest point projection which may not be freely available for many Riemannian manifolds of interest.
We also provide a variance reduction technique to estimate the Riemannian CT-ELBO objective that leverages importance sampling with respect the time integral, which crucially avoids care- fully designing the noise schedule of the inference process.
Empirically, we validate our proposed models on spherical manifolds towards modelling natural disasters as found in earth science datasets, products of spherical manifolds (tori) for protein and RNA, synthetic densities on hyperbolic spaces and orthogonal groups Our empirical results demonstrate that RDM leads to new state-of-art likelihoods over prior manifold generative models.
2 Background
In this section, we provide the necessary background on diffusion models and key concepts from Riemannian geometry that we utilize to build RDMs. For short review of the latter, see Appendix A or Rateliffe (1994) for a more comprehensive treatment of the subject matter.
2.1 Fuclidean diffusion models
diffusion model can be defined the solution to the (Itô) SDE (Øksendal, 2003),
dX = p dl + a dB.,
with the initial condition Xg following some unstructured prior Po such as the standard normal distribution, where B, is standard Brownian motion, and μ and are the drift and diffusion coef- ficients of the diffusion process, which control the deterministic forces driving the evolution and the amount of noise injected at each time step. This provides us a way to sample from the model, via numerically solving the dynamics from / = 0 to / = T for some fixed termination time T. To train the model via maximum likelihood, require expression for the log marginal density of Xx, denoted by log p(x, I",, which is generally intractable.
The marginal likelihood can be represented using stochastic instantancous change-of-variabl for- mula, by applying the Feynman-Ka theorem to the Bokker-Plance PDE of the density. An applica- tion of Girsanov's theorem followed by an application of Jensen's 's inequality leads to the following variational lower bound (Huang et al., 2021; Song et al., 2021a):
log T) 2 R logpo (Vr) - -f-"
(((6
(ala )) V ((Y..T ds Yo
(2)
where a is the variational degree of freedom, V. denotes the (Euclidean) divergence operator, and Y, follows the inference SDE (the generative coefficients are evaluated in reversed time, i.e. T 8)
dy (-H σa) ds + a dB,
(3)
with B, being another Brownian motion. This known as the continuous-tims evidence lower bound, the CT-ELBO for short.

---

2.2 Riemannian manifolds
We work with a d-dimensional Riemannian manifold (M.g) embedded in higher dimensional ambient space RT for m > d. This assumption does not come with a loss of generality, since any Riemannian manifold can be isometrically embedded into Euclidean space by the Nash embedding theorem (Gunther, 1991). In this case, the metric g coincides with the pullback of the Fuclidean metric by the inclusion map. Now, given coordinate chart yp : M -> I2 and its inverse vi = p we can define E, for = 1.... to be the basis vectors of the tangent space Ta-M at point x E M. The tangent space can be understood as the pushforward of the Euclidean derivation of the patch
space along vh; Le., for any smooth function / E O (M), FJJJ) 2 ʃ 4/.
We denote by Pa the orthogonal projection onto the linear subspace spanned by the column vectors
of the Jacobian Jn - / Specifically, Pa can be constructed via Pa - J -=)-- Note that this subspace is isomorphic to the tangent space M, which itself is subspace of T RT As a result, we identify this subspace with T.M. Lastly, we refer to the action of P₂ as the projection onto the tangentill subspace, and P. itself as the tangential projection.
2.3 SDE on manifolds
Unlike Euclidean spaces, Riemannian manifolds generally do not possess a vector space structure. This prevents the direct application of the usual (stochastic) calculus. We can resolve this by defining the process via test functions. Specifically, let V& be a family of smooth vector fields on M, and let z be a family of semimartingales (Protter, 2005). Symbolically, we write
ax,- = Σ Vx,) o dzy if df(x,) = E (xi)) oZ
(4)
for any E C (Hsu, 2002). The in the second differential equation is to be interpreted in the Stratonovich sense (Protter, 2005). The use of the Stratonovich integral is the first step deviating from the Euclidean diffusion model (I), as the Itô integral does not follow the usual chain rule.
Working with this abstract definition is not always convenient, so instead we work with specific coordinates of M. Let &p be a chart, and let ē = (ozk) be matrix representing the coefficients of Vx in the coordinate basis-iee Vi(()) Σ5 82 DZ, 8- o 1 write (X,, i o dZ. Similarly, suppose M is submanifold embedded in R" and denote by U. Then we can
(=p(x)) (=p(_e" This allows us to
v = (v:x) the coefficients wrt the Euclidean basis. v and ö are related by v = de express the dynamics of X as a regular SDE using the Fuclidean space's s coefficients dX = o dz. Notably, by the relation between and ö, the column vectors of e are required to lie in the span of the column vectors of the Jacobian de
dz -.
which restricts the dynamics move tangentially on M.
da
3 Riemannian diffusion models
We now develop d variational framework estimate the likelihood of a diffusion model defined on a Riemannian manifold (M,9). Let X2 E -M be process solving the following SDE:
Generative SDE:
dX = Vo dl + V o dBr,
X Po
(5)
where Vo and the columns of the diffusion matrix V == Vi,..., Vir) smooth vector fields on M, and B, is a w-dimensional Brownian motion. The law of the random variable X, can be written as p(a, t) p(dz). where p(x,1 f) is the probability density function and j is the d-dimensional Hausdorff measure on the manifold associated with the Riemannian volume density. Let V V be a differential operator defined by (V- (V.V)) V,)U Ee_ Us) where V, Un denotes the Riemannian divergence of the vector field Uz:
Σ a az, (G)
V U. G
(6)
Our first result is a stochastic instantaneous change-of-variabe formula for the Riemannian SDE by applying the Peynman-Kac theorem to the Fokker Planck PDE of the density p(z, t).
'hee multiplication is interpreted similarly matrix-vector multiplication i.c. VodB; En Va Vod.
3

---

Theorem 1 (Marginal Density). The density p(x, 1) of the SIDE (5) can be written as
p(x,t) = po (Y)) exp (-L²vo: (100 -Vo- (VVVV,VV ds Yo x
1
v,
where the expectation is taken wrt the following process induced by Brownian motion B3
dy (-Vo + (V v, ds + VodB'
8
For effective likelihood maximization, we require access log p and its gradient. Towards this goal. we prove the following Riemannian CT-ELBO which serves as our training objective and follows from an application of change of measure (Girsanov's theorem) and Jensen's inequality.
Theorem 2 (Riemannian CT-ELBO). Let B, be w-dimensional Brownian motion, and let a process solving the following
Inference SDE:
dY (-Vo+ C.V+Vo ds + d
(9)
where a R x [0,T] + R the variational degree of freedom. Then we have
log log,(r, T) 2 EE loggo(YT) - [-1 Ila(Y.. 3)12 V (vo (vo-2 volV) ds
Yo = æ
(10)
where all the generative degree freedoms V& evaluated in the reversed time direction.
3.1 Computing Riemannian divergence
Similar to the Euclidean case, computing the Riemannian CT-ELBO requires computing the diver- gence "V," " of vector field, which can be achieved by applying the following identity.
Proposition 1 (Riemannian divergence identity). Let (M.g) be d-dimensional Rieman- nian manifold. For any smooth vector field Vo ((M),, the following identity holds:
V. V/2
V E Vh Vv.
(11)
Furthermore, if the manifold a submanifold embedded the ambient space R" equipped with the induced metric g t" g. then
(V, (a) =tr (P= dus dr
Pa
(12)
where Us (Vik,... Vank) are the ambient space coefficients V EEe Usk Da, and Pa the orthogonal projection onto the langent space.
Intrinsic coordinates. The patch-space formula (6) can be used to compute the Riemannian diver- gence. This view adopted by Mathieu & Nickel (2020), where they combined the Hutchinson trace identity and the internal coordinate formula to estimate the divergence. The drawbacks of this framework include: (1) obtaining local coordinates may be difficult for some manifolds, hindering generality in practice; (2) we might need to change patches, which complicates implementations; and (3) the inverse scaling of /IG might result in numerical instability and high variance.
Closest-poine projection. The coordinate-fice expression (11) leads to the closest- ppint projection method proposed by Rozen et al. (2021). Concretely, define the closest-point projection by ((x) := arg min y ||æ yll. where |-|ll is the Euclidean norm. Let Vx(x) be the derivation corresponding to the ambient space vector 2n (z) (x)))) for some unconstrainted TM - Rh Rozen
et al. (2021) showed that V, Vi V (x()), since Us is infinitesimally constant in the normal

---

direction to M.M This allows us to compute the divergence directly in the ambient space. However, the closest-point projection map π may not always be easily obtained.
QR decomposition. An alternative to the closest-point projection is to instead search for an or- thogonal basis for T.M Let Q = [ei,".. ed, P,,".. (mmaa be orthogonal matrix whose first d columns span the TaM,, and the remaining m d vectors span its orthogonal complement TaM-. To construct Q we can simply sample d vectors-e.g from (N, 1)-in the ambient space and orthogonally project them to M using P.. These vectors, although not orthogonal yet, form a basis for T.M. Next we concatenate them with m d random vectors and apply a simple QR decomposition to retrieve an orthogonal basis. Using Q we may rewrite equation (12) as follows:
(Vg. : Vo)(()) tr (QQT Ps
dek dx Ps pa)= = ((PPO)) duk PoQ dx
Σ
dek dx Eg
(13)
where we used (1) the orthogonality of Q. (2) the cyclic property of trace, (3) and the fact that Pzej = Ej and P.n, = O. practice, concatenation with the remaining m vectors is not needed as they are effectively not used in computing the divergence, speeding up computation when m >> Moreover, the vector-Jacobian product can be computed in ((m) time using reverse-mode autograd and importantly does not require the closest-point projection
Projected Hutchinson. When QR too expensive for higher dimensional problems, the Hutchin- son trace estimator (Hutchinson, 1989) can be employed within the extrinsic view representa- tion (12). For example, let x be a standard normal vector (or Rademacher vector), we have (V, (Vx)(z) ==',,'===P.d [2'T dz z'J. Different from a direct application of the trace estimator to the closest-point method, we directly project the random vector to the tangent subspace. Therefore, the closest-point projection is again not needed.
3.2 Fixed-inference parameterization
Following prior work (Sohl-Dickstein et al., 2015; Ho et al., 2020; Huang et al., 2021), we let the inference SDE (9) be defined as a simple noise process taking observed data to unstructured noise:
dY = Un dt + V Vo dB.,
(14)
where Up V logpo and V is the tangential projection matrix; that is, (z(
Ej zi (Pa)sk Bay sion (Girolami & Calderhead, 2011). As long as Po satisfies log-Soboler inequality, the marginal distribution of Y. (i.e. the aggregated posterior) converges to Po at a linear rate in the KL divergence (Wang et al., 2020). For compact manifolds,
af for any smooth function f. This is known as the Riemannian Langevin diffu-
set po to be the uniform density, which means Un = O, and (14) is reduced to the extrinsic construction of Brownian motion on M (Hsu, 2002, Section 1.2). The benefits of this fixedinferennee parameterization are the following:
Stable and Efficient Training. With the fixed-inference parameterization we do not need to opti- mize the vector fields that generate Y., and the Riemannian CT-ELBO can be rewritten as:
1 (V v, Ellog Po (Yr)) - /."" Ey, |a((Y,, s)1 V. Vo (v
Yo ds. (15)
where the first term is constant wrt the model parameters (or it can be optimized separately if we want to refine the prior), and the time integral of the second term can be estimated via importance sampling (see Section 3.3). A sample of Y. can be drawn cheaply by numerically integrating (14), without requiring a stringent error tolerance (see Section 5.2 for an empirical analysis), which allows us estimate the time integral in (15) by evaluating aYY,, s) at single time step s only.
Simplified Riemannian CT-ELBO. The CT-ELBO can be simplified as the differential operator V. V, applied to V yields zero vector when V the tangential projection.
Proposition 2. VfV is the tangential projection matrix, then (V vv)V 0.
This means that we can express the generative SDE Vo using the variational parameter a via dx = (Va(X,T - t) - UX(X x, T - t)) dt - V dB,
(16)
5

---

with the corresponding Riemannian CT-ELBB:
Ellog (o / T
1 + v. (Va [6) Yo
By,
Iall 2
ds.
(17)
3.3 Variance reduction
The inference process can be more generally defined to account for time reparameterization. In fact, this leads to an equivalent model if one can find an invariant representation of the temporal variable. Learning this time rescaling can help to reduce variance (Kingma et al., 2021).
In principle, we can adopt the same methodology, but this would further complicate the parameteri- zation of the model. Alternatively, we opt for simpler view for variance reduction via importance sampling. We estimate the time integral "f... ds" in (17) using the following estimator:
1 I g(3) (2 [al +V, (Va Uo)
where ~ g(s) and Y, q(2, Yo),
(18)
where (()) is proposal density supported on [0, T7. We parameterize ((s) using a ID monotone flow (Huang et al., 2018). As the expected value of this estimator is the same the time integral in (17), it is unbiased. However, this means we cannot train the proposal distribution g(8) by max- imizing this objective, since the gradient wrt the parameters of g(s) is zero in expectation. Instead, we minimize the variance of the estimator by following the stochastic gradient wrt g(s)
(19) The latter can be optimized using the reparameterization trick (Kingma & Welling, 2014) and is well-known variance reduction method in a multitude of settings (Luo et al., 2020; Tucker et al., 2017). can be seen as minimizing the =d-dieergence from density proportional to the magnitude of Ey, [Z] (Dieng al., 2017; Müller et al., 2019).
Va(+) Var(I) V (g(
E[T2]
3.4 Connection to score matching
In the Euclidean case, it can be shown that maximizing the variational lower bound of the fixed- inference diffusion model (16) is equivalent to score matching (Ho et al., 2020; Huang et al., 2021; Song et al., 2021a). In this section, we extend this connection its Riemannian counterpart.
Let q s) be the density of Y, following (14), marginalizing out the data distribution q(3., 0). The score function is the Riemannian gradient of the log-density V, log 9. The following theorem tells us that we can create a family of inference and generative SDEs that induce the same marginal distributions over Y, and XT-, (16) if we have access to its score.
Theorem 3 (Marginally equivalent SDEs). For X S 1, the marginal distributions of XT and Y. of the processes defined below
X Vg logg ds -1 oV
dY Uo
Yo 9(:50)
(20)
2
dx (11-
A
V logg -U)) dt dt+ -I do Vdß,
Xo q(:,T) T)
(21)
both have the density q(: s). In particular, X gives rise equivalent ODE.
This suggests if we can approximate the score function, and plug into the reverse process (21), we obtain a time-reversed process that induces approximately the same marginals.
Theorem (Score matching equivalency). For X 1, let Ex denote the Riemannian CT- ELBO of the generative process (21), with V log 9 replaced by an approximate score Sp, and with (20) being the inference SDE, Assume Sq a compactly supported smooth vector. Then
1x,1811 -Ci 1'' Ey, ||8o - v. gog ds Cz
(22)

---

Earthquake
Fire
>4 1.0
>2 1.0
0.8
0.8
likelhood
likllhoodd
0.6
0.6
0.4
0.4
0.2
0.2
Flood
Volcano
Bkelhood
0.6
0.6
022
0.2
Figure 1: Density of models trained on earth datasets. Red dots are samples from the test set.
where C1 > and C2 are constants wrt
The first implication of the theorem that maximizing the Riemannian CT-ELBO of the plug-in reverse process is equivalent to minimizing the Riemannian score-min more-matching loss. Second, if we set 0, from (135) (in the appendix), we have Va = Sg. which is exactly the fixed-inference training in $3.2. That is, the vector Va a trained using equation (17) is actually approximate score, allowing us to extract an equivalent ODE by substituting Va for V, log q in (20,21) by setting =.
Related work
Diffusion models. Diffusion models can be viewed from two different but ultimately complimentary perspectives. The first approach leverages score based generative models (Song & Ermon, 2019; Song et al., 2021b), while the second approach treats generative modeling as inverting a fixed noise- injecting process (Sohl-Dickstein et al., 2015; Ho et al., 2020). Finally, continuous models can also be embedded within maximum likelihood framework (Huang al., 2021: Song et al., 2021a), which represents the special case of prescribing a flat geometry-ia. Euclidean to the generative model and completely generalized by the theory developed in this work.
diffusion
Riemannian Generative Models. Generative models beyond Euclidean manifolds have recently risen to prominence with early efforts focusing on constant curvature manifolds (Bose al., 2020; Rezende et al.. 2020). Another line of work extends continuous-time flows (Chen et al., 201 8a) to more general Riemannian manifolds (Lou al., 2020; Mathieu & Nickel, 2020; Falorsi & Forré, 2020). To avoid explicitly solving an ODE during training, Rozen et al. (2021) propose Moser Flow whose objective involves computing the Riemannian divergence of parametrized vector field. Concurrent to our work, De Bortoli et al. (2022) develop Riemannian score-based generative models for compact manifolds like the Sphere. While similar in endeavor, RDMs are couched within the the maximum likelihood framework. As result our approach is directly amenable to variance reduction techniques via importance sampling and likelihood estimation. Moreover, our approach is also applicable to non-compact manifolds such as hyperbolic spaces, and we demonstrate this in our experiments on a larger variety manifolds including the orthogonal group and toroids.
Experiments
We investigate the empirical caliber of RDMs on a range of manifolds. We instantiate RDMs by parametrizing in (16) using an MLP and maximize the CT-ELBO (17). We report our detailed training procedure-including selected hyperparameters-Bae all models in $D.
5.1 Sphere
For spherical manifolds, we model the datasets compiled by Mathieu & Nickel (2020), which consist of earth and climate science events on the surface of the earth such as volcanoes (NGDC/WDS, 2022b), carthquakes (NGDC/WDS, 2022a), floods (Brakenridge. 2017), and fires (EOSDIS, 2020).
7

---

Volcano
Earthquake
Flood
Fire
Mixture of Kent
-0.80 -0.80
0.33
0.73
-1.1840.00
Riemannian (Mathieu Nickel, 2020)
+0.16
0.194
0.90=n.08
-066
Moser Flow (Rozon et al., 2021)
-2.02
0.099
0.62=0.04
-0.03
Stereographic Score-Based
4.18 40.30
-0.044
1.31
0.28 2..
Riemannian Score-Based (De Bortoli 2022)
5.00 g0000...0 a6
0.21
0.52go.00
-1.2440.0r
20.97
0.40 20.05
0.43
-1.88.000 20.05
Datusel size
827
6120
4875
12809
Table NLL scores each method on earth datasets. Bold shows best results (up statistical significance).
Means and standard deviations are calculated over 5 runs. Baselines taken from De Bortoli et al. (2022).
uniform q(s)
diectsomplin
diectsor
leameed q(s)
2
direct loss
5
-1
100
200
300
0
10
20
30
10 thousands of iterations
20 30
thousands of iterations
thousands of iterations
Figure 2: Variance reduction with importance sampling.
Figure 3: Direct sampling numerical integration of Brownian tion. Numbers legends indicate number time steps.
Gentral ml
2.
Glycine 4.5
Pre Pro
2.1
Proline
100
0.
J00
E30
00
23 0
103
40
200L
1:.3
13.0
-100
100
-10
100
-10
102
11.4
100
10.1
0
Figure 4: Ramachandran contour plot of the model density for protein datasets. Red dots are set test samples.
In Table for each dataset we report average and standard deviation of test negative log likelihood on 5 different with different splits of the dataset. In Figure we plot the model density in blue while the test data is depicted with red dots.
Variance reduction. We demonstrate the effect of applying variance reduction on optimizing the Riemannian CT-ELBO (17) using the earthquake dataset. As shown in Figure learning an impor- tance sampling proposal effectively lowers the variance and speeds up training.
5.2 Tori
For tori, we use the list of 500 high-resolution proteins compiled in Lovell et al. (2003) and select 113 RNA sequences listed in Murray et al. (2003). Each macromolecule is divided into multiple monomers, and the joint structure is discarded-we model the lower dimensional density of the backbone conformation of the monomer. For the protein data, this corresponds of the amino acid. As one of the angles is normally 180", we also discard it, and model the density nucleotide described by torsion angles in
torsion angles
over the 2D torus. For the RNA data, the monomer the backbone, represented by 7D torus. For protein, we divide the dataset by the type of side chain attached the amino acid, resulting in 4 datasets, and we discard the nucleobases of the RNA.
In Table 2 we report the NLL of our model. Our baseline is a mixture 4, 096 power spherical distributions (De Cao & Aziz, 2020, MoPS). We observe that RDM outperforms the baseline across the board, and the difference most noticeable for the RNA data, which has higher dimensionality.
Numerical integration ablation. We estimate the loss (17) by integrating the inference SDE M. To study the effect integration error, we experiment with various numbers time steps evenly spaced between [0, s] on Glycine. Also, can directly sample the Brownian motion tori

---

General
Glycine
Proline
Pre-Pro
RNA
MoPS RDM
1.15:0.002
2.08
0.27 100000 000
10.019
4.08
1.04 10.012
1.97 10.012
0. 0.12
1.24 1.24
-3.7
Dataset size
138208
13283
7634
6910
9478
Table 2: Negative test likeliiiel for each method Tori datasets. Bold shows best results (up statistical significance). Means and standard deviations are caleulated over runs.
without numerical integration, we use as reference (termed direct loss) for comparison. Figure 3 shows while fewer time steps tend underestimate the loss, the model trained with 100 time steps is already indistinguishable from the one trained with direct sampling. We also find numerical integration is not significant overhead each experiment takes approximately the same wall-clock time with identical setups. This is because the inference path does not involve the neural module a.
5.3 Hyperbolic Manifolds
Hyperbolic manifolds provide example whose closest- point projection is not cheap to obtain, and a claimed closest-point projection in recent literature is in fact not the closest Euclidean projection (Skopek et al., 2019) (see $C for more details). To demonstrate the general- ity of our framework, we model the synthetic datasets in Figure 5, first introduced by Bose al. (2020); Lou et al. (2020). Since hyperbolic manifolds are not compact, we need non-zero drift to ensure the inference processs is not dissipative. We define the prior as the standard normal distribution on the ye-plane and let Ug be 1 Vg log Po. so that Y, will revert back to the origin.
Figure 5: Hyperbolic Manifold. Top: data. Bottom: learned Density
5.4 Special Orthogonal Group
Another example whose closest-point projection ex- pensive to compute is the orthogonal group, as it requires performing the singular value decomposition. To evalu- our framework on this matrix group, we generate data using the synthetic multimodal density defined on SO(3) from Brofos et al. (2021). We view as a submanifold embedded in 3x3 therefore = 3 and m = 9. We use the projected Hutchinson to estimate the Riemannian divergence. Since the data are 3D rotational matrices, can visualize them using the Euler an- gles. We plotted the data density and the learned model density Figure 6, where each coordinate represents the rotation around that particular axis.
Figure 6: SO(3). Left: synthetic multi- modal density. Right: learned density.
6 Conclusion
In this paper, we introduce RDMs that extend continuous diffusion models to arbitrary Riemannian manifolds-incluune challenging non-compact manifolds like hyperbolic spaces. We provide a variational framework to train RDMs by optimizing a novel objective, the Riemannian TTmee ELBO. To enable efficient and stable training we provide several key tools such as a fixed-inference paramterization of the SDE in the ambient space, new methodological techniques compute the Riemannian divergence, well an importance sampling procedure with respect to the time integral to reduce the variance of the loss. On a theoretical front, we also show deep connections between our proposed variational framework and Riemannian score matching through the construction of marginally equivalent SDEs. Finally, we complement our theory by constructing RDMs that achieve state-of-the-an art performance on density estimation on geoscience datasets, protein/RNa data on toroidal, and synthetic data on hyperbolic and orthogonal-group manifolds.

---

References
Boomsma, W., Mardia, K. V. Taylor, C. C., Ferkinghoff- Baggg. J., Krogh, A., and Hamelryck, A generative, probabilisticc model of local protein structure. Proceedings the National Academy of Sciences, 105(26):8932-89223
2008.
Bose, J., Smofsky, A., Liao, R., Panangaden, P., and Hamilton, W. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning, pp. 1045-1055. PMLR, 2020.
Brakenridge, G. Global active archive of large flood events, http://ftocdobberattoryy colorado edu/ Archives/ index.idxxi./. htt.iieeeeddee..mmm., 2017.
Brehmer, J. and Cranmer, K. Flows for simultaneous manifold learning and density estimation. Advances in Neural Information Processing Systems, 33:442-455, 2020.
Brofos, J. A., Brubaker, M. A., and Lederman, R. R. Manifold density estimation via generalized dequantization. arXiv preprint arXiv:2102.07145, 2021.
Burda, Y., Grosse, R., and Salakhutdinov, R. Importance weighted autoencoders. arXiv preprint arXiv: 1509.00599, 2015.
Burrage, K., Burrage, P., and Tian, T. Numerical methods for strong solutions of stochastic differen- tial equations: an overview. Proceedings of the Royal Society of London. Series A: Mathematical, Physical and Engineering Sciences, 46072041) 2004.
Chen, R. T. Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 6572-6583. 2018a.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equa- tions. Advances in Neural Information Processing Systems, 2018b.
Chen, R. T. Q., Amos, B., and Nickel, M. Learning neural event functions for ordinary differential equations. International Conference on Learning Representations, 2021.
Chirikjian, G. S. Stochastic Models, Information Theory, and Lie Groups, Volume I: Classical Results and Geometric Methods. Springer Science & Business Media, 2009.
De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, W., and Doucet, A. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.027700 0200020.222....... 2022.
De Cao, N. and Aziz, W. The power spherical distribution, 2020.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. arXiv preprint 105.v32105.05233, 2021.
Dieng, A. B., Tran, D., Ranganath, R., Paisley, J., and Blei, D. Variational inference via 'cci upper bound minimization. Advances in Neural Information Processing Systems, 30, 2017.
EOSDIS. Land, atmosphere near real-time capability for eos (lance) system operated by nasa's earth science data and information system (esdis). https://earthdata. nasa nasa.gov/ Marth-observation wwwervvtionddtttaaaaea
2020.
Falorsi, L. and Forré, P. Neural ordinary differential equations on manifolds. arXiv preprint
arXiv:2006.06666 06663, 2020.
Frellsen, J., Moltke, I., Thimm, M., Mardia, K. V., Ferkinghoff Borg, J., and Hamelryck, T. A probabilistic model of ma conformational space. PLoS computational biology, 100)040 2009.
Girolami, M. and Calderhead, B. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society: Series (Statistical Methodology). 73(2): 2011.
Gunther, M. Isometric embeddings of riemannian manifolds, kyoto, 1990. In Proc. Intern. Congr. Math., pp. 1137-1143, Math. Soc. Japan, 1991.
10

---

Ho. J., Jain, A.. and Abbeel, P. Denoising diffusion probabilistic models. In Advances in neural information processing systems, 2020.
Hsu, E. P. Stochastic analysis on manifolds. Number 38. American Mathemattcal Soc., 2002.
Huang, C.-W., Krueger, D., Lacoste, A., and Courville, A. Neural autoregressive flows. In Interna- tional Conference on Machine Learning, pp. 2078-2087, 2018.
Huang, C.-W., Lim, J. H., and Courville, A. A variational perspective on diffusion-based generative models and score matching. arXiv preprint vrXiv22166.02800 2021.
Hunter, J. D. Matplotlib: 2d graphics environment. Computing in science & engineering, 9(3): 90, 2007.
Hutchinson, M. F. stochastic estimator of the trace of the influence matrix for laplacian smoothing and Computation, 18(3):1059-1076 1989.
splines. Communications in
Inc., P. T. Collaborative data science, 2015. URL https://plot..t::
Kazhdan, M., Bolitho, M., and Hoppe, H. Poisson surface reconstruction. Proceedings of the fourth Eurographics symposium Geometry processing, volume 7, 2006.
Kidger, P., Foster, J., Li, X., Oberhauser, H., and Lyons, Neural SDEs as limmnte-Dimensionall
GANs. International Conference on Machine Learning. 2021.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible Ixl convolutions. Advances in neural information processing systems, 31, 2018.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Kingma, D. P., Salimans, T., Poole, B., and Ho, Variational diffusion models. arXiv preprint arXiv:2107.006300 2021.
Lee, M. Introduction to Smooth Manifolds. Springer, 2013.
Lee, J. M. Introduction to Riemannian manifolds. Springer, 2018.
Li, X., Wong, L., Chen, R. T., and Duvenaud, D. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870-3882. PMLR, 2020.
Lou, A., Lim, D.,, Katsman, I., Huang, L., Jiang, Q., Lim, N., and De Sa, C. M. Neural manifold ordinary differential equations. Advances in Neural Information Processing Systems, 33:17545 17558, 2020.
Lovell, S. C., Davis, W., Arendall III, W. B., De Bakker, I., Word, J. M., Prisant, M. G., Richardson, S., and Richardson, D. C. Structure validation by сa geometry: Φ, s and cB deviation. Proteins: Structure, Function, and Bioinformatics, 50(3):437-450, 2003.
Luo, Y., Beatson, A., Norouzi, M., Zhu, J., Duvenaud, D., Adams, R. P., and Chen, R. T. Sumo: Unbiased estimation log marginal probability for latent variable models. arXiv preprint 001::2004400353, 2020.
Mardia, K. and Jupp. P. Directional statistics, volume 494. John Wiley & Sons, 2009.
Mardia, K. V., Hughes, G., Taylor, C. C., and Singh, H. A multivariate von mises distribution with applications bioinformatics. Canadian Journal of Statistics, 36(1):99-109 2008
Mathieu, and Nickel, M. Riemannian continuous normalizing flows.
arXiv preprint
10600 2020.

---

Met Office. Cartopy: cartographic python library with Matplotlib interface. Exeter, Devon, 2010 2015. URL https://scitooll..oo org. uk/ cartopy.
Müller, T., McWilliams, B.. Rousselle, F., Gross, M., and Novák, J. Neural importance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19 2019.
Murray, L. J., Arendall, W. B., Richardson, D. C., and Richardson, S. Rna backbone is rotameric. Proceedings of the National Academy of Sciences, 10024))13904-1390, 2003.
NGDC/WDS. Ncei/wds global significant earthquake database. https:/ /wtp://wwww.... www..nnei....................... noa.. gov/ aatees///etaatttllandda http://///ictcccaaaaaon/n//////a 2022a.
nowa. nww..ngo...nn......ooo... mgg
NGDC/WDS. Ncei/wds global significant volcanic eruptions database.
https:
ncei. noaa. www.anww..aaaaeee. httpa/////cccccce/eaaaaaaaa/aaaa ngdc.mgg ha10rd 2022b.
htwpa/iinggaage///ia//////iieiee noaa.
Øksendal, B. Stochastic differential equations. In Stochastic differential equations, pp. 65.
Springer, 2003.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, high-performance deep learning library. In Advances in neural information processing systems, pp. 8026 8037, 2019.
Protter, P. E. Stochastic Integration and Differential Equations. volume 21 of Stochastic Mod- elling and Applied Probability. Springer Berlin Heidelberg, Berlin, Heidelberg, 2005. ISBN 9783642055607 9783662 100615. doi: 000077978-3-6622 0-17778--6662-100111 10061-5.
Ratcliffe, J. G. Foundations of Hyperbolic Manifolds. Number 149 in Graduate Texts in Mathemat- ics. Springer- Verlag, 1994.
Rezende, D. J., Papamakarios, G., Racaniere, S., Albergo, M., Kanwar, G., Shanahan, P., and Cran- mer, K. Normalizing flows on tori and spheres. In International Conference Machine Learning. pp. 8083-8092. PMLR, 2020.
Rozen, N., Grover, A., Nickel, M., and Lipman, Y. Moser flow:
Divergence-based generative modeling manifolds. Advances in Neural Information Processing Systems, 34, 2021.
Rudin, W. Real and Complex Analysis. McGraw-Hii Hill, 1987.
Rudin, W. al. Principles of mathematical analysis, volume 3. McGraw-hill New York, 1976.
Skopek, O., Ganea, O.-E., and Bécigneul, G. Mixed-curvature variational autoencoders. arXiv
preprint arXiv:1911.08411 2019.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp- 2256-2265, PMLR, 2015.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances in neural information processing systems, 2019.
Song, Y.. Durkan, C., Murray, I,, and Ermon, S. Maximum likelihood training of score-based diffusion models. Advances in Newral Information Processing Systems, 34, 2021a.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Thalmaier, A. Sde and pde: Solving pde by running a brownian motion. 2021.
Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and Sohl-Dickstein, J. Rebar: Low-variance,
unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
Wang, X., Lei. Q., and Panageas, Fast convergence of langevin dynamies on manifold: Geodesics meet log-sobolev. Advances in Neural Information Processing Systems, 33:1889; 118888888904444 2020.
12


Reviews:
[{"review_id": "yolitQQRx3d", "reviewer": "Reviewer_KTty", "summary": "The paper derives a variational framework for likelihood estimation of diffusion models on Riemannian manifolds. This is based on prior work on continuous time ELBOs and score matching. The authors generalize this to Riemannian manifolds and address computational issues including fast computation of the Riemannian divergence. The methods are experimentally demonstrated on several manifolds.", "questions": "* line 172-173: I didn't get what is meant by \"is evaluated only once as opposed to estimating the ELBO by integrating over the entire path of Y_s\". What is evaluated only once? The entire sample path is still used?", "limitations": "yes", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* well-written and technically sound paper\n* addresses an important class of problems\n* the presented techniques are mathematically interesting and non-trivial\n* the paper couples stochastic process theory, variational inference and geometry\n\nWeaknesses:\n* the CT-ELBO was derived in earlier paper in the Euclidean setting. To some extend, the extension to the geometric context is not groundbreaking. However, it is in no way trivial either\n\nIn general, I found the paper a very interesting read.\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": "UkMdjXk_IPW", "reviewer": "Reviewer_6m1R", "summary": "The paper proposes a general diffusion model for Riemannian manifolds. It introduces a variational framework based on Riemannian Feynman-Kac representation and Giransov’s theorem, with which a Rimannian continuous-time ELBO is derived. To compute the resulting ELBO, this paper further suggests a QR decomposition-based method for scalable unbiased estimation. Besides, it also provides a variance reduction technique to approximate the Remannian ELBO objective. \n", "questions": "In Eq. (27) of Appendix, what’s the meaning of  \\frac{\\partial \\varphi_{j}^{-1}}{\\partial \\tilde{x}_{i}}? If my understanding is right, \\varphi_{j}^{-1} should be understood as (\\iota \\circ \\varphi)_j^{-1}\n\nRegarding the spherical dataset, it is not clear to me why the used data lie in a sphere? Maybe some explanation or references will help interested readers.\n\nSmall typos should be corrected. For example, Eq. (24) in Appendix, the v in LHS should be capitalized.\n", "limitations": "It seems this paper does not discuss its limitations. It will be highly appreciated if the authors can address my concerns/questions mentioned above.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\nBased on SDE, this paper generalized the Euclidean diffusion model into a Riemannian one. The generalization seems to enjoy a solid mathematical background.\n\nThe experiments demonstrate the effectiveness of the proposed method on the data that are on spherical manifolds, product of spherical manifolds, hyperbolic spaces and orthogonal groups.\n \nWeakness:\n\nThis paper requests a lot of effort to understand its background and the proposed idea.\n\nThe presentation of background (Sec.2) is also very challenging to understand. For example, regarding Euclidean diffusion models, it might be better to further clarify how the diffusion model can be formulated as a fixed Markov chain that is presented in Line 24-26. Besides, it should be necessary to indicate the meanings of \\mu and \\sigma after Eq. (1). The same suggestion could be applied to other equations like Eq.(2). \n\nAs for the computation of the derived CT-ELBO, the suggested QR decomposition based method is an efficient approach. I cannot really understand why the paper further introduces a variance reduction technique to approximate the CT-ELBO. What is the major benefit of the used variance reduction over the QR decomposition based method? It seems that this paper does not study this benefit empirically.\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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "ExhqQgmDCL4", "reviewer": "Reviewer_JCSo", "summary": "This paper extends the application of continuous-time diffusion models into Riemannian manifold. The paper proposes a new method of calculating the Riemannian divergence for computing the Riemannian lower variational bound objective function.  In particular, the authors also makes an insightful connection with Riemannian score matching.", "questions": "- Why is there only one baseline model comparison in the tori results specifically not having the comparison with the Riemannian Score-Based Model?\n- In some tasks, SGM have performed better in certain image generation than diffusion models (e.g. CIFAR-10, CelebA 256x256). In this task, it seems that this model performs better than its score-based counter part. Is there any reason to believe that diffusion models should theoretically perform better in the Riemann space?\n- Projection can be a difficult task, especially projecting onto a high dimensional complex plane. Does this have any detrimental effects on the training/sampling time?", "limitations": "The authors address their limitations (i.e. algorithmic decisions).", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Overall, the paper was well organized and clearly written.\n- The shadowed boxes on the salient equations help follow the derivations.\n- The idea is a novel application/expansion of diffusion models as they have (to the extent of my knowledge) not have been explored on non-Euclidean spaces. \n- The connection between score-based models (in the Riemannian case) is a welcoming and insightful addition.\n\nWeaknesses:\n- The experiments show SOTA results and seem very promising, however, there seems to be a lack of analysis on the results. \n- Limited information on training/sampling complexity.\n \nAdditional Remarks:\n- Pg. 4 Line 136 \"showed that that...\" typo", "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