🤖 machine learning

Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

This paper proposes a general framework for Riemannian diffusion models that overcomes the challenge of unavailable closed-form heat kernels by using physics-informed neural networks to solve the manifold heat equation, thereby enabling effective score-based generative modeling on diverse manifolds such as $S^2$ , $SO(3)$ , and $\mathrm{SPD}(n)$ .

Original authors: Gyeonghoon Ko, Juho Lee

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Gyeonghoon Ko, Juho Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching AI to Walk on Curved Surfaces

Imagine you are teaching a robot to draw pictures. In the world of standard AI (called "Euclidean diffusion"), the robot lives on a flat, infinite sheet of paper. It learns to turn random scribbles into clear images by slowly "denoising" them. This works great for flat things like photos of cats or faces.

But what if the data isn't flat? What if the data lives on a sphere (like earthquake locations on Earth), a twisted loop (like 3D rotations of a robot arm), or a complex shape (like the structure of a molecule)?

If you try to force the robot to draw on a sphere using flat-paper rules, it gets confused. It might stretch the sphere, tear it, or create impossible shapes. To do this right, the robot needs to understand the curved geometry of the surface it's walking on.

The Problem: The "Heat Map" is Missing

To teach the robot how to move on these curved surfaces, the researchers need a special map called a Heat Kernel.

Think of the Heat Kernel as a "smoke spreader." If you drop a single drop of ink (or a puff of smoke) on a specific spot on a curved surface, the Heat Kernel tells you exactly how that smoke will spread out over time.

Why do we need it? The AI uses this "smoke spread" to figure out how to reverse the process: starting with a cloud of smoke and turning it back into a single drop of ink (the original data).
The Catch: For simple shapes (like a flat plane), we have a perfect formula for this smoke spread. But for complex, curved shapes (like the surface of a sphere or a twisted donut), we don't have a formula. The math is too hard. Previous methods tried to guess the formula using shortcuts, but those shortcuts only worked for very specific, simple shapes.

The Solution: The "Physics Detective" (PINN)

The authors propose a clever new way to solve this. Instead of trying to find a magic formula for the smoke spread, they treat it like a physics puzzle.

They know the rules of how smoke spreads (a famous physics equation called the Heat Equation). They also know what the smoke looks like at the very beginning (a tiny, concentrated dot).

So, they use a Physics-Informed Neural Network (PINN). Think of this as a super-smart detective that learns by solving the puzzle directly:

The Setup: They tell the detective, "Here is the rulebook (the Heat Equation). Here is the starting point (a tiny dot of smoke)."
The Training: The detective tries to guess what the smoke looks like at every moment in between. If the guess breaks the rules of physics, the detective gets a "penalty" and tries again.
The Result: Eventually, the detective learns to predict exactly how the smoke spreads on any curved surface, even if no one has ever written down the formula for it before.

How It Works in Practice

The paper shows this "detective" working on several different "worlds":

The Earth (Sphere $S^2$ ): They modeled the spread of natural disasters like volcanoes and earthquakes. The AI learned to generate realistic maps of where these events might happen, respecting the round shape of the Earth.
The 3D Rotator ($SO(3)$): They modeled how objects rotate in 3D space. This is crucial for robotics and 3D modeling.
The Traffic Flow ($SPD$ matrices): They modeled traffic patterns in New York City. The data here isn't just points; it's complex relationships between different city zones. The AI learned to generate realistic traffic flow patterns.
The Molecular Puzzle ( $R^{k \times n}/S_n$ ): This is the most complex part. Molecules are made of atoms, but the order you list the atoms doesn't matter (a molecule is the same whether you list atom A then B, or B then A). The AI learned to generate new molecules by treating them as a "cloud of points" where the order doesn't matter, effectively solving a permutation puzzle.

The Results: Does It Work?

The researchers tested their "detective" against other methods:

Accuracy: On shapes where we do know the answer (like the sphere), their method was just as accurate as the best existing methods, but it didn't need a pre-written formula.
Versatility: It worked on shapes where no other method could work because the math was too hard.
Quality: The generated data (molecules, traffic maps, earthquake locations) looked realistic and followed the correct geometric rules.

The Catch (Limitations)

The paper is honest about the downsides:

It's Heavy: Training this "detective" takes a lot of computer power and time. It's like solving a massive physics simulation every time you want to generate data.
It's Complicated: Setting up the "rulebook" for a new shape requires some math expertise. It's not yet a "plug-and-play" button you can click for any shape.
Trade-offs: In the molecule experiment, the AI made very valid molecules (chemically correct) but they weren't as "new" or unique as some other methods. It's a trade-off between being safe/correct and being creative.

Summary

In short, this paper gives AI a new tool to walk on curved surfaces. Instead of needing a pre-written map (which doesn't exist for most shapes), the AI learns the map on the fly by solving the physics equations of how things spread. This allows AI to generate realistic data for everything from global climate events to 3D molecules, without breaking the rules of geometry.

Technical Summary: Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

1. Problem Statement

Riemannian diffusion models generalize score-based generative modeling to data residing on smooth manifolds by defining forward diffusion processes via Riemannian Brownian motion. A critical bottleneck in training and sampling these models is the requirement to access the manifold heat kernel (the transition density of the Brownian motion) and its gradients. Specifically, the training objective (Denoising Score Matching) requires evaluating the conditional score $\nabla_x \log p_{t|0}(x|x_0)$ , and forward sampling requires drawing from the heat kernel.

While the heat kernel is known in closed form for highly symmetric manifolds (e.g., Euclidean space, spheres, tori), it is rarely available analytically for general Riemannian manifolds. Existing approximations, such as short-time asymptotics (Varadhan's formula) or spectral expansions, are often restricted to specific regimes (small time $t$ ) or require strong structural assumptions (e.g., compact Lie groups), limiting their practical applicability to diverse geometric settings.

2. Methodology

The authors propose a general framework that approximates the heat kernel by directly solving the manifold heat equation using Physics-Informed Neural Networks (PINNs). The method proceeds through four main stages:

A. Coordinate System Selection and Embedding
Instead of patching local coordinate charts, the method embeds the target manifold $M$ into a larger ambient manifold $\tilde{M}$ (typically $\mathbb{R}^N$ ) as a submanifold or quotient manifold.

Submanifolds: Defined by constraints $f_i(x)=0$ . The metric and differential operators are derived using a projection matrix $P(x)$ that maps the ambient space to the tangent space of $M$ .
Quotient Manifolds: Defined by identifying points under a group action (e.g., permutation symmetry). The method treats orbit points as identical states, often using canonical representatives.

B. Derivation of the Heat Equation
Using the ambient coordinates and the projection matrix, the authors derive the explicit form of the Laplace-Beltrami operator $\Delta_M$ and the corresponding Fokker-Planck (heat) equation:
$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$
This equation is expressed entirely in terms of the ambient metric and the projection matrix, allowing for coordinate-based computation without requiring intrinsic coordinate charts.

C. Initial Condition Approximation
Since the Dirac delta initial condition $p_{0|0} = \delta_{x_0}$ cannot be imposed numerically, the authors initialize the PINN at a small time $t_0 > 0$ using a principled short-time asymptotic approximation. They utilize the parametrix method (Minakshisundaram–Pleijel recursion) to obtain a higher-order expansion of the heat kernel:
$p_{t_0}(x|x_0) \approx G_t^{(n)}(x, x_0) = G_t(x, x_0) \left( u_0 + u_1 t + \dots + u_{n-1} t^{n-1} \right)$
where $G_t$ is the Gaussian-like term and $u_k$ are coefficients derived from the geometry (e.g., volume distortion and curvature).

D. PINN Training
A neural network $\phi_\theta(t, x)$ is trained to approximate the log-heat kernel $\log p_{t|0}$ . The training objective minimizes a combination of:

Initial/Boundary Loss: The difference between the network output and the asymptotic initializer at $t=t_0$ .
PDE Residual Loss: The residual of the heat equation $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$ (rewritten for the log-density).

The authors employ adaptive loss weighting and curriculum learning (progressively expanding the time interval) to ensure stability.

E. Integration into Diffusion Models
The trained PINN serves as a surrogate for the heat kernel, enabling two key operations:

Forward Sampling: Samples are generated via Markov Chain Monte Carlo (MCMC) methods (e.g., geodesic random walks) using the PINN-evaluated density.
Score Evaluation: The conditional score $\nabla_x \log p_{t|0}$ is computed by differentiating the PINN, which is then used to train the reverse-time diffusion model via Denoising Score Matching (DSM).

3. Key Contributions

General PDE-Based Surrogate: The paper introduces a method to approximate heat kernels on arbitrary manifolds by solving the heat equation numerically, bypassing the need for closed-form solutions or restrictive spectral assumptions.
Coordinate-Free Derivation: The framework provides a systematic way to derive the heat equation for manifolds defined as submanifolds or quotients of Euclidean spaces, making it applicable to complex geometries like $S^2$ , $SO(3)$, $SPD(n)$, and permutation-quotiented point clouds.
Stability Analysis: The authors provide a theoretical justification (Theorem 4.1) showing that solving the heat equation from an approximate initial condition yields a solution with bounded error for all $t \ge t_0$ , ensuring the reliability of the PINN surrogate.
Permutation-Invariant Diffusion: The method is applied to the quotient manifold $\mathbb{R}^{k \times n}/S_n$ , enabling diffusion processes that explicitly mod out atom permutations for molecule generation, a setting where standard Euclidean diffusion fails to respect structural symmetries.

4. Experimental Results

The method was evaluated on diverse manifolds and datasets:

Spherical Data ( $S^2$ ): Density estimation of geophysical events (volcanoes, earthquakes, floods, wildfires). The proposed method achieved competitive or superior test log-likelihoods compared to Riemannian Score-based Generative Models (RSGM) which rely on Varadhan approximations and spectral methods.
Rotation Data ($SO(3)$): Synthetic mixture modeling showed improved log-likelihoods over RSGM baselines.
Positive Definite Matrices ($SPD(n)$):
- Traffic Analysis ($SPD(10)$): Conditional generation of traffic flow matrices from NYC taxi data. The method reduced the Frobenius distance to ground truth compared to SPD-DDPM.
- Brain Connectivity ($SPD(n)$): Conditional generation of EEG connectivity matrices. The method achieved strong precision and F1 scores, competing with DiffeoCFM (which uses diffeomorphic mappings to Euclidean space) and outperforming LowTriDDPM.
Molecule Generation ( $\mathbb{R}^{k \times n}/S_n$ ): Generation of molecules on the permutation-quotiented space. The method produced molecules with higher validity and stability than a variance-exploding (VE) EDM baseline, though with lower novelty, reflecting a known validity-novelty trade-off.

Validation: The authors validated the PINN surrogates by comparing learned log-densities and scores against reference approximations (where available) and by reporting normalized PDE residuals and boundary errors, confirming the surrogates satisfy the heat equation.

5. Significance and Claims

The paper claims that PDE-based heat-kernel approximation offers a practical and general alternative to existing methods for Riemannian diffusion. By treating the heat kernel as the solution to a known PDE, the approach extends drift-free Riemannian diffusion to manifolds where analytic kernels are unavailable.

The authors position this work as a "solve-the-heat-equation" viewpoint that broadens the applicability of manifold diffusion models beyond simple geometries. However, they modestly acknowledge limitations:

Computational Cost: PINN training requires significant resources due to high-order autodifferentiation and collocation evaluations.
Stability: Training can be unstable in high dimensions or for complex geometries, potentially requiring stronger inductive biases.
Complexity: The pipeline is more complex than alternatives like Flow Matching, which bypass heat kernel estimation entirely.

The work does not claim to replace flow matching but rather to provide a necessary backend for score-based models on manifolds where flow matching's reliance on geodesic interpolation or exponential maps might be insufficient or where the specific structure of the heat kernel is required.