Scaling Riemannian Diffusion Models
Authors: Aaron Lou
Minkai Xu
Adam Farris
Stefano Ermon
Venue: NeurIPS 2023 poster (2023)
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
Aaron Lou, Minkai Xu, Adam Farris, Stefano Ermon Stanford University
faaronlou, minkai, adfarris, ermon)@stanfordd. edu
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 high dimensions. In this work, we reexamine these approximations and pro- pose 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, quantities to high precision. On low dimensional datasets, our correction produces noticeable improvement, allowing diffusion to compete with other methods. Ad- ditionally, we show that our method enables us to scale to high dimensional tasks nontrivial manifolds. In particular, model QCD densities SU(n) lattices and contrastively learned embeddings on high dimensional hyperspheres.
quickly compute relevant
Introduction
By learning to faithfully capture high-dimensional probability distributions, modern deep generative models have transformed countless fields such as computer vision [99 and natural language pro- cessing [41]. However, these models are built primarily for geometrically simple data spaces, such as Euclidean space for images and discrete space for text. many applications such as protein structure prediction K contrastive learning [12. and high energy physics [6] the support of the data distribution is instead a Riemannian manifold such as sphere or torus. Here, naïvely applying standard generative model on the ambient space results in poor performance it doesn't 't properly incorporate the geometric inductive bias and can suffer from singularities (7].
As such, longstanding goal within Geometric Deep Learning has been the development of prin- cipled, general, and scalable generative models on manifolds [ [8] One promising method is the Riemannian Diffusion Model 4
the natural generalization of standard Euclidean space 48 These learn reverse diffusion process on a manifold-in in Riemannian score matching methods. While this approach is
25
score-based diffusion models 46
particular, the heat equation-throuu principled and general, it not scalable. In partieular, 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 -iwddimnsiooal examples.
In our work] we propose several improvements Riemannian Diffusion Models stabilize their performance and enable scaling to high dimensions. In particular, we reexamine the heat kemel [2]] which is the core building block for the denoising score matching objective. To enable denoising
Code found at wttp:://wttt.ccoooaooaoi//iieeed
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 kemnel is Gaussian distribution, but doing this is effectively intractable for general manifolds. By restricting our analysis to Riemannian symmetric spaces 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 mask-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 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 partieular, 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.
show how our exact heat kernel computation enables fast simulation
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 SU) 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 IR are defined through stochastic differential equations 24 45 48 Given
an initial data distribution po on Rdd samples Xo RRd are perturbed with stochastic differential
equation 33
dx; = f(x, t) + (t)
where f and are fixed drift and diffusion coefficients, respectively. The time varying distributions Pe (defined by xx) evolves according the Pokker-Planck Equation
0t Px (x) = - (x) t)) + (()) (xp(
2
and approaches a limiting distribution ग z pr. which is normally a simple distribution like Gaussian (O T) through carefully chosen f and g. Our SDE has a corresponding reversed SDE
dx, (f(x,, t) (())
(t)
(3)
which maps pr back to Pp- Motivated by this correspondence, diffusion models approximate the score function V. log pt (x) using a neural network so(x, 1). To do this, one minimizes the score matching loss 277, which is weighted by constants X,:
Exx,, de [xs, t) - V log
Since this loss is intractable due to the unknown V, log Pa (x,), we instead use an alternative form of the loss. One such loss is the implicit score matching loss 27]:
Exxxropi de div(s))(x,,,)
(x, 1)
(5)
which normally estimated using sliced score Hatching/Hutchinoon" 's trace estimator 26 47
Ex,exe de e (x, t)e + 2 ||5e( t))
(6)

---

where e 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 [oss[5]
E Pxxx A ||so xx, () - Vx (xx(x)))
where Pr (x, (x)) is derived from the SDE in Equation ] and is normally tractable. Once Sa (X,, t) learned, we can construct a generative model by first sampling XT 2 pr and solving the generative SDE from / = T to / 0:
dx, = (f(x,, t) - 92 (t)ss (x,, t)) + gt)
(8)
Furthermore, there exists a corresponding "probability flow ODE" 48
dx; = (f(x,, t) g(t)²
V log (xx))dt
that has the same evolution of P; the SDE in Equation This can be approximated using our score network Sy to get Neural ODE
g(t)²
dxt
(f(x;, 1)
(x t))dt
(10)
2
which can be used to evaluate exact likelihoods the data 20
2.2 Riemannian Diffussin 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 V, and Brownian motion B, 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:
dx; g())
(11)
Unlike in the Euclidean case, pr approaches the uniform distribution UM convergence is fast, getting within numerical precision for / Rs 5).
(in practice, this
The transition density has closed form. Despite the fact that work with the most simple SDE, the transition kernel defined by the manifold heat equation pt (xt (20) has no closed form. This transition kernel is known the heat kernel, which satisfies Equation [ with the additional condition that, as 1 - O, the kernel approaches 6zp- We will denote this by (2, (zo, t), and we highlight that, when M IR4, 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 KMu( (-12o. t), one must simulate a Geodesic Random Walk with the Riemannian exponential map exp:
(12) Additionally, to calculate KM (MM-l2o.F t) or V log KM (z [0 t). one must use cigenfuncions fa which satisfy Âf: fA.. These fa form an orthonormal basis for all L function on M, allowing us to write the heat kernel with an infinite sum:
T++i exp,, At2)
N(O. E T.M
K K (z/20. 1) = Σ e -d (((o))))
(13)
Additionally, previous work has also explored the use of the Varadhan approximation for small values of (which uses the Riemannian logarithmic map log) 43):
K ( t) NO, (2t))ddstt.
log K KM
zo. (tot)) t)
log (x0)
(14)
2t
3

---

3 Method
The key problem with applying denoising score matching in practice is that the heat kernel computa- tion is too expensive and inaccurate. Notably, simulating 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 - O. Worse still, these eigenfunctions formulas are not well-known for most manifolds, and, even when explicit formulas exist, they be numerically unstable (like in the case of S*)). One possible way to alleviate this is to use Varadhan's ' approximation for small f, 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, we do not lose applicability by restricting our attention here. Furthermore, for surfaces (which have appeared as generative modeling test tasks in the literature [2),, one can always define generative model by mapping the data points to 5², learning a generative model there, and mapping back [4]. We empathize that, outside of these two examples, 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 diffcomorphic 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 provides a general intuition with which to build our calculations. We urge interested readers to consult book [22 monograph [9] for a full treatment of the subject.
Definition 3.1. A Riemannian Symmetric Space is Riemannian manifold such that, for all points M, there exists a local isometry s.t. s(x) and D -idT.M.
Intuitivelyy this condition means that, 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 conse- quence: 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/te) (where {e} is the trivial Lie Group), the sphere S S0 (n)) and hyperbolic space H = 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 a compact, connected. and abelian subgroup of G that is not contained in any other such subgroup. For symmetric spaces G/ G, the maximal torus is defined by quotienting oul the maximal torus of G. Importantly, the maximal torus is isomorphic to standard torus T (51)).
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 Sa maximal torus is any great circle. For the Lie group of unitary matrix (n). the maximal torus is the set of all diagonal matrices {diag(e')) "...,,,,, "
0₂ E [0, 2m)}.
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 Maximal Torus). We define root be a vector (really covector) that corresponds with a character X of s.t.
-))
(15)
We define the set R+ to be the set of all positive roots. The multiplicity ma of a root a is the dimension of the eigenspace over the Lie algebra B
{Z E B [H VH)
(16)

---

Examples. For all spheres dimension n. there is only one root of multiplicity n n2-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 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-Beltne Beltrami operalor on M (the manifold generalization of the standard Laplacian) induces the "rudial" Laplacian on T:
8
=
Σ ma cot(a h)) 8u aERt
(17)
where Dr is the standard Laplacian on the torus. Here, h is the "flat" coordinate of on the maximal torus (e.g. for sphere this is the angle between and zo).
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 dimension for spheres and (n²) dimensions to (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 cigenfunction expansion of the induced Laplacian on the maximal torus. This has implicitly appeared in previous work when defining Riemannian Diffusion Models on S² and SO(3) 4 34 allowing one to rewrite the summation respectively
K (o/ t) = 4T
_ee (21 + 1)))((() (1++1) 1=0
(18)
where P/ are the Legendre Polynomials and
2)( sin (21 + 1)0/2
K- (50(3) (x|zo,
0 is the angle between æ, z0
(19)
ST
sin(0/2)
t=
By making this relationship explicit, we can directly generalize this formulation to other symmetric spaces, e.g. SU(3)
pq(p + 4) (x (00 x.P ())e
160 ug 2
EF K
(U))) (x) t)
(20)
p-p=1
where X are the (real) irreducible representations and 0, are the induced angles between x and zo. By reducing the dimensionality, this strategy generally makes the summation more tractable. Further- more, for cases such as the hypersphere 5"h the eigenfunctions 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 f the heat kernel scaling terms e Aut remain large as one increases i, A; + As result, must evaluate many terms (often in the thousands), and the result can be very numerically unstable even with double precision.
To address this problem, examine several refined versions of the Varadhan approximation that use the fact that the manifold symmetric. These approximations allow us control the number of eigenfunctions required and, in some cases, completely obviate the need for them entirely.
5

---

The Schwinger-Dewitt Approximation
The Varadhan approximation is constructed by approximating the heat kernel with Gaussian distribution (with respect to the Riemannian distance function). However, doing this does not account for the curvature of the manifold. Incorporating this premise, get the Schwinger-Dewe Dewitt approximation [
Auo (1))1/2 "Melmgue +
(21)
(4Tl)dim M1/2
det(Dao expzo (logz, (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 s approximation as it better accounts for the curvature of the manifold through the A term, retaining accuracy up to moderate time values.
Here, =
appears to be rather computationally demanding term since it requires evaluating the determinant of Jacobian. Indeed, this naively takes O(dim M evaluations which is completely inaccessible in higher dimensions. However, we again emphasize the fact that we are working with symmetric spaces: A has a particularly simple formula defined by our flat coordinate h from above:
ma
x
II oERt
a
(22)
sin h)
Sum Over Paths
The fact that we can derive a better approximation using a different power of A points to deeper connections between the heat kernel and the "Gaussian" with respect to distance. We draw inspiration from several Fuclidean case examples, such as the flat torus (28 or the unit interval (36. For those cases, the heat kernel is derived by summing Gaussian over all possible paths connecting Zo and a. While this formula does not exactly lift over to Riemannian symmetric spaces, there exists an analogue for Lie Groups [9:
KSOP (xr 1) (4m M/2
e0p
Σ
II
a' (h + 2mT) 2sin(a h/2)
(k+ 2 4t
(23)
2aneT aeRt
Here, is infinite tangent space grid that corresponds with loops of the torus (i.c. 2mn circular intervals), and p² is manifold specific constant. We note that the product over R+ is exactly the A change 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 an exact (albeit infinite) representation for the heat kernel. Compared the eigenfunction expansion Equation [3] the Sum Over Paths represen- tation is accurate for small i, which nicely complements the fact that the eigenfunction representation accurate for large This formula does generalize to split-rank Riemannian symmetric spaces like odd dimensional spheres (more details in Appendix A. and noncompact spaces (Appendix A.2
3.2.3 A Unified Heat Kernel Estimator
We unify these approximations into single heat kernel estimator. Our computation method splits up the heat kernel evaluation based on the value and applies an eigenfunction summation or an improved small time approximation accordingly. This allows to effectively control the errors at both the small and large time steps while significantly reducing the number of function evaluations for cach. Our full algorithm is outlined in Algorithm

---

Algorithm 1: Heat Kernel Computation
Hyperparameters: Riemannian symmetric space M, number of eigenfunctions Ne, time value cutoff т, (optional, depending on if M is Lie Group) number of paths nip
Input: source 20, time i, query value Compute
it then
iM is Lie Group then
return KSOP (a/ao, t) truncated to |na] < Up in the summation over
else
return RM (x)
end
else
return KM (1r/0o, () truncated to |m] <
end
remark V x log KM can be computed with autodifferentiation.
We ablate the accuracy of our various heat kernel (score) approximations in Figure for S 1 and (3)3). In general, we found that computing with standard eigenfunctions was too costly and too prone numerical blowup, and Varadhan's 's approximation was simply too inaccurate. For some spaces like S127, combination of preexisting methods (like is briefly explored in [4] would not work since KEF NaNs out before Varadhan becomes accurate.
Ligenfunction
30"
Eipenfunction
Ligenfunction
Varadhan
Varachan
Varadhan
Schwinger-Dewitt
Schwvinger-Dwitt
35"
Sum Over Paths
30
30
32
10"
18"
13
100
66-
13
30
16
(a) 2
(b) S127
(c) SO(3)
Figure 1: We compare the various heat kernel estimators on a variety of manifolds. We plot the relative error compared with t. Our improved small time asymptotics allow us to control the numerical error, while baseline Varadhan is insufficient. For 5127, Varadhan will still produce an error of 10% when the eigenfunction expansion NaNs out, so we need to use Schwinger-Dewitt
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 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 have a closed form density, easy to sample from, and must not deviate from the heat kernel too much. For large time steps the natural prior distribution is uniform. Conversely, for small time steps, we instead use the wrapped Gaussian distribution, which can be sampled by passing tangent space Gaussian through the exponential map
7

---

(a) O FM Path.
(b) f 0.5 FM Path.
(c) Diffusion SM path.
(d) Closeup of Figure 2c
Figure 2: We visualize the vector fields generated by the flow matching geodesic path and our score matching diffusion path. These are done on S (a) The flow matching path has discontinuity at the pole. (b) The marginal densities of the flow matching path are not smooth and transition sharply at the boundary. (c) Our score matching path has smooth density and smooth vectors. (d) At the pole, our score matching path anneals to 0 to maintain continuity.
and has density
(a sin(a- - (k+2))))
Pwrap (o, (4Tt)dim
Σ II
(a- (h (h+ 2mm)
(24)
2nnET aeR-
In both cases, the ratio can be trivially bounded by examining their behavior numerically.
Heat Kernel ODE Sampling
As alternative, notice that can apply the probability flow ODE sample from the heat kernel exactly. In particular, we draw a sample Fr ~ UM, where is large enough s.t. ( (- T) close to the uniform distribution (within numerical precision). We then solve the ODE de 2T - V. log K (M ( s) from s = T to t. As the sampling procedure follows the probability flow ODE, this guaranteed produce samples from (- t).
This is solvable as a manifold ODE, as previous works have already developed adaptive manifold restrict vector field to the maximal
ODE solvers (37]. Furthermore, by Proposition] 3.4 torus and solve there. Note that this allows us to use preexisting Euclidean space solvers since the torus is effectively Euclidean space. Lastly, scheme like a variance-exploding schedule) to stabilize the numerical values.
scale the time schedule of the ODE (with
Related Work
Our work exists in the established framework of differential equation- based Riemannian generative models. Early methods generalized Neural ODEs to manifolds [5] 37 (38, enabling training with maximal likelihood. More recent methods attempt to remove the simulation components [ 42) but this results in unscalable or biased objectives. We instead work with diffusion models, which are based scores and SDEs and do not have any of these issues. In particular, aim to resolve the main gap that prevents Riemannian Diffusion Models from scaling to high dimensions.
Riemannian Flow Matching (RFM) []0 is 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 data point. RFM introduces several geodesic-base these break the smoothness assumption (see Figure 2). which found to be detrimental for some of the densitics we considered (c.g. Figure 33). However, similar to the Euclidean case, RFM with the diffusion path corresponds to matching with the probability flow ODE, work provides the numerical computation necessary to construct such flow.
vector fields, but
Experiments
5.1 Simple Test Tasks
We start by comparing our accurate denoising score matching objective with the inaccurate approxi- mation scheme suggested by [4] based on 50 eigenfunctions. We test the compiled Earth science

---

Method (5²)
Volcano
Earthquake
Flood
Fire
Sliced Score Matching
4.92 0.25
-0.19 0.07
0.45 0.17
-1.33 0.06
Denoising Score Matching (inaccurate)
-1.28 0.25
0.13 # aas
0.73 * ao4
-0.60 0.18
Denoising Score Matching (accurate)
-4.69
-0.27
0.44
-1.51 0.13
Table I:
We measure the improvement of our improved heat kernel estimator downstream climate science tasks and report negative log likelihood (!). Without our accurate heat kernel estimator, the denoising score matching loss produces substandard results.
Score Matching Flow Matching
Figure 3: We compare Riemannian score matching and flow matching on increasingly complex checkerboard patterns on the torus. Flow matching learns suboptimal distributions with noticeable artifacts (like blurriness and spurious peaks) for simpler distributions and fails outright for more complex checkerboards. Conversely, Riemannian score matching learns accurate densities.
datasets from 38 detailing results are in Table Generally, our accurate heat kernel results in substantial improvement and matches sliced score matching. Note that we do not expect our method outperform sliced score matching since these datasets low dimensional and the objectives equivalent (up to small additional variance).
We also compare directly with RFM series of increasingly complex checkerboard datasets the flat torus. These datasets have appeared in prior work to measure model quality B 37]. As result of the non-smooth vector field dynamies, find that RFM degrades in performance the checkerboard increases in complexity, and is unable to learn past certain point. Our visualized results given in Figure
5.2 Learning the Wilson Action on SU(3) Lattices.
apply our method learn densities on (3) We consider the data density defined by Wilson Action
[33
8 Σ Re tr(KK, K 1 K K
p(K) S(K) S(K)
(25)
For our experiments, we take 3 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 not based on score matching(6 [50]. However, to showcase our score matching techniques,
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 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 Hy Hyperspherical Embeddings
Finally, we examine contrastively learning 12] Standard contrastive losses optimize embeddings of the data that lie the hypersphere. Paradoxically, these embeddings unsuitable for most

---

Method
SVHN
Places365
LSUN
iSUN
Texture
SSD+ 44
31.19
77.74
79.39
80,85
66.63
KNN+ 49
39.23
80.74
48.99
74.99
57.15
CIDER (penultimate layer)
23.09
79.63
16.16
71.68
43.87
CIDER (hypersphere)
93.53
89.92
93.68
89.92
92.41
CIDER (hypersphere) + Diffusion
28.95
76.54
35.78
74.17
62.87
Table 2:
We compare contrastive learning OOD detection methods on CIFAR-100. We report false positive rates (4) for 0.05 false negative rate. The hyperspherical embeddings produce very bad results, but with Riemannian Diffusion Model, it is competitive with surpass the state of the art.
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 fcature-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 the properties of the hyperspherical embedding, the fact that our Riemannian Diffusion Models can extract comparable representation can lead to more principled improvements for future work.
Conclusion
We have 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 the broader adoption of Riemannian generative modeling techniques.
7 Acknowledgements
This project was supported by NSF (#1651565), ARO (W911NF-21-1-0122)) ONR (N00014-23- 1-2159), and CZ Biohub. AL is supported by NSF Graduate Research Fellowship and MX is supported by Stanford Graduate Fellowship.
References
[1] Belal E. Baaquie. Path-integral derivation the heat kernel on (U))) space. Phys. Rev. 12:1007-1010. Aug 1985. doi: (0.1103/hyy.eev.332... 1007. URL http::tt::///ink aps http://boo//o.1110//yeeeeee/
32. 1007
[2] 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. cookbook self-supervised learning. ArXiv, abs/2304. 122110 12210, 2023.
[3] Heli nen-Hamu, Samuel Cohen, Jocy 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.
[4] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson. James Thomton, Yee Whye Teh, and A. Doucet. Riemannian score-based generative modeling. ArXiv, abs/2202.00766: 2022.
A. Bose, Ariella Smofsky, Renjie Liao, Panangaden, and William L. Hamilton. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning, 2020.
[6] 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.05 005466 2020.
10

---

[7] Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estimation. ArXiv, abs/2003.13911 2020.
[8] Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur D. Szlam, and Pierre Vandergheynst. Geometric deep learning: Going heyond euclidean data. IEEE Signal Processing Magazine, 34:18-42 2016.
[9] Roberto Camporesi. Harmonic analysis and propagators homogeneous spaces. Physicx Reports, 196: 1-134, 1990.
[10] Ricky T. Q. Chen and Yaron Lipman. Riemannian flow matching on general geometries. ArXiv, abs/2302.03660, 2023.
[11] Tian Chen, Yulia Rubanova, Jesse Bettencourt, and David Kristjanson Duvenaud. Neural ordinary differential equations. Neural Information Processing Systems, 2018.
[12] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. simple framework for contrastive learning of visual representations. ArXiv, 0bb/2002.055003 2020.
[13] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. ArXiv, abs/2210.01776, 2022.
[14] Victor D. Dorobantu, Charlotte Borcherds, and Yisong Yue. Conformal generative modeling on triangulated surfaces. ArXiv, abs/2303. 10251, 2023.
[15] Luca Falorsi and Patrick Forr'e. Neural ordinary differential equations on manifolds. Abs/2006.06661. 2020.
ArXiv,
[16] H. D. Fegan. The heat equation compact lie group. Transuctions of the American Mathematicul Society, 246:33993555 1978. ISSN 00029947. URL http:///t:///w http:////wwjjj.r.r.oo otggstaalee/1999999
[17] Tomas Geffner, George Papamakarios, and Andriy Mnih. Score modeling for simulation-based inference. ArXiv, 0bb/2209..1249, 2022.
[18] Mevlana Gemici, Danilo Jimenez Rezende, and Shakir Mohamed. Normalizing flows riemannian manifolds. ArXiv, abs/161 2016.
[19] Ian Goodfellow, Jean Pbudet-Abadie. Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In NIPS, 2014.
[20] Will Grathwohl, Ricky Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Kristjanson Duvenaud. Fijord: Free-form continuous dynamics for scalable reversible generative models. ArXiv, abs//8800011333 2018.
[21] Alexander Grigor'yan Spectral theory and geometry: Estimates of heat kernels on riemannian manifolds. 1999.
[22] Sigurdur Helgason. Differential geometry, lie groups, and symmetric spaces. 1978.
[23] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv: Learning, 2016.
[24] Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusiin probabilistic models. ArXiv, abs/2000.11239, 2020.
[25] Chin-Wei Huang, Milad Aghajohari, A. Bose, P. Panangaden, and Aaron C. Courville. Riemannian diffusion models. ArXiv, 00992208.07949, 2022.
[26] Michael E. 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.
[27] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. J. Mach. Learn. Res., 6:695-709, 2005.
[28] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and T. Jaakkola. Torsional diffusion for molecular conformer generation. ArXiv, abs/2206.01729, 2022.
[29] Jurgen Jost. Riemannian geometry and geometric analysis. 1995.
[30] Isay Katsman, Aaron Lou, Derek Lim, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Equivariant manifold flows. In Neural Information Processing Systems, 2021.

---

[31] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6980, 2014.
[32] Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. ArXiv, abs/2107.00630 2021.
[33] Peter E. Kloeden and Eckhard Platen. The numerical solution of stochastic differential equations. The Journal the Australian Mathematical Society. Seriey B. Applied Mathematics, 20:8 12, 1977.
[34] 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 Tope- logical Representation Learning. 2022. URL httpstt/oppeeeeeiieeee wtt.:e/.w.eeeeeoe///ee//eeeeaaaa
[35] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. ArXiv, abs/2210.02747, 2022.
[36] Aaron Lou and Stefano Ermon. Reflected diffusion models. ArXiv, abs/2304.04740, 2023.
[37] 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.
[38] Emile Mathicu and Maximilian Nickel. Riemannian continuous normalizing flows. ArXiv, abs/2006.10655, 2020.
[39] Yifei Ming. Yiyou Sun, Ousmane Dia, and Yixuan Li. How to exploit hyperspherical embeddings for out-of-distribution detection? The Eleventh International Conference on Learning Representations, aet:/foruum?d--aaPPaoWopppp
2023. URL http:tttoppeerrrrriii
[40] Boris Polyak and Anatoli B. Juditsky. Acceleration of stochastic approximation by averaging. Siam Journal on Contral and Optimization, 30:838-855, 1992.
[41] Alec Radford, Jeff Wa, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
[42] Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence-based generative modeling on manifolds. Neural Information Processing Systems, 2021.
[43] Laurent Salo Salo-Coste. The heat kernel and its estimates. 2009.
[44] Vikash Sehwag, Mung Chiang. and Prateek Mittal. Ssd: A unified framework for self- self-supervised outlier detection. ArXiv, abs/2103. 12051, 2021.
[45] Jascha Narain vised learning using nonequilibrium thermodynamics, ArXiv, abs/1503.03585 2015.
Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper-
[46] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. ArXiv, 0tss/907.006000 2019.
[47] Yang Song. Sahaj Garg. Jiaxin Shi, and Stefano Ermon. Sliced score matching: scalable approach density and score estimation. In Conference on Uncertainty in Artificial Intelligence. 2019.
[48] Yang Song. Jascha Sohl-Dickstedi 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 http:ttt/////eeree///
[49] Yiyou Sun, Yifei Ming. Xiaojin Zhu, and Yixuan Li. Out-of-distribtiion detection with deep nearest neighbors. In International Conference on Machine Learning, 2022.
[50] Francisco Vargas, Will Grathwohl, and A. Doucet. Denoising diffusion samplers, ArXiv, abs/2302.13834 2023. htt https://app...... www.mmtnticccchoooo arg/CorpuslD::2//1 257220165
[51] Pascal Vincent. connection between score matching and denoising autoencoders. Neural Computation,
23:
2011.
[52] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. ArXiv, abs/2005.10242 2020.
[53] Kenneth G. Wilson. Confinement of quarks. Physical Review D, 10:2445-2459, 1974.
12

---

A Additional Heat Kernel Computational Details
A.1 Sum Paths for Spheres
For spheres, one has the following sum over paths formula for odd-dimensional spheres S":
Kg (2)1x, e (n-1)21/2
1 a 2m sin
1
Ks (y)2,,
(26)
00
where 0 = 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, the intertwining operator appears which makes the evaluation integral, although a similar recurrent can be derived.
A.2 Extension to Noncompact Symmetric Spaces
For noncompact symmetric spaces, heat kernel evaluations 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 the following formula for noncompact lie groups
Ka(ylæ,t) (4rt)dim M/2
txp
II
a + 2
4
(27)
2 sinnh(a h/2)
oc R
and for hyperbolic space, one has the formula for odd dimensions (cquivalent to the one for odd dimensional spheres):
3
(m-1)
K (y|x,t) =
(28)
2n
2 nt
sinh p 0p.
where p the hyperbolic distance between and
Generally, there the notion of 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 z (Sty" can be realized flat torus [0, 2mr) where each coordinate represents the angular component. Under this construction, we can compute the kernel for each coordinate in S (0,) and then take the product. The eigenfunction expansion is
2
(y
2t (cos(k:)) + sin(ke) sin sin(ky)
(29)
2
The heat kernel also admits sum paths representation. In particular, this agrees with the wrapped probability since the change volume term is 1:
x
Ks (y(0,,
(30)
/4rt
A 00
Spheres. The sphere S is the set {ar Ra+1, |||||| = 1}. We use the formulas (eigenfunction and Schwinger-Dewitt) given the main text, noting that the maximal torus value can be extracted with the equation 0 arccos(æ y) (ie the geodesic distance).
13

---

(O(3). S0(3) is the Lie Group {X E R XXX I, dett(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 and on the manifold. The sum over paths formula given by
SOP K 50(3) x,)
+ 2 sin(O)
e
(042xn)/t
(31)
(2Tt)
SU(3). (3) is the (real) Lie Group {X C3x3 - 1, det((X)] 1). The eigenfunction expansion can be derived from the character classes be derived directly [1] as K5u(3) (v/a, t)SOP
[66 and the sum over paths representation can
Σ 1=-0c
Σ (A(I) +B) + B(m))
m())-mmm)22-m + A(1)
(32)
where A, B are the maximal torus values of x x- (these are log(A)/i for eigenvalues of x x = 8sin A-B sin 2 4+28
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.
2 We use 10000 eigenfunctions for the ground truth and 10 for our comparison.
122 We use 50000 eigenfunctions evaluated at double precision for our ground truth and 100 for comparison.
(3(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 full hyperparameter search. We use a very similar architecture to the one used in Bortoli et al. 4 except the SiLU activation function without learnable parameter [23] and learning rate of 10-4
B.3 2D Torus
We use standard MLP with hidden layers and the SiLU activation function and learn with the Adam optimizer with learning rate set to 16-3 [31]. However, we transform the input x - sin(kæ), cos(Azr) where l ranges from to 6. This was done 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, method. In particular, score matching works with significantly fewer/no Fourier coefficients. We train for 100000 gradient updates with batch size of 100 (each batch is randomly sampled from the checkerboard).
the score matching
B.4 SU(3) Lattice
We generate our 20000 ground truth samples using Riemannian Langevin dynamics with a step size of le - 3 for 10000 update iterations. Our model is similar to the model used in Kingma et al. [32]. except we circular pad convnet and use layers for each up-down block instead. We input compressed version of the 3 x SU (n) matrix, making the input size 18. We train with a learning rate 5 10-1 and perform 1000000 updates with a batch size of 512. evaluate, we use an 0.999 exponential moving average [40] and sample using the manifold ODE sampler [37].
14

---

B.5 Contrastive Learning
We use the pretrained networks given by Ming et [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 learning rate of 5 10 performing 0.999 EMA before using the manifold ODE solver to evaluate likelihoods.
15


Reviews:
[{"review_id": "BWOXLkVliX", "reviewer": "Reviewer_NaGf", "summary": "In 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)$.", "questions": "- 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)? ", "limitations": "- 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.\n\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed review and insightful comments. We are glad that the reviewer found the paper enjoyable to read, and we hope to address any concerns the reviewer has.\n\n* On the generalization to spheres: we will include the generalization to split-rank symmetric spaces e.g. odd dimensional spheres. Effectively, the summation in Equation (18) must be altered to include an additional function of $n$ (see [8]). The computation of this function takes $O(n)$ extra compute but can be vectorized, and it is best to think of it as an extension to the Schwinger-Dewitt approximation.\n\n## Addressing weaknesses\n\n* On section 3.1: we agree with the reviewer that our presentation of these topics was a bit rushed. We will update our final paper to better present these concepts.\n* On [22] compared to [3] in Table 1: [22] uses a very similar training method to [3], but adds in extra components such as variance minimization. For our comparison, we compare only against [3] because it is barebones (and thus has the fewest variables).\n* On Flow Matching: we don’t include RFM results in Table 1 as we are not aiming for SOTA on these datasets; rather, we want to conduct a comparison with minimal confounding variables. RFM also has the discontinuity on the sphere, but this doesn’t seem to matter since the data points have fewer modes and a smaller support. However, we have found that RFM tends to quickly degrade when the task becomes harder (through dimension or data), which hampers applications in real world examples.\n\n## To answer questions \n\n* Equation 11: From what we can see, both [3] and [22] both address the noncompact case using an OU process of the form $dx = -\\log_0(x) dt + dB_t$ (resulting in a limiting distribution of wrapped Gaussian). The issue with these processes is that their transition kernels are not known (and aren’t even heat kernels like the standard OU process). Alternatively, we could apply a noncompact heat kernel computation method, but this would require using a VESDE which is known to be suboptimal for likelihoods since the distribution does not mix.\n* Line 105: We disagree and believe that closed form is an accurate description, since the infinite summation violates the requirements for “closed form”.\n* Line 181: We definitely agree! We wanted to emphasize in the paper that, for the Riemannian symmetric spaces that we consider, this trick of collapsing eigenfunction summations is broadly generalizable, and we will update our paper to better reflect this.\n* Section 3.2.2.: We agree and will update our paper.\n* Eq 20: Apologies, this was supposed to be $h$.\n* Figure 1: We think that Schwinger-Dewitt approximation error is exponential (which accounts for the large jump), but smooths out as the heat kernel stabilizes for large time steps.\n* RFM vector field: certainly! The geodesic path conditional vector field is given by $\\exp_x(t \\log_x(y))$ where $x$ and $y$ are the noise and the data point respectively. The issue is that $\\log_x(y)$ is not defined when $y$ is the cut locus of $x$, and in particular it is highly nonsmooth near these points as well. For example, for spheres we can imagine that when $y$ is opposite $x$ then the vector field $\\log_x(\\cdot)$ is a bunch of large diverging tangent vectors when evaluated around $y$.\n* Section 3.3: The ratio can be trivially bounded for most manifolds, and we found that the main benefit for our simulation strategies was a noticeable speed increase when compared with the many SDE Riemannian exponential map simulation steps that previous methods have to use to get good results [3, 22], which matters when working with many dimensions.\n\n## Addressing Limitations\n\n* On mapping to Euclidean space for noncompact manifolds: we believe that the additional computational benefits we get from working on a simpler space should outweigh potential metric issues, especially since our neural network models are particularly powerful.\n* On analogues for noncompact manifolds: we will update our manuscript with the analogous heat kernel derivations for noncompact manifolds. The maximal torus is replaced with the maximal abelian subgroup, and we can derive a “sum over paths” representation for the heat kernel when the manifold is dual to a valid compact case. In particular, this is a singular sum (as there is only one path between any two points). This works for the SPD manifold (as it is dual to SO(n)) and odd-dimensional hyperbolic space (which is dual to odd-dimensional spheres).\n* On section 5.2: we agree, but this may be a bit out of scope for our current paper. Doing so requires a more principled architectural design as well as a significantly different training procedure [Cite1]. Our current results are there to mainly show that our method can learn the complex data on the more exotic $SU(n)$ manifold at these scales.\n\n[Cite1] https://arxiv.org/pdf/2211.01364.pdf\n"}, {"review_id": "FDW811GC1O", "reviewer": "Reviewer_1Bu5", "summary": "The 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.", "questions": "-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.", "limitations": "No potential negative societal impact beyond that of other typical NeurIPS submissions that I can detect.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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", "weaknesses": "-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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their thorough and interesting comments. We are glad that the reviewer appreciated our presentation, and we hope to address any concerns.\n\n## Addressing weaknesses:\n* On results: we emphatically agree that showing real-world applications of geometric machine learning methods is highly important, and this was a large motivation behind our work (since previous Riemannian diffusion methods do not currently scale to these real world tasks). We also agree that, while our $SU(n)$ example could be useful for lattice QCD systems down the line, it is a bit niche for the ML community. However, we argue that our contrastive embedding experiment is highly relevant to the broader ML community and relies on Riemannian geometry. These embeddings, which include systems like CLIP, are very mainstream and use the hypersphere as an embedding space (and thus is a naturally Riemannian problem) since it produces significantly better results. Much prior work has shown that these embeddings are not usable for many tasks (see [1] Section 3.2), which we argue is due to the weaker machine learning toolkit available on the sphere. Our experiment tests this phenomena on embedding-based OOD detection (which is a common application area) and shows that we can extract a meaningful representation to get nontrivial numerical results.\n* On definition of symmetric space: we agree that our definition of symmetric spaces (as well as some of the technical constructions that we use for our method) was a bit lackluster, and we will definitely update our presentation for the final version of our paper. However, we do note that hyperbolic space is symmetric.\n* On Euclidean mapping of compact spaces: For non-diffeomorphic manifolds any such mapping results in an extremely high distortion even if we could disregard the discontinuity. Previous work [3, 33] has shown that this tends to greatly hamper results. By contrast, [32] has shown that mapping noncompact Riemannian symmetric spaces to Euclidean space produces good results since the distortion is manageable. For diffusion models, the added computational tractability greatly outweighs this minor distortion which can be overcome with the power of the neural network.\n\n## Answering questions:\n* On the Riemannian exponential map: certainly! We will update our paper to include formulas for the Riemannian exponential map as well as a discussion of computational details.\n"}, {"review_id": "fKEu1Srndb", "reviewer": "Reviewer_611J", "summary": "Developing 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.", "questions": "The 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", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Developing 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", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their positive comments!\n\n* On transition kernel: we will update our final version of the paper to include a more detailed explanation for how the transition kernel is derived from the SDE. In particular, we will emphasize that it comes from considering the probability distribution starting from $x_0$ and following the SDE up to time $t$.\n* On experiment 5.2: the purpose of experiment 5.2 is to show that our model can accurately model complex data on more exotic manifolds $SU(n)$ in high dimensions. Our ESS score is competitive with previous results [5] but we are forced to use an incomparable training algorithm due to the scope of the paper. As mentioned in our response to reviewer NaGF, we believe that combining our results with other advances in diffusion modeling and a more principled architecture design is a promising avenue of further research.\n* On mapping from Euclidean space to noncompact Riemannian symmetric spaces: previous work has shown that learning a generative model on the Euclidean space and mapping back produces very good results and can even outperform learning on the manifold directly since both spaces are diffeomorphic [32]. Generally, we believe that the power of the generative model should be enough to overcome most metric distortion issues that arise from such a mapping procedure. Furthermore, since simpler spaces give much more power for computation and modeling capabilities, this is a trade we should take for our scaling concerns.\n"}, {"review_id": "mcVF4VpBj1", "reviewer": "Reviewer_VM6C", "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.", "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?\n\nCan you compare the two techniques in section 4? Which one is more convenient to use in practice? ", "limitations": ".", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "They compare several techniques to approximate the heat kernel in Riemannian symmetric space.\n\nThey propose two techniques to obtain samples from the approximated heat kernel.\n\nThey 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.\n\nFrom Figure 1, these several heat kernel estimators could be very sensitive, so the hyper-parameter $\\tau$ in Algorithm 1 needs very careful tuning.\n\nThe experimental parts consider only few baselines. For example, in figure 3 and table 1, can you compare with more previous Riemannian diffusion model works?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and hope to address any concerns the reviewer may have.\n\n* On novelty of section 3 techniques: we do not claim novelty for the heat kernel computation methods, as these have existed in the physics and applied mathematics fields for decades. Rather, as mentioned by reviewer 1Bu5, our primary contribution is a method to adapt these methods for use in the machine learning community. In particular, our method of combining the approximations is uniquely required for diffusion model training, as it needs fast, accurate, and repeated computation while physicists only need to work with approximate ansatze like the Schwinger-Dewitt approximation. In doing so, we aim to alleviate pre-existing issues in Riemannian diffusion model training, which is an important problem.\n* On tuning $\\tau$: we agree that tuning $\\tau$ is a sensitive task. However, we emphasize that this needs to be done **once** per manifold. Additionally, we will release our heat kernel computation code (and thus our tuned heat kernel values), which will make adoption straightforward.\n* On baselines: for table 1, we compare with [3] only as we want to study the effects of our heat kernel computation on training. For figure 3, we are comparing directly with [22] to show that continuity in the flow path matters. We are not aware of any other Riemannian diffusion works which we can compare against, as the underlying machinery should effectively be the same as [3].\n* On rejection sampling: for our tasks (and these types of tasks in general), the maximal torus dimension is rather low, which makes rejection sampling quite tractable. So, we did not have problems applying it to many of our use cases, although $S^{127}$ has some other issues despite having a one dimensional maximal torus.\n* On comparing techniques: rejection sampling is more convenient and quick to use in practice, while ODE sampling is more cumbersome and slow but works on places where rejection sampling fails (such as $S^{127}$ and in theory manifolds with higher dimensional maximal tori).\n"}]

Author rebuttal:

Meta-review: This paper presents several practical improvements for scaling up Riemannian Diffusion Models on Riemannian symmetric spaces, leading in particular to an improved heat kernel estimator. The proposed method is accompanied by numerical experiments.

Reviewers generally agree that the paper is technically sound, well-written, and of practical interest.

The authors should take the comments from the reviewers into account to improve the presentation of the manuscript.
Decision: Accept (poster)