Riemannian Score-Based Generative Modelling

TL;DR

The paper introduces Riemannian Score-Based Generative Models (RSGMs), extending diffusion-based generative modeling from Euclidean spaces to arbitrary Riemannian manifolds. It develops an intrinsic forward noising process on manifolds, proves a time-reversal formula in the Riemannian setting, and implements approximate sampling via Geodesic Random Walks together with score estimation on manifolds using DSM/ISM losses. The approach is demonstrated on a range of manifolds, including spheres and SO(3), with theoretical convergence guarantees in the compact case and empirical results showing competitive or superior performance to baselines, plus favorable scalability to higher dimensions. Extensions to conditional sampling, Schrödinger bridges, and invariant distributions are discussed, highlighting broad applicability to domains with manifold-valued data such as earth and climate science, robotics, and physics.

Abstract

Score-based generative models (SGMs) are a powerful class of generative models that exhibit remarkable empirical performance. Score-based generative modelling (SGM) consists of a ``noising'' stage, whereby a diffusion is used to gradually add Gaussian noise to data, and a generative model, which entails a ``denoising'' process defined by approximating the time-reversal of the diffusion. Existing SGMs assume that data is supported on a Euclidean space, i.e. a manifold with flat geometry. In many domains such as robotics, geoscience or protein modelling, data is often naturally described by distributions living on Riemannian manifolds and current SGM techniques are not appropriate. We introduce here Riemannian Score-based Generative Models (RSGMs), a class of generative models extending SGMs to Riemannian manifolds. We demonstrate our approach on a variety of manifolds, and in particular with earth and climate science spherical data.

Figures (10)

• Figure 1: Geodesic Random Walks can be used to approximate Brownian motion and more generally SDEs on manifolds. (a) At each step, tangential noise is sampled (red), which is added the drift term (not pictured). This tangent vector is then pushed through the exponential map to produce a geodesics step on the manifold (blue). (b) Iterating this procedure yield approximate sample paths from the process.
• Figure 1: Slice of heat kernel $p_{t|0}(x_t|x_0)$ on $\mathbb{S}^2$ for different approximations.
• Figure 2: Trained score-based generative models on earth sciences data. The learned density is colored green-blue. Blue and red dots represent training and testing datapoints, respectively.
• Figure 2: Illustration of the effect of the corrector step on RSGM. The black line corresponds to the dynamics of the noising process $(p_t)_{t \in \left[0,T\right]}$. The blue dashed lines correspond to the predictor step (going backward in time) and the red dashed lines correspond to the corrector step (projecting back onto the initial dynamics). Note that $\mathcal{L}(\mathbf{Y}^s_\gamma) \approx p_{T- \gamma}$ and $\mathcal{L}(\mathbf{Y}^s_{2\gamma}) \approx p_{T-2\gamma}$.
• Figure 3: Comparison of Moser flows and RSGMs training speed and performance on the synthetic high-dimension torus task. Moser flows trained with $\lambda_{\min}=1$. We report two likelihoods, the 'Moser' closed form density---not guaranteed to be normalized---and the 'ODE' likelihood given by solving an augmented ODE (as in CNFs) with the vector field induced by the Moser flow density---which is guaranteed to have unit volume.

Theorems & Definitions (39)

• theorem 1
• definition 1
• proposition 1
• theorem 2
• definition 2
• proposition 2
• proposition 3
• definition 3
• theorem 3
• proposition 4