Manifold Diffusion Fields

TL;DR

MDF presents a diffusion model over functions defined on general Riemannian manifolds by leveraging an intrinsic coordinate system built from the Laplace-Beltrami eigenfunctions. It encodes fields as coordinate-signal pairs and performs diffusion in function space with a score network operating on context/query pairs via a PerceiverIO backbone, achieving robustness to rigid/isometric transformations and cross-manifold generalization. Empirical results across synthetic and scientific domains (climate on spheres, PDEs on surfaces, and molecular conformers) show improved diversity and fidelity over Euclidean baselines and related manifold methods, with competitive performance even at low spectral dimensionality. The approach highlights the practicality of spectral-geometric priors for scientific problems, enabling diffusion-based priors on complex geometries and multiple manifolds. Potential impact includes improved forward/inverse PDE solving, climate modeling, and molecular design, with avenues for faster inference and broader geometric contexts in future work.

Abstract

We present Manifold Diffusion Fields (MDF), an approach that unlocks learning of diffusion models of data in general non-Euclidean geometries. Leveraging insights from spectral geometry analysis, we define an intrinsic coordinate system on the manifold via the eigen-functions of the Laplace-Beltrami Operator. MDF represents functions using an explicit parametrization formed by a set of multiple input-output pairs. Our approach allows to sample continuous functions on manifolds and is invariant with respect to rigid and isometric transformations of the manifold. In addition, we show that MDF generalizes to the case where the training set contains functions on different manifolds. Empirical results on multiple datasets and manifolds including challenging scientific problems like weather prediction or molecular conformation show that MDF can capture distributions of such functions with better diversity and fidelity than previous approaches.

Figures (25)

• Figure 1: MDF learns a distribution over a collection of fields $f: \mathcal{M} \rightarrow \mathbb{R}^d$, where each field is defined on a manifold $\mathcal{M}$. We show real samples and MDF's generations on different datasets of fields defined on different manifolds. First row: MNIST digits on the sine wave manifold. Second row Middle: ERA5 climate dataset era5 on the 2D sphere. Third row: GMM dataset on the bunny manifold. Fourth row: molecular conformations in GEOM-QM9 qm9_1 given the molecular graph.
• Figure 1: COV and MMD metrics for different datasets on the wave manifold (mean curvature $|K|=0.004$).
• Figure 2: (a) Generative models of fields in Euclidean space dpfgaspfunctagem learn a distribution $p_\theta$ over functions whose domain is $\mathbb{R}^n$. We show an example where each function is the result of evaluating a Gaussian mixture with 3 random components in 2D. (b) MDF learns a distribution $p_\theta$ from a collection of fields whose domain is a general Riemannian manifold$f \sim q(f) | f: \mathcal{M} \rightarrow \mathcal{Y}$. Similarly, as an illustrative example each function is the result of evaluating a Gaussian mixture with 3 random components on $\mathcal{M}$ (i.e. the Stanford bunny). (c) Riemannian generative models riemannian_1riemannian_2riemannian_3riemannian_4 learn a parametric distribution $p_\theta$ from an empirical observations ${\bm{x}} \sim q({\bm{x}}) | {\bm{x}} \in \mathcal{M}$ of points ${\bm{x}}$ on a Riemannian manifold $\mathcal{M}$, denoted by black dots on the manifold.
• Figure 3: Left: Fourier PE of a point ${\bm{x}}$ in 2D Euclidean space. Generative models of functions in ambient space dpfgaspfunctagem use this representation to encode a function's input. Right: MDF uses the eigen-functions $\varphi_i$ of the Laplace-Beltrami Operator (LBO) $\Delta_{\mathcal{M}}$ evaluated at a point ${\bm{x}} \in \mathcal{M}$.
• Figure 4: Training