A parallel framework for graphical optimal transport

TL;DR

This work addresses the computational challenge of multi-marginal optimal transport (MOT) with graph-structured costs by introducing a parallelizable framework based on local entropy regularization. For tree-structured graphs, the MOT problem is reformulated as a sum of coupled bi-marginal OT problems and solved via a Bipartite Iterative Scaling algorithm, enabling two-partite parallel updates and improved iteration complexity. The authors extend the approach to general graphs through modified junction trees, yielding a clique-wise, parallelizable algorithm with provable complexity bounds and practical rounding to enforce marginals. Numerical experiments on Wasserstein barycenter and Wasserstein least-squares problems showcase substantial iteration reductions and confirm the method’s scalability and versatility for graph-structured OT tasks. The paper also connects the framework to probabilistic graphical models, offering a PGM formulation and an Iterative Scaling Belief Propagation algorithm for broader applicability and GPU-friendly implementation.

Abstract

We study multi-marginal optimal transport (MOT) problems where the underlying cost has a graphical structure. These graphical multi-marginal optimal transport problems have found applications in several domains including traffic flow control, barycenter and regression problems in the Wasserstein space, and Hidden Markov model inference problems. The MOT problem can be approached through two formulations: a single big MOT problem, or coupled minor OT problems. In this paper, we focus on the latter approach and demonstrate its efficiency gain from parallelization. For tree-structured MOT problems, we introduce a novel parallelizable algorithm that significantly reduces computational complexity. Additionally, we adapt this algorithm for general graphs, employing the modified junction trees to enable parallel updates. Our contributions, validated through numerical experiments, offer new avenues for MOT applications and establish benchmarks in computational efficiency.

Figures (10)

• Figure 1: Illustration of graphical structures in some applications. Gray nodes correspond to fixed marginals, and white nodes are estimated in the problem. The dotted edge between $x_1$ and $x_J$ in (a) may represent either a bi-marginal constraint, as in \ref{['eq:Euler1']}-\ref{['eq:costEuler1']}, or a cost interaction, as in \ref{['eq:Euler2']}-\ref{['eq:costEuler2']}.
• Figure 1: Bipartite graph with sets $S_1,$$S_2$.
• Figure 1: Modified junction tree for the graph of Euler flow example from \ref{['fig:euler_graph']}.
• Figure 1: Junction tree for the graph from the Euler flow problem in \ref{['fig:euler_graph']}, with leaf nodes corresponding to the constrained marginals.
• Figure 2: Modified junction tree for the graph of Wasserstein least square example from \ref{['fig:wasserstein_least_square']}.

Theorems & Definitions (35)

• Remark 2.1
• Definition 2.2: Junction tree dawid1992applications
• Theorem 3.1
• Definition 3.2
• Proposition 3.3
• Lemma 3.4
• Lemma 3.5
• Lemma 3.6
• Lemma 3.7
• Theorem 3.8