R IEMANNIAN N EURAL SDE: L EARNING S TOCHASTIC R EPRESENTATIONS ON M ANIFOLDS
A BSTRACT
In recent years, the neural stochastic differential equation (NSDE) has gained attention in modeling stochastic representations, while NSDE brings a great success in various types of applications.
However, it typically loses the expressivity when the data representation is manifold-valued.
To overcome such an issue, we suggest a principled way to express the stochastic representation with the Riemannian neural
SDE (RNSDE), which extends the conventional Euclidean NSDE.
Empirical results on the density estimation on manifolds show that the proposed method significantly outperforms baseline methods.
I NTRODUCTION
Recently, there has been a great success in modeling stochastic dynamical systems for complex data representations, containing spatially high stochasticity.
Especially, the stochastic differential equation (SDE) was adopted by recent studies as a fundamental probabilistic model to express the transition of stochastic states.
For example, Li et al.
(2020) suggested the neural SDE , which was the firstly implemented neural network to train the stochastic dynamics and Song et al.
(2020) adopted the parameterized reverse-SDE (Anderson (1982)) to model the score-based generative models.
Park et al.
(2022) introduced the controlled SDE combined with the stochastic optimal control-based theoretical framework to model the time series.
These works opened a new way to express the time-series data by showing the power of stochastic representations induced by SDE models.
Despite the remarkable recent progress, the major interest in using SDE has been focused on the Euclidean geometry ( i .
e ., R d ).
In other words, conventional approaches inevitably lose the expressivity when the data representation us defined on other geometry such as Riemannian manifolds ( i .
e ., M ).
In this paper, we tackle the above concern by introducing the Riemannian neural SDE (RNSDE) to model the stochastic dynamics on manifolds.
The proposed RNSDE is the natural extension of conventional Euclidean SDEs, which defines the intrinsic stochastic transition in the local sense and fully recognizes the geometric structure.
Thus, it can resolve expressivity problems of conventional
SDEs on a fundamental level.
Specifically, our RNSDE is built upon the Eells-Elworthy-Malliavin interpretation, where the stochastic trajectories are expressed under the frame bundle geometry.
Contribution.
The main contribution of this work is to suggest a general framework for modeling stochastic representations on Riemannian manifolds.
Unlike prior works, our objective function enables the proposed model to learn stochastic processes without requiring prior information on target distributions, and thus applicable for high-dimensional and complex datasets.
Moreover, by handling Fokker-Planck Equations, our method can use the network outputs in a direct manner so that its stochastic representation power is strengthened.
Experimental results demonstrate that the proposed framework outperforms conventional approaches on the density estimation.
R ELATED W ORK
Normalizing Flows on Riemannian Manifolds.
MCNF Lou et al.
(2020) and RODE Mathieu & Nickel (2020) defined the ordinary differential equation (ODE) on manifolds and adopted the 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).
In a similar context, diffeomorphic
ODE-based flows were adopted to Lie structures ( i .
e ., SU ( n ) ) by Katsman et al.
(2021), which can preserve the equivariance/invariance.
Riemannian Optimal Transport.
Cohen* et al.
(2021) directly parameterized the convex potential map on Riemannian manifolds with a neural network, which can solve the optimal transport problem. i
Rozen et al.
(2021) adopted the computationally tractable linear-type normalizing equation to express the evolution of density according to time.
R IEMANNIAN N EURAL S TOCHASTIC D IFFERENTIAL E QUATION
Notations.
Let M be complete and connected n -dimensional Riemannian manifolds equipped with the Riemannian metric expressed as the matrix form G ≜ [ g ij ] 1 ≤ i,j ≤ n .
Similarly, we define the inverse of the Riemannian metric ( i .
e ., co-metric) as G − 1 ≜ [ g ij ] .
The Christoffel symbol with respect to the metric g is denoted as [Γ kjk ] 1 ≤ i,j,k ≤ d .
The orthonormal frame bundle and the tangent bundle are denoted as O M and T M , respectively.
Please refer to Émery (2012); Ikeda & Watanabe (2014); Lee & Lee (2003) for more details on Riemannian geometry and stochastic analysis.
Riemannian Neural SDE.
Our major interest is to learn the parameterized stochastic process X θt with parameter θ , which is defined as a solution to the following Stratonovich SDE, Oksendal (1992):
[Equation]
where B t =
[ B 1 t , · · · , B n t ] denotes the n -dimensional standard Wiener process on R n .
We call the combination of three parameterized terms ( i .
e ., (A), (B) and (C) )
a Riemannian Neural Stochastic
Differential Equation (RNSDE).
The detailed explanations are provided as follows: (A) Neural Potential Field.
In the first term, we call the parameterized vector field W ( t, X t ; θ ) ≜ w j ( t, X t ; θ ) ∂ tj : [0 , T ] ×M× Θ → T X t M as the neural potential field , where the set of orthonormal tangent vectors { ∂ tj } 1 ≤ j ≤ n ≜ { ∂∂x j | X t } 1 ≤ j ≤ n denotes the moving frames defined on the trajectories of stochastic dynamics.
This term defines the driving force of stochastic dynamics X θt regarding the inferred network decisions W ( t, X t ; θ ) on the current state ( t, X t ) .
(B) Self-regressive Potential Field The
second term
, which is called a self-regressive potential of the k th order ), is the main object that characterizes our model.
The core idea of this term is to accumulate the decisions of the neural potential field on the history of observed data { s, y s } s ∈ [ t − k,t ) from s =
[ t − k ] to s = t .
Specifically, network decisions W ( s, y s ; θ ) ∈
T X s M on past observed data at time s ∈ T ≜ [ t − k, t ) are translated to the tangent space of the current state by using parallel transport T ts : T X s M → T X t M , and the translated decisions are accumulated overall.
Without loss of generality, we
denote this term as T ( t ; θ ) in this paper.
(C) Stochastic Development .
The last term, called stochastic development (Hsu (2002) Ikeda &
Watanabe (2014)) controls the diffusive behavior of the proposed SDE on manifolds.
The function π : O M → M is the canonical projection from frame bundle to model manifolds, and U t ≜
π − 1 ( X θt ) ∈
O M is the horizontal lift of our stochastic dynamics X θt ∈ M .
The parameterized scalar-valued function β : Θ → R + controls the diffusivity of process X θt .
It is noteworthy that the solution to (1) is called a Brownian motion on M , when W ≡ 0 and β ≡ 1 2 .
Local Representations.
Although the RNSDE defined in (1) provides the rigorous stochastic representation on manifolds, it is not directly applicable for implementation due to the abstract formulations.
Thus, for clarity, we provide the explicit formulation in local representations as follows:
Proposition 1.
(Local Representations)
Let X t be a stochastic process represented as a local coordinate X t ≜ { x it ( w ) } 1 ≤ i ≤ n ⊂ M .
Then, the proposed RNSDE is defined as follows:
[Equation]
The detailed derivation is provided in Appendix A.1.
The intrinsic local coordinate-based representation in (2) only requires the metric tensor g and the Christoffel symbol { Γ ijk } 1 ≤ i,j,k, ≤ n for every point x ∈ M , where the newly appearing terms ( i .
e ., first and second) in (2) reveal the effect of stochastic development.
It transforms the Euclidean stochasticity induced by standard Brownian motion B t into Riemannian stochasticity on manifolds by considering the curvature effect ( i .
e ., g, Γ ).
ii As the propagation of stochastic particles { X t } 0 ≤ t ≤ T randomly changes the moving frame { ∂ j ( X t ) } 1 ≤ j ≤ d and its local coordinate system, we call the solution to (2), X t , as the stochastic local chart flow (SLCF).
In this paper, we regard that the Riemannian manifolds M are embedded compact sub-manifolds of the ambient space, M ⊆ R D for n ≤ D .
In other terms, the coordinate system of the diffusion process X t can be expressed in R D using the global coordinate as Y t ≜ ψ ◦ X t where ψ is the coordinate function.
L EARNING D ENSITIES FROM S AMPLES
Learning Densities from Samples.
The major drawback of prior works is that full information on target probability densities is required.
For example, the conventional RCNF optimizes loglikelihood (or KL-divergence) between source and target distributions, i .
e ., min θ E [log p θ − log p ν ] .
Unfortunately, this formulation may cause problems in high-dimensional and complex datasets, as one needs to priorly access the accurate estimation of the data density p ν ( e .
g ., KDE).
In contrast, our method utilizes only the set of observed particles { y t } ∼ ν t sampled from the unknown probability measure ν t .
In other words, we assume that no prior information on target distributions is given.
Objective Function.
From the N -number of observed particles { y lt } 1 ≤ l ≤ N at time t , we define the empirical target measure as ν t = 1 N (cid:80) N l δ y lt to approximate.
Then, the source probability measure is defined as the law of solution to the proposed RNSDE: µ θt ≜ P ( X θt ∈ · ) ∼ X θt ∈ M .
Note that the source measure µ θt |{ y s } t − k ≤ s<t is conditioned by the set of conditional observations { y s } t − k ≤ s<t encoded in the self-regressive potential field, i .
e ., (1)- (B) .
Regarding the definition of µ θt and ν , our objective function is posed to minimize the distributional discrepancy ( i .
e ., Wasserstein distance) between source and target measures:
[Equation]
Specifically, our objective function is posed to minimize the 2 -Wasserstein distance at time t .
Then, the problem is accurately calculating the Wasserstein distance by relating the probabilistic structure of RNSDE.
To do so, we are especially interested in the Markovian property:
Markov Diffusive Kantorovich Dual.
By the independent increments of B t , one can show that the proposed RNSDE preserves the Markov property.
Let P θt f ≜ E [ f ( X θt ) | X 0 ] be a Markov semigroup of X θt .
Then, by duality, one can define the probability measure called dual semi-group P θ, ∗ t satisfying the equality, E x ∼ µ
[ P θt f ( x )] =
E x ′ ∼ P
θ, ∗ t µ
[ f ( x ′ )] , which identifies source probability measure as µ θt ≜ P θ, ∗ t .
This theoretical interpretation has been proposed by Park et al.
(2021) to express the temporal path of probability measure for the Markov process.
Inspired by their work, we define the dual Kantorovich formulation with the dual semi-group to define the Wasserstein distance:
Definition 1.
(Markov Diffusive Kantorovich Dual) Let us define two arbitrary functions A, B ∈
C 2 ( M ) / R .
Then, the diffusive Kantorovich dual is defined as follows:
[Equation]
It is well-known that the fixed point [ A ⋆ , B ⋆ ] of Sinkhorn-iteration started at [ A, B ] uniquely determines the equality, lim ϵ
→ 0 J ϵ ([ A ⋆ , B ⋆ ]; θ ) =
W 2 ( µ θt , ν t ) .
Using this relation, our objective function in (3) can be reformulated as min θ J ∆ ϵ
([ A ⋆ , B ⋆ ]; θ ) with pre-determined parameter ϵ ≜ ∆ ϵ ≈ 0 .
To solve this new minimization problem, we apply the gradient descent as follows:
[Equation]
where τ m is the learning rate.
Please note that the expectation is taken to measure P .
Detailed derivations are provided in Appendix A.3.
4. C OMPARISON TO E XISTING M ETHODS
In conventional Riemannian normalizing flows, the model density p θ is transited by following the propagation rule: ∂ t log p θ ( z t ) =
− div M [ V θ ( z t , t )] , where z t denotes the ODE flow on the manifold, iii and V θ is the parameterized vector field by the neural network.
For the numerical simulation of the equation, they require the computation of Hessian trace.
In our case, the model density p t is defined as the solution to the second-order PDE called a parameterized Fokker-Planck Equation, (pFPE) :
[Equation]
where V θ ( t, · ) ≜ W ( t, · ; θ )+ T ( t ; θ ) .
In contrast to the conventional approach, the proposed RNSDE does not need to simulate the pFPE directly.
As shown in (4), we indirectly access the measure µ θt by calculating the Markovian semi-group.
Thus, our method can avoid the computational burden to calculate the geometric operations for normalizing flows ( e .
g ., divergence, log-determinant), while preserving the representational power.
Katsman et al.
(2021) induce the vector field V ( t, · ) ≜ ∇ log( p ν ( t, · )
/p θ ( t, · ))
∈ T M , which also solves the pFPE.
However, the representational power can be limited as the network outputs ( i .
e ., Neural ODE) are implicitly utilized to calculate model density p θ .
In contrast, our method directly utilizes the network outputs ( e .
g ., W, T ∈ T M ) to model the vector field V in (6), which enables rich stochastic representations.
E XPERIMENTS
In the experiment, we evaluated the performance for estimating densities on the 2 -sphere, including three different densities: 8-shapes, Two moons, and Spiral .
Note that the complexity of these densities is higher compared to those of previous works ( e .
g ., a mixture of von Mises), and the density estimation is more challenging.
To define the target densities for baseline models, including MCNF (Lou et al.
and
EMSRE (Rezende et al.
, we utilized samples from target densities as the anchor points of Gaussian radial functions.
After defining the approximated target densities obtained from the Gaussian
KDE, both models were trained to minimize the KL divergence.
For our model, we initialized the start point of stochastic trajectories as X θ 0 ∼ µ 0 ≜ ψ # N (0 , I n ) , where N denotes the Gaussian distribution and ψ is coordinate function.
For self-regressive behavior, the conditional observations are defined as y s ∼ ψ # N ( E [ X θs ] , I n ) .
As a consequence, the information of past trajectories X θ { s<t } is consecutively encoded in the current state.
For the RNODE, which is a deterministic version of our
RNSDE, we utilized only the potential field term in (1)-(A), where other terms were set to 0 .
[Table]
Experimental Results.
As shown in Table 1, the proposed RNSDE significantly outperforms other baselines by a large margin.
This highlights the representation power of the proposed RNSDEs.
The RNODE shows an inferior outcome compared to the RNSDE, which shows the effectiveness of providing additional information ( i .
e ., conditional observations in (1)-(B) and (C)).
Figs 1 anddisplay samples from learned densities ( i .
e ., black dots) and target densities ( i .
e ., red dots) for the baseline and our model, respectively.
While the MCNF fails to approximate the target densities due to the complex geometric shapes, our method can fit the target densities accurately.
C ONCLUSION
In this paper, we suggested a principled way to express the stochastic representations on manifolds.
Specifically, we introduced the Riemannian neural SDE with stochastic development defined on an orthonormal frame bundle.
In future work, we plan to theoretically extend the current work by analyzing the stability of Sinkhorn-iteration and show the learnability of the RNSDE.
iv
A A PPENDIX A. L OCAL E XPRESSION OF S TOCHASTIC D EVELOPMENT
In the orthonormal frame bundle O M , the local coordinate is defined as U t =
( X t , E kj ) ∈
O M , where the tangent of frame bundle can be decomposed as direct sum of two subspaces: TO M =
V O M ⊕ HO M .
[Equation]
i,θ
[Equation]
where V θ ( t, · ) ≜ W ( t, · ; θ ) + T ( t ; θ ) .
In full description, the Stratonovich SDE in the orthonormal frame bundle is written as the semi-martingale:
[Equation]
where ˆ V θ is a horizontal lift of vector field V θ .
By applying Itó’s lemma to the SDE in (8) with the function f
=
R ◦ π for arbitrary R : M → R
[Equation]
one can obtain the following representation of SDE.
[Equation]
In the local coordinate presentation
[Equation]
we can written (9) as follows:
[Equation]
As the lifted function f = R ◦ π uniquely determines the orthonormal basis of tangent space ue i ∈ T M by the property of the fundamental horizontal vector field H i ( i .
e ., there exists a unique relation π ⋆ H
i ( U t ) =
U t e i ), we can express H i f
( U t ) =
U t e i .
This fact leads the second equality in the following formulation:
[Equation]
To understand the above relation precisely, M t can be expressed as local coordinate.
For this, we denote the basis of horizontal vector space HO M as { ∂ j ≜ ∂∂x j , ¯ ∂ km ≜
[Equation]
In this case, the horizontal curve U t in orthonormal frame bundle O M can be represented in local coordinates, U t =
[ X it , e ij ( t )] as follows:
[Equation]
where H i ( U t ) ≜ e ji ∂ j − e ji e lm Γ kjl ( X t ) ¯
∂ km is the fundamental horizontal vector field.
The third equality holds as ¯ ∂ km ◦ f
([ X it , e ij ( t )]) = ¯
∂ km R ◦ π ([ X it , e ij ( t )]) =
¯ ∂ km R ( X t ) = 0 .
In the last equality, we define the vector field T ∈ T M as ue
i
=
e ji ∂ j ≜ T ( X t ) .
To estimate the vector field T ( X t ) , we need to find out the explicit numeric of the orthogonal matrix e ji by solving the following equation:
[Equation]
13) vi
To solve the equation, we take the Riemannian inner product between Ue l and Ue m , as follows:
[Equation]
where we denote ⟨ X, Y ⟩ g as the inner product between vector fields X, Y ∈ T M and δ lm is the coordinate delta.
Using the relation in (14), the following identity can be easily obtained:
[Equation]
One may express the identity in (15) as the matrix form, as follows:
[Equation]
where we denote E = { e ik } , G − 1 = g ij .
To obtain the explicit form of matrix E , we apply the
Cholesky decomposition to the co-metric matrix.
( i .
e ., E =
Ch ◦
[ G − 1 ] ).
Finally, the derivation form of horizontal vector field to f , H i f ( U t ) can be written in local coordinate as follows:
[Equation]
As the metric matrix is semi-definite positive in our case, the following equality holds by the elementary algebraic property:
[Equation]
Finally, the local martingale term M t is written in local coordinate as follows:
[Equation]
[Equation]
Let us calculate the quadratic variation of the process M t .
[Equation]
20) where [ B, B ] t = dt .
Thus, M t is the local martingale.
[Equation]
where ∆ O M = (cid:80) i H 2 i denotes the Bochner’s horizontal Laplacian operator.
The last equality is known as the Bochner’s Laplacian identity.
The last term V θt =
V ( X t ; θ ) R ( X t ) corresponds to the anti-development of vector field ˆ V θ .
By collecting the defined stochastic representations, M , N , and V , the local expressions suggested in (2) is derived.
A. L OCAL COORDINATE SYSTEM
Among the various definitions of the local chart ψ , we can select the following two types of local coordinate system: 1) Parameterization. Let
( ϑ , φ ) be a local coordinate of the 2 -sphere.
Then, the global expression is parameterized with spherical coordinate in the following form:
[Equation]
2) Normal Coordinate.
While
this coordinate system induces the vanishing
metric
tensor and the
Christoffel symbol at X t ( i .
e .
, g | X t , Γ | X t = 0 ), the diffusion process X t in the ambient coordinate system on R k is expressed as the following form:
[Equation]
In the literature, the deterministic version of the proposed differential equation in (23) is known as the exponential-map flows , Rezende et al. vii
A. M ARKOV D IFFUSIVE K ANTOROVICH D UAL F ORMULATION
In this section, we introduce the intuition behind the gradient descent suggested in (5).
The propositions introduced in this section can be found in Berman (2020) and Vialard (2019).
For each function A, B ∈ C 2 ( M ) / R , we define two functional operators on C 2 ( M ) / R , as follows:
[Equation]
Then, the composition of two operators ( i .
e .
, H ϵµ t , H ϵν ) is called a log-Sinkhorn iteration S ≜
H ϵµ t ◦ H ϵν .
Let us define A l is the transformed shape of A after l -th iteration:
A l = S l ( A ) =
S
[Equation]
( A ) .
Let us define the functional F ( A l ) :
C 2 ( M ) / R → R , as follows:
[Equation]
Then, the log-Sinkhorn iteration uniquely minimizes the functional F by the following proposition.
Proposition 2.
The log-Sinkhorn iteration S has a fixed point
A l →∞ in C 2 ( M ) / R .
This fixed point is determined up to additive constant, and minimize the functional F uniformly:
[Equation]
Let us assume that, for the large enough l > L with small enough ϵ ≈ 0 , the log-Sinkhorn iteration is converged S ◦
A l = A l , and the functional F is minimized.
Then, the function A l ∨ L is approximated to the d 2 -transformation (Villani (2008)) of the function B m ∨ M .
[Equation]
28) where χ V ( x ) = 0 , if x ∈ V , and χ V ( x ) = ∞ if x / ∈ V for the compact subset V ⊂ M , and [ f ] c denotes the d 2 -transformation of f .
By the Kantorovich duality theorem, this theoretical characteristic of log-Sinkhorn iteration induces the optimal transport between µ t and ν t .
Theorem 1.
(Unique fixed-point of S and Optimal transport) Assume that ( µ ( ϵ ) t , ν ( ϵ ) t ) → ( µ t , ν t ) in P ( M ) -weak
sense.
If A is the fixed point of the log-Sinkhorn operator S on C 2 ( M ) / R , then B converges uniformly to a d 2 -convex function such that there exists a unique and optimal transport map Φ satisfying: [Equation]
[Equation]
where exp and ∇ g are Riemannian exponential and gradient, respectively.
In short, the proposed iteration approximates the d 2 -convex function B ≈ A c and solves the optimal transport problem, i .
e ., 2 -Wasserstein distance (McCann (2001)).
Parameterized Generator.
Our interest is to investigate the stochastic and geometric effect on evaluating the functional F regarding the probabilistic property of X θt .
To do so, we first define the second-order partial differential operator called infinitesimal generator as follows:
[Equation]
where ∆ M ≜ div M ◦ ∇ M is the Laplace-Beltrami operator on M .
Note that, due to the timedependent behavior, the proposed generator induces the time-inhomogeneous Markov process.
Lemma 1.
The solution of the proposed RNSDE, X θt , is the Markov process for any θ ∈ Θ .
As the proposed process X θt is the Markov diffusion process regarding the definition of generator in (30), one can apply the geometric version of Itó’s lemma (Wang (2014)) to obtain the following equality:
[Equation]
viii where U − 1 s : T X θs M
→ R d is the inverse of frame U s at X s .
Note that the notation for A m is rewritten in LHS of (31), A θ ∈ C ( M , R )
× Θ to emphasize the parameterized generator L θ in RHS.
[Equation]
Notice.
In the first equality, the parameter in the second term (cid:82) H ϵµ θ
[Equation]
t [ A m ]( y ) dν
( y ) is considered as fixed θ
= θ 0 during the update.
This trick makes us to avoid calculating ∂
[Equation]
( y
[Equation]
dν ( y ) , which leads unstable results.
Gradient flow (descent) of functional F with respect to Parameter θ .
In dual formulation, the gradient of functional F can be written as follows:
[Equation]
In local coordinate, the gradient with respect to parameter θ is written as follows:
[Equation]
where A is the fixed point of log-Sinkhorn iteration.
Finally, the gradient flow of θ to minimize the functional F is defined in the following form:
[Equation]
The proposed gradient descent in (5) is temporal discretized version of (35).
A. G EOMETRIC C ALCULATIONS
In this section, we provide the detailed geometric calculation on the 2 -Sphere, S 2 .
Riemannian Metric Tensor, Christoffel Symbol.
The following calculation is conducted upon the spherical coordinate system defined in (22).
Then, the Riemannian metric and Christoffel symbol is calculated as follows:
[Equation] [Equation]
where the Christoffel symbol is defined as Γ ijk = 12 g ke ( ∂ j g ei + ∂ i g ej − ∂ e g ij ) and ∂ i denotes the (spatial) partial derivative with respect to the i -th component.
Parallel Transport, Exponential Map.
Let U be a tangent vector on T X s M at time s , and let X t be a another point at time t .
Then, the parallel transport of tangent vector U from T X s M to T X t M is defined as follows:
[Equation]
Note that the stochastic process X s is recognize as global representations in the evaluation.
For the
Riemannian exponential map, we define the following normalized version:
[Equation]
(38) ix