Sampling On Metric Graphs

Rajat Vadiraj Dwaraknath, Lexing Ying

Keywords: sampling, metric graphs, Langevin diffusion, Brownian motion, Euler-Maruyama, GPUs, CUDA, Monte Carlo, parallel algorithms

TL;DR: We present a parallelizable timestep splitting Euler-Maruyama algorithm for sampling on metric graphs with theoretical guarantees and a fast, custom CUDA kernel implementation.

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 a custom CUDA kernel that is up to ~8000x faster than a GPU implementation using PyTorch. We corroborate our theoretical results with numerical experiments applying our implementation to star metric graphs. In terms of accuracy and efficiency, our scheme significantly outperforms a baseline finite volume scheme.

Submission Number: 28