Title: Scaling Riemannian Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion Models

Geometric Deep Learning

Manifolds

Numerical Algorithms
Abstract: Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ans\"{a}tze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement and is competitive with other techniques. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds, including $SU(n)$ lattices in the context of lattice quantum chromodynamics (QCD). Finally, we apply our models to contrastively learned hyperspherical embeddings, curbing the representation collapse problem in the projection head and closing the gap between theory and practice.

Scaling Riemannian Diffusion Models

Abstract

Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on SU ( n ) lattices and contrastively learned embeddings on high dimensional hyperspheres.

1 Introduction

By learning to faithfully capture high-dimensional probability distributions, modern deep generative models have transformed countless fields such as computer vision [19] and natural language processing [41]. However, these models are built primarily for geometrically simple data spaces, such as Euclidean space for images and discrete space for text. For many applications such as protein structure prediction [13], contrastive learning [12], and high energy physics [6], the support of the data distribution is instead a Riemannian manifold such as the sphere or torus. Here, naïvely applying a standard generative model on the ambient space results in poor performance as it doesn't properly incorporate the geometric inductive bias and can suffer from singularities [7].

As such, a longstanding goal within Geometric Deep Learning has been the development of principled, general, and scalable generative models on manifolds [8, 18]. One promising method is the Riemannian Diffusion Model [4, 25], the natural generalization of standard Euclidean space score-based diffusion models [46, 48]. These learn to reverse a diffusion process on a manifold-in particular, the heat equation-through Riemannian score matching methods. While this approach is principled and general, it is not scalable. In particular, the additional geometric complexity renders the denoising score matching loss intractable. Because of this, previous work resorts to inaccurate approximations or sliced score matching [47], but these degrade performance and can't be scaled to high dimensions. We emphasize that this fundamental problem causes Riemannian Diffusion Models to fail for even trivial distributions on high dimensional manifolds, which limits their applicability to relatively simple low-dimensional examples.

In our work, we propose several improvements to Riemannian Diffusion Models to stabilize their performance and enable scaling to high dimensions. In particular, we reexamine the heat kernel [21], which is the core building block for the denoising score matching objective. To enable denoising

score matching, one needs to be able to sample from and compute the gradient of the logarithm of the heat kernel efficiently. This can be done trivially in Euclidean space as the heat kernel is a Gaussian distribution, but doing this is effectively intractable for general manifolds. By restricting our analysis to Riemannian symmetric spaces [9, 29], which are a class of a manifold with a special structure, we can make substantial improvements. We leverage this additional structure to quickly and precisely compute heat kernel quantities, allowing us to scale up Riemannian Diffusion Models to high dimensional real-world tasks. Furthermore, since almost all manifolds that practitioners work with are (or are diffeomorphic to) Riemannian symmetric spaces, our improvements are generalizable and not task-specific. Concretely our contributions are:

We present a generalized strategy for numerically computing the heat kernel on Riemannian symmetric spaces in the context of denoising score matching. In particular, we adapt known heat kernel techniques to our more specific problem, allowing us to quickly, accurately, and stably train with denoising score matching.

We show how to exactly sample from the heat kernel using our quick computations above. In particular, we show how our exact heat kernel computation enables fast simulation free techniques on simple manifolds. Furthermore, we develop a sampling method based on the probability flow ODE and show that this can be quickly computed with standard ODE solvers on the maximal torus of the manifold.

We empirically demonstrate that our improved Riemannian Diffusion Models improve performance and scale to high dimensional real world tasks. For example, we can faithfully learn Wilson action on 4 × 4 SU (3) lattices ( 128 dimensions). Furthermore, when applied to contrastively learned hyperspherical embeddings ( 127 dimensions), our method enables better model interpretability by recovering the collapsed projection head representations. To the best of our knowledge, this is the first example where differential equation-based manifold generative models have scaled to real world tasks with hundreds of dimensions.

2 Background

2.1 Diffusion Models

Diffusion models on R d are defined through stochastic differential equations [24, 45, 48]. Given an initial data distribution p 0 on R d , samples x 0 ∼ R d are perturbed with a stochastic differential equation[33]

where f and g are fixed drift and diffusion coefficients, respectively. The time varying distributions p t (defined by x t ) evolves according to the Fokker-Planck Equation

and approaches a limiting distribution π ≈ p T , which is normally a simple distribution like a Gaussian N (0 , σ 2 T I ) through carefully chosen f and g . Our SDE has a corresponding reversed SDE

which maps p T back to p 0 . Motivated by this correspondence, diffusion models approximate the score function ∇ x log p t ( x ) using a neural network s θ ( x , t ) . To do this, one minimizes the score matching loss [27], which is weighted by constants λ t :

Since this loss is intractable due to the unknown ∇ x log p t ( x t ) , we instead use an alternative form of the loss. One such loss is the implicit score matching loss[27]:

which normally is estimated using sliced score matching/Hutchinson's trace estimator[26, 47]:

where ϵ is drawn over some 0 mean and identity covariance distribution like the standard normal distribution or the Rademacher distribution. Unfortunately, the added variance from the ϵ normally renders this loss unworkable in high dimensions[46], so practitioners instead use the denoising score matching loss[51]

where p t ( x t | x 0 ) is derived from the SDE in Equation 1 and is normally tractable. Once s θ ( x t , t ) is learned, we can construct a generative model by first sampling x T ∼ π ≈ p T and solving the generative SDE from t = T to t = 0 :

Furthermore, there exists a corresponding 'probability flow ODE' [48]

that has the same evolution of p t as the SDE in Equation 1. This can be approximated using our score network s θ to get a Neural ODE [11]

which can be used to evaluate exact likelihoods of the data [20].

2.2 Riemannian Diffusion Models

To generalize diffusion models to d -dimensional Riemannian manifolds M , which we assume to be compact, connected, and isometrically embedded in Euclidean space, one adapts the existing machinery to the geometrically more complex space [4]. Riemannian manifolds are deeply analytic constructs, so Euclidean space operations like vector fields v , gradients ∇ , and Brownian motion B t have natural analogues for M . This allows one to mostly port over the diffusion model machinery from Euclidean space. Here, we highlight some of the core differences.

The forward SDE is the heat equation. The particle dynamics typically follow a Brownian motion:

Unlike in the Euclidean case, p t approaches the uniform distribution U M as t →∞ (in practice, this convergence is fast, getting within numerical precision for t ≈ 5 ).

The transition density has no closed form. Despite the fact that we work with the most simple SDE, the transition kernel defined by the manifold heat equation p t ( x t | x 0 ) has no closed form. This transition kernel is known as the heat kernel, which satisfies Equation 2 with the additional condition that, as t → 0 , the kernel approaches δ x 0 . We will denote this by K M ( x t | x 0 , t ) , and we highlight that, when M is R d , this corresponds to a Gaussian and is easy to work with.

This has several major consequences which cause prior work to favor sliced score matching over denoising score matching. First, to sample a point x ∼ K M ( ·| x 0 , t ) , one must simulate a Geodesic Random Walk with the Riemannian exponential map exp :

Additionally, to calculate K M ( ·| x 0 , t ) or ∇ log K M ( x | x 0 , t ) , one must use eigenfuncions f i which satisfy ∆ f i = -λ i f i . These f i form an orthonormal basis for all L 2 function on M , allowing us to write the heat kernel with an infinite sum:

Additionally, previous work has also explored the use of the Varadhan approximation for small values of t (which uses the Riemannian logarithmic map log ) [43]:

3 Method

The key problem with applying denoising score matching in practice is that the heat kernel computation is too expensive and inaccurate. Notably, simulating a geodesic random walk is expensive since it requires many exponential maps. Furthermore, the eigenfunction expansion in Equation 13 requires more and more eigenfunctions (numbering in the tens of thousands) as t → 0 . Worse still, these eigenfunctions formulas are not well-known for most manifolds, and, even when explicit formulas exist, they can be numerically unstable (like in the case of S n ). One possible way to alleviate this is to use Varadhan's approximation for small t , but this is also unreliable except for very small t .

To remedy this issue, we instead consider the case of Riemannian symmetric spaces [22]. We emphasize that most manifold generative modeling applications already model on Riemannian Symmetric Spaces like the sphere, torus, or Lie Groups, so we do not lose applicability by restricting our attention here. Furthermore, for surfaces (which have appeared as generative modeling test tasks in the literature [42]), one can always define a generative model by mapping the data points to S 2 , learning a generative model there, and mapping back [14]. We empathize that, outside of these two examples, we are unaware of any other manifolds which have been used for Riemannian generative modeling tasks. For this very reason, we will also focus primarily on the compact case, since the irreducible (base case) noncompact symmetric spaces are all diffeomorphic to Euclidean space, although we discuss the extensions of our method in Appendix A.2.

3.1 Heat Kernels on Riemannian Symmetric Spaces

In this section, we will define Riemannian Symmetric Spaces and showcase their relationship with the heat kernel. We empathize that our exposition is neither rigorous nor fully defines all terms, although it provides a general intuition with which to build our calculations. We urge interested readers to consult a book [22] or monograph [9] for a full treatment of the subject.

Definition 3.1. A Riemannian Symmetric Space is a Riemannian manifold such that, for all points x ∈ M , there exists a local isometry s s.t. s ( x ) = x and D x s = -id T x M .

Intuitively, this condition means that, at each point, the manifold 'looks the same" in all directions, which simplifies the analysis. While this is defined as a local condition, it also has a global consequence: all Riemannian symmetric spaces are quotients G/K for Lie Groups G and compact isotropic subgroups K . Importantly, most Riemannian manifolds that are used in practice are symmetric:

Examples. Lie Groups G ∼ = G/ { e } (where { e } is the trivial Lie Group), the sphere S n ∼ = SO ( n + 1) /SO ( n ) , and hyperbolic space H n ∼ = SO ( n, 1) /O ( n ) are all symmetric spaces.

On symmetric spaces, one can define a special structure called the maximal torus:

Definition 3.2 (Maximal Torus) . A maximal torus T is a compact, connected, and abelian subgroup of G that is not contained in any other such subgroup. For symmetric spaces G/K , the maximal torus is defined by quotienting out the maximal torus of G . Importantly, the maximal torus is isomorphic to a standard torus T n ≈ ( S 1 ) n .

Intuitively, the maximal torus forms the 'most stable" part of the symmetric space, and it shows up as a natural object when performing analysis due to its isometry with the standard flat torus.

Examples. For the sphere S n , a maximal torus is any great circle. For the Lie group of unitary matrix U ( n ) , the maximal torus is the set of all diagonal matrices { diag( e iθ 1 , . . . , e iθ n ) : θ k ∈ [0 , 2 π ) } .

To work with the maximal torus algebraically (as we shall do in our definition of the heat kernel), we instead analyze the root system.

Definition 3.3 (Roots of a Maximal Torus) . We define a root α be a vector (really a covector) that corresponds with a character χ of T s.t.

We define the set R + to be the set of all positive roots. The multiplicity m α of a root α is the dimension of the eigenspace over the Lie algebra G

Examples. For all spheres dimension n , there is only one root of multiplicity n -1 , which corresponds (roughly) to measuring the angle on the maximal torus/great circle.

Notably, the heat kernel is invariant on the maximal torus. This allows us to write out the formula for the heat kernel on the general manifold as a simplified formula on the maximal torus

Proposition 3.4 (Heat Kernel Reduces on Maximal Torus) . The Laplace-Beltrami operator on M (the manifold generalization of the standard Laplacian) induces the 'radial' Laplacian on T :

where ∆ T is the standard Laplacian on the torus. Here, h is the 'flat" coordinate of x on the maximal torus (e.g. for sphere this is the angle between x and x 0 ).

One can readily see that the above formula allows one to reduce computations over M to computations over T , which greatly reduces dimensionality ( n dimensions to 1 dimension for spheres and O ( n 2 ) dimensions to O ( n ) for Matrix Lie Groups).

3.2 Improved Heat Kernel Estimation

We now use this maximal torus perspective to improve computations related to the heat kernel. In particular, we can greatly improve both the speed and fidelity of the numerical evaluation during our training process.

3.2.1 Eigenfunction Expansion Restricted to the Maximal Torus

We note that the maximal torus relationship in Proposition 3.4 reduces the eigenfunction expansion in Equation 13 to an eigenfunction expansion of the induced Laplacian on the maximal torus. This has implicitly appeared in previous work when defining Riemannian Diffusion Models on S 2 and SO (3) [4, 34], allowing one to rewrite the summation as, respectively

where P l are the Legendre Polynomials and

By making this relationship explicit, we can directly generalize this formulation to other symmetric spaces, e.g. SU (3)

where χ are the (real) irreducible representations and θ, ϕ are the induced angles between x and x 0 .

By reducing the dimensionality, this strategy generally makes the summation more tractable. Furthermore, for cases such as the hypersphere S n , the eigenfunctions are numerically unstable (rendering them useless for our purposes), while the maximal torus setup stable.

3.2.2 Controlling Errors for Small Time Values

While the torus eigenfunction expansion greatly reduces the computational cost (particularly for higher dimensions), for small values of t the heat kernel scaling terms e -λ i t remain large as one increases i, λ i →∞ . As a result, one must evaluate many more terms (often in the thousands), and the result can be very numerically unstable even with double precision.

To address this problem, we examine several refined versions of the Varadhan approximation that use the fact that the manifold is symmetric. These approximations allow us to control the number of eigenfunctions required and, in some cases, completely obviate the need for them entirely.

The Schwinger-Dewitt Approximation

The Varadhan approximation is constructed by approximating the heat kernel with a Gaussian distribution (with respect to the Riemannian distance function). However, doing this does not account for the curvature of the manifold. Incorporating this premise, we get the Schwinger-Dewitt approximation [9]:

Here, ∆ x 0 ( x ) = det( D x 0 exp x 0 (log x 0 ( x ))) is the (unnormalized) change of volume term introduced by the exponential map, and R is the scalar curvature of the manifold. Generally, this is much more accurate than Varadhan's approximation as it better accounts for the curvature of the manifold through the ∆ term, retaining accuracy up to moderate time values.

∆ appears to be a rather computationally demanding term since it requires evaluating the determinant of a Jacobian. Indeed, this naively takes O (dim M 3 ) evaluations which is completely inaccessible in higher dimensions. However, we again emphasize the fact that we are working with symmetric spaces: ∆ has a particularly simple formula defined by our flat coordinate h from above:

Sum Over Paths

The fact that we can derive a better approximation using a different power of ∆ points to deeper connections between the heat kernel and the 'Gaussian' with respect to distance. We draw inspiration from several Euclidean case examples, such as the flat torus [28] or the unit interval [36]. For those cases, the heat kernel is derived by summing a Gaussian over all possible paths connecting x 0 and x . While this formula does not exactly lift over to Riemannian symmetric spaces, there exists an analogue for Lie Groups [9]:

Here, Γ is infinite tangent space grid that corresponds with loops of the torus (i.e. 2 πn circular intervals), and ρ 2 is a manifold specific constant. We note that the product over R + is exactly the ∆ change in variables as given above since m α = 1 , but we simply extend this to every other root.

Generally, this formula is rather powerful as it gives us an exact (albeit infinite) representation for the heat kernel. Compared to the eigenfunction expansion in Equation 13, the Sum Over Paths representation is accurate for small t , which nicely complements the fact that the eigenfunction representation is accurate for large t . This formula does generalize to split-rank Riemannian symmetric spaces like odd dimensional spheres (more details in Appendix A.1) and noncompact spaces (Appendix A.2).

3.2.3 A Unified Heat Kernel Estimator

We unify these approximations into a single heat kernel estimator. Our computation method splits up the heat kernel evaluation based on the value of t and applies an eigenfunction summation or an improved small time approximation accordingly. This allows us to effectively control the errors at both the small and large time steps while significantly reducing the number of function evaluations for each. Our full algorithm is outlined in Algorithm 1.

Algorithm 1: Heat Kernel Computation

We ablate the accuracy of our various heat kernel (score) approximations in Figure 1 for S 2 , S 127 , and SO (3) . In general, we found that computing with standard eigenfunctions was too costly and too prone to numerical blowup, and Varadhan's approximation was simply too inaccurate. For some spaces like S 127 , a combination of preexisting methods (like is briefly explored in [4]) would not work since K EF M NaNs out before Varadhan becomes accurate.

(a)

S

2

(b)

S

127

(c)

SO

(3)

106

Eigenfunction

1023

Eigenfunction

106

Eigenfunction

Varadhan

Varadhan

Varadhan

104

Schwinger-Dewitt

1019

Schwinger-Dewitt

104

Sum OverPaths

102

1015

10²

1011

100

100

107

102

10-2

103

10~4

10-1

10-4

?27

10-6

105

10-6

M

103

10-2

10-1

100

10-2

10-1

100

103

10-2

10-1

100

3.3 Exact Heat Kernel Sampling

To train with the denoising score matching objective, one must produce heat kernel samples to monte carlo samples the loss. Typically, Riemannian Diffusion Models sample by discretizing a Brownian motion through a geodesic random walk [29]. However, this can be slow as it requires taking many computationally expensive exponential map steps on M and can drift off the manifold due to compounding numerical error [25]. In this section, we discuss strategies to sample from the heat kernel quickly and exactly.

Cheap Rejection Sampling Methods

When our dimension is low enough, we can sample using rejection sampling using our fast heat kernel evaluator. The key detail is the prior distribution, which needs to have a closed form density, be easy to sample from, and must not deviate from the heat kernel too much. For large time steps t the natural prior distribution is uniform. Conversely, for small time steps, we instead use the wrapped Gaussian distribution, which can be sampled by passing a tangent space Gaussian through the exponential map

(a)

t

= 0

FM Path.

(b)

t

= 0

.

5

FM Path.

(c) Diffusion SM path.

(d) Closeup of Figure 2c

and has a density

In both cases, the ratio can be trivially bounded by examining their behavior numerically.

Heat Kernel ODE Sampling

As an alternative, we notice that we can apply the probability flow ODE to sample from the heat kernel exactly. In particular, we draw a sample x T ∼ U M , where T is large enough s.t. K M ( ·| x 0 , T ) is close to the uniform distribution (within numerical precision). We then solve the ODE d ds x T = -1 2 ∇ x log K M ( x s | x 0 , s ) from s = T to s = t . As the sampling procedure follows the probability flow ODE, this is guaranteed to produce samples from K M ( ·| x 0 , t ) .

This is solvable as a manifold ODE, as previous works have already developed adaptive manifold ODE solvers [37]. Furthermore, by Proposition 3.4, we can restrict our vector field to the maximal torus and solve it there. Note that this allows us to use preexisting Euclidean space solvers since the torus is effectively Euclidean space. Lastly, we can scale the time schedule of the ODE (with a scheme like a variance-exploding schedule) to stabilize the numerical values.

4 Related Work

Our work exists in the established framework of differential equation-based Riemannian generative models. Early methods generalized Neural ODEs to manifolds[15, 37, 38], enabling training with maximal likelihood. More recent methods attempt to remove the simulation components[3, 42], but this results in unscalable or biased objectives. We instead work with diffusion models, which are based on scores and SDEs and do not have any of these issues. In particular, we aim to resolve the main gap that prevents Riemannian Diffusion Models from scaling to high dimensions.

Riemannian Flow Matching (RFM) [10] is a concurrent work that attempts to achieve similar goals (i.e. scaling to high dimensions) by generalizing flow matching [35] to Riemannian manifolds. The fundamental difficulty is that one must design smooth vector fields that flows from a base distribution (ie the uniform distribution) to a data point. RFM introduces several geodesic-based vector fields, but these break the smoothness assumption (see Figure 2), which we found to be detrimental for some of the densities we considered (e.g. Figure 3). However, similar to the Euclidean case, RFM with the diffusion path corresponds to score matching with the probability flow ODE, so our work provides the numerical computation necessary to construct such a flow.

5 Experiments

5.1 Simple Test Tasks

We start by comparing our accurate denoising score matching objective with the inaccurate approximation scheme suggested by [4] based on 50 eigenfunctions. We test on the compiled Earth science

Flow Matching

Score Matching

datasets from [38], detailing results are in Table 1. Generally, our accurate heat kernel results in a substantial improvement and matches sliced score matching. Note that we do not expect our method to outperform sliced score matching since these datasets are low dimensional and the objectives are equivalent (up to a small additional variance).

We also compare directly with RFM on a series of increasingly complex checkerboard datasets on the flat torus. These datasets have appeared in prior work to measure model quality [3, 5, 37]. As a result of the non-smooth vector field dynamics, we find that RFM degrades in performance as the checkerboard increases in complexity, and is unable to learn past a certain point. Our visualized results are given in Figure 3.

5.2 Learning the Wilson Action on SU (3) Lattices.

We apply our method to learn densities on SU (3) 4 × 4 . We consider the data density defined by the Wilson Action [53]:

For our experiments, we take β = 9 as a standard hyperparameter and learn this distribution, achieving an effective sample size of 0 . 62 . Generally, learning to sample from p ( K ) directly requires a separate objective that is not based on score matching[6, 50]. However, to showcase our score matching techniques, we instead learn the density through precomputed samples, and concurrent work has shown that this type of training can be coupled with diffusion guidance to improve variational inference methods [17]. As such, our model has the potential to improve SU ( n ) lattice QCD samplers, although we leave this task for future work. We also note that our model can likely be further improved by building data symmetries into the score network [30].

5.3 Contrastively Learned Hyperspherical Embeddings

Finally, we examine contrastively learning [12]. Standard contrastive losses optimize embeddings of the data that lie on the hypersphere. Paradoxically, these embeddings are unsuitable for most

downstream tasks, so practitioners instead use the penultimate layer [2]. This degrades interpretability, since the theoretical analyses in this field work on the hyperspherical embeddings [52].

We investigate this issue further in the context of out of distribution (OOD) detection. We use the pretrained embedding network from CIDER[39]. Using the hyperspherical representation for OOD detection produces very subpar results. However, using likelihoods from our Riemannian Diffusion Model stabilizes performance and achieves comparable results with other penultimate feature-based methods (see Figure 2). We emphasize that the embedding network has been tuned to optimize the performance using the penultimate layer. Since our established theory exclusively focuses on the properties of the hyperspherical embedding, the fact that our Riemannian Diffusion Models can extract a comparable representation can lead to more principled improvements for future work.

6 Conclusion

Wehave introduced several practical improvements for Riemannian Diffusion Models that leverage the fact that most relevant manifolds are Riemannian symmetric spaces. Our improved capabilities allow us, for the first time, to scale differential equation manifold models to hundreds of dimensions, where we showcase applications in lattice QCD and constrastive learning. We hope that our improvements help open the door to the broader adoption of Riemannian generative modeling techniques.

References

Belal E. Baaquie. Path-integral derivation of the heat kernel on SU( n ) space. Phys. Rev. D , 32:1007-1010, Aug 1985. doi: 10.1103/PhysRevD.32.1007. URL https://link.aps.org/doi/10.1103/PhysRevD. 32.1007 .

Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari S. Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Grégoire Mialon, Yuandong Tian, Avi Schwarzschild, Andrew Gordon Wilson, Jonas Geiping, Quentin Garrido, Pierre Fernandez, Amir Bar, Hamed Pirsiavash, Yann LeCun, and Micah Goldblum. A cookbook of self-supervised learning. ArXiv , abs/2304.12210, 2023.

Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Maximilian Nickel, Aditya Grover, Ricky T. Q. Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. In International Conference on Machine Learning , 2022.

Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and A. Doucet. Riemannian score-based generative modeling. ArXiv , abs/2202.02763, 2022.

A. Bose, Ariella Smofsky, Renjie Liao, P. Panangaden, and William L. Hamilton. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning , 2020.

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. ArXiv , abs/2008.05456, 2020.

Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estimation. ArXiv , abs/2003.13913, 2020.

Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur D. Szlam, and Pierre Vandergheynst. Geometric deep learning: Going beyond euclidean data. IEEE Signal Processing Magazine , 34:18-42, 2016.

Roberto Camporesi. Harmonic analysis and propagators on homogeneous spaces. Physics Reports , 196: 1-134, 1990.

Ricky T. Q. Chen and Yaron Lipman. Riemannian flow matching on general geometries. ArXiv , abs/2302.03660, 2023.

Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Kristjanson Duvenaud. Neural ordinary differential equations. In Neural Information Processing Systems , 2018.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. ArXiv , abs/2002.05709, 2020.

Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and T. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. ArXiv , abs/2210.01776, 2022.

Victor D. Dorobantu, Charlotte Borcherds, and Yisong Yue. Conformal generative modeling on triangulated surfaces. ArXiv , abs/2303.10251, 2023.

Luca Falorsi and Patrick Forr'e. Neural ordinary differential equations on manifolds. ArXiv , abs/2006.06663, 2020.

H. D. Fegan. The heat equation on a compact lie group. Transactions of the American Mathematical Society , 246:339-357, 1978. ISSN 00029947. URL http://www.jstor.org/stable/1997977 .

Tomas Geffner, George Papamakarios, and Andriy Mnih. Score modeling for simulation-based inference. ArXiv , abs/2209.14249, 2022.

Mevlana Gemici, Danilo Jimenez Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds. ArXiv , abs/1611.02304, 2016.

Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In NIPS , 2014.

Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Kristjanson Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. ArXiv , abs/1810.01367, 2018.

Alexander Grigor'yan. Spectral theory and geometry: Estimates of heat kernels on riemannian manifolds. 1999.

Sigurdur Helgason. Differential geometry, lie groups, and symmetric spaces. 1978.

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv: Learning , 2016.

Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv , abs/2006.11239, 2020.

Chin-Wei Huang, Milad Aghajohari, A. Bose, P. Panangaden, and Aaron C. Courville. Riemannian diffusion models. ArXiv , abs/2208.07949, 2022.

Michael F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics - Simulation and Computation , 18:1059-1076, 1989.

Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. J. Mach. Learn. Res. , 6:695-709, 2005.

Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and T. Jaakkola. Torsional diffusion for molecular conformer generation. ArXiv , abs/2206.01729, 2022.

Jurgen Jost. Riemannian geometry and geometric analysis. 1995.

Isay Katsman, Aaron Lou, Derek Lim, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Equivariant manifold flows. In Neural Information Processing Systems , 2021.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR , abs/1412.6980, 2014.

Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. ArXiv , abs/2107.00630, 2021.

Peter E. Kloeden and Eckhard Platen. The numerical solution of stochastic differential equations. The Journal of the Australian Mathematical Society. Series B. Applied Mathematics , 20:8 - 12, 1977.

Adam Leach, Sebastian M Schmon, Matteo T. Degiacomi, and Chris G. Willcocks. Denoising diffusion probabilistic models on SO(3) for rotational alignment. In ICLR 2022 Workshop on Geometrical and Topological Representation Learning , 2022. URL https://openreview.net/forum?id=BY88eBbkpe5 .

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. ArXiv , abs/2210.02747, 2022.

Aaron Lou and Stefano Ermon. Reflected diffusion models. ArXiv , abs/2304.04740, 2023.

Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Neural manifold ordinary differential equations. ArXiv , abs/2006.10254, 2020.

Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. ArXiv , abs/2006.10605, 2020.

Yifei Ming, Yiyou Sun, Ousmane Dia, and Yixuan Li. How to exploit hyperspherical embeddings for out-of-distribution detection? In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=aEFaE0W5pAd .

Boris Polyak and Anatoli B. Juditsky. Acceleration of stochastic approximation by averaging. Siam Journal on Control and Optimization , 30:838-855, 1992.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence-based generative modeling on manifolds. In Neural Information Processing Systems , 2021.

Laurent Salo-Coste. The heat kernel and its estimates. 2009.

Vikash Sehwag, Mung Chiang, and Prateek Mittal. Ssd: A unified framework for self-supervised outlier detection. ArXiv , abs/2103.12051, 2021.

Jascha Narain Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. ArXiv , abs/1503.03585, 2015.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. ArXiv , abs/1907.05600, 2019.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Conference on Uncertainty in Artificial Intelligence , 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=PxTIG12RRHS .

Yiyou Sun, Yifei Ming, Xiaojin Zhu, and Yixuan Li. Out-of-distribution detection with deep nearest neighbors. In International Conference on Machine Learning , 2022.

Francisco Vargas, Will Grathwohl, and A. Doucet. Denoising diffusion samplers. ArXiv , abs/2302.13834, 2023. URL https://api.semanticscholar.org/CorpusID:257220165 .

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 23:1661-1674, 2011.

Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. ArXiv , abs/2005.10242, 2020.

Kenneth G. Wilson. Confinement of quarks. Physical Review D , 10:2445-2459, 1974.

A Additional Heat Kernel Computational Details

A.1 Sum over Paths for Spheres

For spheres, one has the following sum over paths formula for odd-dimensional spheres S n :

where θ = arccos( x · y ) is the distance between two points. One can derive the sum over paths formulation by applying the derivative to each summand. This formula does not generalize to even dimensional spheres, as the intertwining operator appears which makes the evaluation an integral, although a similar recurrent can be derived.

A.2 Extension to Noncompact Symmetric Spaces

For noncompact symmetric spaces, our heat kernel evaluations are only given by the sum over paths representation. Luckily, there is only 'one" path since the space is not compact, meaning we don't need to do an infinite sum. In particular, this gives us the following formula for noncompact lie groups

and for hyperbolic space, one has the formula for odd dimensions (equivalent to the one for odd dimensional spheres):

where ρ is the hyperbolic distance between y and x .

Generally, there is the notion of a OU process on a manifold [4, 25] which limits towards the wrapped Gaussian distribution. We note that the intermediate densities are not heat kernels (unlike in the Euclidean case).

A.3 Explicit Heat Kernel Formulas

In this section, we highlight the formulas we used for computing the heat kernel.

Torus. The Torus T n ∼ = ( S 1 ) n can be realized as a flat torus [0 , 2 π ) n , where each coordinate represents the angular component. Under this construction, we can compute the kernel for each coordinate in S 1 ∼ = [0 , 2 π ) and then take the product. The eigenfunction expansion is

The heat kernel also admits a sum over paths representation. In particular, this agrees with the wrapped probability since the change of volume term is 1 :

Spheres. The sphere S n is the set { x ∈ R n +1 : ∥ x ∥ = 1 } . We use the formulas (eigenfunction and Schwinger-Dewitt) given in the main text, noting that the maximal torus value can be extracted with the equation θ = arccos( x · y ) (ie the geodesic distance).

SO ( 3 ) . SO (3) is the Lie Group { X ∈ R 3 × 3 : XX ⊤ = I, det( X ) = 1 } . We have already given the eigenfunction expansion in the main text, but the sum over paths method can be derived from the fact that the maximal torus value θ is the distance between x and y on the manifold. The sum over paths formula is given by

SU ( 3 ) . SU (3) is the (real) Lie Group { X ∈ C 3 × 3 : XX H = I, det( X ) = 1 } . The eigenfunction expansion can be derived from the character classes [16], and the sum over paths representation can be derived directly [1] as K SU (3) ( y | x, t ) SOP =

where A,B are the maximal torus values of x -1 y (these are log( λ ) /i for eigenvalues λ of x -1 y ), S = 8sin A -B 2 sin 2 A + B 2 A +2 B 2

B Experimental Details

B.1 Heat Kernel Estimates

We sample a random point (based off of the heat kernel probability) for each time step to compute. S 2 . We use 10000 eigenfunctions for the ground truth and 10 for our comparison.

S 127 . We use 50000 eigenfunctions evaluated at double precision for our ground truth and 100 for our comparison.

SO ( 3 ) . We use 10000 eigenfunctions for our ground truth and 50 for our comparison. We sum over 10 paths.

B.2 Earth Science Datasets

We do not perform a full hyperparameter search. We use a very similar architecture to the one used in Bortoli et al. [4] except we use the SiLU activation function without a learnable parameter [23] and a learning rate of 5 · 10 -4 .

B.3 2D Torus

Weuse a standard MLP with 4 hidden layers and the SiLU activation function and learn with the Adam optimizer with learning rate set to 1 e -3 [31]. However, we transform the input x → sin( kx ) , cos( kx ) where k ranges from 1 to 6 . This was done to ensure that the input respects the boundary constraints. We note that this architecture is generally quite powerful, as the Fourier coefficients can capture finer grain features, but this was optimized for the flow matching baseline, not the score matching method. In particular, score matching works with significantly fewer/no Fourier coefficients. We train for 100000 gradient updates with a batch size of 100 (each batch is randomly sampled from the checkerboard).

B.4 SU(3) Lattice

We generate our 20000 ground truth samples using Riemannian Langevin dynamics with a step size of 1 e -3 for 10000 update iterations. Our model is similar to the model used in Kingma et al. [32], except we circular pad the convnet and use 3 layers for each up-down block instead. We input a compressed version of the 3 × 3 SU ( n ) matrix, making the input size 18 . We train with a learning rate of 5 · 10 -4 and perform 1000000 updates with a batch size of 512 . To evaluate, we use an 0 . 999 exponential moving average [40] and sample using the manifold ODE sampler [37].

B.5 Contrastive Learning

We use the pretrained networks given by Ming et al. [39] to construct our hyperspherical embeddings. Our Riemannian diffusion model is similar to the simplex diffusion model given by [36], although we use 3 layers instead of 4 . We train using the Adam optimizer with a learning rate of 5 · 10 -4 , performing a 0 . 999 EMA before using the manifold ODE solver to evaluate likelihoods.

Peer reviews:
comments: SUMMARY:
This paper considers the diffusion model on Riemann space. They propose several strategies for numerically computing the heat kernel on Riemannian symmetric spaces. They also propose two techniques to sample from their proposed heat kernel approximations, including a rejection sampling method and a probability flow method.

STRENGTHS:
They compare several techniques to approximate the heat kernel in Riemannian symmetric space.

They propose two techniques to obtain samples from the approximated heat kernel.

They scale the models to the task of 128 dimension, which seems to be promising.

WEAKNESSES:
The proposed several techniques in section 3 seem quite standard, and do not involve much novelty.

From Figure 1, these several heat kernel estimators could be very sensitive, so the hyper-parameter $\tau$ in Algorithm 1 needs very careful tuning.

The experimental parts consider only few baselines. For example, in figure 3 and table 1, can you compare with more previous Riemannian diffusion model works?

QUESTIONS:
Rejection sampling is known to be non-scalable (low acceptance rate.) in high dimension. Do you have this problem when you apply Cheap Rejection Sampling Methods in section 4?

Can you compare the two techniques in section 4? Which one is more convenient to use in practice?

LIMITATIONS:
.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDeveloping diffusion model on Riemannian manifolds has received attention recently. However, the current proposals did not scale well to high dimension as they used sliced-score matching or denoising score matching with bad approximations of the heat kernel. This work focuses on using better approximations of the heat kernel on (compact) symmetric Riemannian manifolds. Then, they show the benefit of their proposals on several manifolds such as the torus, $S^2$, $S^{127}$ or $SU(n)$, and apply it to different problems such as density estimation of earth data or on the torus, or for contrastively learned hypersperical embeddings.\n\nSTRENGTHS:\nDeveloping efficient generative models on Riemannian manifolds is of much interest, and this paper proposes to use efficient approximations of the transition kernel, which is pretty appealing.\n\n- Different kernel approximations with an analysis on the setting in which they work or not\n- Very competitive results compared to previous methods, especially on high-dimensional manifolds\n- Good Applications\n\nWEAKNESSES:\n- I find the introduction of the transition kernel a bit quick and it is thus not very clear how to relate it to the equation (7,8) in the Euclidean diffusion. Maybe it would be more clear rewrite the equations with the kernel.\n- The experiment in Section 5.2 is not very clear.\n\nQUESTIONS:\nThe results of the experiment in Section 5.2 are not clear to me as it only provides an ESS of 0.62, but without any comparison or value of any other method. Thus, I do not understand what it brings.\n\nNot very related to the point of the paper, but it is stated that for non compact Riemannian manifolds, one can \"map points to surfaces and learn the diffusion on $\\mathbb{R}^d$. Is it not a bit naive? I would expect such method to be less performant than by directly applying the diffusion on the manifold.\n\n\nTypos:\n- Equation 10: lack closing parenthesis\n- Line 109: \"prior to work to\"\n- The notation \"cot\" in equation (15) is not defined.\n- Line 288: the abreviation ESS is not defined\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe theory of stochastic differential equations on vector spaces forms the basis for recent diffusion generative models, and recent efforts have attempted to extend this framework to stochastic differential equations on Riemannian manifolds (Riemannian diffusion models). From a practical computational perspective this extension is less than straightforward as pointed out in this paper, e.g., computing things like \"straight lines\" (minimal geodesics) become very difficult two-point boundary value problems; quantities that admitted closed-form analytic formulas in the vector space setting do not translate to closed-form formulas in the Riemannian setting. Past attempts at practical work-arounds and computational approximations are limited to low-dimensional cases; these methods do not scale well to higher-dimensional spaces. By restricting the focus to Riemannian manifolds that are also symmetric spaces, the authors present a collection of computation and approximation techniques that, taken together, demonstrate both computational tractability and scalability.  Experiments validating these claims are presented.\n\nSTRENGTHS:\nThe paper is technically sound and well-written, and the contributions are laid out in a clear and concise manner without claiming too much. Rather than making claims about the significance of a particular result, I appreciate that the authors claim that it is the sum of a set of smaller contributions that, taken together, lead to a computationally practical, scalable set of techniques to make Riemannian diffusion models workable in practice. The focus on Riemannian symmetric spaces seems reasonable, and the detailed computation formulas and approximations are laid out in enough detail that they can be implemented by a reader familiar with stochastic differential equations and some Riemannian geometry. The examples and experiments appear to support the claims of the authors about computational tractability and scalability to higher-dimensional problems.\n\nWEAKNESSES:\n-The following is a general comment applicable not only to this paper, but to all papers that present a “geometric” version of an existing algorithm or method in machine learning. In the past, developing a geometric, manifold version of a vector space algorithm could be regarded as a meaningful contribution in itself, and I would argue that the Riemannian diffusion model is somewhat in this spirit (equivariant models I would argue are not in this spirit however). Of course, it’s hard to argue against the claim that the Riemannian diffusion model is needed for problems in which the underlying data are manifold-valued; in that case I would be much more convinced by examples and case studies drawn from mainstream applications rather than narrow ones.  As geometric methods have become more mainstream, the threshold for what constitutes a meaningful contribution is justifiably higher: given the much more difficult computations involved in computing, e.g., derivatives, gradients, Laplacians, minimal geodesics, kernels, etc., the considerable increase in computation needs to be justified by results. For this paper, the authors recognize and point out the difficulty of computing the geometric quantities and propose more efficient approximations, which is worthwhile (but whether they deserve to be published in NeurIPS is another matter). It would be helpful if the authors can address this question more directly – are the added computational difficulties justified by a commensurate improvement in the results? The experiments do not strike me as mainstream problems, for one thing. (Note: This comment could also have been placed in the \"questions\" section, but I place it here as it could be a potential weakness).\n-The definition of symmetric space could be sharpened, as this is an important underlying assumption in the paper, e.g., clarify the isometry requirement (isometry between what spaces?), make the distinction between local vs global symmetric space, provide brief intuition (reflexive symmetry about a point), and most importantly, provide examples of spaces arising in ML applications that are symmetric (the authors list some -- a few more would be helpful, e.g., the space of symmetric positive-definite matrices, which is nonimpact but a space that arises constantly in ML) and also examples of non-symmetric spaces (e.g., hyperbolic manifolds).\n-The assumption of the Riemannian manifold being embedded in a higher-dimensional ambient Euclidean space in itself is not restrictive thanks to Whitney and Nash, and most practical problems that I've encountered usually admit some natural embedding. The compactness assumption, however, could be somewhat restrictive for certain problems in which the underlying data manifold is potentially unbounded. The authors brush aside this case by asserting that noncompact Riemannian manifolds are diffeomorphic to Euclidean space, and therefore the standard way of using the embedding in R^n to model diffusion, then mapping back to the surface, is sufficient. One could then ask why this approach doesn't work for any manifold embedded in Euclidean space; it would be helpful if the authors could clarify this point.\n\nQUESTIONS:\n-Explaining in more detail the Riemannian exponential map would be helpful, so that the reader has a better idea of what the computation entails.\n-The examples are for SU(3) and the hypersphere, which are rather specific manifolds. Particularly in the case of the hypersphere, I suspect that the heavy Riemannian diffusion machinery may not be needed to arrive at the result, since Brownian motion on the sphere is well-characterized and projections to the hypersphere are trivial.\n\nLIMITATIONS:\nNo potential negative societal impact beyond that of other typical NeurIPS submissions that I can detect.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work the authors leverage properties of symmetric spaces to get better approximation of the heat kernel and therefore of the denoising score matching loss for training diffusion models on compact manifolds.\nIn particular, using the concept of maximal torus and roots of the symmetric space, the Laplacian can be written in terms of the radial Laplacian which itself depends on these roots and their multiplicity.\nMost interestingly, for (compact) Lie groups, the heat kernel can also be written as a summation over all possible paths between the origin and evaluation points, similarly to the heat kernel expression on the torus.\nAdditionally, for small time scale, the authors propose to use the Schwinger-Dewitt approximation of the heat kernel which can be seen (I think) as a higher order Taylor approximation than the Varadhan one, and therefore is accurate for much higher values of $t$ while staying computationally fast.\nThe authors then show that these approximations are crucial for training diffusion models on 'high' dimensional manifolds, but still improves performance on smaller ones such as $\\mathcal{S}^2$ and $\\mathrm{SU}(3)$.\n\nSTRENGTHS:\n- I enjoyed reading this paper, as it is well organised, sounds and generally well written.\n\n- I was not aware of the result from Equation 21 'Sum Over Paths', and have been wondering in the past why does the approximation from a truncation of Eq 13 was getting worse as $t$ gets smaller, while it is the opposite for the heat kernel of the torus. This is really interesting, and would like to know more about the generalisation mentioned at line 219.\nFigure 1 is quite illustrative in what the Schwinger-Dewitt approximation brings compared the Varadhan!\n\n- The exact sampling from the (marginal) Brownian motion is also neat.\n\nWEAKNESSES:\n- Section 3.1: Would be useful to give some intuition to the Definition 3.1 of symmetric manifolds. Generally, the concept of maximal torus and Equation 15 deserve a better presentation. The 'positive roots of our symmetric space' are not defined. Would be useful to say how the 'radial' Laplacian relates to the Laplacian of this symmetric space, and similarly for its heat kernel to the heat kernel of the symmetric space.\n- Table 1: Would be worth including [22] as they overperform [3] on several datasets.\n- Flow matching: Why not including the results from Flow Matching [9]? They seem to report better results on the Volcano and Fire datasets. Why does flow matching fails on the chequerboard on the torus but not on the sphere? Isn't the vector field also discontinuous there?\n\nQUESTIONS:\n- Equation 11: As stated in the paragraph below this noising process is indeed for compact manifolds, yet [3] more generally proposed an 'Ornstein-Uhlenbeck' like process for non-compact manifolds, which simplifies to Eq 11 when compact. Would be worth having this more general setting in Section 2.2?\n- line 105: Perhaps slightly misleading to say that the heat kernel has 'has no closed form' since it is given by the Sturm–Liouville decomposition of Equation 13, although it is a converging series, and has such not 'computable' (since there is generally no finite time algorithm).\n- Line 181: It is indeed good to stress that one can bypass directly computing the eigenfunctions, although I don't think that in the past people have done this (naive approach)? E.g. equation 18 can be found in appendix F of [3].\n- Section 3.2.2: Worth linking this with the $\\exp(-\\lambda_i t)$ term in Eq. 13, so for small values of $t$, terms in the series would tend to contribute more and more equally.\n- Eq 20: What is $v$? Where does this come from?\n- Figure 1: Would you know why the Schwinger-Dewitt approximation has this very fast change in accuracy?\n- 'RFM introduces several geodesic-based vector fields, but these break the smoothness assumption and the theoretical framework'. Could you please expand on that? This is using which (conditional) vector field?\n- Section 3.3: This is a really neat idea. Can the ratio be trivially bounded? Worth adding this in the paper! Does this help during training (for the Monte Carlo estimator of the denoising score matching loss)?\n\nLIMITATIONS:\n- This paper 'do not consider the case of noncompact Riemannian symmetric spaces, as these are diffeomorphic to Euclidean space'. I agree that since the topology is Euclidean it is perhaps less crucial to parameterise a probabilistic model directly on that space, yet the geometry is still different and only locally linear, so the geometry is more and more distorted further away, e.g. at boundary for the Poincare model of hyperbolic geometry. Is there some equivalent notion of the maximal Torus for non compact manifolds?\n- Would be worth expanding the evaluations in Section 5.2 which is a really interesting setting.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}