1 Introduction↩︎

Diffusion models have become a powerful and widely used framework for generative modeling. They define a forward process that gradually corrupts data with noise, and learn a reverse process that transforms noise back into samples, enabling high-quality generation across many domains [1], [2]. However, a growing number of applications involve data that do not naturally live in \(\mathbb{R}^n\): typical examples include spheres, special orthogonal groups \(SO(n)\) and spaces of symmetric positive definite (SPD) matrices. Treating such objects in Euclidean space via some embeddings or projections can break intrinsic constraints, distort distances and volumes, and create coordinate-dependent artifacts. In contrast, modeling directly on the underlying manifold respects the geometry by construction, aligns the noise process with geodesic structure, and yields a principled inductive bias that can improve both sample quality and generalization.

On a Riemannian manifold \((\mathcal{M},g)\), a canonical drift-free forward process is Riemannian Brownian motion, whose time-marginal density evolves according to the heat equation \(\partial_t p_t=\Delta^{\mathcal{M}}p_t\). Following the drift-free formulation of Riemannian score-based generative modeling of [3], training and sampling require access to the manifold heat kernel (the transition density of Brownian motion), in particular to sample from it and evaluate its log-density/gradients. While several approximations have been proposed—most notably short-time asymptotics and spectral expansions—these can be accurate only in restricted regimes or rely on strong structural assumptions, limiting their practical scope [4].

Instead, we adopt a simple viewpoint: the heat kernel is the solution of a known partial differential equation (PDE) with a known initial condition. For each fixed \(x_0\), it satisfies \[\label{eq:fp-eqn-intro} \partial_t p(t,x)=\Delta_x^{\mathcal{M}}p(t,x),\qquad p(t,\cdot)\to \delta_{x_0}\;\text{as }t\downarrow 0.\tag{1}\] We approximate this object by learning a numerical surrogate of the heat equation. Concretely, we (i) express \(\Delta_{\mathcal{M}}\) in explicit coordinates obtained from an embedding or quotient representation, (ii) use a principled short-time asymptotic to provide an approximate initializer at a small \(t_0>0\), and (iii) train a physics-informed neural network (PINN) to minimize the PDE residual and match the initializer [5], [6]. The resulting surrogate provides both heat-kernel evaluation (and its gradient) and supports heat-kernel-based sampling, enabling drift-free Riemannian diffusion on manifolds beyond the few simple ones.

We evaluate our method across diverse geometries and applications: spherical climate-event density estimation on \(S^2\) (volcanoes, earthquakes, floods, wildfires), synthetic mixtures on \(SO(3)\), conditional traffic-flow generation on \(\mathop{\mathrm{SPD}}(10)\) from the NYC taxi dataset, and class-conditional EEG brain-connectivity generation on \(\mathop{\mathrm{SPD}}(n)\). Moreover, for molecule generation, we go beyond the standard \({\mathbb{R}}^{k\times n}\) point-cloud view and formulate the task on the quotient manifold \({\mathbb{R}}^{k\times n}/S_n\), explicitly modding out atom permutations. This leads to a novel diffusion process that evolves over permutation-equivalence classes, so atom identities are shared across molecules. Overall, the results show that PDE-based heat-kernel approximation enables practical drift-free Riemannian diffusion beyond simple geometries.

2 Related Work↩︎

2.0.0.1 Riemannian diffusion models.

Riemannian diffusion models generalize Euclidean diffusion/score-based generative modeling to data on a smooth manifold \({\mathcal{M}}\) by defining a forward diffusion SDE on \({\mathcal{M}}\) and learning a reverse-time SDE. RSGM [3] provides a general formulation and proposes practical training objectives, including Denoising Score Matching (DSM) and Implicit Score Matching (ISM). Complementary likelihood-oriented formulations [7] derive a variational framework for an evidence lower bound (ELBO) on manifolds, which requires computing quantities such as the Riemannian divergence. Recent work improves scalability by exploiting additional geometric structure (e.g., symmetric spaces), yielding more accurate and efficient evaluation and approximation of manifold diffusion quantities, including heat-kernel-related terms, in higher dimensions [4], [8]. Relatedly, [9] propose Riemannian Diffusion Schrödinger Bridge (RDSB), extending diffusion Schrödinger bridges to compact Riemannian manifolds for generative modeling.

2.0.0.2 Flow matching on Riemannian manifolds.

Flow matching is an alternative strategy for building generative models on Riemannian manifolds. It trains a continuous-time generative model by learning a time-dependent vector field whose flow transports a simple base distribution to the data distribution, avoiding likelihood computation and divergence estimation. [10] extend conditional flow matching to general geometries via Riemannian Flow Matching (RFM), where the training targets are constructed by geodesic interpolation. This formulation requires only \(\exp\) and \(\log\) maps, yielding a scalable and simulation-free objective. This framework has been adopted in geometric generative modeling, including \(\mathrm{SE}(3)\)-equivariant protein modeling and torus-valued torsion-angle generation [11]–[14].

2.0.0.3 Applications of generative modeling on Riemannian manifolds.

Manifold-aware generative models are motivated by geometric constraints and symmetries, especially in scientific and 3D data. Typical examples include distributions on \(S^2\) (e.g., geophysical event locations), rotations on \(SO(3)\) for 3D orientation modeling, and rigid-body frames on \(SE(3)\) in robotics and protein structure generation [3], [11], [15], [16]. In molecular and materials modeling, periodic degrees of freedom such as internal angles and torsions live on products of circles (tori), motivating manifold generative modeling on tori for conformer generation and related tasks [17], [18]. Structured matrix manifolds such as \(\mathrm{SPD}(n)\) arise in domains including neuroimaging and traffic analytics, motivating diffusion-style generative modeling directly in SPD geometry [19]–[21]. These applications motivate scalable manifold generative methods that remain applicable when heat kernels are not analytically tractable.

3 Background↩︎

3.0.0.1 Definitions and notations.

Let \({\mathcal{M}}\) be a complete, connected and boundaryless smooth Riemannian manifold equipped with a Riemannian metric \(g\). The metric \(g\) is a smooth rank-2 tensor bundle; for each \(x \in {\mathcal{M}}\), the metric \(g_x:T_x{\mathcal{M}}\times T_x{\mathcal{M}}\rightarrow {\mathbb{R}}\) is a positive definite bilinear function. In a local coordinate chart \((x^1,\dots,x^n)\), the metric is represented by the matrix \(g=(g_{ij})\), and we write \(|g| := \det(g_{ij})\). At \(x \in {\mathcal{M}}\), the exponential map \(\exp_x:T_x{\mathcal{M}}\rightarrow {\mathcal{M}}\) is defined by transporting the tangent vector in \(T_x {\mathcal{M}}\) along the geodesic starting from \(x\). The inverse of the exponential map \(\log_x = \exp_x^{-1}\) is also defined, although it may not be defined globally. A vector field on \({\mathcal{M}}\) can be written in coordinates with the local frame \(\{\partial_i\}_{i=1,\cdots,n}\). Throughout the paper, we use the Einstein summation convention without raising or lowering indices, i.e. the indices appearing twice in the tensor expression are regarded as being summed, e.g., a standard matrix multiplication is written as \((AB)_{ik} = A_{ij}B_{jk}\).

The metric \(g\) induces the gradient \(\nabla f\) of a smooth function \(f:{\mathcal{M}}\to{\mathbb{R}}\) via \(g(\nabla f, V) = df(V)\) for all vector fields \(V\). The divergence of a vector field \(V = V_i\partial_i\) is given in local coordinates by \(\mathrm{div}(V) = \frac{1}{\sqrt{|g|}}\,\partial_i\!\big(\sqrt{|g|}\,V_i\big)\), or identically, for the Levi-Civita connection \(\nabla\), \(\mathrm{div}(V) = \mathop{\mathrm{tr}}(\nabla V)\). The Laplace–Beltrami operator is defined by \[\label{eq:lap-def} \Delta^{{\mathcal{M}}} f = \mathrm{div}(\nabla f) = \frac{1}{\sqrt{|g|}}\,\partial_i\!\Big(\sqrt{|g|}\,g^{-1}_{ij}\partial_j f\Big).\tag{2}\] The (Riemannian) Brownian motion \(\mathbf{B}_t^{{\mathcal{M}}}\) is the diffusion on \({\mathcal{M}}\) whose infinitesimal generator is \(\tfrac12\Delta^{{\mathcal{M}}}\). More detailed background on Riemannian geometry can be found in, e.g., [22], [23], and for stochastic processes on Riemannian manifolds, see, e.g., [3], [24], [25].

3.0.0.2 Riemannian diffusion models.

Let \(p_0\) be a distribution on \({\mathcal{M}}\). We consider the drift-free Riemannian diffusion \[\label{eq:brownian-def} dX_t \;=\; \sqrt{2} d\mathbf{B}_t^{{\mathcal{M}}}, \qquad X_0\sim p_0.\tag{3}\] A time-dependent scalar noise schedule \(dX_t=\sigma_t\,d\mathbf{B}_t^{{\mathcal{M}}}\) can be introduced via a time reparameterization, so we present this formulation for simplicity. Let \(p_t\) be the marginal density of \((X_t)\) with respect to the Riemannian volume measure of \({\mathcal{M}}\), and \(p_{t|0}(x_t | x_0)\) be the transition density (or the heat kernel, since the SDE is simple Brownian). The heat kernel \(p_{t|0}\) satisfies the Fokker-Planck equation (or the heat equation) \[\label{eq:fp-eqn-background} \begin{align} \partial_t p_{t|0}(x\mid x_0) \;&=\; \,\Delta_x^{{\mathcal{M}}} p_{t|0}(x\mid x_0),\\ \lim_{t\rightarrow 0} p_{t|0}(\cdot\mid x_0) \;&=\; \delta_{x_0}(\cdot). \end{align}\tag{4}\] Under mild regularity assumptions, the reverse-time process \((X_{T-t})_{t\in [0,T]}\) satisfies the reverse-time SDE [3] \[\label{eq:reverse-sde} \begin{align} dX_{T-t} \;&=\; 2s_{T-t}(X_{T-t})\,dt + \sqrt{2}d\widetilde{\mathbf{B}}_t^{{\mathcal{M}}},\\ X_T \;&\sim\; p_T, \end{align}\tag{5}\] where \(\widetilde{\mathbf{B}}_t^{{\mathcal{M}}}\) is a (reverse-time) Brownian motion on \({\mathcal{M}}\) and \[\label{eq:score-fn} s_t(x) \;:=\; \nabla_x \log p_t(x)\tag{6}\] is the score function.

Given a neural network \(s_\theta: [0, T] \times {\mathcal{M}}\to T{\mathcal{M}}\) with \(s_\theta(t,x) \in T_x{\mathcal{M}}\), one can approximate the score function \(s\) by \(s_\theta\) by learning \(\theta\) with denoising score matching (DSM) objective [3], [26], \[\label{eq:dsm-loss} \begin{align} &\mathcal{L}_{\text{DSM}}(\theta) :=\\ & \mathbb{E}\Big[ \lambda(t)\, \big\|\, s_\theta(t,X_t) -\nabla_{X_t}\log p_{t|0}(X_t\mid X_0)\,\big\|_{g_{X_t}}^2 \Big]. \end{align}\tag{7}\]

When \({\mathcal{M}}\) is Euclidean, this formulation is identical to the standard Euclidean diffusion models [1], [2]. In this case, the heat kernel \(p_{t|0}\) becomes the isotropic Gaussian, so both sampling the forward diffusion \(p_{t|0} (\cdot,x_0)\) and computing the conditional score function \(\nabla_{x_t}\log p_{t|0}(x_t\mid x_0)\) are tractable. However, when working with a general Riemannian manifold \({\mathcal{M}}\), the heat kernel is not known in closed form in general. Previous works [3], [4] introduce a few tricks to overcome this challenge.

3.0.0.3 Asymptotic methods.

When \(t\) is sufficiently small, \(p_{t|0}\) can be approximated by a Gaussian-like form: \[\label{eq:approx-warped} \begin{align} p_{t|0}(x_t|x_0) &\approx G_t(x_t, x_0) \\ &:= \frac{1}{(4 \pi t)^{n/2}} \exp \bigg(-\frac{d(x_t,x_0)^2}{4t}\bigg) \end{align}\tag{8}\] where \(n=\dim {\mathcal{M}}\) and \(d(\cdot, \cdot)\) is the geodesic distance. This justifies sampling from \(p_{t|0}(\cdot|x_0)\) by sampling a Gaussian in a tangent space \(T_{x_0}{\mathcal{M}}\) and push-forwarding by the exponential map \(\exp_{x_0}:T_{x_0}{\mathcal{M}}\rightarrow {\mathcal{M}}\). This distribution is sometimes called the warped Gaussian. When \(t\) is large, sampling from \(p_{t|0}(\cdot|x_0)\) can be approximated by the geodesic random walk algorithm [3], which is a variant of the random walk algorithm where the Gaussian is replaced by the warped Gaussian. The score function \(\nabla \log p_{t|0}\) also admits the asymptotic form: \[\label{eq:approx-varadhan} \begin{align} \\ &\approx \nabla_{x_t} \log G_t(x_t, x_0) = \frac{1}{2t}\log_{x_t}(x_0), \end{align}\tag{9}\] called the Varadhan’s asymptotics [27]. However, this approximation does not incorporate the volume element distorted by the curvature of the manifold. To account for this, [4], [28] introduced the Schwinger-DeWitt approximation: \[\label{eq:approx-sd} \begin{align}\\ &\,\,\,:= \frac{D_{x_0}(x_t)^{-1/2}}{(4 \pi t)^{n/2}} \exp \bigg(-\frac{d(x_t,x_0)^2}{4t} + \frac{tR}{6}\bigg) \end{align}\tag{10}\] where \(D_{x_0}(x_t) = |\det (d( \exp_{x_0})_{\log_{x_0}(x_t)})|\) measures the volume distorted by the exponential map, and \(R\) is the scalar curvature, assuming \({\mathcal{M}}\) has a constant scalar curvature. However, this also only works with sufficiently small time \(t\).

3.0.0.4 Spectral methods.

When the manifold \({\mathcal{M}}\) is compact, the Laplace-Beltrami operator \(\Delta^{{\mathcal{M}}}\) is a self-adjoint operator on the Hilbert space \(L^2({\mathcal{M}})\) (with respect to the Riemannian volume measure). Hence it admits a spectral decomposition \(\{(\lambda_i,\phi_i)\}_{i=0}^\infty\) satisfying \(\Delta^{{\mathcal{M}}}\phi_i \;=\; -\lambda_i \phi_i\), with eigenvalues \(0=\lambda_0<\lambda_1\le\cdots\). The heat kernel then has the spectral expansion \[\label{eq:approx-spectral} p_{t|0}(x\mid x_0) \;=\; \sum_{i=0}^\infty e^{-\lambda_i t}\,\phi_i(x)\,\phi_i(x_0),\tag{11}\] which can be approximated by truncating the series. In practice, however, this only works with compact manifolds and requires access to the eigenfunctions of \(\Delta^{{\mathcal{M}}}\), which are rarely available beyond a few highly symmetric examples. Recent work [4] stabilizes and accelerates such computations by reformulating the expansion using representation theory on Lie groups, enabling efficient evaluation on certain groups and homogeneous/quotient spaces. Nevertheless, these methods rely on strong algebraic structure and are only applicable to (quotients of) compact Lie groups or Euclidean groups.

3.0.0.5 Mapping into Euclidean spaces.

When the exact computation on manifolds is challenging, using their Euclidean counterparts can be a useful trick. [3] used stereographic projection to map data on 2D sphere into \({\mathbb{R}}^2\). [19] mimicked the logic of Euclidean diffusion on the space of positive definite matrices, using operations defined on its tangent space. Similar tricks are common for Lie groups: one can perform diffusion updates in the associated Lie algebra (a vector space) via the \(\log\) and \(\exp\) maps, e.g., for \(SO(3)\) or \(SE(3)\) [15], [29].

4 Methods↩︎

Given an explicit description of a Riemannian manifold \(({\mathcal{M}},g)\), we approximate the heat kernel \(p_{t|0}(\cdot| x_0)\) by numerically solving the heat equation \(\partial_t p=\Delta^{{\mathcal{M}}}p\) with the initial condition \(p_{0|0} = \delta(\cdot)\). For this work, we first locate the manifold \({\mathcal{M}}\) into an easily-computable coordinate system, in which we derive the heat equation. Then we choose a properly approximated initial condition, since the original initial condition \(p_{0|0} = \delta(\cdot)\) cannot be imposed numerically. Finally, we train a PINN that approximates the heat kernel. The PINN is used for both (i) forward sampling of Brownian motion and (ii) computing the conditional score function for the DSM loss.

4.1 Choosing a coordinate system↩︎

When working with manifolds mathematically, it is common to choose multiple local coordinate systems and patch them to cover the entire manifold. However, when working with manifolds numerically, using multiple coordinate systems would be computationally inconvenient. Instead, we choose to work with a global coordinate system, by embedding the manifold \({\mathcal{M}}\) into a larger space \(\tilde{{\mathcal{M}}}\) (called an ambient manifold; typically \({\mathbb{R}}^N\), possibly with a non-Euclidean metric \(\tilde{g}\)) as a submanifold or quotient manifold.

For \(\tilde{{\mathcal{M}}} = {\mathbb{R}}^N\), a submanifold \({\mathcal{M}}\subseteq {\mathbb{R}}^N\) can be practically defined by \[\label{eq:submanifold-def} \begin{align} {\mathcal{M}} &=\Bigl\{\,x\in{\mathbb{R}}^N \;\Big|\;f_i(x)=0 \;\text{for } i=1,\ldots,k,\\ &\qquad\qquad\;\;f_j(x)>0 \;\text{for } j=k+1,\ldots,k+l\,\Bigr\}. \end{align}\tag{12}\] for functions \(f_1,\cdots,f_{k+l} : {\mathbb{R}}^N \rightarrow {\mathbb{R}}\). The manifold inherits the metric \(\tilde{g}\) of \({\mathbb{R}}^N\). For \(x \in {\mathcal{M}}\), the projection map \(P(x):T_x {\mathbb{R}}^N \rightarrow T_x {\mathcal{M}}\subseteq T_x {\mathbb{R}}^N\) is given by \[\label{eq:projection-def} P(x)=I-\tilde{g}_x^{-1}\,J(x)^{T}\,\bigl(J(x)\,\tilde{g}_x^{-1}\,J(x)^{T}\bigr)^{-1}\,J(x),\tag{13}\] where \(\tilde{g}_x\) is the matrix representation of the metric at \(x\), and \(J(x)=\bigl(J_{ij}(x)\bigr):=\bigl(\partial_j f_i(x)\bigr)\in{\mathbb{R}}^{k\times N}\) is the Jacobian. The ambient metric \(\tilde{g}\) and the projection matrix \(P(x)\) are sufficient to express various differential-geometric operations on \({\mathcal{M}}\) in coordinates.

A quotient manifold \({\mathcal{M}}=\tilde{{\mathcal{M}}}/G\) is obtained by identifying points related by an isometry group action of \(G\). For simplicity, we restrict to a finite or discrete group \(G\), so that \({\mathcal{M}}\) inherits a smooth structure from \(\tilde{{\mathcal{M}}}\); numerically, quotienting can be implemented by treating all points in an orbit as the same state, e.g., via a canonical representative. We do not explicitly address freeness or properness of the group action; although non-free orbits may introduce lower-dimensional singularities in the quotient, these are negligible for the continuous distributions considered in our deep-learning applications.

4.2 Fokker–Planck equation in coordinates↩︎

Using the coordinates on the ambient manifold \(\tilde{{\mathcal{M}}}\) with the projection matrix \(P(x)\), we can express all the required differential operators in coordinates. Let \(\tilde{\nabla}\) denote the Levi–Civita connection of \((\tilde{{\mathcal{M}}},\tilde{g})\). Although we present the full expressions for generality, in practice, some terms (e.g., the Levi-Civita connection \(\tilde{\nabla}\)) simplify by choosing a simple ambient manifold \(\tilde{{\mathcal{M}}}\); for instance, when \(\tilde{{\mathcal{M}}}\) is a Euclidean space, \(\tilde{\nabla}\) reduces to partial derivatives. The Fokker-Planck equation (or the heat equation) on \({\mathcal{M}}\) can be explicitly written in coordinates, with expressions depending only on the ambient metric \(\tilde{g}\) and the projection matrix \(P\): \[\label{eq:fp-coords} \begin{align} &\partial_t p_{t|0} (x | x_0) = \Delta^{{\mathcal{M}}}p_{t|0}(x|x_0) \\ &\quad = P_{ij}(x)\,\tilde{\nabla}_k\!\Bigl(P_{j\ell}(x)\,\tilde{g}^{-1}_{\ell m}(x)\,\partial_m \bar p_{t|0}(x|x_0)\Bigr)\,P_{ki}(x). \end{align}\tag{14}\] where \(\bar p_{t|0}\) denotes any smooth extension of \(p_{t|0}\) on \(\tilde{{\mathcal{M}}}\). The detailed derivation is in 8.1.

4.2.0.1 Example: a 2D unit sphere.

Let \(S^2 = \{x \in {\mathbb{R}}^3 ~|~ \lVert x\rVert = 1\}\) be the unit sphere. Although it’s standard to use a local coordinate system such as the spherical coordinates, here we directly embed the 2D sphere in the Euclidean manifold \({\mathbb{R}}^3\). For \(x \in S^2\), the projection matrix \(P(x):T_x{\mathbb{R}}^3 \rightarrow T_xS^2 \subseteq {\mathbb{R}}^3\) is given as \[\label{eq:s2-proj} P(x) = I - xx^T; \quad P_{ij} = \delta_{ij}-x_ix_j\tag{15}\] where \(\delta_{ij}\) is the Kronecker delta. In the ambient manifold \({\mathbb{R}}^3\), the metric \(\tilde{g}\) is the identity matrix, and the Levi-Civita connection \(\tilde{\nabla}_i\) simplifies to the partials \(\partial_i\). Now we can work out the Laplace-Beltrami operator: \[\label{eq:s2-lap} \begin{align} \Delta^{S^2}f &= P_{ij}\partial_i\partial_jf + P_{ki}(\partial_kP_{ij})\partial_jf \\ &= (\delta_{ij} - x_ix_j) \partial_i \partial_j f - 2x_i \partial_i f. \end{align}\tag{16}\]

4.3 Initial condition for the Fokker–Planck equation↩︎

Since it is not possible to impose the initial condition \(p_{0|0}(x \mid x_0)=\delta_{x_0}(x)\) directly, we instead initialize the PDE using the small-time asymptotics of the heat kernel. Let \(B_{x_0}(r_{\max})=\{x\in{\mathcal{M}}\mid d(x,x_0)\le r_{\max}\}\) be the geodesic ball of radius \(r_{\max}\) centered at \(x_0\). On \(x \in B_{x_0}(r_{\max})\), the heat kernel admits the zeroth-order expansion \[\label{eq:approx-0th} \begin{align} p_t(x\mid x_0) &= G_t^{(0)}(x,x_0) + O(t),\\ G_t^{(0)}(x,x_0) &:= D_{x_0}(x)^{-1/2}\,G_t(x,x_0)\\ &= \frac{D_{x_0}(x)^{-1/2}}{(4\pi t)^{n/2}} \exp\!\bigg(-\frac{d(x,x_0)^2}{4t}\bigg). \end{align}\tag{17}\] A principled way to obtain tighter approximations is the parametrix method (also known as the Minakshisundaram–Pleijel recursion) [30], which yields an order-\(n\) expansion of the form \[\label{eq:approx-series} \begin{align} & p_t(x|x_0) = G_t^{(n)}(x, x_0) + O(t^n), \\ &G_t^{(n)}(x, x_0) := G_t(x,x_0)\Big(u_0(x,x_0) \\ & \quad +u_1(x, x_0)t + \cdots + u_{n-1}(x,x_0)t^{n-1}\Big). \end{align}\tag{18}\] Here \(u_0(x,x_0)=D_{x_0}(x)^{-1/2}\), and \(u_1,\ldots,u_{n-1}\) can be computed recursively; we defer the details to 8.2. The Schwinger-DeWitt approximation 10 is also a simplified version of the first-order parametrix expansion.

4.3.0.1 Stability analysis of PDE.

In general, similarity in initial condition for PDEs does not guarantee similarity of the corresponding solutions at later times. However, due to various stability properties of the heat equation, solving the heat equation from an approximate initial condition at \(t = t_0\) yields a solution that remains reliable for all \(t \geq t_0\).

2 is a rigorous version of this theorem. Moreover, in 10.2 we extend this result to learned PINNs with small log-PDE residual, where the log-density error is bounded by the boundary error plus the accumulated residual.

4.4 Training PINN↩︎

A PINN is a deep-learning approach for approximating solutions to PDEs by embedding the governing equation directly into the training objective [5]. While many PINN variants also incorporate observations at interior collocation (quadrature) points, we describe here a basic formulation that uses only the initial/boundary condition and the PDE residual. Let the space–time domain be \([t_0, t_{\max}] \times U,\) where \(U \subset {\mathbb{R}}^N\) is closed, and consider a well-posed evolution equation written in residual form \[\label{eq:pde-def} {\mathcal{R}}[f](t, x) := {\mathcal{R}}(t, x, \partial_t f, \partial_x f, \partial_x^2 f,\cdots) = 0\tag{19}\] with initial condition \(f(t_0,x) = f_0(x)\). In the PINN framework, we parameterize the solution by a neural network \(f_{\theta}:[t_0, t_{\max}] \times U \rightarrow {\mathbb{R}}\) and train \(\theta\) minimizing a combination of the initial/boundary loss: \[\label{eq:pinn-loss} \begin{align} {\mathcal{L}}_B(f_\theta) &:= \mathbb{E}_x [\lVert f_\theta(x, t_0) - f_0(x)] \rVert^2 ] \\ {\mathcal{L}}_{I}(f_\theta) &:= \mathbb{E}_{t,x} [\lVert {\mathcal{R}}[f_\theta](x, t)\rVert^2 ] \end{align}\tag{20}\] where \(x\in U\) and \((t,x)\in [t_0,t_{\max}] \times U\) are sampled on some adequate distributions on their domain.

This PINN formulation is well suited to approximating the heat kernel. For a given manifold \({\mathcal{M}}\), we derive the heat equation written in coordinates and train the PINN with the approximate initial condition \(\tilde{p}_{t_0|0}\). For numerical stability, we let \(p_{t|0} = e^{\phi_{t|0}}\) and rewrite the equation in terms of \(\phi_{t|0}\), and train a PINN \(\phi^{(\theta)}_{t|0}\) approximating it. The trained PINN is used jointly with the short-time approximation, to approximate the heat kernel: \[\label{eq:logp-approx} \log p^{(\text{approx})}_{t|0}(x| x_0)= \begin{cases} \begin{array}{@{}l@{}} \log \tilde{p}_{t|0}(x| x_0), \quad t<t_0,\\ \text{\small(short-time asymptotics)} \end{array} \\[6pt] \begin{array}{@{}l@{}} \phi^{(\theta)}_{t|0}(x| x_0), \quad t\ge t_0.\\ \text{\small(PINN approximation)} \end{array} \end{cases}\tag{21}\] For stable training of PINN, we adopt several techniques of [6], including modified MLP structures with embeddings, adaptive weight scaling of two losses, and time-scheduled learning method.

4.4.0.1 Representational power of PINN.

Quantifying whether a PINN can accurately represent the heat kernel on a given (potentially large and geometrically complex) manifold \({\mathcal{M}}\) is challenging, since the approximation error depends strongly on modeling choices such as the network architecture. Nevertheless, the heat equation has a smoothing property: for any \(t>t_0\), the evolved log-density \(\log p_{t|0}\) is guaranteed to be more regular than the initial log-density \(\log p_{t_0|0}\), e.g., in terms of Sobolev norms [31]. This suggests that, once a reasonable approximation is available at \(t=t_0\), learning the forward solution for \(t\ge t_0\) may be easier than fitting the initial condition itself. In high dimensions, PINNs can be difficult to train and may require more computational and memory resources [32], [33], often necessitating additional inductive biases (e.g., incorporating known symmetries and using expressive input embeddings) to improve optimization behavior.

4.5 Riemannian diffusion model with PINN↩︎

All ingredients for building Riemannian diffusion models are now set. We train a Riemannian diffusion model as described in 3. Having knowledge of the density \(\log p^{(\text{approx})}_{t|0}(x| x_0)\), the forward sampling \(x_t \sim p^{(\text{approx})}_{t|0}(\cdot| x_0)\) can be done by a Markov chain Monte Carlo (MCMC) method. Since the warped Gaussian discussed in 3 is an effective approximation of \(p_{t|0}\), we use it for the initial distribution of the MCMC, then run a geodesic random walk for a fixed step size \(\varepsilon>0\) for proposal. After \(x_t\) is drawn, the conditional score \(\nabla \log p^{(\text{approx})}_{t|0}\) is computed by differentiating the PINN, and train the diffusion model under the DSM loss.

The architecture of a diffusion model on \({\mathcal{M}}\) depends on the geometry of \({\mathcal{M}}\). For example, if \({\mathcal{M}}\) is (a quotient of) a submanifold of an ambient space \(\tilde{{\mathcal{M}}}={\mathbb{R}}^N\), then an extrinsic network \(s_{\theta}:[0,T]\times{\mathbb{R}}^N \rightarrow {\mathbb{R}}^N\) can be used, possibly with appropriate symmetry constraints induced by the quotient. Alternatively, one may use an intrinsic architecture tailored to \({\mathcal{M}}\), e.g., \(s_{\theta}:[0,T]\times{\mathcal{M}}\rightarrow {\mathcal{M}}\). In this case, since the denoiser corresponds to a tangent vector \(\nabla_{x} \log p_{t|0}(x\mid x_0)\), we adopt a standard trick from Euclidean diffusion models that reparameterizes prediction in terms of the denoised data [1]. Concretely, instead of directly learning \(\nabla_{x} \log p_{t|0}(x\mid x_0)\), we train the "denoised input" \(\exp_x\!\big(2t\nabla_x \log p_{t|0}(x\mid x_0)\big)\in {\mathcal{M}}\).

For the reverse sampling, we also require the perturbed (or prior) distribution \(p_T \approx p(\cdot | x_0)\). When \({\mathcal{M}}\) is compact, we use the uniform distribution \(p_T \approx \mathop{\mathrm{Unif}}({\mathcal{M}})\) assuming \(T\) is sufficiently large. When \({\mathcal{M}}\) is not compact, we set a rough mean of the data distribution \(x_{\text{mean}}\) and take \(p_T \approx p_T(\cdot,x_{\text{mean}})\). This setting is common in recent state-of-the-art Euclidean diffusion models [34].

5 Worked examples↩︎

5.1 2D sphere \(S^2\)↩︎

As derived in 4.2, the heat equation on \(S^2 \subset {\mathbb{R}}^3\) is given as \(\partial_t p = (\delta_{ij} - x_ix_j) \partial_i \partial_j p - 2x_i \partial_i p\), and for \(\phi = \log p\), \[\label{eq:logheat-s2} \begin{align} \partial_t \phi &= [(\delta_{ij}-x_ix_j)(\partial_i\partial_j\phi + \partial_i\phi\partial_j\phi)-2x_i\partial_i\phi] \\ &= \mathop{\mathrm{tr}}(\nabla^2\phi) + \lVert\nabla \phi\rVert^2 -x^T(\nabla^2\phi) x \\ &\quad -(x^T\nabla\phi)^2 - 2x^T\nabla\phi \end{align}\tag{22}\] The initial condition can be computed using the parametrix method as described in 4.3 up to third order. The recursion integrals are evaluated via series expansion using a computer algebra system (CAS). Detailed procedure is in 8.3.

5.2 Special orthogonal group \(SO(3)\)↩︎

The special orthogonal group \(SO(3)\) is isometric to \(S^3 / \{\pm I\}\), where \(S^3 \subset {\mathbb{R}}^4\) is the 3D unit sphere and \(\cdot / \{\pm I\}\) denotes quotienting by identifying the opposite points on sphere, i.e. \(\pm x \in S^3\) are regarded as an identical point. This isometry can be explicitly written using quaternions. Once the mapping is constructed, the heat equation is similar to the case of the 2D sphere: \[\label{eq:logheat-so3} \begin{align} \partial_t \phi &= [(\delta_{ij}-x_ix_j)(\partial_i\partial_j\phi + \partial_i\phi\partial_j\phi)-3x_i\partial_i\phi] \\ &= \mathop{\mathrm{tr}}(\nabla^2\phi) + \lVert\nabla \phi\rVert^2 -x^T(\nabla^2\phi) x \\ &\quad -(x^T\nabla\phi)^2 - 3x^T\nabla\phi \end{align}\tag{23}\] and the initial condition is again computed using the parametrix method, similar to 5.1.

5.3 Space of positive-definite matrices \(\mathop{\mathrm{SPD}}(n)\)↩︎

Generative modeling on the space of positive definite matrices \(\mathop{\mathrm{SPD}}(n) = \{X \in {\mathbb{R}}^{n \times n} ~|~ X \text{ is positive definite}\}\) has drawn interest in recent research, e.g. for traffic analysis or fMRI data modeling [19], [20]. The hyperbolic geometry of \(\mathop{\mathrm{SPD}}(n)\) is well suited to the affine-invariant metric \(g_X(U,V) = \mathop{\mathrm{tr}}(X^{-1}UX^{-1}V)\), written in infinitesimal form. The derivation of the heat equation is in 8.4. For any initial point \(X_0\in \mathop{\mathrm{SPD}}(n)\), since the map \(X \mapsto X_0^{-1/2}XX_0^{-1/2}\) is an isometry under the affine-invariant metric that maps \(X_0\) to \(I \in \mathop{\mathrm{SPD}}(n)\), it suffices to consider the heat equation starting at \(I\). The heat equation starting at \(I\) only depends on the eigenvalues of \(X \in \mathop{\mathrm{SPD}}(n)\), so we parametrize the matrices by their log-eigenvalues \((r_1,\cdots,r_n)\) where \(X = Q \mathop{\mathrm{diag}}(e^{r_1},\cdots,e^{r_n})Q^T\) for some \(Q \in O(n)\), then the heat kernel is reformulated by an equation in \(r_1,\cdots,r_n\): \[\label{eq:heat-spd} \partial_tp_{t|0} = \sum_i \partial_i^2p_{t|0} + \frac{1}{2}\sum_{i < j} \coth\Big(\frac{r_i-r_j}{2}\Big)(\partial_i p_{t|0}-\partial_j p_{t|0})\tag{24}\] where the partials are in the \(r\)-coordinates, i.e. \(\partial_i = \frac{\partial}{\partial r_i}\). For the initial condition, we use the 0-th order approximation in 4.3: \[\label{eq:initial-spd} \tilde{p}_{t_0|0}(X | I) = \frac{D(r)^{-1/2}}{(4 \pi t_0)^{n(n+1)/4}} \exp \bigg(-\frac{\lVert r\rVert^2}{4t_0}\bigg)\tag{25}\] where \(D(r) =\prod_{i < j} \frac{\sinh((r_i-r_j)/2)}{(r_i-r_j)/2}\).

5.4 Space of point clouds quotiented by permutation symmetry \({\mathbb{R}}^{k \times n} / S_n\)↩︎

A point cloud \(X = [x^{(1)},\cdots,x^{(n)}] \in {\mathbb{R}}^{k \times n}\) usually has the permutation symmetry \(S_n\) by reordering the points, and it is common to incorporate permutation symmetry for e.g. architectural designs [35]–[37] or probabilistic formulation [38] in deep learning. Although existing works for generative modeling of point clouds (e.g., molecules) mostly use permutation symmetric neural networks for denoisers [39], [40], their diffusion process is formulated in \({\mathbb{R}}^{k \times n}\). To model a diffusion on the quotient manifold \({\mathbb{R}}^{k \times n} / S_n\), one requires a heat kernel on this space, given as: \[\label{eq:heat-perm} \begin{align} p_{t|0}(x\mid x_0) &= \sum_{\sigma\in S_n} \frac{1}{(4\pi t)^{kn/2}} \\ &\quad \cdot \exp\!\bigg(-\frac{1}{4t}\sum_{i=1}^n \lVert x^{(i)}-x_0^{(\sigma(i))}\rVert^2\bigg). \end{align}\tag{26}\] Intuitively, in this setting, all the points \(x^{(1)},\cdots,x^{(n)}\) share the identity. Although the heat kernel is in closed-form, it is intractable to compute due to the \(n!\) elements of \(S_n\). We use a permutation-symmetric neural network as a PINN to approximate this kernel, with the equation and (approximate) initial condition same as the Euclidean one: \[\label{eq:initial-perm} \begin{align} \partial_t\tilde{p}_{t|0} &= \sum_{i=1}^n\sum_{j=1}^k \frac{\partial^2}{{\partial x_j^{(i)}}^2}\tilde{p}_{t|0}, \\ \tilde{p}_{t_0|0}(x\mid x_0)&= \frac{1}{(4\pi t_0)^{\tfrac{kn}{2}}}\exp\bigg(-\frac{\sum_{i=1}^n \lVert x^{(i)}-x^{(i)}_0\rVert^2}{4t_0}\bigg). \end{align}\tag{27}\]

6 Experiments↩︎

In this section, we demonstrate how our method applies across diverse settings and a range of manifolds. Our experiments are designed to highlight the generality of the framework, while minimizing task-specific engineering aimed solely at boosting performance. Our implementation is available at https://github.com/kogyeonghoon/riem-diff-pinn.git.

6.1 Training PINN↩︎

We largely follow the PINN training protocol of [6], adapting its architectural and optimization strategies to our manifold setting. Specifically, we employ a modified MLP architecture in which each hidden layer is fused with manifold-aware coordinate embeddings. This design enriches the coordinate-based neural representation and improves its ability to approximate nonlinear and geometrically complex surrogate solutions. Since the accuracy of later-time dynamics depends critically on the quality of the solution learned at earlier times, we adopt a curriculum strategy that progressively expands the training interval to \([t_0,t_{\max}]\). In addition, we use an adaptive loss-balancing scheme based on the exponential moving average of gradient norms, together with time-dependent residual weights that encourage the PINN to learn the solution in a temporally causal manner.

The input embeddings are chosen according to the geometry of each manifold. For \(S^2\) and \(SO(3)\), we use Fourier embeddings following [6]. For \(\mathrm{SPD}(n)\), we use the symmetric and antisymmetric embeddings described in 9, motivated by the permutation-invariant structure of the radial heat kernel. For \({\mathbb{R}}^{k \times n}/S_n\), we replace the modified MLP with a permutation-symmetric GNN architecture, reflecting the permutation symmetry of the quotient heat-kernel surrogate. The spatial and temporal training domains of the PINN, specified by the radius cutoff \(r_{\max}\) and the time interval \([t_0,t_{\max}]\), are selected separately for each manifold and downstream dataset. We choose \(r_{\max}\) to cover the region where the heat-kernel surrogate is queried, while \(t_{\max}\) is set large enough so that the forward heat diffusion sufficiently corrupts the data by the terminal time.

To directly validate the learned heat-kernel surrogate, we evaluate it separately from the downstream diffusion-model performance. On \(S^2\) and \(SO(3)\), where accurate reference heat-kernel approximations are available via [4], we compare both the learned log-density \(\log p_{t|0}^{(\theta)}(x|x_0)\) and the conditional score \(\nabla_x\log p_{t|0}^{(\theta)}(x|x_0)\) against the reference quantities across several noise levels in 1. For the remaining manifolds, such reference kernels are not tractable, so we instead assess whether the trained PINN satisfies the defining heat-kernel problem. Specifically, for all manifolds we report the normalized initial/boundary-condition error and the normalized PDE residual of the learned log-density in 2. The PDE residual is normalized by the largest-magnitude term appearing in the equation; for example, for a PDE written in the form \(A=B+C\), we report \(|A-B-C|/\max{|A|,|B|,|C|}\). These diagnostics measure, respectively, how well the surrogate matches the prescribed short-time initial condition and how accurately it satisfies the manifold heat equation after training. Thus, the quality of the surrogate is evaluated both by direct comparison to existing heat-kernel approximations when available, and by geometry-agnostic PINN consistency checks across all considered manifolds. Finally, we report the wall-clock runtime of the surrogate in 9.6; since the surrogate is implemented in JAX and evaluated after JIT compilation, log-density and score evaluations incur only a small computational overhead.

3pt

2pt

6.2 Climate science datasets on \(S^2\)↩︎

Following the experimental setup of RSGM [3], we validate our method on the spherical density-estimation benchmarks comprising geolocated events on Earth’s surface—volcanic eruptions [41], earthquakes [42], floods [43], and wildfires [44]. Each event is represented as a point on the unit sphere \(S^2\). We report test log-likelihoods computed with the same manifold-adapted likelihood estimator for score-based models used in RSGM, and we compare against their method. The results are summarized in 3. In RSGM, the computation of the heat kernel is approximated using a combination of the Varadhan approximation and spectral decomposition.

2pt

6.3 Synthetic data on \(SO(3)\)↩︎

We again follow the experimental setting of RSGM and evaluate our method on \(SO(3)\) using synthetic data drawn from a mixture of wrapped Gaussians. In RSGM, the heat kernel is approximated using the Varadhan approximation. Results are reported in 3.

6.4 Traffic analysis on \(\mathop{\mathrm{SPD}}(10)\)↩︎

[19] is one of the early works to formulate diffusion on the space of positive definite matrices. They use the NYC taxi dataset [45], where New York City is divided into 10 regions and each \(\mathop{\mathrm{SPD}}(10)\) matrix encodes traffic flows between pairs of regions. The dataset includes 13 predictor features, such as weather conditions on the observed day. Following their experimental setting, we train a conditional diffusion model that takes predictors \(y\) as input and generates an SPD matrix \(X \in \mathop{\mathrm{SPD}}(10)\). We evaluate using the Frobenius distance between generated matrices and the ground-truth data. Results are in 4.

3pt

6.5 Brain connectivity (EEG) analysis on \(\mathop{\mathrm{SPD}}(n)\)↩︎

DiffeoCFM [20] introduces conditional flow matching using a diffeomorphic mapping \(\mathop{\mathrm{SPD}}(n)\rightarrow {\mathbb{R}}^{n(n+1)/2}\), enabling SPD matrices to be modeled in a Euclidean space. Following their experimental protocol, we evaluate our method on the EEG datasets BNCI2014-002 [46] and BNCI2015-001 [47]. Each SPD matrix is paired with a binary motor-imagery label, and we generate SPD samples conditionally on label. We evaluate the generated samples using distributional quality metrics (\(\alpha\)-precision, \(\beta\)-recall and F1) based on [48], and downstream classification accuracy (AUC and F1), obtained by training a classifier on generated samples and evaluating it on the test set.

In addition to DiffeoCFM, we compare against the diffusion-style baseline LowTriDDPM [20], which models diffusion over the lower-triangular entries of SPD matrices. As shown in 5 and 6, DiffeoCFM performs best on downstream AUC/CAS-F1, while our method is competitive and achieves strong precision/Q-F1 on some settings. We hypothesize that the data-centric diffeomorphic parameterization (via a \(\log\) map around a central reference point) captures global structure more effectively than operating directly on \(\mathop{\mathrm{SPD}}(n)\) under the affine-invariant metric.

2pt

2pt

6.6 Molecule generation on \({\mathbb{R}}^{k \times n} / S_n\)↩︎

Finally, we study molecule generation on the permutation-quotiented space \({\mathbb{R}}^{k \times n} / S_n\), using the E(3)-Equivariant Diffusion Model (EDM) [39] as a baseline. QM9 [49] contains molecules with up to 29 atoms, along with atomic positions and molecular properties. Following EDM, we concatenate atomic coordinates with one-hot atom types to form per-atom feature vectors, representing a molecule with \(n\) atoms as a point cloud in \({\mathbb{R}}^{k \times n}\). To enable a more direct comparison with our method, we train an EDM variant based on the variance-exploding (VE) SDE [2]. We report molecule stability, validity, uniqueness, and novelty in 7. We find that our method produces more valid molecules but lower novelty, reflecting a common validity–novelty trade-off in molecule generation. We also note that our method is not compared to the state-of-the-art baselines [40], [50] which use more parameters and incorporate edge features.

2pt

7 Conclusion↩︎

In this work, we addressed a central practical bottleneck in drift-free Riemannian diffusion models: training and sampling typically require access to the manifold heat kernel and its gradients, which are rarely available beyond a few simple manifolds. We proposed a general, PDE-based alternative: given an explicit specification of a manifold, we (i) select a global coordinate representation, (ii) derive the corresponding heat (Fokker–Planck) equation in coordinates, (iii) replace the singular Dirac-delta initial condition with a principled short-time asymptotic initializer, and (iv) train a physics-informed neural network (PINN) to approximate the log heat kernel. The resulting surrogate heat kernel supports both forward noising and conditional-score evaluation, enabling denoising score matching and reverse-time sampling on manifolds where closed-form kernels are unavailable.

Empirically, we demonstrated that this solve-the-heat-equation viewpoint can serve as a practical backend for manifold diffusion modeling across diverse geometries and data modalities, including density estimation on \(S^2\), synthetic modeling on \(SO(3)\), conditional generation on \(\mathop{\mathrm{SPD}}(n)\), and quotient-manifold settings such as permutation-invariant point clouds for molecule generation. Overall, the experiments support the thesis that PDE-based heat-kernel approximation can broaden the applicability of Riemannian diffusion models beyond cases where spectral methods or analytic kernels are readily available.

7.0.0.1 Limitations.

First, training a PINN for the heat kernel can become nontrivial as the effective dimension grows or when the manifold geometry/topology induces a highly structured solution; in such regimes, optimization can be unstable and may require stronger inductive biases (e.g., symmetry-aware architectures or better feature embeddings). Second, the overall pipeline is computationally heavy in practice: PINN training requires many collocation evaluations with higher-order autodifferentiation, and downstream usage may further require repeated evaluations (and sometimes MCMC-style forward sampling), leading to substantial runtime and engineering overhead. Finally, while the method is general, it is also more complicated than alternatives such as (Riemannian) flow matching, which typically bypass heat-kernel estimation entirely and can be simpler to implement when accurate exponential/log maps are available.

7.0.0.2 Future work.

A key direction is to automate and systematize the full workflow so that applying the method to a new manifold becomes closer to “plug-and-play.” Concretely, this includes more automatic identification/derivation of the governing PDE (e.g., Laplace–Beltrami in chosen coordinates), robust construction of the short-time initializer and boundary/domain specification, and more reliable PINN training protocols. Beyond improving usability and scalability, it is also promising to extend the PDE-surrogate idea to a broader class of generative modeling frameworks—e.g., manifold Schrödinger bridges and controlled diffusions, or hybrids that combine the geometric convenience of flow matching with stochasticity—where learned transition-density surrogates (or their score analogues) could serve as modular building blocks.

Acknowledgements↩︎

This work was partly supported by Institute of Information & communications Technology Planning & Evaluation(IITP) grant funded by the Korea government(MSIT) (No.RS-2019-II190075, Artificial Intelligence Graduate School Program(KAIST), No.RS-2024-00509279, Global AI Frontier Lab, and No.RS-2022-II220713, Meta-learning Applicable to Real-world Problems).

Impact Statement↩︎

This work advances foundational methods for generative modeling on Riemannian manifolds. We do not anticipate direct societal impact from the results presented. While the techniques could be used in downstream applications, they do not introduce new application domains or materially change the known risk profile of modern generative models; no specific ethical concerns beyond standard responsible-use considerations are expected.

8 Mathematical Details↩︎

8.1 Fokker-Planck equation in coordinates↩︎

The notation follows from 4.2.

8.1.0.1 Gradient, divergence and Laplace-Beltrami.

For a smooth function \(f:{\mathcal{M}}\to{\mathbb{R}}\), to work in the ambient manifold \(\tilde{{\mathcal{M}}}\), we use any smooth extension \(\bar f\) to a neighborhood of \({\mathcal{M}}\) in \(\tilde{{\mathcal{M}}}\). The gradient of \(\bar f\) in the ambient manifold \(\tilde{{\mathcal{M}}}\) is given by \[(\widetilde{\nabla}\,\bar f)_i=\tilde{g}^{-1}_{ij}\,\partial_j \bar f,\] and the gradient of \(f\) in the submanifold \({\mathcal{M}}\) is obtained by projection: \[(\nabla f)(x)=P(x)\,\widetilde{\nabla}\,\bar f(x),\qquad (\nabla f)_i=P_{ij}(x)\,\tilde{g}^{-1}_{jk}(x)\,\partial_k \bar f(x).\] For a smooth tangent vector field \(X\) on \({\mathcal{M}}\) (so \(PX=X\)), the divergence is the trace over tangent directions: \[\mathop{\mathrm{div}}X=\mathop{\mathrm{tr}}\!\bigl(P\,(\tilde{\nabla} X)\,P\bigr) = P_{ij}(x)\,(\tilde{\nabla}_k X)_j\,P_{ki}(x).\] The Laplace–Beltrami operator is \(\Delta^{{\mathcal{M}}}f:=\mathop{\mathrm{div}}(\nabla f)\), hence \[\Delta^{{\mathcal{M}}}f =\mathop{\mathrm{tr}}\!\Bigl(P\,\tilde{\nabla}\bigl(P\,\widetilde{\nabla}\,\bar f\bigr)\,P\Bigr) = P_{ij}(x)\,\tilde{\nabla}_k\!\Bigl(P_{j\ell}(x)\,\tilde{g}^{-1}_{\ell m}(x)\,\partial_m \bar f(x)\Bigr)\,P_{ki}(x).\]

8.1.0.2 Fokker–Planck equation.

Therefore, the Fokker–Planck equation of the simple Brownian motion is given by \[\partial_t p_{t|0}(x\mid x_0) =\Delta^{{\mathcal{M}}}p_{t|0}(x\mid x_0) = P_{ij}(x)\,\tilde{\nabla}_k\!\Bigl(P_{j\ell}(x)\,\tilde{g}^{-1}_{\ell m}(x)\,\partial_m \bar p_{t|0}(x\mid x_0)\Bigr)\,P_{ki}(x).\] Note that, although \(p_{t|0}\) is defined only on \({\mathcal{M}}\), we evaluate \(\Delta^{{\mathcal{M}}}p_{t|0}\) by choosing any smooth extension \(\bar p_{t|0}\) to a neighborhood of \({\mathcal{M}}\) in \(\tilde{{\mathcal{M}}}\). Any such extension can be used since the projection matrix \(P(x)\) projects vectors into the tangent bundle of \({\mathcal{M}}\).

8.2 Minakshisundaram–Pleijel recursion formula↩︎

We follow the notation in 4.3. The heat kernel has an \(n\)-th order parametrix expansion [30]: \[\begin{align} p_t(x|x_0) &= G_t^{(n)}(x, x_0) + O(t^n), \\ G_t^{(n)}(x, x_0) &:= G_t(x,x_0) \Big(u_0(x,x_0)+u_1(x, x_0)t + \cdots, u_{n-1}(x,x_0)t^{n-1}\Big). \end{align}\] where \(u_0(x,x_0) = D_{x_0}(x)^{-1/2}\) and \(u_1,\cdots,u_{n-1}\) can be computed recursively from earlier ones: \[\label{eq:parametrix-recursion} u_i(x,x_0) = r^{-i}D^{-1/2}(x)\int_0^r D^{1/2}(z(s))\Delta_x^{{\mathcal{M}}}u_{i-1}(x,z(s))s^{i-1}ds\tag{28}\] where \(r = d(x,x_0)\) and \(z:[0,r] \rightarrow {\mathcal{M}}\) is the unit-speed geodesic from \(x_0\) to \(x\).

8.3 Worked example of Minakshisundaram–Pleijel recursion formula on \(S^2\)↩︎

Since the recursion formula 28 has an integral, it cannot be analytically computed in general. Instead, if the dimension of \({\mathcal{M}}\) is sufficiently small or the heat kernel only depends on a small number of variables, we can Taylor-expand the terms \(u_i\) and compute the recursion formula, using a Computer Algebra System (CAS), e.g., sympy [51].

Take \(x_0 = [0,0,1]^{T}\) and for \(x \in S^2 \subseteq {\mathbb{R}}^3\), let \(r = \arccos(x^Tx_0)\) be the geodesic distance between \(x_0\) and \(x\). Due to the symmetry of \(p_{t|0}\), the heat equation only depends on the distance \(r\). For a function \(f = f(r):S^2 \rightarrow {\mathbb{R}}\) that only depends on \(r\), the Laplace-Beltrami operator reduces to: \[\Delta^{S^2} f = f'' + (\frac{1}{r} + \frac{D'}{D})f'\] where \(D(r) = \frac{\sin r}{r}\) is the volume distortion. The recursion formula 28 reduces to: \[u_i = r^{-i}D(r)^{-1/2}\int_0^rD^{1/2}(s)\Delta u_{i-1}(s)s^{i-1}ds\] where \(u_0 = D(r)^{-1/2}\). In the sympy environment, we conduct series expansion before computing the integral in \(s\). Below we list some examples of computed terms \(u_i\) using sympy: \[\begin{align} u_1 &= \frac{9689 r^{8}}{319334400} + \frac{17 r^{6}}{56700} + \frac{31 r^{4}}{10080} + \frac{r^{2}}{30} + \frac{1}{3} \\ u_2 &= - \frac{1433 r^{8}}{139708800} + \frac{3277 r^{6}}{19958400} + \frac{19 r^{4}}{15120} + \frac{r^{2}}{105} + \frac{1}{15} \\ u_3 &= - \frac{9619 r^{8}}{1026432000} - \frac{5713 r^{6}}{52390800} + \frac{71 r^{4}}{110880} + \frac{r^{2}}{315} + \frac{4}{315} \end{align}\] where the series are truncated at 10-th order. The computed parametrix approximation up to third order is used for the initial condition of heat kernel, for \({\mathcal{M}}= S^2\) and \(SO(3)\).

8.4 Heat equation on \(SPD(n)\) in log-eigenvalue coordinates↩︎

Let \(\mathop{\mathrm{SPD}}(n)=\{X\in\mathbb{R}^{n\times n} : X=X^\top,\;X\succ 0\}\) equipped with the affine-invariant metric \[g_X(U,V) \;=\; \mathop{\mathrm{tr}}\!\big(X^{-1}UX^{-1}V\big), \qquad X\in\mathop{\mathrm{SPD}}(n),\;U,V\in T_X\mathop{\mathrm{SPD}}(n)\cong \mathop{\mathrm{Sym}}(n). \label{eq:spd-metric-app}\tag{29}\] The congruence action of \(GL(n)\), \(X\mapsto AXA^\top\), is isometric under 29 . In particular, the heat kernel started at \(I\) is conjugation-invariant: \(p_t(QXQ^\top\mid I)=p_t(X\mid I)\) for all \(Q\in O(n)\), hence it depends only on the eigenvalues of \(X\).

8.4.0.1 Log-eigenvalue coordinates.

Write the spectral decomposition \[X \;=\; Q\,\mathop{\mathrm{diag}}(\lambda_1,\dots,\lambda_n)\,Q^\top,\qquad Q\in O(n),\;\lambda_i>0,\] and define log-eigenvalues \(r_i:=\log\lambda_i\), so that \[X \;=\; Q\,\mathop{\mathrm{diag}}(e^{r_1},\dots,e^{r_n})\,Q^\top. \label{eq:spd-logeig-coords}\tag{30}\] For a conjugation-invariant function \(f(X)=\bar f(r)\) (symmetric in \(r\)), the Laplace–Beltrami operator reduces to its radial part.

8.4.0.2 Radial volume density.

In the coordinates \((Q,r)\), the Riemannian volume element factorizes as \[d\mu_{\mathop{\mathrm{SPD}}}(X) \;=\; C\,J(r)\,dr\,d\mu_{\mathrm{Haar}}(Q), \qquad J(r)\;=\;\prod_{1\le i<j\le n}\sinh\!\Big(\frac{r_i-r_j}{2}\Big), \label{eq:spd-jacobian}\tag{31}\] where \(C\) is a constant. It is often convenient to write \[J(r) \;=\; \Big(\prod_{i<j}\frac{r_i-r_j}{2}\Big)\, D(r), \qquad D(r):=\prod_{i<j}\frac{\sinh\!\big(\frac{r_i-r_j}{2}\big)}{\frac{r_i-r_j}{2}}. \label{eq:spd-D-def-app}\tag{32}\]

8.4.0.3 Radial Laplacian.

For \(f(X)=\bar f(r)\), the radial Laplace–Beltrami operator is \[(\Delta f)(X) \;=\; (\Delta_{\mathrm{rad}}\bar f)(r) \;:=\; \frac{1}{J(r)}\sum_{i=1}^n \partial_{r_i}\!\Big(J(r)\,\partial_{r_i}\bar f(r)\Big). \label{eq:spd-radial-laplacian-div-form}\tag{33}\] Expanding 33 yields \[\Delta_{\mathrm{rad}}\bar f(r) \;=\; \sum_{i=1}^n \partial_{r_i}^2\bar f(r) \;+\; \sum_{i=1}^n \big(\partial_{r_i}\log J(r)\big)\,\partial_{r_i}\bar f(r). \label{eq:spd-radial-laplacian-expand}\tag{34}\] Using 31 , \[\log J(r)=\sum_{i<j}\log\sinh\!\Big(\frac{r_i-r_j}{2}\Big) \quad\Longrightarrow\quad \partial_{r_i}\log J(r) = \frac{1}{2}\sum_{j\neq i}\coth\!\Big(\frac{r_i-r_j}{2}\Big).\] Plugging this into 34 gives the explicit radial operator \[\Delta_{\mathrm{rad}}\bar f(r) = \sum_{i=1}^n \partial_{r_i}^2\bar f(r) + \frac{1}{2}\sum_{i\neq j}\coth\!\Big(\frac{r_i-r_j}{2}\Big)\,\partial_{r_i}\bar f(r). \label{eq:spd-radial-laplacian-final}\tag{35}\] Thus, the heat equation \(\partial_t p_t=\Delta p_t\) becomes \[\partial_t p_t(r) = \sum_{i=1}^n \partial_{r_i}^2 p_t(r) + \frac{1}{2}\sum_{i< j}\coth\!\Big(\frac{r_i-r_j}{2}\Big)\, \big(\partial_{r_i}-\partial_{r_j}\big)p_t(r). \label{eq:spd-heat-r}\tag{36}\] after grouping the symmetric terms.

9 Experiment Details↩︎

This section describes the experimental settings.

9.0.0.1 Architectures and optimization.

Unless noted otherwise, for PINN training we use the modified MLP of [6] with manifold-dependent embeddings. We implement the PINN in JAX [52], which enables computing higher-order derivatives while staying within practical GPU memory limits. For the denoiser, we use the same architecture as in the baseline experiments. All models are trained with the Adam optimizer and a cosine learning-rate schedule. The initial learning rate is tuned over \(\{3\cdot10^{-3}, 1\cdot10^{-3}, 3\cdot10^{-4}, 1\cdot10^{-4}\}\).

9.0.0.2 Hardware and training speed.

All models are trained on a single NVIDIA GeForce RTX 3090 Ti GPU. All experiments finish within one day, except training the diffusion model on QM9, which takes about 3 days.

9.1 Climate science datasets on \(S^2\)↩︎

For PINN training, we fix the initial point \(x_0 = [0,0,1]^T \in S^2\) and learn the heat kernel conditioned on \(x_0\). For an arbitrary \(x_0 \in S^2\), the heat kernel \(p_{t|0}(\cdot \mid x_0)\) is recovered via an isometry that maps \(x_0 \mapsto [0,0,1]^T\). The PINN is a 4-layer modified MLP with 256 hidden units and sinusoidal embeddings. We train the PINN on \([t_0, t_{\max}] = [0.1, 5]\), imposing the initial condition at \(t = t_0\). The denoiser is a 5-layer MLP with 512 hidden units.

9.2 Synthetic data on \(SO(3)\)↩︎

We view \(SO(3)\) as the quotient manifold \(S^3 / \{\pm I\}\), where \(S^3 \subset {\mathbb{R}}^4\) is the unit 3-sphere. The PINN and denoiser architectures are identical to 9.1, except for input/output dimensions.

9.3 Traffic analysis on \(SPD(10)\)↩︎

Motivated by the symmetric and antisymmetric terms in 24 , we introduce two permutation-invariant embeddings of \(r\): \[S_k(r) = \Big(\sum_i |r_i|^k\Big)^{1/k}, \quad A_k(r) = \Big(\sum_{i<j} |r_i-r_j|^k\Big)^{1/k},\] where \(k\) is sampled uniformly from \([1,5]\). Since \(p_{t|0}\) is symmetric in \(r = (r_1,\cdots,r_n)\), these embeddings provide a permutation-invariant representation. We combine them with a sinusoidal embedding of \(t\) and feed the resulting features into the modified MLP PINN. We train the PINN on \([t_0, t_{\max}] = [0.1, 1]\), since the hyperbolic geometry of \(SPD(n)\) can cause diffusion across scales even for small \(t\). For the denoiser, we use SPD-Net introduced in [19]. This architecture mimics the shape of U-Net [53] with operations on \(SPD(n)\).

9.4 Brain connectivity (EEG) analysis on \(SPD(n)\)↩︎

The PINN architecture is identical to 9.3. For the denoiser, we use a 4-layer MLP with 256 hidden units that takes the upper-triangular entries (including the diagonal) of an SPD matrix as input and outputs an upper-triangular matrix \(U\) (including the diagonal). We then map \(U \mapsto U^T U\) to ensure positive-definiteness.

9.5 Molecule generation on \({\mathbb{R}}^{k \times n} / S_n\)↩︎

Since the heat kernel is permutation-symmetric, we use permutation symmetry as an inductive bias for the PINN. For the initial point \(X_0 = [x_0^{(1)},\cdots,x_0^{(n)}]^{T}\) and a perturbed configuration \(X = [x^{(1)},\cdots,x^{(n)}]^{T}\), the PINN input is the matrix of relative distances \(R = (R_{ij})_{1 \leq i,j \leq n}\), with \(R_{ij} = |x^{(i)} - x_0^{(j)}|\). To encode permutation symmetry, we use a Graph Neural Network (GNN) [54], treating diagonal entries as node features and off-diagonal entries as edge features. Node and edge features are updated by MLPs with two hidden layers and 64 hidden units, and edge features are aggregated by summation. We stack two such GNN layers as the PINN.

For the denoiser, following the baseline [39], we use the EGNN architecture [55] with 9 layers and 256 hidden units.

9.6 Runtime analysis↩︎

We report the runtime overhead of the learned PINN heat-kernel surrogate. All timings are measured on a single NVIDIA GeForce RTX 3090 Ti GPU using the same implementation as in the main experiments. Since the PINN is implemented in JAX, all per-call timings are measured after JIT compilation. We evaluate the average wall-clock time for computing the learned log heat kernel \(\log p^{(\theta)}_{t|0}\), its conditional score \(\nabla\log p^{(\theta)}_{t|0}\), and one MCMC-based forward sampling step with batch size \(128\).

3pt

The cost of evaluating the learned log heat kernel and its score is small for all considered manifolds. The main additional overhead comes from MCMC-based forward sampling, especially on \(\mathrm{SPD}(10)\), where each proposal involves matrix operations under the affine-invariant geometry. The permutation-quotiented point-cloud experiment does not use the same MCMC forward-sampling procedure, so we omit its MCMC timing.

10 Statements and proofs↩︎

10.1 Stability under a zero PINN residual, with an approximate initializer↩︎

10.2 Stability under a nonzero PINN residual↩︎

2 analyzes the idealized setting where the approximate initial condition is evolved exactly under the heat equation. In practice, however, the learned PINN is only an approximate solver: it is trained to match the initial/boundary condition and to make the PDE residual small, but the residual is not identically zero. We therefore give a complementary stability statement showing that a uniformly small residual for the log heat equation leads to a controlled log-density error.

References↩︎