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

2026-05-29

AI Summary

This work addresses the challenge of training Riemannian diffusion models on general manifolds, which typically relies on intractable heat kernel sampling and differentiation. The authors propose a novel approach based on physics-informed neural networks (PINNs) to directly solve the heat equation on manifolds, yielding a differentiable approximation of the log heat kernel. This enables both forward noising and conditional score estimation without requiring closed-form heat kernels. Notably, the method is the first to employ PINNs for approximating heat kernels on arbitrarily defined explicit manifolds, overcoming the limitation of prior techniques that only apply to highly symmetric spaces. Experiments demonstrate effective diffusion-based generative modeling across diverse manifolds, including the sphere \(S^2\), the rotation group \(SO(3)\), the manifold of symmetric positive-definite matrices \(SPD(n)\), and permutation quotient point clouds.

Abstract

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.

Problem

• Riemannian diffusion models
• manifold heat kernel
• score-based generative modeling
• stochastic diffusion equations
• manifold-supported data

Innovation

• Riemannian diffusion models
• physics-informed neural networks
• heat kernel approximation
• manifold generative modeling
• score-based generative models