SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aaron Lou, Minkai Xu, Adam Farris, Stefano Ermon
{aaronlou, minkai, adfarris, ermon}@stanford.edu
Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on  $SU(n)$  lattices and contrastively learned embeddings on high dimensional hyperspheres.
By learning to faithfully capture high-dimensional probability distributions, modern deep generative models have transformed countless fields such as computer vision [19] and natural language processing [41]. However, these models are built primarily for geometrically simple data spaces, such as Euclidean space for images and discrete space for text. For many applications such as protein structure prediction [13], contrastive learning [12], and high energy physics [6], the support of the data distribution is instead a Riemannian manifold such as the sphere or torus. Here, naively applying a standard generative model on the ambient space results in poor performance as it doesn't properly incorporate the geometric inductive bias and can suffer from singularities [7].
As such, a longstanding goal within Geometric Deep Learning has been the development of principled, general, and scalable generative models on manifolds [8, 18]. One promising method is the Riemannian Diffusion Model [4, 25], the natural generalization of standard Euclidean space score-based diffusion models [46, 48]. These learn to reverse a diffusion process on a manifold-in particular, the heat equation-through Riemannian score matching methods. While this approach is principled and general, it is not scalable. In particular, the additional geometric complexity renders the denoising score matching loss intractable. Because of this, previous work resorts to inaccurate approximations or sliced score matching [47], but these degrade performance and can't be scaled to high dimensions. We emphasize that this fundamental problem causes Riemannian Diffusion Models to fail for even trivial distributions on high dimensional manifolds, which limits their applicability to relatively simple low-dimensional examples.
In our work, we propose several improvements to Riemannian Diffusion Models to stabilize their performance and enable scaling to high dimensions. In particular, we reexamine the heat kernel [21], which is the core building block for the denoising score matching objective. To enable denoising
Code found at https://github.com/louaaron/Scaling-Riemannian-Diffusion
score matching, one needs to be able to sample from and compute the gradient of the logarithm of the heat kernel efficiently. This can be done trivially in Euclidean space as the heat kernel is a Gaussian distribution, but doing this is effectively intractable for general manifolds. By restricting our analysis to Riemannian symmetric spaces [9, 29], which are a class of a manifold with a special structure, we can make substantial improvements. We leverage this additional structure to quickly and precisely compute heat kernel quantities, allowing us to scale up Riemannian Diffusion Models to high dimensional real-world tasks. Furthermore, since almost all manifolds that practitioners work with are (or are diffeomorphic to) Riemannian symmetric spaces, our improvements are generalizable and not task-specific. Concretely our contributions are:
- We present a generalized strategy for numerically computing the heat kernel on Riemannian symmetric spaces in the context of denoising score matching. In particular, we adapt known heat kernel techniques to our more specific problem, allowing us to quickly, accurately, and stably train with denoising score matching.
- We show how to exactly sample from the heat kernel using our quick computations above. In particular, we show how our exact heat kernel computation enables fast simulation free techniques on simple manifolds. Furthermore, we develop a sampling method based on the probability flow ODE and show that this can be quickly computed with standard ODE solvers on the maximal torus of the manifold.
- We empirically demonstrate that our improved Riemannian Diffusion Models improve performance and scale to high dimensional real world tasks. For example, we can faithfully learn Wilson action on  $4 \times 4$  SU(3) lattices (128 dimensions). Furthermore, when applied to contrastively learned hyperspherical embeddings (127 dimensions), our method enables better model interpretability by recovering the collapsed projection head representations. To the best of our knowledge, this is the first example where differential equation-based manifold generative models have scaled to real world tasks with hundreds of dimensions.
Diffusion models on  $\mathbb{R}^d$  are defined through stochastic differential equations [24, 45, 48]. Given an initial data distribution  $p_0$  on  $\mathbb{R}^d$ , samples  $\mathbf{x}_0 \sim \mathbb{R}^d$  are perturbed with a stochastic differential equation[33]
$$
\mathrm {d} \mathbf {x} _ {t} = \mathbf {f} \left(\mathbf {x} _ {t}, t\right) \mathrm {d} t + g (t) \mathrm {d} \mathbf {B} _ {t} \tag {1}
$$
where  $\mathbf{f}$  and  $g$  are fixed drift and diffusion coefficients, respectively. The time varying distributions  $p_t$  (defined by  $\mathbf{x}_t$ ) evolves according to the Fokker-Planck Equation
$$
\frac {\partial}{\partial t} p _ {t} (\mathbf {x}) = - \operatorname {d i v} \left(p _ {t} (\mathbf {x}) \mathbf {f} (\mathbf {x}, t)\right) + \frac {g (t) ^ {2}}{2} \Delta_ {x} p _ {t} (\mathbf {x}) \tag {2}
$$
and approaches a limiting distribution  $\pi \approx p_T$ , which is normally a simple distribution like a Gaussian  $\mathcal{N}(0, \sigma_T^2 I)$  through carefully chosen  $\mathbf{f}$  and  $g$ . Our SDE has a corresponding reversed SDE
$$
\mathrm {d} \mathbf {x} _ {t} = \left(\mathbf {f} \left(\mathbf {x} _ {t}, t\right) - g (t) ^ {2} \nabla_ {x} \log p _ {t} \left(\mathbf {x} _ {t}\right)\right) \mathrm {d} t + g (t) \mathrm {d} \overline {{\mathbf {B}}} _ {t} \tag {3}
$$
which maps  $p_T$  back to  $p_0$ . Motivated by this correspondence, diffusion models approximate the score function  $\nabla_x \log p_t(\mathbf{x})$  using a neural network  $\mathbf{s}_{\theta}(\mathbf{x}, t)$ . To do this, one minimizes the score matching loss [27], which is weighted by constants  $\lambda_t$ :
$$
\mathbb {E} _ {t, \mathbf {x} _ {t} \sim p _ {t}} \lambda_ {t} \| \mathbf {s} _ {\theta} (\mathbf {x} _ {t}, t) - \nabla_ {x} \log p _ {t} (\mathbf {x} _ {t}) \| ^ {2} \tag {4}
$$
Since this loss is intractable due to the unknown  $\nabla_{x}\log p_{t}(\mathbf{x}_{t})$ , we instead use an alternative form of the loss. One such loss is the implicit score matching loss[27]:
$$
\mathbb {E} _ {t, \mathbf {x} _ {t} \sim p _ {t}} \lambda_ {t} \left[ \operatorname {d i v} \left(\mathbf {s} _ {\theta}\right) \left(\mathbf {x} _ {t}, t\right) + \frac {1}{2} \| \mathbf {s} _ {\theta} \left(\mathbf {x} _ {t}, t\right) \| ^ {2} \right] \tag {5}
$$
which normally is estimated using sliced score matching/Hutchinson's trace estimator[26, 47]:
$$
\mathbb {E} _ {t, \epsilon , \mathbf {x} _ {t} \sim p _ {t}} \lambda_ {t} \left[ \epsilon^ {\top} D _ {x} \mathbf {s} _ {\theta} (\mathbf {x} _ {t}, t) \epsilon + \frac {1}{2} \| \mathbf {s} _ {\theta} (\mathbf {x} _ {t}, t) \| ^ {2} \right] \tag {6}
$$
where  $\epsilon$  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  $\epsilon$  normally renders this loss unworkable in high dimensions[46], so practitioners instead use the denoising score matching loss[51]
$$
\mathbb {E} _ {t, \mathbf {x} _ {0} \sim p _ {0}, \mathbf {x} _ {t} \sim p _ {t} (\cdot | \mathbf {x} _ {0})} \lambda_ {t} \| \mathbf {s} _ {\theta} (\mathbf {x} _ {t}, t) - \nabla_ {x} \log p _ {t} (\mathbf {x} _ {t} | \mathbf {x} _ {0}) \| ^ {2} \tag {7}
$$
where  $p_t(\mathbf{x}_t|\mathbf{x}_0)$  is derived from the SDE in Equation 1 and is normally tractable. Once  $\mathbf{s}_{\theta}(\mathbf{x}_t,t)$  is learned, we can construct a generative model by first sampling  $\mathbf{x}_T \sim \pi \approx p_T$  and solving the generative SDE from  $t = T$  to  $t = 0$ :
$$
\mathrm {d} \mathbf {x} _ {t} = (\mathbf {f} (\mathbf {x} _ {t}, t) - g ^ {2} (t) \mathbf {s} _ {\theta} (\mathbf {x} _ {t}, t)) \mathrm {d} t + g (t) \mathrm {d} \overline {{\mathbf {B}}} _ {t} \tag {8}
$$
Furthermore, there exists a corresponding "probability flow ODE" [48]
$$
\mathrm {d} \mathbf {x} _ {t} = (\mathbf {f} (\mathbf {x} _ {t}, t) - \frac {g (t) ^ {2}}{2} \nabla_ {x} \log p _ {t} (\mathbf {x} _ {t})) \mathrm {d} t \tag {9}
$$
that has the same evolution of  $p_t$  as the SDE in Equation 1. This can be approximated using our score network  $\mathbf{s}_{\theta}$  to get a Neural ODE [11]
$$
\mathrm {d} \mathbf {x} _ {t} = \left(\mathbf {f} \left(\mathbf {x} _ {t}, t\right) - \frac {g (t) ^ {2}}{2} \mathbf {s} _ {\theta} \left(\mathbf {x} _ {t}, t\right)\right) \mathrm {d} t \tag {10}
$$
which can be used to evaluate exact likelihoods of the data [20].
2.2 Riemannian Diffusion Models
To generalize diffusion models to  $d$ -dimensional Riemannian manifolds  $\mathcal{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  $\mathbf{v}$ , gradients  $\nabla$ , and Brownian motion  $\mathbf{B}_t$  have natural analogues for  $\mathcal{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:
$$
\mathrm {d} \mathbf {x} _ {t} = g (t) \mathrm {d} \mathbf {B} _ {t} \tag {11}
$$
Unlike in the Euclidean case,  $p_t$  approaches the uniform distribution  $\mathcal{U}_{\mathcal{M}}$  as  $t \to \infty$  (in practice, this convergence is fast, getting within numerical precision for  $t \approx 5$ ).
The transition density has no closed form. Despite the fact that we work with the most simple SDE, the transition kernel defined by the manifold heat equation  $p_t(x_t|x_0)$  has no closed form. This transition kernel is known as the heat kernel, which satisfies Equation 2 with the additional condition that, as  $t \to 0$ , the kernel approaches  $\delta_{x_0}$ . We will denote this by  $K_{\mathcal{M}}(x_t|x_0,t)$ , and we highlight that, when  $\mathcal{M}$  is  $\mathbb{R}^d$ , this corresponds to a Gaussian and is easy to work with.
This has several major consequences which cause prior work to favor sliced score matching over denoising score matching. First, to sample a point  $x \sim K_{\mathcal{M}}(\cdot |x_0,t)$ , one must simulate a Geodesic Random Walk with the Riemannian exponential map exp:
$$
x _ {t + \Delta t} = \exp_ {x} (\sqrt {\Delta t} z) \quad z \sim \mathcal {N} (0, I _ {d}) \in T _ {x} \mathcal {M} \tag {12}
$$
Additionally, to calculate  $K_{\mathcal{M}}(\cdot |x_0,t)$  or  $\nabla \log K_{\mathcal{M}}(x|x_0,t)$ , one must use eigenfunctions  $f_{i}$  which satisfy  $\Delta f_{i} = -\lambda_{i}f_{i}$ . These  $f_{i}$  form an orthonormal basis for all  $L^2$  function on  $\mathcal{M}$ , allowing us to write the heat kernel with an infinite sum:
$$
K _ {\mathcal {M}} ^ {\mathrm {E F}} (x \mid x _ {0}, t) = \sum_ {i = 0} ^ {\infty} e ^ {- \lambda_ {i} t} f _ {i} \left(x _ {0}\right) f _ {i} (x) \tag {13}
$$
Additionally, previous work has also explored the use of the Varadhan approximation for small values of  $t$  (which uses the Riemannian logarithmic map log) [43]:
$$
K _ {\mathcal {M}} (x | x _ {0}, t) \approx \mathcal {N} (0, \sqrt {2 t}) (\operatorname {d i s t} _ {\mathcal {M}} (x _ {0}, x)) \Rightarrow \nabla_ {x} \log K _ {\mathcal {M}} (x | x _ {0}, t) \approx \frac {1}{2 t} \log_ {x} (x _ {0}) \tag {14}
$$
The key problem with applying denoising score matching in practice is that the heat kernel computation is too expensive and inaccurate. Notably, simulating a geodesic random walk is expensive since it requires many exponential maps. Furthermore, the eigenfunction expansion in Equation 13 requires more and more eigenfunctions (numbering in the tens of thousands) as  $t \to 0$ . Worse still, these eigenfunctions formulas are not well-known for most manifolds, and, even when explicit formulas exist, they can be numerically unstable (like in the case of  $S^n$ ). One possible way to alleviate this is to use Varadhan's approximation for small  $t$ , but this is also unreliable except for very small  $t$ .
To remedy this issue, we instead consider the case of Riemannian symmetric spaces [22]. We emphasize that most manifold generative modeling applications already model on Riemannian Symmetric Spaces like the sphere, torus, or Lie Groups, so we do not lose applicability by restricting our attention here. Furthermore, for surfaces (which have appeared as generative modeling test tasks in the literature [42]), one can always define a generative model by mapping the data points to  $S^2$ , learning a generative model there, and mapping back [14]. We empathize that, outside of these two examples, we are unaware of any other manifolds which have been used for Riemannian generative modeling tasks. For this very reason, we will also focus primarily on the compact case, since the irreducible (base case) noncompact symmetric spaces are all diffeomorphic to Euclidean space, although we discuss the extensions of our method in Appendix A.2.
3.1 Heat Kernels on Riemannian Symmetric Spaces
In this section, we will define Riemannian Symmetric Spaces and showcase their relationship with the heat kernel. We empathize that our exposition is neither rigorous nor fully defines all terms, although it provides a general intuition with which to build our calculations. We urge interested readers to consult a book [22] or monograph [9] for a full treatment of the subject.
Definition 3.1. A Riemannian Symmetric Space is a Riemannian manifold such that, for all points  $x \in \mathcal{M}$ , there exists a local isometry  $s$  s.t.  $s(x) = x$  and  $D_x s = -\mathrm{id}_{T_x \mathcal{M}}$ .
Intuitively, this condition means that, at each point, the manifold "looks the same" in all directions, which simplifies the analysis. While this is defined as a local condition, it also has a global consequence: all Riemannian symmetric spaces are quotients  $G / K$  for Lie Groups  $G$  and compact isotropic subgroups  $K$ . Importantly, most Riemannian manifolds that are used in practice are symmetric:
Examples. Lie Groups  $G \cong G / \{e\}$  (where  $\{e\}$  is the trivial Lie Group), the sphere  $S^n \cong SO(n + 1) / SO(n)$ , and hyperbolic space  $H^n \cong SO(n,1) / O(n)$  are all symmetric spaces.
On symmetric spaces, one can define a special structure called the maximal torus:
Definition 3.2 (Maximal Torus). A maximal torus  $T$  is a compact, connected, and abelian subgroup of  $G$  that is not contained in any other such subgroup. For symmetric spaces  $G / K$ , the maximal torus is defined by quotienting out the maximal torus of  $G$ . Importantly, the maximal torus is isomorphic to a standard torus  $T^n \approx (S^1)^n$ .
Intuitively, the maximal torus forms the "most stable" part of the symmetric space, and it shows up as a natural object when performing analysis due to its isometry with the standard flat torus.
Examples. For the sphere  $S^n$ , a maximal torus is any great circle. For the Lie group of unitary matrix  $U(n)$ , the maximal torus is the set of all diagonal matrices  $\{\mathrm{diag}(e^{i\theta_1},\ldots ,e^{i\theta_n}):\theta_k\in [0,2\pi)\}$ .
To work with the maximal torus algebraically (as we shall do in our definition of the heat kernel), we instead analyze the root system.
Definition 3.3 (Roots of a Maximal Torus). We define a root  $\alpha$  be a vector (really a covector) that corresponds with a character  $\chi$  of  $T$  s.t.
$$
\chi (\exp X) = \exp (2 \pi i (\alpha \cdot X)) \tag {15}
$$
We define the set  $R^{+}$  to be the set of all positive roots. The multiplicity  $m_{\alpha}$  of a root  $\alpha$  is the dimension of the eigenspace over the Lie algebra  $\mathfrak{G}$
$$
\{Z \in \mathfrak {G}: [ H, Z ] = \pi i \lambda (H) Z \forall H \} \tag {16}
$$
Examples. For all spheres dimension  $n$ , there is only one root of multiplicity  $n - 1$ , which corresponds (roughly) to measuring the angle on the maximal torus/great circle.
Notably, the heat kernel is invariant on the maximal torus. This allows us to write out the formula for the heat kernel on the general manifold as a simplified formula on the maximal torus
Proposition 3.4 (Heat Kernel Reduces on Maximal Torus). The Laplace-Beltrami operator on  $\mathcal{M}$  (the manifold generalization of the standard Laplacian) induces the "radial" Laplacian on  $T$ :
$$
L _ {r} = \Delta_ {T} + \sum_ {\alpha \in R ^ {+}} m _ {\alpha} \cot (\alpha \cdot h) \frac {\partial}{\partial \alpha} \tag {17}
$$
where  $\Delta_T$  is the standard Laplacian on the torus. Here,  $h$  is the "flat" coordinate of  $x$  on the maximal torus (e.g. for sphere this is the angle between  $x$  and  $x_0$ ).
One can readily see that the above formula allows one to reduce computations over  $\mathcal{M}$  to computations over  $T$ , which greatly reduces dimensionality ( $n$  dimensions to 1 dimension for spheres and  $O(n^2)$  dimensions to  $O(n)$  for Matrix Lie Groups).
3.2 Improved Heat Kernel Estimation
We now use this maximal torus perspective to improve computations related to the heat kernel. In particular, we can greatly improve both the speed and fidelity of the numerical evaluation during our training process.
3.2.1 Eigenfunction Expansion Restricted to the Maximal Torus
We note that the maximal torus relationship in Proposition 3.4 reduces the eigenfunction expansion in Equation 13 to an eigenfunction expansion of the induced Laplacian on the maximal torus. This has implicitly appeared in previous work when defining Riemannian Diffusion Models on  $S^2$  and  $SO(3)$  [4, 34], allowing one to rewrite the summation as, respectively
$$
K _ {S ^ {2}} ^ {E F} (x \mid x _ {0}, t) = \frac {1}{4 \pi} \sum_ {l = 0} ^ {\infty} (2 l + 1) P _ {l} (\langle x, x _ {0} \rangle) e ^ {- l (l + 1) t} \tag {18}
$$
where  $P_{l}$  are the Legendre Polynomials and
$$
K _ {S O (3)} ^ {E F} \left(x \mid x _ {0}, t\right) = \frac {1}{8 \pi^ {2}} \sum_ {l = 0} ^ {\infty} e ^ {- 2 l (l + 1) t} \frac {\sin (2 l + 1) \theta / 2}{\sin (\theta / 2)} \quad \theta \text {i s} \tag {19}
$$
By making this relationship explicit, we can directly generalize this formulation to other symmetric spaces, e.g.  $SU(3)$
$$
K _ {S U (3)} ^ {E F} \left(x \mid x _ {0}, t\right) = \sum_ {p, q = 1} ^ {\infty} \frac {p q (p + q)}{4} \left(\chi^ {p, q} (\theta , \phi) + \chi^ {q, p} (\theta , \phi)\right) e ^ {- \frac {t \left(p ^ {2} + q ^ {2} + p q\right)}{6}} \tag {20}
$$
where  $\chi$  are the (real) irreducible representations and  $\theta, \phi$  are the induced angles between  $x$  and  $x_0$ .
By reducing the dimensionality, this strategy generally makes the summation more tractable. Furthermore, for cases such as the hypersphere  $S^n$ , the eigenfunctions are numerically unstable (rendering them useless for our purposes), while the maximal torus setup stable.
3.2.2 Controlling Errors for Small Time Values
While the torus eigenfunction expansion greatly reduces the computational cost (particularly for higher dimensions), for small values of  $t$  the heat kernel scaling terms  $e^{-\lambda_it}$  remain large as one increases  $i, \lambda_i \to \infty$ . As a result, one must evaluate many more terms (often in the thousands), and the result can be very numerically unstable even with double precision.
To address this problem, we examine several refined versions of the Varadhan approximation that use the fact that the manifold is symmetric. These approximations allow us to control the number of eigenfunctions required and, in some cases, completely obviate the need for them entirely.
The Schwinger-Dewitt Approximation
The Varadhan approximation is constructed by approximating the heat kernel with a Gaussian distribution (with respect to the Riemannian distance function). However, doing this does not account for the curvature of the manifold. Incorporating this premise, we get the Schwinger-Dewitt approximation [9]:
$$
K _ {\mathcal {M}} ^ {\mathrm {S D}} (x \mid x _ {0}, t) = \frac {\overline {{\Delta}} _ {x _ {0}} (x) ^ {1 / 2} e ^ {- \frac {d _ {\mathcal {M}} (x _ {0} , x) ^ {2}}{4 t} + \frac {t R}{6}}}{(4 \pi t) ^ {\dim \mathcal {M} / 2}} \tag {21}
$$
Here,  $\overline{\Delta}_{x_0}(x) = \operatorname*{det}(D_{x_0}\exp_{x_0}(\log_{x_0}(x)))$  is the (unnormized) change of volume term introduced by the exponential map, and  $R$  is the scalar curvature of the manifold. Generally, this is much more accurate than Varadhan's approximation as it better accounts for the curvature of the manifold through the  $\Delta$  term, retaining accuracy up to moderate time values.
$\Delta$  appears to be a rather computationally demanding term since it requires evaluating the determinant of a Jacobian. Indeed, this naively takes  $O(\dim \mathcal{M}^3)$  evaluations which is completely inaccessible in higher dimensions. However, we again emphasize the fact that we are working with symmetric spaces:  $\Delta$  has a particularly simple formula defined by our flat coordinate  $h$  from above:
$$
\bar {\Delta} _ {x _ {0}} (x) = \prod_ {\alpha \in R ^ {+}} \left(\frac {\alpha \cdot h}{\sin (\alpha \cdot h)}\right) ^ {m _ {\alpha}} \tag {22}
$$
Sum Over Paths
The fact that we can derive a better approximation using a different power of  $\Delta$  points to deeper connections between the heat kernel and the "Gaussian" with respect to distance. We draw inspiration from several Euclidean case examples, such as the flat torus [28] or the unit interval [36]. For those cases, the heat kernel is derived by summing a Gaussian over all possible paths connecting  $x_0$  and  $x$ . While this formula does not exactly lift over to Riemannian symmetric spaces, there exists an analogue for Lie Groups [9]:
$$
K _ {\mathcal {M}} ^ {\mathrm {S O P}} \left(x \mid x _ {0}, t\right) = \frac {e ^ {t \rho^ {2}}}{(4 \pi t) ^ {\dim \mathcal {M} / 2}} \sum_ {2 \pi n \in \Gamma} \prod_ {\alpha \in R ^ {+}} \frac {\alpha \cdot (h + 2 \pi n)}{2 \sin (\alpha \cdot h / 2)} e ^ {- \frac {(h + 2 \pi n) ^ {2}}{4 t}} \tag {23}
$$
Here,  $\Gamma$  is infinite tangent space grid that corresponds with loops of the torus (i.e.  $2\pi n$  circular intervals), and  $\rho^2$  is a manifold specific constant. We note that the product over  $R^{+}$  is exactly the  $\Delta$  change in variables as given above since  $m_{\alpha} = 1$ , but we simply extend this to every other root.
Generally, this formula is rather powerful as it gives us an exact (albeit infinite) representation for the heat kernel. Compared to the eigenfunction expansion in Equation 13, the Sum Over Paths representation is accurate for small  $t$ , which nicely complements the fact that the eigenfunction representation is accurate for large  $t$ . This formula does generalize to split-rank Riemannian symmetric spaces like odd dimensional spheres (more details in Appendix A.1) and noncompact spaces (Appendix A.2).
3.2.3 A Unified Heat Kernel Estimator