Sampling on Metric Graphs

TL;DR

This work develops the first practical algorithm for simulating Brownian motions on metric graphs by a timestep-splitting Euler-Maruyama scheme, and extends it to Langevin diffusion for sampling on graphs. It provides theoretical guarantees on finiteness of vertex crossings and convergence of exit probabilities to vertex-edge jump probabilities as the timestep vanishes. A fast, memory-aware CUDA implementation enables large-scale, GPU-accelerated simulations, achieving substantial speedups over baseline finite-volume methods on both toy star graphs and a real cortical vascular network. The results demonstrate improved stability at larger timesteps and significant computational efficiency, highlighting the method's potential for real-world diffusion and transport problems on complex networks.

Abstract

Metric graphs are structures obtained by associating edges in a standard graph with segments of the real line and gluing these segments at the vertices of the graph. The resulting structure has a natural metric that allows for the study of differential operators and stochastic processes on the graph. Brownian motions in these domains have been extensively studied theoretically using their generators. However, less work has been done on practical algorithms for simulating these processes. We introduce the first algorithm for simulating Brownian motions on metric graphs through a timestep splitting Euler-Maruyama-based discretization of their corresponding stochastic differential equation. By applying this scheme to Langevin diffusions on metric graphs, we also obtain the first algorithm for sampling on metric graphs. We provide theoretical guarantees on the number of timestep splittings required for the algorithm to converge to the underlying stochastic process. We also show that the exit probabilities of the simulated particle converge to the vertex-edge jump probabilities of the underlying stochastic differential equation as the timestep goes to zero. Finally, since this method is highly parallelizable, we provide fast, memory-aware implementations of our algorithm in the form of custom CUDA kernels that are up to ~8000x faster than a GPU implementation using PyTorch on simple star metric graphs. Beyond simple star graphs, we benchmark our algorithm on a real cortical vascular network extracted from a DuMuX tissue-perfusion model for tracer transport. Our algorithm is able to run stable simulations with timesteps significantly larger than the stable limit of the finite volume method used in DuMuX while also achieving speedups of up to ~1500x.

Figures (6)

• Figure 1: An example metric graph $\mathbf{\Gamma}$ and its associated spaces.
• Figure 2: Visual depiction of the inward derivative $\partial_{e_1}f\left( v \right)$ along an edge $e_1$ at a vertex $v$. Its sign is independent of the orientation of the edge.
• Figure 3: Error in density estimation for linear and quadratic potentials. The FVM scheme directly solves the Fokker-Planck equation to obtain the steady-state density. We compare the best case error (over discretization parameters) of this scheme with the error obtained by running \ref{['alg:euler_maruyama_metric_graph_bothsplit']} for multiple particle counts and values of the timestep. We estimate the density using a simple histogram with a bin size equal to the discretization of the FVM scheme. The error is computed as the empirical L2 distance between the estimated density and the analytical steady-state density. We observe that \ref{['alg:euler_maruyama_metric_graph_bothsplit']} results in significantly lower error compared to the FVM scheme for the same level of spatial and temporal discretizations.
• Figure 4: Normalized runtimes per step, aggregated over different discretization parameters for \ref{['alg:euler_maruyama_metric_graph_bothsplit']} for linear and quadratic potentials compared with the best runtime for the FVM scheme. We observe that the FVM scheme has a significantly higher runtime compared to \ref{['alg:euler_maruyama_metric_graph_bothsplit']} for the same level of spatial and temporal discretizations. Additionally, our custom CUDA kernel for \ref{['alg:euler_maruyama_metric_graph_bothsplit']} is significantly faster (up to $\sim$8000x speedup) than the PyTorch implementation (speedups indicated on the bars). We observe slightly higher runtimes for the linear potential, which is expected due to the increased likelihood of vertex crossings per timestep. All experiments were run on an NVIDIA RTX A6000 GPU.
• Figure 5: Simulation of tracer transport as advection--diffusion on a real cortical vascular network, with runtime scaling (left) and steady state tracer density visualization on the network (right). We observe that \ref{['alg:euler_maruyama_general']} achieves up to $\sim$1500x speedup over the DuMuX FVM scheme, while remaining stable beyond the FVM stability limit for large timesteps where FVM fails to converge. This demonstrates the practical utility of our algorithm on real-world metric graphs. These experiments were run on an NVIDIA RTX 4090 GPU.

Theorems & Definitions (11)

• Definition 1: Metric Graph
• Definition 2: Star Metric Graph
• Definition 3
• Definition 4: Standard Boundary Conditions
• Theorem 1: Lemma 2.2 and Corollary 2.4 in freidlin_diffusion_2000
• Theorem 2: Finite vertex crossings with high probability
• Theorem 3: Number of crossings is 1 with high probability
• Corollary 1: Jump probabilities converge to $b_v$
• proof : Proof of \ref{['thm:finite_bounces_high_prob']}
• proof : Proof of \ref{['thm:em_one_crossing']}