Score matching for sub-Riemannian bridge sampling

TL;DR

This work addresses the challenge of simulating diffusion bridges on sub-Riemannian manifolds, where hypoellipticity prevents explicit score computations. It develops a neural score model $S^\theta$ trained via sub-Riemannian score matching using divergence and denoising losses and leverages time-reversed bridge dynamics to enable conditioning. The authors derive local-coordinate descriptions, short-time approximations, and specialized procedures for the Heisenberg group, including adapted coordinates and stochastic Taylor expansions. Experiments on the Heisenberg group illustrate successful bridge sampling and concentration effects for small conditioning times, highlighting the approach's potential for inference and geometric statistics in hypoelliptic settings.

Abstract

Simulation of conditioned diffusion processes is an essential tool in inference for stochastic processes, data imputation, generative modelling, and geometric statistics. Whilst simulating diffusion bridge processes is already difficult on Euclidean spaces, when considering diffusion processes on Riemannian manifolds the geometry brings in further complications. In even higher generality, advancing from Riemannian to sub-Riemannian geometries introduces hypoellipticity, and the possibility of finding appropriate explicit approximations for the score of the diffusion process is removed. We handle these challenges and construct a method for bridge simulation on sub-Riemannian manifolds by demonstrating how recent progress in machine learning can be modified to allow for training of score approximators on sub-Riemannian manifolds. Since gradients dependent on the horizontal distribution, we generalise the usual notion of denoising loss to work with non-holonomic frames using a stochastic Taylor expansion, and we demonstrate the resulting scheme both explicitly on the Heisenberg group and more generally using adapted coordinates. We perform numerical experiments exemplifying samples from the bridge process on the Heisenberg group and the concentration of this process for small time.

Figures (4)

• Figure 1: In a sub-Riemannian geometry, the increments of a horizontal stochastic process $(X_t)_{t\geq 0}$ lie in the horizontal bundle $E$. When approximating the score for short time intervals by taking derivatives of the density of the steps, a common approximation takes the increment at the $i$th step to be normally distributed in the distribution $E_{X_{t_i}}$ at $X_{t_i}$ (sketched dark grey). However, this distribution is not differentiable with horizontal derivatives in the distribution $E_{X_{t_{i+1}}}$ at step $X_{t_{i+1}}$. The hypoellipticity of $(X_t)_{t\geq 0}$ implies that the distribution has a vertical component in $V_{X_{t_i}}$ (sketched light grey). We rely on Taylor expansion for the stochastic integral to approximate this component, which we then exploit to derive a sub-Riemannian denoising loss for use in the training of neural network score approximators.
• Figure 2: Left: Sample path from bridge starting at $(0.5,0,0.8)$ conditioned on hitting $(0,0,0)$ at $T=0.1$ (blue curve). Geodesic between the same points corresponding to the limit $T\to 0$ (red curve). Right: Norm of $(x,y)$-component (red) and $z$-component (blue) for the sample path as a function of $t\in[0,0.1]$.
• Figure 3: Mean $(x,y)$-component and $z$-component norms over 100 samples with quartiles for $T=0.1$, $T=0.2$, $T=0.5$ and $T=1$, respectively.
• Figure 4: Left: Estimated score field for $x,y,z\in[-1,1]$. Right: Estimated score for $x,y\in[-1,1]$ and $z=0.2$. Top row: Training with denoising loss \ref{['eq:denoising_loss_Heisenberg']}. Bottom row: Training with divergence loss \ref{['eq:loss_samples']}. Arrow lengths scaled for visualisation, and arrows colored according to lengths.

Theorems & Definitions (16)

• Remark 2.1: Choice of measure
• Remark 2.2
• Lemma 3.1
• proof
• Lemma 3.2
• proof
• Corollary 3.3
• proof
• Remark 3.4: Generative models
• Theorem 4.1