Riemannian Diffusion Models
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 lowerbound 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.
Introduction
By learning to transmute noise, generative models seek to uncover the underlying generative factors that give rise to observed data.
These factors can often be cast as inherently geometric quantities as the data itself need not lie on a ﬂat Euclidean space.
Indeed, in many scientiﬁc domains such as high-energy physics (Brehmer & Cranmer, 2020), directional statistics (Mardia &
Jupp, 2009), geoscience (Mathieu & Nickel, 2020), computer graphics (Kazhdan et 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 a Riemannian manifold with a 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 a learning framework.
Recently diffusion based generative models have emerged as an attractive model class that not only achieve likelihoods comparable to state-of-the-art autogressive models (Kingma et al., 2021) but match the sample quality of GANs without the pains of adversarial optimization (Dhariwal &
Nichol, 2021).
Succinctly, a diffusion model consists of a ﬁxed Markov chain that progressively transforms data to a prior deﬁned by the inference path, and a generative model which is another
Markov chain that is learned to invert the inference process (Ho et al., 2020; Song et al., 2021b).
While conceptually simple, the learning framework can have a variety of perspectives and goals.
For example, Huang et al.
(2021) provide a 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 an 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 speciﬁcally, image data.
It does not easily translate to general Riemannian manifolds.
In this paper, we introduce Riemannian Diffusion Models (RDM)—generalizing 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 conventional chain rule of calculus holds, which, as we demonstrate in section §3, can be exploited to deﬁne diffusion on a Riemannian manifold.
Furthermore, we take an extrinsic view of geometry by deﬁning the Riemannian manifold of interest as an embedded sub-manifold within a higher dimensional (Euclidean) ambient space.
Such a choice enables us to deﬁne 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 a variational framework built on the Riemannian Feynman-Kac representation and Giransov’s theorem.
In Theorem 2 we derive a Riemannian continuous-time ELBO, strictly generalizing the CT-ELBO in Huang et al.
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 to compute the Riemannian divergence of our parametrized vector ﬁeld, for which we introduce a QR-decomposition-based method that is computationally efﬁcient for low dimensional manifolds as well a 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 to the time integral, which crucially avoids carefully 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.
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 a short review of the latter, see
Appendix A or Ratcliffe (1994) for a more comprehensive treatment of the subject matter.
2.1 Euclidean diffusion models
A diffusion model can be deﬁned as the solution to the (Itô) SDE (Øksendal, 2003),
[Equation]
with the initial condition X 0 following some unstructured prior p 0 such as the standard normal distribution, where B t is a standard Brownian motion, and µ and σ are the drift and diffusion coefﬁcients 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 t = 0 to t =
T for some ﬁxed termination time T .
To train the model via maximum likelihood, we require an expression for the log marginal density of X T , denoted by log p ( x, T ) , which is generally intractable.
The marginal likelihood can be represented using a stochastic instantaneous change-of-variable formula, by applying the Feynman-Kac theorem to the Fokker-Planck PDE of the density.
An application of Girsanov’s theorem followed by an application of Jensen’s inequality leads to the following variational lower bound (Huang et al., 2021;
Song et al., 2021a):
[Equation]
where a is the variational degree of freedom, ∇· denotes the (Euclidean) divergence operator, and
Y s follows the inference SDE (the generative coefﬁcients are evaluated in reversed time, i.e. T − s )
[Equation]
with ˆ B s
2.2 Riemannian manifolds
We work with a d -dimensional Riemannian manifold ( M , g ) embedded in a higher dimensional ambient space R m , for m > d .
This assumption does not come with a loss of generality, since any
Riemannian manifold can be isometrically embedded into a Euclidean space by the Nash embedding theorem (Gunther, 1991).
In this case, the metric g coincides with the pullback of the Euclidean metric by the inclusion map.
Now, given a coordinate chart ϕ : M → R d and its inverse ψ = ϕ − 1 , we can deﬁne ˜
E j for j = 1 , · · · , d to be the basis vectors of the tangent space T x M at point x ∈ M .
The tangent space can be understood as the pushforward of the Euclidean derivation of the patch space along ψ ; i.e. , for any smooth function f ∈ C ∞ ( M ) ,
˜ E j ( f ) =
∂∂ ˜ x
j f ◦ ψ .
We denote by P x the orthogonal projection onto the linear subspace spanned by the column vectors of the Jacobian J x = d ψ/ d˜ x .
Speciﬁcally, P x can be constructed via P x =
J x ( J Tx J x ) − 1 J Tx .
Note that this subspace is isomorphic to the tangent space T x M , which itself is a subspace of T x R m .
As a result, we identify this subspace with T x M .
Lastly, we refer to the action of P x as the projection onto the tangential subspace, and P x 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 deﬁning the process via test functions.
Speciﬁcally, let V k be a family of smooth vector ﬁelds on M , and let
Z k be a family of semimartingales (Protter, 2005).
Symbolically, we write
[Equation]
for any f ∈ C ∞ ( M ) (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 ﬁrst step deviating from the Euclidean diffusion model (1), as the Itô integral does not follow the usual chain rule.
Working with this abstract deﬁnition is not always convenient, so instead we work with speciﬁc coordinates of M .
Let ϕ be a chart, and let ˜ v =
(˜ v jk ) be a matrix representing the coefﬁcients of V k in the coordinate basis—i.e. V k ( f )
= (cid:80) dj
=1 ˜ v jk ∂∂ ˜ x
j f ◦ ϕ − 1 (cid:12)(cid:12) ˜ x
= ϕ ( x ) .
This allows us to write d ϕ ( X t ) =
˜ v ◦ d Z .
Similarly, suppose M is a submanifold embedded in R m , and denote by v = ( v ik ) the coefﬁcients wrt the Euclidean basis.
v and ˜ v are related by v =
d
ϕ − 1 d˜ x ˜ v .
Then we can express the dynamics of X as a regular SDE using the Euclidean space’s coefﬁcients d X = v ◦ d Z .
Notably, by the relation between v and ˜ v , the column vectors of v are required to lie in the span of the column vectors of the Jacobian d ϕ − 1 d˜ x which restricts the dynamics to move tangentially on M .
Riemannian diffusion models
We now develop a variational framework to estimate the likelihood of a diffusion model deﬁned on a Riemannian manifold ( M , g ) .
Let X t ∈ M be a process solving the following SDE:
[Equation]
where V 0 and the columns of the diffusion matrix 1 V :
=
[ V 1 , · · · , V w ] are smooth vector ﬁelds on M , and B t is a w -dimensional Brownian motion.
The law of the random variable X t can be written as p ( x, t ) µ (d x ) , where p ( x, t ) is the probability density function and µ is the d -dimensional
Hausdorff measure on the manifold associated with the Riemannian volume density.
Let V · ∇ be a differential operator deﬁned by ( V · ∇ g ) U := (cid:80) wk
=1 ( ∇ g · U k )
V k , where ∇ g · U k denotes the
Riemannian divergence of the vector ﬁeld U k :
[Equation]
Our ﬁrst result is a stochastic instantaneous change-of-variable formula for the Riemannian SDE by applying the Feynman-Kac theorem to the Fokker Planck PDE of the density p ( x, t ) .
Theorem 1 ( Marginal Density ) .
The density p ( x, t ) of the SDE (5) can be written as
[Equation]
where the expectation is taken wrt the following process induced by a Brownian motion B (cid:48)
s
[Equation]
For effective likelihood maximization, we require access to 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 s be a w -dimensional Brownian motion, and let
Y s be a process solving the following
Inference SDE
[Equation]
where a : R m × [0 , T ] → R m is the variational degree of freedom.
Then we have
[Equation]
where all the generative degree of freedoms V k are evaluated in the reversed time direction.
3.1 Computing Riemannian divergence
Similar to the Euclidean case, computing the Riemannian CT-ELBO requires computing the divergence “ ∇ g · ” of a vector ﬁeld, which can be achieved by applying the following identity.
Proposition 1 ( Riemannian divergence identity ) .
Let ( M, g ) be a d -dimensional Riemannian manifold.
For any smooth vector ﬁeld V k ∈ X ( M ) , the following identity holds:
[Equation]
Furthermore, if the manifold is a submanifold embedded in the ambient space R m equipped with the induced metric g = ι ∗ ¯ g , then
[Equation]
where v k = ( v 1 k , · · · , v mk ) are the ambient space coefﬁcients
V k = (cid:80) mi
=1 v ik ∂∂x i and P x is the orthogonal projection onto the tangent space.
Intrinsic coordinates .
The patch-space formula (6) can be used to compute the Riemannian divergence.
This view was 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 difﬁcult for some manifolds, hindering generality in practice; (2) we might need to change patches, which complicates implementations; and (3) the inverse scaling of (cid:112) | G | might result in numerical instability and high variance.
Closest-point projection .
The coordinate-free expression (11) leads to the closest-point projection method proposed by Rozen et al.
Concretely, deﬁne the closest-point projection by π ( x ) :=
arg min y ∈M (cid:107) x
− y (cid:107) , where (cid:107)·(cid:107) is the Euclidean norm.
Let V k ( x ) be the derivation corresponding to the ambient space vector v k ( x ) =
P π ( x ) u ( π ( x )) for some unconstrainted u : R m → R m .
Rozen et al.
(2021) showed that ∇ g · V k ( x ) = ∇ · v k ( x ) , since v k is inﬁnitesimally constant in the normal direction to T x 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 orthogonal basis for T x M .
Let Q =
[ e 1 , · · · , e d , n 1 , · · · , n m − d ] be an orthogonal matrix whose ﬁrst d columns span the T x M , and the remaining m − d vectors span its orthogonal complement
T x M ⊥ .
To construct Q we can simply sample d vectors— e.g. from N (0 , 1) –in the ambient space and orthogonally project them to T x M using P x .
These vectors, although not orthogonal yet, form a basis for T x 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:
[Equation]
where we used (1) the orthogonality of Q , (2) the cyclic property of trace, (3) and the fact that
P x e j = e j and P x n j = 0 .
In practice, concatenation with the remaining m − d vectors is not needed as they are effectively not used in computing the divergence, speeding up computation when m (cid:29) d .
Moreover, the vector-Jacobian product can be computed in O ( m ) time using reverse-mode autograd and importantly does not require the closest-point projection π .
Projected Hutchinson .
When QR is too expensive for higher dimensional problems, the Hutchinson trace estimator (Hutchinson, 1989) can be employed within the extrinsic view representation (12).
For example, let z be a standard normal vector (or a Rademacher vector), we have ( ∇ g · V k )( x ) =
E z ∼N ,z (cid:48) =
P x z [ z (cid:48)(cid:62) d v k d x z (cid:48) ] .
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 deﬁned as a simple noise process taking observed data to unstructured noise:
[Equation]
14) where U
= 12 ∇ g log p
and V is the tangential projection matrix; that is, V k ( f )( x ) = (cid:80) mj
=1 ( P x ) jk ∂f∂x j for any smooth function f .
This is known as the Riemannian Langevin diffusion (Girolami & Calderhead, 2011).
As long as p 0 satisﬁes a log-Sobolev inequality, the marginal distribution of Y s ( i.e. the aggregated posterior) converges to p 0 at a linear rate in the KL divergence (Wang et al., 2020).
For compact manifolds, we set p 0 to be the uniform density, which means
U 0 = 0 , and (14) is reduced to the extrinsic construction of Brownian motion on M (Hsu, 2002,
Section 1.2).
The beneﬁts of this ﬁxed-inference parameterization are the following:
Stable and Efﬁcient Training .
With the ﬁxed-inference parameterization we do not need to optimize the vector ﬁelds that generate Y s , and the Riemannian CT-ELBO can be rewritten as:
[Equation]
where the ﬁrst term is a constant wrt the model parameters (or it can be optimized separately if we want to reﬁne the prior), and the time integral of the second term can be estimated via importance sampling (see Section 3.3).
A sample of Y s 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 to estimate the time integral in (15) by evaluating a ( Y s , s ) at a single time step s only.
Simpliﬁed Riemannian CT-ELBO .
The CT-ELBO can be simpliﬁed as the differential operator
V · ∇ g applied to V yields a zero vector when V is the tangential projection.
Proposition 2.
If V is the tangential projection matrix, then ( V · ∇ g ) V = 0 .
This means that we can express the generative SDE V 0 using the variational parameter a via
[Equation]
with the corresponding Riemannian CT-ELBO:
[Equation]
3.3 Variance reduction
The inference process can be more generally deﬁned to account for a time reparameterization.
In fact, this leads to an equivalent model if one can ﬁnd 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 parameterization of the model.
Alternatively, we opt for a simpler view for variance reduction via importance sampling.
We estimate the time integral “ (cid:82) . . .
d s ” in (17) using the following estimator:
[Equation]
where s ∼ q ( s ) and Y s ∼ q
( Y s | Y 0 ) , (18) where q ( s ) is a proposal density supported on [0 , T ] .
We parameterize q ( s ) using a 1D monotone ﬂow (Huang et al., 2018).
As the expected value of this estimator is the same as the time integral in (17), it is unbiased.
However, this means we cannot train the proposal distribution q ( s ) by maximizing this objective, since the gradient wrt the parameters of q ( s ) is zero in expectation.
Instead, we minimize the variance of the estimator by following the stochastic gradient wrt q ( s )
[Equation]
The latter can be optimized using the reparameterization trick (Kingma & Welling, 2014) and is a well-known variance reduction method in a multitude of settings (Luo et al., 2020; Tucker et al., 2017).
It can be seen as minimizing the χ 2 -divergence from a density proportional to the magnitude of E Y s [ I ] (Dieng et al., 2017; Müller et al., 2019).
3.4 Connection to score matching
In the Euclidean case, it can be shown that maximizing the variational lower bound of the ﬁxedinference 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 to its Riemannian counterpart.
Let q ( y s , s ) be the density of Y s following (14), marginalizing out the data distribution q ( y 0 , 0) .
The score function is the Riemannian gradient of the log-density ∇ g log q .
The following theorem tells us that we can create a family of inference and generative SDEs that induce the same marginal distributions over Y s and X T − s as (16) if we have access to its score.
Theorem 3 ( Marginally equivalent SDEs ) .
For λ ≤ 1 , the marginal distributions of X T − s and Y s of the processes deﬁned as below
[Equation]
both have the density q ( · , s ) .
In particular, λ = 1 gives rise to an equivalent ODE.
This suggests if we can approximate the score function, and plug it into the reverse process (21), we obtain a time-reversed process that induces approximately the same marginals.
Theorem 4 ( Score matching equivalency ) .
For λ < 1 , let E ∞ λ denote the Riemannian CTELBO of the generative process (21), with ∇ g log q replaced by an approximate score S θ , and with (20) being the inference SDE.
Assume S θ is a compactly supported smooth vector.
Then
[Equation]
where C 1 > 0 and C 2 are constants wrt θ .
The ﬁrst implication of the theorem is that maximizing the Riemannian CT-ELBO of the plug-in reverse process is equivalent to minimizing the Riemannian score-matching loss .
Second, if we set λ = 0 , from (135) (in the appendix), we have V a = S θ , which is exactly the ﬁxed-inference training in §3.2.
That is, the vector V a trained using equation (17) is actually an approximate score, allowing us to extract an equivalent ODE by substituting V a for ∇ g log q in (20,21) by setting λ = 1 .
Related work
Diffusion models .
Diffusion models can be viewed from two different but ultimately complimentary perspectives.
The ﬁrst approach leverages score based generative models (Song & Ermon, 2019;
Song et al., 2021b), while the second approach treats generative modeling as inverting a ﬁxed noiseinjecting process (Sohl-Dickstein et al., 2015; Ho et al., 2020).
Finally, continuous-time diffusion models can also be embedded within a maximum likelihood framework (Huang et al., 2021; Song et al., 2021a), which represents the special case of prescribing a ﬂat geometry— i.e. Euclidean—to the generative model and is completely generalized by the theory developed in this work.
Riemannian Generative Models .
Generative models beyond Euclidean manifolds have recently risen to prominence with early efforts focusing on constant curvature manifolds (Bose et al., 2020;
Rezende et al., 2020).
Another line of work extends continuous-time ﬂows (Chen et al., 2018a) to more general Riemannian manifolds (Lou et 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 a parametrized vector ﬁeld.
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 a 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 of 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 a in (16) using an MLP and maximize the CT-ELBO (17).
We report our detailed training procedure—including selected hyperparameters—for 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), earthquakes (NGDC/WDS, 2022a), ﬂoods (Brakenridge, 2017), and ﬁres (EOSDIS, 2020).
[Table]
[Table]
In Table 1 for each dataset we report average and standard deviation of test negative log likelihood on 5 different runs with different splits of the dataset.
In Figure 1 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 2, learning an importance 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.
and select 113 RNA sequences listed in Murray et al.
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 to 3 torsion angles of the amino acid.
As one of the angles is normally 180°, we also discard it, and model the density over the 2D torus.
For the RNA data, the monomer is a nucleotide described by 7 torsion angles in the backbone, represented by a 7D torus.
For protein, we divide the dataset by the type of side chain attached to 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 of 4 , 096 power spherical distributions (De Cao & Aziz, 2020, MoPS).
We observe that RDM outperforms the baseline across the board, and the difference is most noticeable for the RNA data, which has a higher dimensionality.
Numerical integration ablation .
We estimate the loss (17) by integrating the inference SDE on M .
To study the effect of integration error, we experiment with various numbers of time steps evenly spaced between [0 , s ] on Glycine.
Also, as we can directly sample the Brownian motion on tori [Table] without numerical integration, we use it as a reference (termed direct loss) for comparison.
Figure 3 shows while fewer time steps tend to underestimate the loss, the model trained with 100 time steps is already indistinguishable from the one trained with direct sampling.
We also ﬁnd numerical integration is not a signiﬁcant overhead as 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 an example whose closestpoint 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 generality of our framework, we model the synthetic datasets in
Figure 5, ﬁrst introduced by Bose et al.
Lou et al.
Since hyperbolic manifolds are not compact, we need a non-zero drift to ensure the inference processs is not dissipative.
We deﬁne the prior as the standard normal distribution on the yz -plane and let U 0 be 12 ∇ g log p 0 , so that Y s will revert back to the origin.
5.4 Special Orthogonal Group
Another example whose closest-point projection is expensive to compute is the orthogonal group, as it requires performing the singular value decomposition.
To evaluate our framework on this matrix group, we generate data using the synthetic multimodal density deﬁned on SO (3) from Brofos et al.
We view it as a submanifold embedded in R 3 × 3 , therefore d = 3 and m = 9 .
We use the projected Hutchinson to estimate the Riemannian divergence.
Since the data are 3D rotational matrices, we can visualize them using the Euler angles .
We plotted the data density and the learned model density in Figure 6, where each coordinate represents the rotation around that particular axis.
Conclusion
In this paper, we introduce RDMs that extend continuous-time diffusion models to arbitrary
Riemannian manifolds—including challenging non-compact manifolds like hyperbolic spaces.
We provide a variational framework to train RDMs by optimizing a novel objective, the Riemannian
Continuous-Time ELBO.
To enable efﬁcient and stable training we provide several key tools such as a ﬁxed-inference paramterization of the SDE in the ambient space, new methodological techniques to compute the Riemannian divergence, as well as 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-art performance on density estimation on geoscience datasets, protein/RNA data on toroidal, and synthetic data on hyperbolic and orthogonal-group manifolds.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] There are a few algorithmic choices we made (limited by the framework and the fact we have a dynamical model) such as numerical integration, and divergence computation / estimation.
Our ablation studies indicate that it does not pose a big problem at least for the problems we are looking at.
However, that does not mean it can be applied to other settings (like higher dimensional problems) out of the box though.
(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?
[Yes] 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 ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (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 to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.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?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[No] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[No] 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, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to 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]
A Riemannian manifolds
Notation Convention .
There are a number of different notations used in differential geometry and all have their place.
The most abstract level is with tensors (of which forms and vectors are special cases) and it is the best for establishing general properties.
We used an index-free notation in the main paper.
A coordinate-based, but still intrinsic, description that uses local charts which describe explicit coordinate systems for patches of the manifold: for the most part we use intrinsic coordinates in the main paper.
For computational purposes it is convenient to view manifolds as hypersurfaces embedded in R m even though this obscures the geometric meaning; these are called extrinsic coordinates which we use for actual implementations.
We use capital letters to denote vectors, and tilded letters to denote vectors and variables deﬁned on the local patch.
A.1 Smooth manifolds and tangent vectors
We recall some preliminaries of smooth manifolds.
See Lee (2013) for a more detailed and comprehensive account.
A smooth d -manifold is a topological space M (assumed to be paracompact, Hausdorff and second countable) and a family of pairs { ( U i , ϕ i ) } , where the U i are open sets that together cover all of M and each ϕ
i is a homeomorphism from U i to an open set in R d ; these pairs are called charts .
They are required to satisfy a compatibility condition: if U i and U j have non-empty intersection, say V , then ϕ
i ◦ ϕ − 1 j | V has to be an inﬁnitely differentiable map from ϕ j ( V ) ⊂
R d to ϕ i ( V ) ⊂ R d .
The use of charts allows one to talk about differentiability of functions or vectors ﬁelds, by moving to R d as needed.
A smooth function f on M has type M → R and is such that for any chart ( U, ϕ ) the map f ◦ ϕ − 1 : R d → R is smooth 2 .
The set of smooth functions on M is denoted C ∞ ( M ) .
Let M be a smooth manifold, and ﬁx a point x in M .
A derivation at x is a linear operator
D : C ∞ ( M ) → R satisfying the product rule
[Equation]
for all f, g ∈ C ∞ ( M ) .
The set of all derivations at x is a d -dimensional real vector space called the tangent space T x M , and the elements of T x M are called the tangent vectors (or tangents) at x .
For the Euclidean space M = R d , we have that T x R d
= span { ∂∂x 1 , · · · , ∂∂x d } .
We now see how to use the Euclidean derivations to induce the tangent space of arbitrary Riemannian manifolds.
Let N be another smooth manifold.
For any tangent V ∈ T x M and smooth map ϕ : M → N , the differential dϕ x : T x M → T ϕ ( x ) N is deﬁned as the pushforward of V acting on f ∈ C ∞ ( N ) :
[Equation]
Note that, if ϕ is a diffeomorphism, dϕ x is an isomorphism between T x M and T ϕ ( x ) N , and the inverse map satisﬁes ( dϕ x ) − 1 = d ( ϕ − 1 ) ϕ ( x ) .
Furthermore, differentials follow the chain rule, i.e. the differential of a composite is the composite of the differentials.
Let ˜ x
= (˜ x 1 , · · · , ˜ x d ) = ϕ ( x ) be a local coordinate.
Since dϕ x : T x M → T ϕ ( x ) R d is an isomorphism, we can characterize T x M via inversion.
We deﬁne the basis vector ˜ E i of T x M by
[Equation]
which means
[Equation]
The tangent space T x M of M at x is spanned by (cid:110)
[Equation]
This means any tangent vector V can be represented by (cid:80) di =1 ˜ v i ˜
E
i for some coordinate-dependent coefﬁcients ˜ v i .
A manifold M is said to be embedded in R m if there is an inclusion map ι : M → R m such that
M is homeomorphic to ι ( M ) and the differential at every point is injective.
Every smooth manifold can be embedded in some R m with m > d for some suitably chosen m .
When M is embedded in R m , we can view T x M as a linear subspace of T x R m ; note that this map has trivial kernel.
Let ι : M → R m denote the inclusion map, i.e. ι ( x ) = x ∈ R m for x ∈ M .
Then
[Equation]
This means we can rewrite a tangent vector using the ambient space’s basis
[Equation]
28) where ¯ v j = (cid:80) di =1 ˜ v i ∂ϕ − 1 j ∂ ˜ x
i is the coefﬁcient corresponding to the j ’th ambient space coordinate.
What exactly is ϕ − 1 j ?
Note that ι ◦ ( ϕ − 1 ) is a map from R d to R m
and it takes ϕ ( x ) to ι ( x ) .
It is this that we are writing as ϕ − 1 j .
In matrix-vector form, we can write ¯ v =
d
ϕ − 1 d˜ x ˜ v , where ¯ v is a vector that represents the mdimensional coefﬁcients in the ambient space.
This also means ¯ v lies in the linear subspace spanned by the column vectors of the Jacobian ∂ϕ − 1 ∂ ˜ x i .
This linear subspace is isomorphic to T x M , which itself is a subspace of T x R m .
We refer to this linear subspace as the tangential linear subspace
Intuitively, this means a particle traveling at speed ¯ v and position x can only move tangentially on the surface.
Therefore it is restricted to move on the manifold.
A vector ﬁeld V is a continuous map that assigns a tangent vector to each point on the manifold; that is V ( x ) ∈ T x M .
We abuse the notation a bit and use capital letters to denote both vector ﬁelds and vectors.
It should be clear in the context whether it it meant to be a function of points on the manifold or not.
Such a vector ﬁeld can also map a smooth function to a function, via the assignment x ∈ M
(cid:55)→ V ( x )( f ) ∈ R .
If it maps smooth functions to smooth functions we say that the vector ﬁeld is smooth.
The space of smooth vector ﬁelds on M is denoted by by X ( M ) .
A.2 Riemannian metric
A Riemannian manifold ( M , g ) is a d -dimensional smooth manifold M equipped with an inner product g x : T x M × T x M
→ R on the tangent space of each x ∈ M (Lee, 2018).
g x is called the metric tensor at x .
A metric tensor ﬁeld is an assignment of a metric tensor to each point x of M ; we denote it by g .
The metric tensor ﬁeld g is said to be smooth if for any smooth vector ﬁelds u and v , g ( U, V )( x )
=
g x
( U ( x ) , V ( x )) is a smooth function of x .
When it is clear from the context, we suppress the subscript for simplicity.
Since g is an inner product, we also write g ( u, v ) =
(cid:104) u, v (cid:105) g .
The Euclidean metric ¯ g for R m is deﬁned as the Euclidean inner product, characterized by the delta function
[Equation]
[Equation]
which is equal to 1 if i =
j ; otherwise it is equal to 0 .
This means for any U, V ∈ T x M ,
[Equation]
Generally, given a set of basis vectors, such as ˜ E i , the metric tensor can be represented in a matrix form, via
[Equation]
This allows us to write the metric using the patch coordinates
[Equation]
where G is a matrix whose i ’th row and j ’th column corresponds to g ij .
Using the components of the metric tensor, we can deﬁne the dual basis ˜
E i = (cid:80) j g ij ˜ E j , where g ij stands for the ( i, j ) ’
th entry of the inverse matrix G − 1 .
( ˜ E 1 , · · · , ˜ E d ) is called the dual basis for ( ˜ E 1 , · · · , ˜ E d ) since they form a bi-orthogonal system, meaning
[Equation]
If M is a submanifold, e.g. if it is embedded in an ambient space [Equation] it automatically inherits the ambient manifold’s metric.
Suppose M ⊂ R m , where m > d is the dimensionality of the ambient space.
Then g =
ι ∗ ¯ g is a metric induced by the inclusion map, deﬁned by
[Equation]
Unwinding the deﬁnitions, we have
[Equation]
k , m (cid:88) [Equation] (cid:43) [Equation]
= m
[Equation] ∂ϕ [Equation]
∂
[Equation]
∂ [Equation]
That is, if ψ = ϕ − 1 is the inverse map of ϕ , we can write G = d ψ d˜ x (cid:62) d ψ d˜ x , which can be equivalently deduced from equating (30) and (32).
An important use of the metric is to deﬁne a measure over measurable subsets of the manifold.
Let ( U, ϕ ) be a chart and consider all functions smooth functions f supported in U .
Then
[Equation]
is a positive linear functional.
Since M is Hausdorff and locally compact, by the Riesz representation theorem (Rudin, 1987, Theorem 2.14), there exists a unique Borel measure µ g (over U ) such that (cid:82) U f
d µ g is equal to the evaluation of the functional above.
We can then apply a partition of unity (Lee, 2013, Theorem 2.23) to extend this construction of µ g to be deﬁned over the entire
M , which says that for any open cover { U i } of M , there exists a set of continuous functions Φ i satisfying the following properties:
≤ Φ i ( x ) for all x ∈ M . 2.
supp Φ i ⊆
U i .
(cid:80) i Φ i ( x )
= 1 for all x ∈ M . 4.
Any x ∈ M has a neighborhood that intersects with only ﬁnitely many supp Φ i .
By means of the partition, we can consider the following positive linear functional instead:
[Equation] [Equation]
which is always well-deﬁned since f is compactly supported in M (only ﬁnitely many summands are non-zero).
(cid:112) | det G | is called the volume density .
We write | G | = | det G | for short.
A probability density p over M can be thought of as a non-negative integrable function satisfying (cid:82)
M p d µ g = 1 .
A.3 Riemannian gradient and divergence
Riemannian gradient Another crucial structure closely related to the metric is the Riemannian gradient .
The deﬁnition of Riemannian gradient ∇ g : f ∈ C ∞ ( M ) (cid:55)→ ∇ [Equation] f ∈ X ( M ) is motivated by the directional derivative in Euclidean space, satisfying
[Equation]
36) for any V ∈ X ( M ) .
To obtain an explicit formula for the Riemannian gradient, we expand both sides of (36):
[Equation]
37) where we let ˜ u i and ˜ v j denote the coefﬁcients of the gradient and V respectively.
And,
[Equation]
[Equation]
Since v is arbitrary, this means for all j
[Equation]
Riemannian divergence Recall that we deﬁne the Riemannian divergence using the patch coordinates in (6), which we later show has a coordinate-free form (11) and can be computed in the ambient space (12) if the manifold is embedded.
The following theorem extends the Stokes theorem to Riemannian manifolds.
Theorem 5 ( Divergence theorem ) .
For any compactly supported f ∈ X ( M ) , (cid:82) M ∇ g · f d µ g = 0 .
Proof.
Let { (Ψ i , U i ) } be a partition of unity.
By compactness, we can choose a ﬁnite subcover over the support of f , so the index set of i is ﬁnite.
[Equation]
All of the ﬁnitely many summands equal 0 by an application of Stokes’ theorem in R d (Rudin et al., 1976, Theorem 10.33).
This is because the support of Ψ i ◦ ϕ − 1 i is contained in ϕ
i ( U i )) ; therefore at the boundary of ϕ
i ( U i ) , Ψ i ◦ ϕ − 1 i is equal to 0 .
The Riemannian divergence satisﬁes the following product rule.
Proposition 3 ( Product rule ) .
Assume V ∈ X
( M ) and f ∈ C ∞ ( M ) .
Then
[Equation]
Proof.
Using (11), the product rule of the Afﬁne connection (see Appendix A.4),
[Equation]
Proposition 4 ( Expanding Riemannian gradient ) .
Let V denote the tangential projection matrix in the sense of Proposition 2.
Then for any f ∈ C ∞ ( M )
[Equation]
A.4 Covariant derivative
An afﬁne connection allows us to compare values of a vector ﬁeld at nearby points.
It is a differential operator denoted by ∇ : X ( M ) × X ( M ) → X ( M ) and written as U, V (cid:55)→ ∇ U V for
U, V ∈ X ( M ) , satisfying the following deﬁning properties:
Linearity in U : ∇ fU 1 + gU 2 V = f ∇ U 1 V + g ∇
U
2 V for f, g ∈ C ∞ ( M ) and U 1 , U 2 , V ∈ X ( M ) .
Linearity in V : ∇ U ( aV 1 + bV 2 ) =
a ∇ U V 1 + b ∇ U V 2 for a, b ∈ R and U, V 1 , V 2 ∈ X ( M ) .
Product rule: ∇ U ( fV ) =
f ∇ U V
+ U ( f ) V for f ∈ C ∞ ( M ) and U,
V ∈ X ( M ) .
∇ U V is called the covariant derivative of V in the U -direction.
If U, V ∈ X ( R m ) , the Euclidean connection ∇ is deﬁned as
[Equation]
It can be veriﬁed that the Euclidean connection is indeed an afﬁne connection.
We can express a connection internally in terms of a coordinate system ˜
E i .
For any pair of indices i and j , we deﬁne the connection coefﬁcients of ∇ , denoted by Γ , as d 3 smooth functions satisfying
[Equation]
Then for any U, V ∈ X ( M ) , we have
[Equation]
Now given a metric tensor, we say that ∇ is a Levi-Civita connection of g if it is
Compatible with g : U ( g ( V, W ))
=
g ( ∇ U V, W ) + g ( V, ∇ U W ) .
Symmetric: ∇ u v − ∇ V U =
[ U, V ] , where [ U, V ] :
= (cid:80) di
=1 U ( V i )
˜
E i − V
( U i ) ˜
E i is the Lie bracket.
The ﬁrst condition looks messy but it essentially says that the Levi-Civita connection leaves the metric invariant.
It is equivalent to saying that the covariant derivative of g in any direction is zero.
Theorem 6 ( Fundamental Theorem of Riemannian Geometry ) .
Let ( M , g ) be a Riemannian manifold.
There exists a unique Levi-Civita connection of g .
See Lee (2018, Theorem 5.10) for proof.
The connection coefﬁcients of the Levi-Civita connection are called the Christoffel symbols of g .
They are symmetric in the lower indices, i.e. Γ kij = Γ kji .
A by-product of the proof of the fundamental theorem is the following identity, which will turn out to be useful in deriving the identity for the Riemannian divergence:
[Equation]
An example of a Levi-Civita connection is the Euclidean connection of ( R d , ¯ g ) .
It can be checked that ∇ is both symmetric and compatible with ¯ g .
Furthermore, for any d -submanifold M embedded in R m for m > d , we can deﬁne a tangential connection
[Equation]
58) for U, V ∈ X ( M ) , where U and V are any 3 smooth extensions of U and V to R m .
P is the tangential projection deﬁned as
[Equation]
59) for any V ∈ X ( R m ) .
Recall that P x is the orthogonal projection onto the tangent space spanned by ∂ψ∂ ˜ x i .
The tangential connection ∇ (cid:62) is the Levi-Civita connection on the embedded submanifold
M (Lee, 2018, Proposition 5.12).
B Proofs
Theorem 1 ( Marginal Density ) .
The density p ( x, t ) of the SDE (5) can be written as
[Equation]
where the expectation is taken wrt the following process induced by a Brownian motion B (cid:48)
s
[Equation]
Proof.
Our ﬁrst step is to express the time derivative of the density using derivations (spatial derivatives); this gives us a partial differential equation (PDE) on the manifold.
Second, we apply the
Feynman-Kac formula (Thalmaier, 2021, Proposition 3.1) to the solution of the PDE.
We denote by d ˜ X t = ˜ v 0 d t + ˜ v ◦ d B t the Stratonovich SDE deﬁned on the patch.
The density p of the process satisﬁes the Fokker-Planck equation (Chirikjian, 2009, Equation (8.16)):
[Equation]
[Equation]
We would like to re-express the RHS using the abstract vectors V 0 and V .
Note the ﬁrst term can be written as −∇ g · ( pV 0 ) .
We now show that we can also rewrite the second term in terms of the
Riemannian divergence.
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
Summing these two terms give us
[Equation]
Next
[Equation]
we expand the above formula using the product rule (43):
[Equation] [Equation]
( ∇
[Equation] [Equation]
[Equation]
(cid:33) (78) =
[Equation]
[Equation] [Equation]
(cid:88)
[Equation]
(cid:32)
[Equation]
( p )
[Equation] [Equation]
· ( V
[Equation]
[Equation] [Equation]
· ( V
[Equation]
))(cid:33) (79) =
[Equation]
( p
[Equation] [Equation]
V
[Equation] [Equation]
(cid:0) ( V ·
[Equation]
) V (cid:1) ( p )
[Equation] [Equation]
(cid:88)
[Equation]
= (cid:32)
V
[Equation]
( p )
[Equation]
p
[Equation]
g · (( ∇ g · V k ) V k ) (cid:33) (80) = − V
( p
[Equation]
p
[Equation]
· ( V 0 )
[Equation]
(cid:0) ( V ·
[Equation]
) V (cid:1) ( p )
[Equation]
1w (cid:88) k
= (cid:32)
V
[Equation]
( p )
[Equation]
p
[Equation]
g · (( ∇ g · V k ) V k ) (cid:33) (81) = − V
( p ) − p
[Equation]
g · ( V 0 ) + (cid:0) ( V ·
[Equation]
g ) V (cid:1) ( p ) + 1w (cid:88) k
=
V
[Equation]
( p ) +
[Equation]
(cid:88) k
= p ∇ g · (( ∇ g · V k ) V k ) (82) = −
[Equation]
0 ( p )
[Equation]
p
[Equation]
g · ( V 0 )
[Equation]
(cid:0) ( V ·
[Equation]
g ) V (cid:1) ( p ) +
[Equation]
(cid:88) k
=
V
[Equation]
( p ) + 1 2 p
[Equation]
g · (( V · ∇ g ) V ) (83)
In order to apply the Feynman-Kac formula, we group all the terms by the order of differentiation (of p ), which gives us
[Equation]
Now the above is a parabolic PDE, which can be solved using the Feynman-Kac formula (Thalmaier, 2021, Proposition 3.1).
Let Y be induced (8) restated below
[Equation]
Then p ( x, t ) is given by
[Equation]
where p 0 = p ( x, 0) is the prior distribution.
Theorem 2 ( Riemannian CT-ELBO ) .
Let ˆ B s be a w -dimensional Brownian motion, and let
Y s be a process solving the following
Inference SDE
[Equation]
where a : R m × [0 , T ] → R m is the variational degree of freedom.
Then we have log p
[Equation]
x, T
[Equation]
≥ E
[Equation]
where all the generative degree of freedoms V k are evaluated in the reversed time direction.
Proof.
Let P be the probability measure under which B (cid:48) is a Brownian motion.
Let
[Equation]
where a is the variational degree of freedom.
Let Q be deﬁned as
[Equation]
Note that the ﬁrst term is an Itô integral.
Then by the Girsanov theorem (Øksendal, 2003, Theorem 8.6.3), ˆ B is a Brownian motion wrt Q .
Therefore, changing the measure from P to Q to the expression in Theorem 1 yields
[Equation]
x
[Equation]
which by Jensen’s inequality, is lower bounded by
E Q
[Equation]
Now under the expectation, the Radon-Nikodym derivative can be simpliﬁed:
E Q
[Equation]
where we used the deﬁnition of Q (89), the deﬁnition of d ˆ B (88), and the Martingale property of the Itô integral (Øksendal, 2003, Corollary 3.2.6).
This concludes the proof.
Proposition 1 ( Riemannian divergence identity ) .
Let ( M, g ) be a d -dimensional Riemannian manifold.
For any smooth vector ﬁeld V k ∈ X ( M ) , the following identity holds:
[Equation]
Furthermore, if the manifold is a submanifold embedded in the ambient space R m equipped with the induced metric g = ι ∗ ¯ g , then
[Equation]
where v k = ( v 1 k , · · · , v mk ) are the ambient space coefﬁcients
V k = (cid:80) mi
=1 v ik ∂∂x i and P x is the orthogonal projection onto the tangent space.
Proof.
We drop the index on k (since the statement is for any smooth vector).
Using product rule, the LHS of (11) is equal to
[Equation]
Using the chain rule, Jacobi’s formula, and the identity (57), we have
[Equation]
[Equation] [Equation]
[Equation] [Equation]
Therefore, the LHS reduces to
[Equation]
Now we express the covariant derivative on the RHS using the connection coefﬁcients (56)
[Equation]
which means
[Equation]
Now summing all terms yields
[Equation]
Relabeling i → j and j → k in the term term shows this is equal to the LHS.
For the second half of the theorem, recall that the Levi-Civita connection is equal to the tangential connection.
Therefore, changing the basis via ˜ E j = (cid:80) mk =1 ∂ψ k ∂ ˜ x
j ∂∂x k , we can rewrite it as
[Equation]
On the other hand,
[Equation]
Since g is the induced metric, the summation over j = 1 , · · · , d is equivalent to the Frobenius inner product (cid:104)· , ·(cid:105) F of the two m × d matrices
[Equation]
Proposition 2.
If V is the tangential projection matrix, then ( V · ∇ g ) V = 0 .
Proof.
By deﬁnition,
[Equation]
Denote by the j ’th column of P x by ( P x ) : j .
Applying the resulting tangent vector to any smooth function f (evaluated at x ) and applying (12) gives
[Equation]
That is, the resulting tangent vector’s coefﬁcients correspond to the tangential projection of the vector
[Equation]
which we claim is orthogonal to the tangential linear subspace
[Equation]
To prove the claim [Equation] we ﬁrst note that we can rewrite P x as
[Equation]
where n x is of type R m × ( m − d ) [Equation] and the column vectors of n x are orthonormal, and orthogonal to the tangential linear subspace; that is to say, P x n x = 0 .
Using this representation [Equation] we can write the
Jacobian as
[Equation]
Now multiplying by the projection matrix from both sides gives
[Equation]
Lastly, let
[Equation]
which means (118) is simply
[Equation]
This implies the claim is true, since this is nothing more than a linear combination of the column vectors of n x , which is orthogonal to the tangential linear subspace.
Theorem 3 ( Marginally equivalent SDEs ) .
For λ ≤ 1 , the marginal distributions of X T − s and Y s of the processes deﬁned as below
[Equation]
both have the density q ( · , s ) .
In particular, λ = 1 gives rise to an equivalent ODE.
Proof.
We work with the derivation version of (14):
[Equation]
That is, U 0 ( f ) =
(cid:80) k
( Pr ) k
∂∂ ˜ x k f ◦ ψ , and V is the tangential projection.
The marginal density q follows the Fokker-Planck PDE
[Equation]
where we have used the product rule, and Proposition 2, the chain rule, and Proposition 4.
For λ ≤ 1 , we can rearrange the Fokker-Planck and get
[Equation]
which is the Fokker-Planck equation of the process (20).
To construct a reverse process inducing the same family of marginal densities, we mirror the diffusion term around 0 :
[Equation]
Now we apply a change of variable of time via p ( x, t )
=
q ( x, T − t ) , which means ∂ t p = − ∂ s q | s = T − t and thus
[Equation]
which is the Fokker-Planck of (21).
Theorem 4 ( Score matching equivalency ) .
For λ < 1 , let E ∞ λ denote the Riemannian CTELBO of the generative process (21), with ∇ g log q replaced by an approximate score S θ , and with (20) being the inference SDE.
Assume S θ is a compactly supported smooth vector.
Then
[Equation]
where C 1 > 0 and C 2 are constants wrt θ .
Proof.
Approximating ∇ g log q in (21) using S θ and plugging it in (5) and (20) in (9), we get [Equation] Also
[Equation]
as we only need to focus on the tangential components of a [Equation] note that
[Equation]
where E j denote the ambient space Euclidean derivation ∂∂x j .
Thus, we have
[Equation]
Summing up these two parts gives us E ∞ λ .
Taking the expectation over q ( · , 0) and applying the divergence theorem give us the desired identity.
C Manifolds
We provide some background on the manifolds used in this paper.
C.1 Spheres and tori
Spheres are deﬁned as submanifolds in an Euclidean space of points with unit Euclidean norm. Precisely, an d -sphere is S d = { x ∈ R d +1 : (cid:107) x (cid:107) 2 = 1 } . Therefore the ambient space dimensionality of a d sphere is m = d + 1 .
Tori are products of 1 -spheres (or
Naturally, we can embed a d -torus in a m = 2 d -dimensional ambient space.
Tangential projection Without loss of generality, we derive the orthogonal projection to the tangent space of spheres.
The tangential projection of tori is just the same linear operator applied to d Rvectors independently.
To derive the tangential project, we note that any any incremental change in x , denoted by d x , will need to leave the norm (cid:107) x (cid:107) 2 unchanged.
That is,
[Equation]
This means x is normal to the tangential linear subspace.
We can ﬁnd the orthogonal projection onto the tangent space by subtracting the normal component, via P x = I − xx (cid:62) (cid:107) x (cid:107) 22 .
Closest-point projection The closest-point projection onto the sphere is π ( x ) = [Equation]
One can verify this is the point on S d that minimizes the Euclidean distance from x ∈ R d +1 \ { 0 } .
C.2 Hyperbolic spaces
We work with the Lorentzian model of the hyperbolic manifold, which, like the d -spheres, is a d -manifold embedded in R d +1 , deﬁned as
[Equation]
142) where K < 0 is the curvature of the manifold, and (cid:104)· ,
·(cid:105) L is the Lorentzian inner product
[Equation]
In our experiments, K = − 1 .
The d + 1 -dimensional Euclidean space endowned with the Lorentzian inner product ( R d +1 , (cid:104)· , ·(cid:105) L ) is known as the Minkowski space.
The Lorentz inner product is in general indeﬁnite.
Therefore, technically it is not an inner product.
But it is positive deﬁnite when restricted to H dK , and as a result induces a valid Riemannian metric g L .
Equation (12), however, relies on the Euclidean geometry of the ambient space.
Therefore, we model the density p E associated with the metric tensor g E induced by the regular Euclidean inner product.
That is, p E is a probability density of the manifold ( H dK , g E ) .
Note that all the data points still lie on the same topological space H dK , and the density can be translated via p E = p L (cid:113) | G L | | G E | , where G L and G E are the components of the matrix g L and g E , and p L is the actual density on the Hyperbolic manifold ( H dK , g L ) .
This change-of-volume relation implies instead of maximizing the likelihood log p L , we can simply maximize log p E .
Alternatively, one can also compute the Riemannian divergence wrt the metric g L using the internal coordinates, as is done in (Lou et al., 2020).
In this case, the learned density will be the actual density p L on the hyperbolic manifold.
Tangential projection Similar to the spheres, we analyze the contribution of the differential d x .
[Equation]
where n x = ( − x 0 , x 1 , . .
, x d ) is the normal vector.
Subtracting the normal contribution gives rise to the tangential projection P x
= I − n
[Equation]
Note that this is different from the usual “Lorentz” orthogonal projection P L x ( u ) =
u − (cid:104) x, [Equation] (cid:105) [Equation]
(cid:104) [Equation] ,x (cid:105) L x (Ratcliffe, 1994); the latter is not orthogonal in the Euclidean inner product.
Closest-point projection
We ﬁrst derive the closest-point projection wrt the Lorentz inner product.
For any x ∈ { x (cid:48) : (cid:104) x (cid:48) , x (cid:48) (cid:105)
L <
[Equation]
where (cid:107) x (cid:107) L :
= (cid:112) (cid:104) x, x (cid:105) L is the Lorentz norm.
To deal with the constraint y ∈ H dK , we can introduce the Lagrange multiplier λ , and ﬁnd the stationary point of the function
[Equation]
Taking the gradient wrt y and setting it to be zero yield
[Equation]
On the other hand, y ∈ H dK , which means λ + 1 = (cid:112) K (cid:107) x
[Equation]
L , and therefore
[Equation]
This projection, however, is not the closest-point projection in Euclidean distance in general, as depicted in Figure 7.
This is contrary to the claim made by Skopek et al.
In fact, following the same derivation (using Euclidean distance in place of the Lorentz norm in (145)), we would end up with a Lagrange multiplier that cannot be analytically solved, as it involves solving a root ﬁnding problem.
This projection, albeit not the shortest one in Euclidean distance, is still a valid projection.
We use it in numerical integration to simulate the dynamics.
C.3 Orthogonal groups
The orthogonal groups are deﬁned as O ( n ) = { X ∈ R n × n : X (cid:62) X
= XX (cid:62) = I } .
The determinant of X is either 1 or − 1 .
The subgroup with determinant 1 is called the special orthogonal group, denoted by SO ( n ) .
Naturally, R n × n is an ambient space of the orthogonal groups.
Tangential projection Following the differential analysis,
[Equation]
That is, d XX (cid:62) is skew-symmetric.
Denote the set of skew-symmetric matrices by Skew n = { X ∈ R n × n : X (cid:62) = − X } .
Let U be an arbitrary matrix in R n × n .
We want to project it orthogonally onto T X O ( n ) .
The orthogonal projection needs to be the closest-point projection onto the subspace.
We can use the
Frobenius norm to induce the Euclidean distance metric over the entries of the matrix.
Then ﬁnding the closest-point projection V of U amounts to ﬁnding the stationary point of
[Equation]
[Equation]
where Λ is the Lagrange multiplier.
Taking the gradient wrt V yields
[Equation]
Equating the last step with 0 yields
[Equation]
Since V needs to satisfy XV (cid:62) + V X (cid:62) = 0 , we have
[Equation]
) which means
[Equation]
Substituting this into (151) yields
[Equation]
That is, P X ( U ) =
U − XU (cid:62) Xfor orthogonal groups.
Closest-point projection
Again, using the Lagrange multiplier Λ for the constraint that the projection M of X should satisfy M (cid:62) M
= I , we try to ﬁnd the stationary point of the following quantity
[Equation]
Equating the gradient wrt M with 0 gives d
[Equation]
which means
[Equation]
Since M is orthogonal, we have
[Equation]
which means
[Equation]
Let X = UDV (cid:62) be the singular value decomposition of X .
Then
[Equation]
Substituting this into (157), we get
[Equation]
That is, π ( X ) =
UV (cid:62) for orthogonal groups, where U, V are the left and right singular matrices of
X .
[Table] [Table]
D Experimental details D.1 Architecture
In our experiments, we parameterize the a network as a multi-layer perceptron (MLP) with either the sinusoidal or the swish activation function.
For the hyperbolic experiments, the ﬁrst layer of the MLP has an additional ActNorm layer (Kingma & Dhariwal, 2018) which we ﬁnd adds extra numerical stability.
The ActNorm layer is initialized before training with one batch such that its output has a mean of zero and a standard deviation of one.
In an analogous manner to training the MLP the ActNorm parameters are updated via backpropagation.
For the orthogonal group experiments we ﬂatten the input matrix into a vector before passing it to the MLP.
The details of our various model are given in Table 3.
For our importance sampler which is used to represent a differentiable distribution over [0 , T ] , we use a deep sigmoidal ﬂow (Huang et al., 2018) (without the ﬁnal logit activation) followed by a ﬁxed scaling ﬂow, which represents the range [0 , T ] .
We disconnect the gradient from the numerical solver to save compute; i.e. Y s is not differentiable.
This would result in slightly biased gradient updates for minimizing the variance of the importance estimator, but we still observe substantial reduction in variance (see Figure 2).
Finally, we use PyTorch (Paszke et al., 2019) as our deep learning framework.
Computational Resources .
We run all of our experiments either on a single NVIDIA Tesla V or a single NVIDIA Quadro RTX 8000 GPU for a maximum of 30 hours.
D.2 Optimization
We use the Adam (Kingma & Ba, 2015) optimizer to train the a network.
The learning rate and momentum parameters used for each manifold is mentioned in the Table 4.
For the sphere experiments, we slowly decrease the learning rate during training using a cosine scheduler.
For optimization of our importance sampler, we use Adam with a ﬁxed learning rate of 0 .
We update the importance sampler every 500 steps of our training loop for the a network.
Lastly, to optimize our mixture of power spherical distributions for the tori experiments we use Adam with a learning rate of 0 .
03 with β 1 = 0 .
9 and β 2 = 0 .
D.3 KELBO
The gap between the exact likelihood of the data given the model, i.e. log p ( x ) , and the Riemannian
CT-ELBO may be large.
This evaluation gap makes empirical validation of the models using the
Riemannian CT-ELBO imprecise.
We acquire a tighter lower bound by using K > 1 samples and importance sampling similar to Burda et al.
In details, we know from (90) that:
[Equation]
[Figure]
[Table]
We rewrite this as:
[Equation]
where L ( Y ) is deﬁned to be:
[Equation]
Then by Jensen’s inequality [Equation] we have that:
[Equation]
where Y i s are i.i
[Equation]
d. trajectories sampled from Q .
We call this new lower bound KELBO.
Note that this is a tighter lower bound because we can write:
[Equation]
In fact, this lower bound increases monotonically to the true likelihood as K → ∞ .
We use KELBO with K = 100 for evaluating all of our models.
We have experimented with the K to be up toand found out the results stop changing much for K > 100 .
D.4 Numerical integration of the SDEs
During training and evaluation, we numerically integrate the SDE on each respective manifold using the Stratonovich-Heun method as described in Burrage et al.
Each iteration is followed by the closest-point projection (in the case of H dK , we use the closest-point project wrt the Lorentz inner product).
The number of integration steps for each manifold during training is reported in Table 5.
During evaluation, as described in D.3, we numerically integrate the data from s = 0 to s = T , and the Itô integral involved in the KELBO is approximated using the Euler-Maruyama scheme (note that the dynamics is still generated using Stratonovich-Heun).
As computing KELBO requires forward passes through the a network, it may not be as smooth as just integrating the inference SDE.
Therefore, we use an adaptive step size for integration.
We adapted the torchsde library (Kidger et al., 2021; Li et al., 2020) to calculate errors and adapt the step size accordingly.
The error tolerance and minimum step size used in integration for all the experiments are reported in Table 6.
Also, for plotting densities we use the exact log likelihood of the equivalent ODE.
To numerically integrate the ODE for computing the exact likelihood, we use the default dopri5 solver from the torchdiffeq library (Chen et al., 2018b, 2021).
Finally, we use cartopy (Met Ofﬁce, 2010 - 2015), matplotlib (Hunter, 2007), and plotly (Inc., 2015) for visualization.