📊 statistics

Generative Modeling on Metric Graphs via Neural Optimal Transport

This paper introduces the first deep generative modeling framework for probability distributions on compact metric graphs by embedding them into smooth ambient spaces and solving entropic optimal transport via neural semiduals, thereby enabling scalable, graph-supported sample generation that converges to valid transport couplings and outperforms discrete baselines on real-world urban mobility data.

Original authors: Alessandro Micheli, Yueqi Cao, Anthea Monod, Samir Bhatt

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alessandro Micheli, Yueqi Cao, Anthea Monod, Samir Bhatt

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to move a crowd of people from one location to another, but there's a catch: everyone must stay strictly on a specific network of roads. They cannot walk through parks, fly over buildings, or cut across fields. They must stay on the asphalt.

This is the problem the paper solves. The authors have created a new "smart traffic system" (a deep generative model) that learns how to move people (or data) along complex road networks (metric graphs) without ever letting them step off the pavement.

Here is how they did it, explained through simple analogies:

1. The Problem: The "Off-Road" Trap

Usually, when computers try to move data from Point A to Point B, they treat the world like a smooth, open field (Euclidean space). They draw a straight line between two points.

The Issue: If your data lives on a winding river or a city street grid, a straight line might cut through a building or a lake.
The Old Way: To fix this, people used to chop the roads into tiny dots (discretization) and solve the problem dot-by-dot. But this is like trying to navigate a city by only looking at a pixelated map; it gets messy and slow if you zoom in too much.

2. The Solution: The "Magic Slide"

The authors' method uses a clever three-step trick to keep the data on the road while doing the hard math in the open air.

Step 1: The Magic Slide (Embedding)
Imagine you take your winding, 2D road map and slide it onto a giant, smooth, 3D trampoline (a smooth ambient space).

The roads are still there, but now they exist inside a smooth, open world where standard math works easily.
The authors use two types of slides:
- The "Map" Slide: Just taking the road coordinates as they are (like a GPS map).
- The "Tropical" Slide: A special mathematical transformation that turns the road network into a donut shape (a torus). This is great because it captures the "loops" and "cycles" of the road network (like a roundabout) perfectly, regardless of how the roads are drawn on a map.

Step 2: The "Ghost" Transport (Neural Optimal Transport)
Now that the roads are on the smooth trampoline, the computer uses a "neural network" (a type of AI) to figure out the best way to move the crowd from the start to the finish.

Instead of forcing the crowd to stay on the road during the calculation, the AI calculates a "ghost" path through the open air.
It uses a technique called Entropic Optimal Transport. Think of this as the AI learning a "heat map" of how people should flow. It doesn't just move one person; it learns the general flow pattern of the whole crowd.

Step 3: The "Snap-Back" (Projection)
Here is the magic part. The AI generates a sample point in the open air (the trampoline). This point might be floating in the sky, off the road.

The system immediately snaps that point back to the nearest spot on the road.
It then translates that road spot back into the original map coordinates.
Result: You get a generated sample that is mathematically optimal but physically guaranteed to be on the road.

3. Why This is a Big Deal

The paper claims this is the first time this has been done for continuous road networks.

Better than the old way: They tested their method against older "dot-based" methods. Their new method was more accurate (less error) and didn't slow down as the maps got bigger.
Real-world proof: They tested this on 1 million real Uber pickup locations in Manhattan.
- Imagine trying to move 1 million people from a uniform distribution (people everywhere) to the actual hotspots where people actually get into cabs.
- Their model successfully learned the complex, winding geometry of Manhattan's streets and generated realistic pickup locations that respected the actual road layout, whereas older methods struggled or failed at that scale.

The Bottom Line

The authors built a system that learns how to move things along a network (like roads or blood vessels) by:

Lifting the network into a smooth, open space to do the math.
Using AI to learn the flow.
Snapping the results back onto the network so nothing ever leaves the "road."

They proved mathematically that as the AI gets smarter, this "snap-back" method becomes a perfect way to transport data along the network, and they showed it works on massive, real-world city maps.

Technical Summary: Generative Modeling on Metric Graphs via Neural Optimal Transport

Problem Statement

Many scientific and geometric datasets, such as road networks, vascular structures, and river systems, are naturally supported on metric graphs. A metric graph is a one-dimensional geodesic space where edges carry lengths and are treated as continuous intervals. Existing generative modeling approaches face a fundamental limitation: standard methods generate points in an ambient Euclidean space, failing to ensure samples remain supported on the graph or respect its intrinsic metric geometry.

Conversely, traditional Optimal Transport (OT) approaches for graphs often rely on discretization, approximating continuous distributions with finite node-based measures. While computationally accessible, this approach introduces mesh-dependent approximations and scales poorly as discretization is refined, losing the continuous nature of the original domain.

The paper addresses the need for a continuous, graph-supported generative framework that can leverage smooth ambient structures for computation without sacrificing the geometric fidelity of the metric graph.

Methodology

The authors propose a deep generative modeling framework that embeds the metric graph into a smooth ambient space, solves an entropic OT problem there, and projects samples back to the graph. The pipeline consists of three main stages:

Embedding: The metric graph $\Gamma$ is embedded into a smooth ambient manifold $M_\Psi$ via a map $\Psi$ . The paper investigates two specific embeddings:
- Extrinsic Euclidean Realization ( $\iota$ ): Maps the graph to $\mathbb{R}^m$ (e.g., planar coordinates of a road network).
- Intrinsic Tropical Abel–Jacobi Embedding ( $\Phi_p$ ): Maps the graph into its Tropical Jacobian torus $J(\Gamma)$ . This embedding depends only on edge lengths and cycle structure, providing an intrinsic geometry independent of a specific Euclidean drawing.
Neural Entropic Optimal Transport:
- The source ( $P_\Gamma$ ) and target ( $Q_\Gamma$ ) measures are pushed forward to the ambient space ( $P_\Psi, Q_\Psi$ ).
- An entropic Kantorovich problem is solved between these embedded measures using a neural semidual parameterization.
- Unlike map-based Monge formulations, this approach learns a coupling (plan). The target-side Schrödinger potential is parameterized by a neural network $g_\theta$ .
- To generate continuous outputs rather than samples restricted to target atoms, the learned conditional law is smoothed using the ambient heat kernel. This produces samples in the ambient space that may not lie exactly on the embedded graph.
Recovery (Projection–Pullback):
- Since heat-smoothed samples may fall off the embedded graph, a nearest-point projection onto the embedded graph $\Psi(\Gamma)$ is performed.
- The inverse embedding $\Psi^{-1}$ is then applied to map the projected point back to the original metric graph $\Gamma$ .
- The final generator is constructed as $\rho_\Gamma^\Psi = \Psi^{-1} \circ \text{proj}_{\Psi(\Gamma)}$ , ensuring all generated samples are strictly supported on the original graph.

Key Contributions

First Deep Generative Framework for Metric Graphs: The paper introduces, to the authors' knowledge, the first deep generative model for probability distributions continuously supported on compact metric graphs.
Dual Embedding Geometries: The framework is developed for two distinct transport geometries:
- An extrinsic Euclidean realization.
- The intrinsic tropical Abel–Jacobi embedding into the Jacobian torus.
Theoretical Convergence: The authors prove that in the joint limit of increasing neural expressivity and vanishing heat scale, the learned heat-smoothed projection–pullback generator converges weakly to a valid transport coupling between the original graph measures. This guarantees that despite learning in an ambient space, the limiting object is a bona fide coupling on the graph.
Bridge Handling: The method includes a canonical augmentation (adding virtual parallel edges) to handle bridges in the graph, ensuring the tropical Abel–Jacobi map remains injective and the graph can be recovered from its image.

Experimental Results

The method was evaluated across synthetic benchmarks and real-world data:

Synthetic Benchmarks: Tested on theta, wheel, grid, and small road graphs. The neural OT variants (both Tropical and Euclidean) consistently reduced error (in terms of intrinsic graph Wasserstein distance $W_1^\Gamma, W_2^\Gamma$ and density/CDF $L_1$ errors) compared to heuristic baselines based on discrete graph OT (node-interpolation and ambient pushforward). The two embedding geometries performed comparably, suggesting the primary gain comes from the amortized entropic transport rule rather than a specific embedding advantage.
Scalability: The neural approach scales favorably with support size. Unlike explicit pairwise graph OT methods that require evaluating intrinsic shortest-path distances on large support pairs, the neural method trains on minibatches and applies the learned rule via batched inference.
Real-World Application: The framework was applied to one million Uber pickup locations in Manhattan, New York City.
- The task was graph-supported spatial density generation (learning to generate pickups from a uniform road source).
- Dense graph OT baselines were infeasible at this scale due to memory constraints ( $10^{12}$ pairwise costs).
- The Neural Euclidean (Gromov) model successfully generated samples that recovered the dominant pickup structure on the Manhattan road graph, outperforming heuristic baselines trained on a subsampled subset (16,384 points).

Significance and Claims

The paper claims to provide a principled solution for generative modeling on metric graphs that avoids the pitfalls of discretization. By solving the transport problem in a smooth ambient space and recovering graph-supported samples via projection, the method retains the continuous geometry of the domain.

The authors emphasize that their approach is plan-based rather than map-based, allowing for the handling of singular measures on the graph without mollifying them into ambient densities. The theoretical guarantee of weak convergence to a valid coupling ensures that the method is not merely a heuristic but has a rigorous foundation in optimal transport theory.

The work demonstrates that neural OT can be adapted to non-Euclidean, one-dimensional geometric domains, offering a scalable alternative to discrete graph OT for applications ranging from urban mobility to biological networks. The authors note limitations, including the asymptotic nature of the convergence proof (no finite-sample rates) and the current focus on unconditional generation rather than temporal or covariate-conditioned forecasting.