Infinite-dimensional Diffusion Bridge Simulation via Operator Learning

TL;DR

The paper addresses the challenge of simulating diffusion bridges in infinite-dimensional function spaces under nonlinear conditioning. It combines score-matching ideas with operator learning to directly learn the time-reversed infinite-dimensional bridge, providing a tractable objective and a time-dependent neural operator (continuous-time U-shaped FNO) that handles resolution-invariant data. The key contributions are the derivation of the infinite-dimensional time-reversed bridge, a KL-divergence-based training objective via a parametric operator, and a practical sampling pipeline using finite truncations; validation on functional Brownian bridges and stochastic shape data demonstrates high fidelity and zero-shot generalization across discretizations. This framework enables efficient, high-resolution diffusion-bridge simulation for continuous data types such as shapes and images, with potential impact on stochastic shape analysis, phylogenetic simulations, and other infinite-dimensional stochastic modeling tasks.

Abstract

The diffusion bridge, which is a diffusion process conditioned on hitting a specific state within a finite period, has found broad applications in various scientific and engineering fields. However, simulating diffusion bridges for modeling natural data can be challenging due to both the intractability of the drift term and continuous representations of the data. Although several methods are available to simulate finite-dimensional diffusion bridges, infinite-dimensional cases remain under explored. This paper presents a method that merges score matching techniques with operator learning, enabling a direct approach to learn the infinite-dimensional bridge and achieving a discretization equivariant bridge simulation. We conduct a series of experiments, ranging from synthetic examples with closed-form solutions to the stochastic nonlinear evolution of real-world biological shape data. Our method demonstrates high efficacy, particularly due to its ability to adapt to any resolution without extra training.

Figures (11)

• Figure 1: Continuous-time U-shaped FNO architecture, acting as a parameterized operator $\mathcal{G}^{(\theta)}:(x, t)\mapsto y$ for $x,y\in \mathcal{H}$. Both $x, y$ are evaluated on a discrete spatial-temporal grid $(\xi, t)$ with $\xi\in\mathbb{R}^{m^d\times d}$ and $t\in\mathbb{R}^n$. The figure only shows the case when $d=1$, but the application to $d>1$ can be achieved through similar architecture, e.g. $x(\xi, t)\in\mathbb{R}^{m\times m\times n\times d_x}$ for $d=2$ and so on.
• Figure 2: Qualitative results for Brownian bridges between two quadratic functions; (a) One sample from the learned reversed Brownian bridge, evaluated at 128 evenly distributed points; (b) One sample from the true reversed Brownian bridge, simulated with the same random seed as (a) for comparison.
• Figure 3: Quantive evaluation of the model on different levels of discretization; (a) RMSE between the model's output and the ground true drift term evaluated on the whole trajectory; (b) RMSE between the end of estimated trajectories and the true target shape; All the statistics are done for 64 independent samplings.
• Figure 4: Visualization of Brownian bridges between ellipse shapes evaluated under different levels of discretizations, the training is done on 16 points. The blue shade represents 64 independent samples of end shape, and only one colored sample of the trajectories is shown with different colors indicating time steps.
• Figure 5: Visualization of the Brownian bridges between two nested spheres with different radii. The top row is the estimated bridge, and the bottom row is the ground true path with the same seed. The intermediate shapes are plotted in light orange. The training is done on $16\times 16$ grid size and evaluated on $48\times 48$ grid size.

Theorems & Definitions (7)

• Theorem 1
• Theorem 2
• Lemma 1
• Theorem 3
• proof
• proof
• proof