Riemannian Neural SDE: Learning Stochastic Representations on Manifolds
Authors: Sung Woo Park
Hyomin Kim
Kyungjae Lee
Junseok Kwon
Venue: NeurIPS 2022 Accept (2022)
Keywords: Stochastic representation on Manifolds
Riemannian neural stochastic differential equation

Abstract: In recent years, the neural stochastic differential equation (NSDE) has gained attention for modeling stochastic representations with great success in various types of applications. However, it typically loses expressivity when the data representation is manifold-valued. To address this issue, we suggest a principled method for expressing the stochastic representation with the Riemannian neural SDE (RNSDE), which extends the conventional Euclidean NSDE. Empirical results for various tasks demonstrate that the proposed method significantly outperforms baseline methods.

Riemannian Neural SDE: Learning Stochastic Representations on Manifolds
Sung Woo Park', Hyomin Kim?, Kyungjae Lee Junseok Kwon
School of Computer Science and Engineering, Chung-Ang University, Korea 1,3,4 Artificial Intelligence Graduate School, Chung-Ang University, Korea
LGG Al Research
ppswkiki@gmill com, *icecream1260cau ac.kr kr
'kyungjae. aatggjaaeeeeeiiii. cau. ac.kr, kr, "jskwonccau.ac. ac.kr
Abstract
In recent years, the neural stochastic differential equation (NSDE) has gained attention for modeling stochastic representations with great success in various types of applications. However, it typically loses expressivity when the data representation is manifold-valued. To address this issue, we suggest a principled method for expressing the stochastic representation with the Riemannian neural SDE (RNSDE), which extends the conventional Euclidean NSDE. Empirical results for various tasks demonstrate that the proposed method significantly outperforms baseline methods.
Introduction
Recently, there has been great success in modeling stochastic dynamical systems for complex data representations with spa- tially high stochasticity. Particularly, recent studies have uti- lized the stochastic differential equation (SDE) as fundamen- tal probabilistic model to express the transition of stochastic states. The pioneering work, Latent SDE (14), introduced the deep-learning framework to utilize SDE for modeling spatio- temporal representations. Following their work, parameterized reverse-SDE (l) was adopted model score-based generative models (27). Recently, a controlled SDE combined with a stochastic optimal control theoretical framework was intro- duced to model time series (21). These studies provided a new method for describing stochastic representations and illustrated the superiority of neural SDE models in various real-world applications. Although remarkable progress has been recorded in recent studies, the primary interest in using SDE has focused Euclidean geometry (i.e., (R.). Thus, conventional approaches inevitably lose their expressivity when data representation is defined in generic geometry, such as Riemannian manifolds (i.e., M). In this paper, we solve the aforementioned problem by introducing the Riemannian neural SDE (RNSDE), which can model the stochastic representation on manifolds. The proposed RNSDE is a natural extension of conventional Euclidean SDEs that defines intrinsic stochastic transitions while fully considering the geometric structure. Specifically, we follow the www-e-lwwowwwlllleee representation is expressed on the frame bundle.
Figure 1: Brownian motions on torus T (black) and R (red).
interpretation of the diffusion Riemannian SDE, where the stochastic
Example. To illustrate the problem mentioned above, we represent the stochastic trajectories of Riemannian (black) and Euclidean (red) Brownian motions on 2- in Fig As shown in the figure, the Buclidean SDE unaware of the underlying geometry (i.e., torus) and its trajectory easily steps out of the curved surface. Accordingly, the use of Euclidean SDE is problematic and lacks geometry-awareen
contrast, the proposed Riemannian SDE respects its underlying geometry and yields a trajectory that clearly lies on the manifold.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Contribution. The main contribution of this work is to suggest general framework for modeling the stochastic representations on Riemannian manifolds. Specifically, we define the diffuusion process on manifolds based on the Bells-Elworthy-Mallln Maaliaviin interpretation. Given well-defined Riemannian SDEs, we introduce a novel Markov diffusive Kantorovich dual formulation and suggest gradient flow-based algorithm to train the neural network.
Notation. Einstein's 's summation conventions are used throughout the paper. The Euclidean and Riemannian norms are denoted as ||u|], IIull, with vectors u E IR", E ToMM, respectively. Os, == a and 0, øt denote spatial and temporal partial derivatives, respectively. We consider the de;
probability space (M, Σ, IP) where the filtration (FiJosesT is augmented by the Brownian motion (B) in the time interval [0, T].
Riemannian geometry. In this paper, we consider complete and oriented n-dimensional Riemannian manifolds (M. g) equipped with Riemannian metric g. In particular, we focus on embedded compact sub-manifolds .M (e.g.. sphere and torus) of the ambient Euclidean space, M C RD,, for n < D < 2r The metric tensor is expressed in the matrix form G = [9t]] (S.jgn Similarly, we define the Riemannian co-metric as the inverse of matrix G as G = [g") The Christoffel symbol with respect metric g is denoted by the tensor form [T Islssis, For the tangent space Tx.M along the sample path, we denote {0) Jsisnn as the moving frame. The Riemannian volume measure is denoted as d'V = v det ((xx)ddz. The orthogonal frame bundle and its corresponding local coordinate function are denoted as OM and &b U C R" M, respectively.
2 Riemannian Neural Stochastic Differential Equation
We introduce novel Riemannian SDE that defines the stochastic representations on Riemannian manifolds, Additionally, explicit local representations and numerical scheme for implementation are presented. We begin by introducing the conventional neural SDE and extend the discussion.
Euclidean Neural SDE. The major interest of neural SDE is train the parameterized stochastic process X defined as solution to the following Stratonovich SDE (20):
ax h (t, x,; 0) dt + a (X²:0) dB.,
where B, = -B, BT] denotes the n-dimenssinal standard Brownian motion on R", and Eu- clidean drift, and diffusion functions (i.e., h [o, T] x Rd x è - IR and Rd x θ - Rdxn respectively) are parameterized by the neural network 0 E O. Euclidean SDE has obvious limitations in handling EunFuulidden datasets, because it produces outputs in ambient space. Therefore, we propose Riemannian neural SDE, which notably respects its underlying geometry.
Riemannian Neural SDE. Similar to Euclidean NSDE, our goal is to suggest a differential equation on manifolds. Specifically, we suggest the Riemannian Neural Stochastic Differential Equation (RNSDE) parameterized by two novel terms (i.e., (A) and (B)), defined as follows:
dx W (1, X);0) di 3(0)= (X?) dB
(2)
(A)
(B)
(A) Neural Potential Field. In the first term, refer to the parameterized vector field W (t; = w Xe: 0)05 [o. T| x M x + M as the neural potential field, where the set of orthonormal tangent vectors h8shsssn =
Wit Nea.
M + M is the moving frame on the trajectories of stochastic dynamics. This term defines the de- terministic drift on manifolds respect to the inferred network decision W (t; 0); based on the current state information (t, X.). The neural network W produces tangent vectors (red arrows in Fig which consecutively drive the stochastic trajectory to determine the next direction on Tx (white planes in F
M
Figure 2: Neural potential fields.
While the Euclidean drift h in (T) limited producing vectors in the ambient space R the neural potential field W in (2) produces vectors in the tangent space, thus respecting the geometric surfaces.
'The Whitney's 's theorem (II) guarantees the existence of smooth embedding on R D=2n, Thus, we set the maximal dimension of D to be bounded by 2n.
2

---

(B) Stochastic Development. The second term, called stochastic development 8 9) controls the diffusive/stochaatic behavior of the proposed SDE on manifolds. The function π OM + M is canonical projection that maps the orthogonal frame bundle to model manifolds, where U (x OM. This term projects the stochastic trajectory defined on the orthogonal frame bundle Ug OM, which locally behaves like flat Euclidean SDE, onto the curved model manifolds. This procedure is called Eells-
interpretation to present the SDE on the manifold. Intuitively, the stochastic development could be considered the procedure for rolling the trajectory on curved surface (e. over flat plane (e.g., Euclidean SDE in (.. Meanwhile, the parameterized scalar-valued function B θ + R+ controls the diffusivity of the process X It is noteworthy that the solution ( is the Brownian motion on M, when W and 3 This
case is illustrated in Fig
Local Representations. Although the proposed RNSDE in ( is mathematically well-defined, it is not directly applicable the numerical al- gorithm because of the abstract formulations of horizontal vector fields. Thus, for clarity, provide explicit local representations expressed as Itô's SDE, follows:
Naural Potensal
0.1o Dr. X
Stocharsi Develosmans
hou
Proposition 1. (Local Representations) Let X be a parameterized stochastic process which is the solution the RNSDE in (2). In local coordinate system, the proposed RNSDE can be rewritten as the Itô's representation:
0,6o X.
Figure 3: Local representations of the RNSDE.
X²- = B(0) 1-" Vor(x()au3 + '." [ -5(00 2 9 (XP)) + (X,: (x
ds. (3)
The detailed derivation is provided in Appendix C.3. The local-coordinate-asaee representation in (3) simultaneously requires geometric terms including metric tensor Christoffel symbol J'jk- and frame 07. along the sample path X .M, that is caleulated during sample propagation. The effect of stochastic development can be inferred from the newly appearing terms (i.e., the first and second terms in ). It transforms the Euclidean diffusion induced by standard Brownian motion B, into Riemannian diffusion on manifolds by considering the curvature effect (i.e.,g.1).
Euler-Me Euler-Maruyane Maruyama Scheme. To approximate the solution to the RNSDE in local coordinate system, we apply the Euler-Maruyama scheme to simulate the SDE in (3). Let t E = {tohi be the total simulation time and A, be priorly fixed time interval. This scheme provides sampled stochastic particles approximated
& Am,,A,, (X where the mean m, - =Bg jk + W considers both the geometrical properties (i.e., g, T) and neural potential fields. It is noteworthy that these two components are computed simultaneously for every temporal states. Fig 3 shows the detailed numerical scheme for the local representation.
3 Learning Representations on Manifolds
In this section, we describe an efficient method for learning densities over manifolds. We begin by analyzing the density estimation methods used in previous studies. Let assume that both the model and target measures admit densities pe (dVV = oh and pod'V = vt, respectively. Conventional methods (e.g., RCNF (17) define the minimization problem of relative entropy H(-I-) between the model probability measure n and the target probability measure Vt, follows:
dup dv:
min H (p. (1)) min log #E
min [log ps PE0
log pul
(4)
To compute the numerical estimation in (4) conventional methods require preprocessing step (e.g., Gaussian KDE) to obtain an explicit and accurate estimation of the target density Pv. Owing to this shorteoming, the induced performance is highly dependent the preprocessing step, which must be correctly tuned for different applications and manifolds. In contrast, the proposed method directly utilizes set of observed particles sampled from the target object. Specifically, it is constructed
"This term also called "rolling without slipping" to emphasize the intuition [6. g
3

---

from the N sampled particles (1/ defined as Σ1 bng Consequently, the proposed approach does not require any distributional information to represent the target object.
Markov Semi-group. Because our main objective is to find a solution X. of the RNSDE, the model measure is defined as a law of the stochastic path 1 = P(X E E ~ X E M, where X is the solution to the proposed RNSDE. In spite of the rigorous definition, the model measure ii is not applicable numerical algorithms owing to its abstract form. To resolve this issue, suggest algebraic trick for evaluating the model measure u.. Specifically, we utilize the Markov property of the proposed RNSDE to identify the probabilistic structures of the model measure. Let f = (X,) [o]] be Markov group (28) of X². Then, by duality, we can define the probability measure called a dual semi-group Pe satisfying the following equality:
E Ex-nn (X( Xo x] Ez [P fxx)] pf - Exp~ut Exp=x;1((Xx)) (5)
where Hi indicates the initial state of the RNSDE. Evidently, the expectation according the measure 14, is identified as the expectation with the dual semi-group μu p. which can be easily computed using the relation in (5). Throughout this paper, the model measure is interchangeably denoted as pf and u
Objective Function. Similar to previous studies, the proposed method aims to minimize the discrepancy between the model and target measures; however, a different approach adopted. Specifically, we utilize the static Schrödinger bridge to define the discrepancy between P. and vy: Definition (Static Schrödinger Bridge ([2)) Let us assume that there exists a coupling no such that To & Ph: 8 vr for fixed time I. E [0, T) where we E <I(PP". ,u)) the set of parameterized couplings. Let R be the reference measure satisfying "(r"""Qu1)) x e- (2 with the Riemannian distance 7,. Then, the (static) Schrödinger bridge x² a unique solution to the following minimization problem:
dR
[ log
dr dR (x, vi) x.)
T = arg min T.EIIP" (M))
(TT = arg min meEl(Pp .n)) MXM
dx(Xt, y4).
(6)
By adopting the Markov property, the proposed Schrödinger bridge problem in 66 enables the transformation of the fixed static bridge suggested in (I2) into a temporally variational formulation. In addition, the marginal measure p..* of the bridge Te propagated over time by the parameterized Fookeer-Pande Equation (pFPE) because of the Markov property of X Consequently, the optimal coupling in (6) can be interpreted as being similar to the dynamic Schrödinger bridge explored in recent study (4) for generative modeling. With the definition of the Schrödinger bridge (6), our objective function is defined as follows:
min min tE rEI(PP" (m))
H (TT) = min (T)(RR. 9C8
The objective function in (7) first connects the model measure to the target measure by searching the Schrödinger bridge म (inner problem), and minimizes the relative entropy using the neural network 0 E ④ according to Tē (outer problem). Unlike the minimization problem suggested in (4) which
requires accurate approximation of the Radon-Nikodym derivative duf function adopts an alternative approach. It is reformulated to identify feasible scheme for searching for the Schrödinger bridge τÿ. the following, we introduce dual formulation of (7), called the Markov diffausive Kantorovich dual to induce the bridge in computationally tractable manner:
di, the proposed objective
Definition 2. (Markov diffusive Kantorovich dual) Let 72 be the squared Riemannian distance on M. Then, we define the functional I C(M) x C(M) x R+ x [0, θ > R. as follows:
I (1,,,), f A(x)dP( / B(y)deruv
/
8 v4). (8)
In the following proposition, we relate the proposed dual functional to the regularized version of the Schrödinger bridge problem in (6) in the following proposition:

---

Proposition 2. Log-Sinkhorr with Dual formulation) Let H' be regularized relative entropy Then, the Schrödinger bridge problem 7 has a dual formulation the following equality:
min E
min PP" Jn)
H (To R) min J(AA", c,, = min J (c, 8E8
where c > 0 is a predetermined constant, [A", B is the fixed points of log- Nog-Sinkhorn operator
S = Hz Hyg and each functional Wi, Hi C(M) + CMM is defined as follows:
H. (gllA())) = e elog SM e Hz(BB(z) = e log SM e --
(10)
Note that < 1 restores the original problem posed in In 22 13), they showed that the evaluation of J on the right-hand side of ( is the entropic regularized 2- Wasserstein distance. This indicates that the proposed objective function approximates the Wasserstein distance to minimize the discrepancy between model and target measures given small z O. The optimal solution to the inner problem obtained Proposition | by combining the Markov diffusive Kantorovich duality and log- Sinkhorn operator. Our next goal to solve the outer problem by optimizing neural networks
(3, W).
Gradient Descent Scheme. In the expectations of Definition 2 the Markov dual semi semi-group
Pf plays central role, explicitly revealing the neural networks 3, and W in the equations by reformulating the Kantorovich functional. To specify the discussion, we introduce the equivalent formulation of the first term in 8
/ ((())dPp )dP.
5 f PP A(Xo) + / ((((((AA)(())) - WaA(X ds, (11)
where AM is the Laplace-Beltramt operator on M, and pr is the law for the initial state of the RNSDE. Owing to the explicit relation above, the gradient of the first term [T] can be expressed by calculating the gradients of the neural networks, 0g,3, Og W. Regarding this relation, we propose the gradient flow-based update rule to minimize the dual formulation in ( with neural networks in the following proposition.
Proposition 3. The following gradient descent scheme (0(m)) minimizes the dual functional in the right-hand side of (
[ TMT
0m+1 = 0m - <R
-KR
0o (g3(0m))g) 0ig Ar (X
[o TMT
0g3(0m)) 0,A ( dss f:rm The
Oew (X?m; (,
dds]
(12)
where i. is the learning rate, and T x C1 is the random stopping time that ensures the avoidance of the gradient explosion induced by geometric effects. Then, there exist numerical constants co, c1, Co, C₁ > 0, such that the convergence speed is bounded as O [n' A AT) ci, Co, Ci)].
Detailed discussions the derivation and relation between 8 and 12) are provided in Appendix C.5. Algorithm summarizes the training procedure for the proposed RNSDE.
Role of Stopping Time As shown in 12), the proposed gradient descent method requires the evaluation of geometric terms to update the parameters the RNSDE. particular, the Riemannian co-metric (i.e., 9 (XP)) and Christoffel symbol (i.e., (X()) are required to be consecutively calculated for the given stochastic trajectory X². As addressed in (17), these geometric quantities are not generally bounded. Thus they may cause the gradient explosion problem during the parameter update of the RNSDE. Regarding the issue, the role of stopping time T is to ensure safe training by imposing constraint the gradient norm. These problems are not problematic in Euclidean geometry (e.g., Euclidean SDE) because the geometric terms are set to constants (i.e., = =1, 1, = 0). Thus, an additional module (e.g., 7) is not required in the Euclidean SDE because the raised problem is purely a geometric consequence depending the non-flat Riemannian structure of manifolds.
Please refer to Appendix C.4 for detailed information.
5

---

Algorithm/ 1 Learning Representations on Manifolds with the proposed RNSDE
Require: Neural parameters 8m-u E ④ learning rate x, time interval A., stopping time threshold e1
for m = 1 to M do
1) Simulate the proposed RNSDE by applying the Euler-Maruyama scheme in for 1 T/ / do
(4 [X!," JisoST/A..
1-1) Update the geometric terms according to the current state (t, x(rr) ).
g 9 (x ), J 3 (X2." ), 0 + 0 (X!x))
(13)
1-2) Sample the mean-zero standard Gaussian random variable N(0, In)-
1-3) Propagate stochastic trajectories:
X = ((0m)) Σ VogA.iz + [- 3(0m)
Σ J.k
T 5k +
Σ u² (XP Pm 0m)) A.
(14)
2
end for
2) Search the fixed point [A", B*] by applying the log- Sinkhorn iteration.
3) Sample random stopping-time
inf infistogy/aL (fo: 1 [T co).
4) Update the neural network 0m.+ + O.. by applying the proposed gradient descent in (12 with stopping
time T, and the obtained fixed point A".
end for
4 Comparison to Existing Methods
In conventional Riemannian continuous normalizing flows, model density po is transited according to the following differential equation:
a (z)) div. t),, Ōt
RCNF:
(15)
where 2, is the ODE flow on a manifold, and Vi is the vector field parameterized by the neural network. As shown in (15), the infinitesimal difference between log-densities is proportional to the divergence of the learnable vector field Vo. where the spatio-temporal latent information need to be priorly encoded in the neural ODE (i.e., 21). While the probability transition rule in (15) is dependent on the temporal derivative of the neural ODE (i.e,, by the chain rule, the differential equation in (15) implicitly expresses the probability transition. In the proposed method, the parameterized Fokker-Planck Equation (pFPE) obeys the transition rule of model density X² ~ ps:
② कe (o(, - 3(0) x) div divo x) (o.,
PFPE:
(16)
where Vo(t. = Wt., : 6) is the neural potential field. Instead of borrowing additional dynamics (i.e., ODE in RCNF) to express the probability transitions of density, the pFPE can explicitly express the infinitesimal evolution of density p(l, xr) with neural networks and W for both spatio- temporal variables (t, æ) and the solution the pFPE follows the law of path trajectory X².
In 26), the FPE has been numerically simulated by linearization to approximate density pt. In contrast, the proposed method can avoid the computational burden of simulating the pFPE, while considering the dual measure 14, = po (t, dV ~ X² as the law of X using the property of the Markovian semi- group to access the probability at time In the Lie-group valued generative flow model proposed in (10), it showed that the specific form of vector field Va((,.) := V log(p> (1, //po((,.)) E TM can also induce the solution to the FPE. However, its representational power may be limited because the network outputs (i.e., Neural ODE) are implicitly utilized to estimate the model density pa- In contrast, the neural networks in the proposed model are directly incorporated into the pFPE vector field, which enables rich stochastic representations.
5 Related Work
MCNF (15) and RODE 18 defined the ordinary differential equation (ODE) on manifolds and adopted continuous normalizing flows (i.e., RCNF) to express the transition of data representations. In their methods, the stochastic transition was expressed as the ODE flow of log-density on manifolds and directly calculated the geometric operations (e.g., divergence). ODE-based diffeomorphic flows for the Lie structure (e.g., nespecial unitary group (n)) that preserves equivariannce/ivariance EMSRE (23) proposed a new set of expressive normalizing flows on complex geometries, such torus and spheres. RCPM (6) directly parameterized the convex
a similar context, (I0) adopted

---

(b) RNSDE (ours) Figure 4: Learned densities 2-sphere S², Green arrows indicate the north poles the sphere 00,0,1) 1) R
(a) MCNF 15
Rs
Figure 5: Learned densities 2- 2-Torus T The target shapes helix coil where the number of coils set
Nc 4, in left and right figures.
potential map on Riemannian manifolds with a neural network, which can solve the optimal transport problem. MF (24) introduced computationally tractable linear-type normalizing equation to express the evolution of density over time. 77),, 4., generative modelings was formulated by adopting the Euclidean dynamic Schrödinger bridge and SDEs were estimated for applications.
6 Experiments
In this section, we compared the proposed RNSDE baseline methods for various tasks, including generative modeling, interpolation, and reconstruction. Although the proposed method can be applied to general compact manifolds, we focused on well-known structures. In particular, the following two model manifolds were considered: 2-sphere S² and 2-torus T = S1 x St. The corresponding Riemannian metrics and detailed information on geometric calculations are provided Appendix
Generative Modeling, In this experiment, we evaluated the performance for density estimation on model manifolds of geometrie shapes: 8-shapes, two moons, spiral, and helix coil. Note that the complexities of the target densities in our experiments are considerably higher than those in previous studies (e.g., mixture of von Mises). Every methods take the initial states of the stochastic trajectories as the mean-zero standard Gaussian in the local coordinate, X Po == (iggNOO, I,). where the coordinate function denoted as vb. The goal is then set to accurately restore the target densities. To define the target densities for baseline models, including MCNF (15), RCPM (6). and EMSRE (23), we utilized samples from target densities as anchor points Gaussian kernels, where the bandwidths were identically set to 10, suggested in their methods. After obtaining the approximated target densities from the Gaussian KDE, each baseline model was trained by minimizing KL divergence. To train EMSRE, the number of transformation modules and radial components were set to Nr = 24 and K = 5, respectively. For RNODE, which is a deterministic version of our RNSDE, we utilized only the potential field term in (A), where the diffusive term was set to 0.
In Tables and we estimated the 2- Wasserstein distance (i.e., W2) between the empirical target measure Ue and the model measure 14 to evaluate the model performance. Table summarizes the performance comparisons of the density estimation tasks. Evidently, the proposed method (RNSDE) outperforms all the basclines by approximately 0.1 8.2% in all four experiments (i)-(iv,, despite incorporating no prior information on the target density. RNODE exhibits an inferior performance compared to the proposed RNSDE, which empirically verifies the effectiveness of stochastic representation (i.e., stochastic development) compared the case of arbitrary manifolds with no definite formulas for extrinsic geometric operations (e.g., global exponential map by vector projection from Rd), flow-based models (e.g.. MCNF) cannot be implemented in such geometry because their methods require explicit computation of geometric operations. Accordingly, MCNF is incapable of generating samples on the torus. Unlike their methods, RNSDE only requires intrinsic and local geometric objects (e.g., Riemannian metric tensor,
deterministic one. In
7

---

Table I: Performance comparison of density estimation synthetic datasets. Each model was evaluated with the 2- 2-Wasserstein distance W22 (x10 ²). The best performance is rewritten bold.
Manifold
Density
MCNF
EMSRE
Moser
RCPM
RNODE
RNSDE
8-shapes
11.26
9.83
5.81
9.40
8.01
5.67
Sphere
(ii) Two moons
14.34
9.11
6.10
9.02
7.68
5.66
(iii) Spiral
15.15
10.13
7.56
9.25
9.32
6.97
Torus
(iv) Helix coil
7.64
3.18
2.86
Table 2: Shape interpolation
Table 3: Volcano dataset
Table 4: Vessel Route dataset
Methods
(i)
(ii)
(i) -> (iii)
Methods
W2
Methods
Ey
W2
EMSRE
15.37
11.40
EMSRE RCPM
5.70
Latent ODE
14.25
7.17
RCPM
9.32
9.24
1.83
ODE RNN
16.37
7.29
RNODE
7.09
9.01
RNSDE
1.08
RNSDE
12.83
6.23
RNSDE
6.09
7.61
Christoffel symbols, and geodesic) to define the stochastic representations: thus, the experimental results can be easily followed.
Figs/4a and 4b show samples from the learned densities (i.e., black dots) and target densities (i.e., red dots) of the baseline and proposed method, respectively. In the figures, our RNSDE accurately restored the complex target shapes, whereas the MCNF failed to capture the geometric patterns. The samples from the learned densities on 2-Torus are shown in Fig where the number of coils is set to N. and N. 6. As shown in the figure, the proposed method correctly approximates the target shapes, regardless of the complexity. These qualitative results highlight the representational power of the proposed RNSDEs for various geometries.
Initial
Targel
p(t, x) ~ X1.
to
4
f
f6
1/20
f/1
1/8
f31
Figure 6: Result of shape interpolation starting from 8-shape to spiral on The red (at fai) and black dots represent target and model densitics, respectively.
The
Shape Interpolation. In the second experiment, we conducted shape interpolation starting from (i) 8-shapes to (ii) two moons and (iii) spiral. The total interpolation time (T) was set to 0.1, and 32 intermediate samples were taken in the sequence. Fig 6 shows sampled shapes with uniform temporal intervals. Although the distributional constraint is only imposed at the end of the sequence (e.g., target), the intermediate samples exhibit smooth changes between increasing temporal states. This example emphasizes the representability of the proposed RNSDE, showing the superiority of the PFPF, which can express the entanglement of spatio-temporal variables. Table 2 shows the estimated performance comparison on scale of 10-2, As shown in the table, the proposed model outperformed baselines by approximately 1.0 ~ 9.3%, which is similar to the performance improvement in the previous experiment. Generally, quantitative and qualitative results of density estimation tasks show the versatility of various geometric spaces (i.e., sphere and torus) compared to the previous models. In the following experiment, we evaluated the expressiveness of the proposed model real-world datasets, including volcano and vessel route.
Volcano Eruption. As the third application, we modeled the density of volcano eruption dataset (19), which has been explored in previous studies works [70 (4).. Table 3 shows the quantitative results of the density estimation the geographic dataset at a scale of 10 Owing to the dependency on hyperparameters for approximating target densities (i.e., optimal bandwidth for KDE), baseline methods produced relatively poor results in learning the target density, shown in Table 3 On the contrary, the proposed model precisely restored the intricate distribution of volcano datasets and demonstrated the superiority of expressivity. Fig 7 highhights the samples from the learned (black) and target densities (red) distributed around the Pacific Rim. In the figure, the proposed model accurately captured the volcano eruptions that occurred on the "Ring of Fire".
8

---

Vessel Route. In the final experiment, we re- constructed the vessel routes in the ocean, ex- pressed as time sequences. We compared the proposed method with neural ODE models for time-series data, including Latent ODE (5) and ODE-RNN (25). To implement the baseline models in the Riemannian setting, projected the Euclidean vectors produced by their meth- ods onto the 2-sphere using the stereographic projection. Then, the baseline models were trained to minimize the mean geodesic errors (MGEs) (i.e., Ey²). Table 4 shows the com- parison results for both MGES and Wasserstein distances. The results show that RNSDE out- performs other latent models (e.g., Latent ODE and ODE-RNN), Appendix E provides the cor- responding figure depicting the sampled trajec- tories of the trained RNSDE and thoroughly discusses the suggested metric.
Figure 7: Density estimation on the geographic volcano dataset. Table shows the numerical results.
7 Discussion
Conclusion. In this paper, we presented a principled approach for expressing stochastic representa- tions on manifolds by suggesting a novel Riemannian neural SDE. To define diffusivity, we followed the Moass--llottyy-alllli
interpretation, where the Brownian motion on manifolds derived from the stochastic trajectories on the orthogonal frame bundle. Theoretically, we combined the Markov property with the static Schrödinger bridge problem and proposed Markov diffusive dual formulation. Then, the corresponding gradient descent scheme was proposed. Empirical results for various tasks including generative modeling, interpolation, and reconstruction, showed that the proposed method surpasses its counterpart.
Limitation and Future Work. In the implementation of the proposed RNSDE, we focused on Euclidean embedded manifolds despite their extensive applicability to general compact manifolds. Regarding this issue, we plan extend our method to other geometries including the non-compact Lie group (e.g., SU) for quantum physics (IO) 3) and the product spheres for human motion data and global weather data as high-level applications.
Acknowledgement. This work was partly supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government by Ministry of Culture, Sports and Tourism and Korea Creative Content Agency (Project Number: R2021040032).
and partly supported
References
[1] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
[2] Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 348. Springer Science & Business Media, 2013.
[3] Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael S.
Albergo, Kyle Cranmer, Daniel C. Hackett, and Phiala E. Shanahan. Sampling using SU(n)
gauge equivariant flows. Phys. Rev. D) 103:074504, Apr 2021.
[4] Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. Recovering stochastic dynamics via gaussian schrödinger bridges. arXiv preprint arXiv:2002.05722, 2022.
[5] Ricky T. Q. Chen, Yulia Rubanova. Jesse Bettencourt, and David Duvenaud. Neural ordinary
differential equations. NewrIPS, 2018.
[6] Samuel Cohen*, Brandon Amos*, and Yaron Lipman. Riemannian Convex Potential Maps. In
ICML, 2021.
[7] Valentin De Bortoli, James Thomton, Jeremy Heng, and Arnaud Doucet. Diffussin schrödinger bridge with applications to score- based generative modeling. NewrIPS, 2021.
Elton P Hsu. Stochastic analysis on manifolds. Number 38. American Mathematical Soc., 2002.
[9] Nobuyuki Ikeda and Shinzo Watanabe. Stochastic differrntial equations and diffusion processes.
Elsevier, 2014.

---

[10] Isay Katsman, Aaron Lou, Derek Lim, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Equivariant manifold flows. In NeurIPS. 2021.
[11] J.M. Lee and J.M. Lee. Introduction to Smooth Manifolds. Graduate Texts in Mathematics. Springer, 2003.
[12] Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014. [13] Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Continuous regularized wasserstein barycenters. In NewrIPS, 2020.
[14] Xuechen Li, Ting- Ting-Kam Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable
gradients for stochastic differential equations. In AISTATS, 2020.
[15] Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Neural manifold ordinary differential equations. In NeurIPS, 2020.
[16] Jonathan H Manton. A primer on stochastic differential geometry for signal processing. IEEE
Journal of Selected Topics in Signal Processing, ((4):681-699 2013.
[17] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In NeurIPS, 2020.
[18] M. F. Mathieu, J. J. Zhao, J. Zhao, A. Ramesh, P. Sprechmann, and Y. LeCun. Disentangling factors of variation in deep representation using adversarial training. In NIPS, 2016.
[19] NOAA. Global significant volcano database. "https://dat http:s:/ddaaaannta noda. noaa gov/cgi-bin/
iso?id=gov
www.ddaooooaa......oo noa.. ngdc. mgg. hazards: C10147 National Geophysical Data Center
World Data Service (NGDC/WDS); NCEI/WDS Global Significant Volcanic Eruptions Database.
NOAA National Centers for Environmental Information.
[20] Bernt Oksendal. Stochastic Differential Equations: An Introduction with Applications. Springer-
Verlag, Berlin, Heidelberg, 1992.
[21] Sung Woo Park, Kyungjae Lee, and Junseok Kwon. Neural markov controlled SDE: Stochastic
optimization for continuous-ting
data. ICLR, 2022.
[22] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trendy@ in Machine Learning, 11(5-6)3355-6077 2019.
[23] Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In ICML, 2020.
[24] Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence-
based generative modeling manifolds. In NeurlPS, 2021.
[25] Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equations
for pvegularly-samplee
time series. NeurIPS, 2019.
[26] Arno Solin, Ella Maija Tamir, and Prakhar Verma. Scalable inference in SDEs by direct matching the fokker- tlkkeer-plankk-.o planker-ooookkkooooooooooo equation. In NeurlPS, 2021.
[27] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based based generative modeling through stochastic differential equations. In ICLR, 2020.
[28] Daniel W Stroock. An introduction to Markov processes, volume 230. Springer Science & Business Media, 2013.
10

---

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]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [N/A]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [No]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
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]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
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? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]


Reviews:
[{"review_id": "8Ivqt_XIt7", "reviewer": "Reviewer_hLQB", "summary": "The submission proposes an approach for modeling stochastic differential equations on Riemannian manifolds using neural networks called Riemannian neural SDE (RNSDE). This approach allows for learning a probability distribution over a manifold and modeling stochastic processes happening on the manifold. RNSDE uses neural networks to model the terms in the  local representation of an SDE and allows for probability density estimation using the dual metrics formulation. Experiments show that RNSDE outperforms some of the related deep learning baselines in a set of illustrative toy tasks and in several real-world benchmarks.", "questions": "1. It would be worth to consider diffusion models in related works, especially [1]\n2. In the experiments with the Vessel route dataset, you choose Latent ODE and ODE RNN as the baselines. As far as I know the main feature of these models is that they can deal with irregular time series, i.e. where measurements are not equidistant wrt time. Does RNSDE capable to take this irregularity into account? If not, this comparison looks like comparing apples with oranges, where it is hard to make any conclusion why one model outperforms the others. I would highly recommend authors to add a comment on this in the text of the paper.\n3. In the introduction section there’s a strong claim that conventional approaches which do not take into account geometry of the task are losing in their expressivity. However, there’s a little of validation for that in the experiments section. I would recommend to add this kind of conventional approaches as baselines to all the experiments to give a clear evidence of the importance of taking into account geometry of the data.", "limitations": "Authors clearly state  that they focus only on Euclidian manifolds while there’re other domains (non-compact Lie groups, product spheres) important for downstream applications to consider. There’re no discussion of the negative social impact provided. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Originality\n\nAccording to the provided literature review, the submission proposes a novel approach to modeling SDEs on manifolds. Authors clearly state the difference between their method and related works, highlighting that RNSDE requires less trainable components to model evolution of the density and claiming that it provides potentially more flexible solution than some other restricted methods. However, the literature review does not cover diffusion models which seem to be very relevant to the proposed method, especially in its application to density estimation. Specifically, there’s a relevant work [1] studying diffusion models on Riemannian manifolds which would be worth to compare with or at least to mention as a competitive work and highlight the key differences.\n\n### Quality\n\nThe submission is technically sound and complete with a solid theoretical basis and good enough experiments. Authors’s claims are mostly supported by experimental results and most of the baseline methods are appropriate. \n\n### Clarity\n\nThe presentation of the method is very decent, authors introduce all the important concepts and provide preliminaries required for understanding the method. They also support some of the involved concepts and derivations with intuitive explanations which are very helpful for a broader machine learning audience. Moreover, informative illustrations helps a lot to make sense of the method and experiments. However, a solid background in manifold theory is required to fully understand the method and its theoretical basis. \n\n### Significance\n\nThe proposed method addresses an important task of modeling SDEs on manifolds in an original way surpassing some of the existing methods. The literature review and experimental comparison suggests that the proposed method is fairly significant.\n\n[1] De Bortoli, Valentin, et al. \"Riemannian score-based generative modeling.\" arXiv preprint arXiv:2202.02763\n (2022).", "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": "7cpUOuIXL8M", "reviewer": "Reviewer_mk19", "summary": "This paper introduces a neural network method to learn and represent stochastic differential equations on manifolds. Applications are presented for density modeling on manifolds and learning latent paths on manifolds.", "questions": "My questions are covered in the weaknesses section. In particular, I would like to see a comparison with Moser Flow and a discussion with the Riemannian Score Based Modeling Paper.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n-----------\n+ The paper is presented very well. I appreciated the clear presentation, precise mathematical language, thorough writing, and informative figures.\n+ The paper is very technically solid, as the core difficulty of extending Neural SDEs to manifolds is aptly handled with the correct Riemannian-geometric constructs.\n+ The experiments are pretty well done, and the Vessel Route dataset is an interesting contribution to this space.\n\nWeaknesses\n--------------\n- There is work that slightly predates this one on extending SDEs to manifolds in the context of scored based generative models and diffusion (https://arxiv.org/abs/2202.02763). While I believe that the two works are different, the above work should be included in the discussion.\n- For the density modeling experiments, some baselines are missing. In particular, I believe there should be a comparison with Moser Flow.\n\nVerdict\n--------\nI generally lean on accepting the paper, as I think that it is a solid contribution to the space. However, I would want some more discussion with the baselines and previous methods mentioned above (as I think they sufficiently predate the work).", "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": "20RnQXosm8P", "reviewer": "Reviewer_rDVp", "summary": "This paper proposes a method to express the stochastic representation over manifolds, using a novelly defined Riemannian Neural stochastic differential equation (RNSDE), the method is theoretically well-analyzed at each step, empirically RNSDE outperforms several baselines on some simple tasks. ", "questions": "Q1: I understand that Euclidean Neural SDE gives results not on the manifold, but a very naive approach is to enforce a hard constraint (e.g. projection to the manifold) after each Euclidean neural SDE step, or something as a first order retraction? One can also move along local frameworks to get a trajectory on the manifold, do you include this baseline to compare with? Both performance and computations\n\nQ2: I’m curious what’s the choice of the stopping time tau could affect in practice? How large it needs to be, how severe the gradient explosion is and how to choose this value? ", "limitations": "Mostly discussed above, the practical usage seems to be rare, and the computations are heavy, limited to a small set of manifolds and low dimensions. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I’m not an expert on neural SDE, so some points maybe stupid :)\nS1: the method is theoretically sound in most steps, except the learning step and some approximations;\nS2: the extension from Euclidean neural SDE to Riemannian Neural SDE makes it easer to understand and follow;\nS3: Empirically the proposed RNSDE outperforms some previous baselines on tasks including generative modeling, interpolation and reconstruction. \n\nW1: I think this paper is not that ‘readable’ for general NeurIPS audience, but only experts in this area, currently it’s heavy math, I would suggest add more explanations/backgrounds and intuitions to the text, for example the Eells-Elworthy-Malliavin interpretation of the diffusion process, and what each step in the Scheme is doing, this could significant help abstract more readers into this area. \n\nW2: RNSDE incurs more computations than the Euclidean case, including metric tensors, christoffee symbol and etc, which has to be computed for every local neighborhood as moving along the path, I would like to see some computation burden analysis or empirical results of RNSDE, plus compared to some baselines, particularly to one I mention in next question part. \n\nW3: The practical usages/applications of this method looks rare to me, what’s more,  the manifolds are also simple, though S^n is claimed to be one model, I only see n=2 case, and with a torus model, is it possible to test on high dimensional manifolds? ", "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