Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks
This work provides a general method for training Riemannian diffusion models on arbitrary manifolds, which is a significant step for researchers working with non-Euclidean data.
This paper addresses the challenge of training Riemannian diffusion models on general manifolds, where the manifold heat kernel is typically unknown. They propose using a physics-informed neural network (PINN) to approximate the heat kernel by solving the manifold heat equation, enabling both forward noising and conditional-score evaluation for denoising score matching on diverse manifolds like S2, SO(3), SPD(n), and permutation-quotiented point clouds.
Riemannian diffusion models generalize score-based generative modeling to manifold-supported data via stochastic diffusion equations on the manifold. However, training requires sampling from and differentiating the manifold heat kernel, which is rarely available in closed form beyond a few highly symmetric manifolds. We propose a general approach that approximates the heat kernel by directly solving the manifold heat equation with a physics-informed neural network (PINN). Given an explicit manifold specification, we choose a coordinate system, derive the corresponding heat (Fokker--Planck) equation and a short-time asymptotic approximation, and then train a PINN to learn the log heat kernel. The resulting surrogate enables both forward noising (heat-kernel sampling) and conditional-score evaluation for denoising score matching. We demonstrate the method on diverse manifolds including $S^2$, $SO(3)$, $\mathrm{SPD}(n)$, and permutation-quotiented point clouds.