SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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
Hyperparameters: Riemannian symmetric space  $\mathcal{M}$ , number of eigenfunctions  $n_e$ , time value cutoff  $\tau$ , (optional, depending on if  $\mathcal{M}$  is a Lie Group) number of paths  $n_p$   
Input: source  $x_0$ , time  $t$ , query value  $x$   
Compute  
if  $t < \tau$  then  
if  $\mathcal{M}$  is a Lie Group then  
| return  $K_{\mathcal{M}}^{\mathrm{SOP}}(x|x_0,t)$  truncated to  $|n| < n_p$  in the summation over  $\Gamma$ .  
else  
| return  $K_{\mathcal{M}}^{\mathrm{SD}}(x|x_0,t)$ .  
end  
else  
return  $K_{\mathcal{M}}^{\mathrm{EF}}(x|x_0,t)$  truncated to  $|n| < n_e$ .  
end  
remark  $\nabla_x \log K_{\mathcal{M}}$  can be computed with autodifferentiation.
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_{\mathcal{M}}^{\mathrm{EF}}$  NaNs out before Varadhan becomes accurate.
(a)  $S^2$
(b)  $S^{127}$
(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  $S^{127}$ , 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 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  $\mathcal{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.
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^1$ . (a) The flow matching path has a 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 a smooth density and smooth vectors. (d) At the pole, our score matching path anneals to 0 to maintain continuity.
(c) Diffusion SM path.
(d) Closeup of Figure 2c
and has a density
$$
p _ {\text {w r a p}} (x \mid x _ {0}, t) = \frac {1}{(4 \pi t) ^ {\dim \mathcal {M} / 2}} \sum_ {2 \pi n \in \Gamma} \prod_ {\alpha \in R ^ {+}} \left(\frac {\sin (\alpha \cdot h)}{\alpha \cdot (h + 2 \pi n)}\right) e ^ {- \frac {(h + 2 \pi n) ^ {2}}{4 t}} \tag {24}
$$
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} \sim \mathcal{U}_{\mathcal{M}}$ , where  $T$  is large enough s.t.  $K_{\mathcal{M}}(\cdot |x_0,T)$  is close to the uniform distribution (within numerical precision). We then solve the ODE  $\frac{d}{ds} x_T = -\frac{1}{2}\nabla_x\log K_{\mathcal{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_{\mathcal{M}}(\cdot |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.
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.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
Method (S2)	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.28	0.13 ± 0.03	0.73 ± 0.04	-0.60 ± 0.18
Denoising Score Matching (accurate)	-4.69 ± 0.29	-0.27 ± 0.05	0.44 ± 0.03	-1.51 ± 0.13
Table 1: We measure the improvement of our improved heat kernel estimator on downstream climate science tasks and report negative log likelihood  $(\downarrow)$ . Without our accurate heat kernel estimator, the denoising score matching loss produces substandard results.
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 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\times 4}$ . We consider the data density defined by the Wilson Action [53]:
$$
p (K) \propto e ^ {- S (K)} \quad S (K) = - \frac {\beta}{3} \sum_ {x, y \in \mathbb {Z} _ {4} ^ {2}} \operatorname {R e} \operatorname {t r} \left(K _ {x, y} K _ {x + 1, y} K _ {x + 1, y + 1} ^ {*} K _ {x, y + 1} ^ {*}\right) \tag {25}
$$
For our experiments, we take  $\beta = 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
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  $(\downarrow)$  for 0.05 false negative rate. The hyperspherical embeddings produce very bad results, but with a Riemannian Diffusion Model, it is competitive with or 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 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.
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 to the broader adoption of Riemannian generative modeling techniques.
This project was supported by NSF (#1651565), ARO (W911NF-21-1-0125), ONR (N00014-23-1-2159), and CZ Biohub. AL is supported by a NSF Graduate Research Fellowship and MX is supported by a Stanford Graduate Fellowship.
[1] Belal E. Baaquie. Path-integral derivation of the heat kernel on  $\mathrm{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.
[2] Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari S. Marcos, 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.
[3] 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.
[4] 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.
[5] 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.
[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  $\mathrm{su}(\mathfrak{n})$  gauge equivariant flows. ArXiv, abs/2008.05456, 2020.
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$ :
$$
K _ {S ^ {n}} (y | x, t) = e ^ {- (n - 1) ^ {2} t / 2} \frac {1}{2 \pi} \left(- \frac {1}{\sin \theta} \frac {\partial}{\partial \theta}\right) ^ {\frac {n - 1}{2}} K _ {S ^ {1}} (y | x, t) \tag {26}
$$
where  $\theta = \arccos (x\cdot 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
$$
K _ {G} (y | x, t) = \frac {e ^ {t \rho^ {2}}}{(4 \pi t) ^ {\dim \mathcal {M} / 2}} \prod_ {\alpha \in R ^ {+}} \frac {\alpha \cdot (h + 2 \pi n)}{2 \sinh (\alpha \cdot h / 2)} e ^ {- \frac {h ^ {2}}{4 t}} \tag {27}
$$
and for hyperbolic space, one has the formula for odd dimensions (equivalent to the one for odd dimensional spheres):
$$
K _ {H ^ {n}} (y | x, t) = \left(- \frac {1}{2 \pi}\right) ^ {\frac {n - 1}{2}} \frac {1}{2 \sqrt {\pi t}} \left(\frac {1}{\sinh \rho} \frac {\partial}{\partial \rho}\right) ^ {\frac {n - 1}{2}} e ^ {- \frac {(n - 1) ^ {2}}{4} t - \frac {\rho^ {2}}{4 t}} \tag {28}
$$
where  $\rho$  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 \cong (S^1)^n$  can be realized as a flat torus  $[0,2\pi)^n$ , where each coordinate represents the angular component. Under this construction, we can compute the kernel for each coordinate in  $S^1 \cong [0,2\pi)$  and then take the product. The eigenfunction expansion is
$$
K _ {S ^ {1}} (y | x, t) = \frac {1}{\pi} \left(\frac {1}{2} + \sum_ {k = 1} ^ {\infty} e ^ {- k ^ {2} t} (\cos (k x) \cos (k y) + \sin (k x) \sin (k y)\right) \tag {29}
$$
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:
$$
K _ {S ^ {1}} (y | x, t) = \frac {1}{\sqrt {4 \pi t}} \sum_ {k = - \infty} ^ {\infty} e ^ {- \frac {(y - x + 2 \pi k) ^ {2}}{4 t}} \tag {30}
$$
Spheres. The sphere  $S^n$  is the set  $\{x \in \mathbb{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  $\theta = \arccos(x \cdot y)$  (ie the geodesic distance).
SO(3).  $SO(3)$  is the Lie Group  $\{X \in \mathbb{R}^{3 \times 3} : XX^{\top} = 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  $\theta$  is the distance between  $x$  and  $y$  on the manifold. The sum over paths formula is given by
$$
K _ {S O (3)} ^ {\mathrm {S O P}} (y | x, t) = \frac {e ^ {\frac {t}{1 6}}}{(2 \pi t) ^ {\frac {3}{2}}} \sum_ {n = - \infty} ^ {\infty} \frac {\theta + 2 \pi n}{\sin (\theta)} e ^ {- (\theta + 2 \pi n) / t} \tag {31}
$$
$\mathbf{SU}(3)$ .  $SU(3)$  is the (real) Lie Group  $\{X \in \mathbb{C}^{3 \times 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)^{\mathrm{SOP}} =$
$$
\frac {1}{S} \sum_ {l = - \infty} ^ {+ \infty} \sum_ {m = - \infty} ^ {\infty} (A (l) - B (m)) (A (l) + 2 B (m)) (2 A (l) + B (m)) e ^ {- \frac {A (l) ^ {2} + B (m) ^ {2} + A (l) B (m)}{t}} \tag {32}
$$
where  $A, B$  are the maximal torus values of  $x^{-1}y$  (these are  $\log (\lambda) / i$  for eigenvalues  $\lambda$  of  $x^{-1}y$ ),  $S = 8\sin{\frac{A - B}{2}}\sin{\frac{2A + B}{2}}\frac{A + 2B}{2}$
B.1 Heat Kernel Estimates
We sample a random point (based off of the heat kernel probability) for each time step to compute.
$\mathbf{S}^2$ . We use 10000 eigenfunctions for the ground truth and 10 for our comparison.
$\mathbf{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 \cdot 10^{-4}$ .
B.3 2D Torus
We use a standard MLP with 4 hidden layers and the SiLU activation function and learn with the Adam optimizer with learning rate set to  $1e - 3$  [31]. However, we transform the input  $x\rightarrow \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  $1e - 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 \times 3 SU(n)$  matrix, making the input size 18. We train with a learning rate of  $5 \cdot 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 \cdot 10^{-4}$ , performing a 0.999 EMA before using the manifold ODE solver to evaluate likelihoods.