Generative Modeling on Metric Graphs via Neural Optimal Transport
2026-06-15 • Machine Learning
Machine Learning
AI summaryⓘ
The authors present a new method that uses deep learning to create probability models on networks (graphs) where points are connected by edges with distances. Their approach embeds these graphs into smooth spaces, solves an optimal transport problem using neural networks, and maps samples back onto the graph itself. They study two ways to embed the graphs and prove their model converges to a valid solution as it gets more expressive. They show their method works well on different types of graphs and is efficient enough to handle large-scale data like Uber pickups in Manhattan.
deep generative modelingmetric graphsoptimal transportentropic Kantorovich problemneural semidual parameterizationembeddingtropical Abel-Jacobi embeddingweak convergenceurban mobility data
Authors
Alessandro Micheli, Yueqi Cao, Anthea Monod, Samir Bhatt
Abstract
We introduce, to our knowledge, the first deep generative modeling framework for probability distributions continuously supported on compact metric graphs. Given source and target measures on a metric graph, our method embeds the graph into a smooth ambient space, solves an entropic Kantorovich problem via a neural semidual parameterization, and projects generated samples back onto the original graph. We study two embedded geometries: an extrinsic Euclidean realization and the intrinsic tropical Abel--Jacobi embedding into the Jacobian torus. In both cases, the resulting generator is graph-supported by construction. We prove that, in the joint limit of increasing neural expressivity, the learned generator converges weakly to a valid transport coupling between the original graph measures. Empirically, across a range of geometrically distinct graphs, our method matches or improves upon heuristic transport baselines based on discrete graph OT, while scaling more favorably. Finally, we demonstrate scalability on real-world urban mobility data by training our model on one million Uber pickup locations in Manhattan, New York City.