Efficient algorithm for Wasserstein-1 distance on graph

316 Views Asked by At

I'm looking for an efficient algorithm to calculate the Wasserstein-1 distance in the following setting:

Let $G = (V, E)$ be a (locally) finite, undirected graph, with associated weight function $w : E \to \mathbb{R}_{\geq 0}$. From the weights we derive the geodesic distance $d : V \times V \to \mathbb{R}_{\geq 0}$ as either $d(x, y) = 1 / w(x, y)$ when $(x, y) \in E$ or the shortest path over existing edges between $x$ and $y$ otherwise.

Also associated to the graph $G$ is a simple random walk, where in each vertex $x \in V$ the SRW has probability \begin{align*} \mathbb{P}_x(y) = \begin{cases} \frac{w(x, y)}{\sum_{z \sim x}w(x, z)} & \text{ if } y \sim x \\ 0 & \text{otherwise} \end{cases} \end{align*} to jump to vertex $y \in V$.

The Wasserstein-1 distance $\mathcal{W}_1$ between probability measures $\mathbb{P}_x$ and $\mathbb{P}_y$ is defined as: \begin{align*} \mathcal{W}_1(\mathbb{P}_x, \mathbb{P}_y) = \inf_{\gamma \in \prod(\mathbb{P}_x, \mathbb{P}_y)} \left\{\sum_{(a, b) \in V \times V} d(a, b) \gamma(a, b)\right\} \end{align*} where $\prod(\mathbb{P}_x, \mathbb{P}_y)$ denotes the space of couplings between $\mathbb{P}_x$ and $\mathbb{P}_y$.

Essentially, this is the problem of modifying one mass configuration to the other, over a cost function $d$. Thus, $\gamma(x, y)$ represents the amount of mass that is sent from vertex $x$ to vertex $y$, and $d(x, y) \gamma(x, y)$ is the cost of transporting this mass.

Since this is a discrete setting, calculating $\mathcal{W}_1$ isn't too difficult, but I'd like to use $\mathcal{W}_1$ for very large graphs, and I have some difficulty designing an efficient, preferably $O(n^2)$ algorithm.

My first implementation was to simply use linear programming, since the coupling properties can be used as constraints and the coupling $\gamma$ can be interpreted as several variables.

The problem here is that the constraints require an enormous amount of memory: take 200 vertices, then in the worst case scenario you need 40,000 variables, and get a 40,000 by 400 matrix of constraints at least.

Another method is to model a minimal cost flow problem: start with a start $s$ and sink $t$, add the vertices where $\mathbb{P}_x$ is not zero on the left and connect to $s$, add the vertices where $\mathbb{P}_y$ is not zero on the right and connect to $t$, and finally form a complete bipartite subgraph between these two groups of vertices.

With potentials equal to $\mathbb{P}_x$ and $\mathbb{P}_y$ on the vertices and cost $d$ on the bipartite subgraph, this is a minimal cost flow problem.

I've made a greedy implementation of this problem, but the problem here is that with sorting all the paths the algorithm isn't very efficient: only $O(n^3)$ or $O(n^4)$.

I realise that my attempts have not been very clever, so I'd be very grateful if someone can help me with a better algorithm or point to a C++ or R package with a good implementation!