category: literaturenote citekey: fanparallelframeworkgraphicaloptimal2024 title: A parallel framework for graphical optimal transport authors: "Fan, Jiaojiao; Haasler, Isabel; Zhang, Qinsheng; Karlsson, Johan; Chen, Yongxin" year: 2024 date: 2024-00-00 2024 url: "https://arxiv.org/abs/2406.10849v1" zotero_key: GQAX6472 zotero_storage: C7423MZ7 collections: transpordi_planeerimine folder: 001_artiklid firstAuthor: "Fan, Jiaojiao"
Jiaojiao Fan∗ Isabel Haasler† Qinsheng Zhang‡ Johan Karlsson§ Yongxin Chen¶
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 and regression problems in the Wasserstein space. MOT problem can be approached through two aspects: a single big MOT problem, or coupled minor OT problems. In this paper, we focus on the latter approach and demonstrate it has efficiency gain from the 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.
Optimal transport (OT) theory offers a fundamental mathematical framework for the comparison of probability distributions (Peyr´e et al. 2017). Central to OT is the quest for the most cost-efficient strategy to morph one distribution into another. With significant strides in the development of computationally efficient methods (Cuturi 2013, Zanetti and Gondzio 2023), OT has burgeoned into a plethora of applications. It has made substantial impacts in biology (Bunne et al. 2022, Demetci et al. 2020), generative modeling (Fan et al. 2021, 2022b, Korotin et al. 2022), image processing (Peyr´e et al. 2017, Fan et al. 2021, Zhou and Parno 2024), domain adaptation (Seguy et al. 2017, Perrot et al. 2016), Bayesian statistics (Fan et al. 2020, Srivastava et al. 2015), and control theory (Chen et al. 2016, 2021a, Singh et al. 2022, Ringh et al. 2021a), showcasing the versatility and the expansive scope of OT as a tool for analytical and computational advancement. Multi-marginal optimal transport (MOT) (Ba and Quellmalz 2022, Beier et al. 2023, Trillos et al. 2023, Yang and Zha 2023), a generalization of OT, involves optimizing over transports among more than two distributions, thereby extending the applicability of OT to a broader range of problems (Pass 2012, 2015, Nenna 2016). Its applications are vast and varied, encompassing
∗Georgia Institute of Technology, jiaojiaofan@gatech.edu. The first two authors contribute equally.
†EPFL, isabel.haasler@epfl.ch
‡NVIDIA, qinshengz@nvidia.com. Work was done when the author was at Georgia Institute of Technology.
§KTH Royal Institute of Technology, johan.karlsson@math.kth.se
¶Georgia Institute of Technology, yongchen@gatech.edu
This work was supported by the Swedish Research Council (VR) under grant 2014-5870, SJTU-KTH cooperation grant, the Knut and Alice Wallenberg foundation under grant KAW 2021.0274, and the NSF under grant 2206576.
areas such as signal processing (Elvander et al. 2020), estimation and control (Ringh et al. 2023, 2021a), and machine learning (Haasler et al. 2021c, Noble et al. 2023).
The MOT problem frequently embodies a graphical structure where each node denotes a distribution, and the edges represent the transport plans. For example, the well-known Wasserstein barycenter problem finds a generalized average of a set of distributions (Agueh and Carlier 2011). This is done by finding a distribution that minimizes the sum of optimal transport distances to the given distributions, which can be represented by a star-shaped graph (Fan et al. 2022a). Originally, the Wasserstein barycenter problem was formulated as a sum of coupled bi-marginal OT problems (Carlier and Ekeland 2010, Rabin et al. 2012). However, it can equivalently be posed as a MOT problem, where the graphical structure is instead invoked through the underlying cost. This equivalence between the MOT formulation and the coupled OT formulation holds for any structured problems, where the underlying graph is a tree, i.e., an acyclic graph (Haasler et al. 2021b). Although the coupled OT methodology often presents a more intuitive formulation, its computational development has not been extensively pursued, and existing numerical solvers for general graphical structures rely predominantly on the MOT formulation (Elvander et al. 2020, Haasler et al. 2021b,c, Fan et al. 2022a). To close this gap in the literature, we develop numerical methods for graph-structured OT problems that rely on the coupled OT formulation.
More precisely, following the powerful standard approach of adding an entropic term to the OT objective (Cuturi 2013), we first develop a novel scheme for a sum of OT problems that are coupled according to a tree-structure. We also prove complexity results for solving the problem to a desired accuracy. It turns out that this algorithm can be easily parallelized, thus significantly improving the runtime of previous algorithms that are based on the equivalent MOT problem. We also extend our algorithm to the setting of general graph-structures that may contain cycles. In this setting, the algorithm is based on factorizing the underlying graph as a modified junction tree and solving an OT problem that is structured according to this modified junction tree. Thus, in this case the problem is posed as a sum of several coupled MOT problems of small size, where each MOT problem corresponds to a clique, i.e., a collection of nodes, in the modified junction tree.
To summarize, our contributions are the following.
regularization. Extending our previous work Fan et al. (2022a), this framework allows for marginal constraints on the joint distributions. Furthermore, it incorporates the capacity for cost functions to depend on multiple marginals, a feature previously implied but now explicitly defined and integrated into our algorithm.
In this section, we review some background material and introduce the problem we consider in this work. In particular, we first define optimal transport and multi-marginal optimal transport and recap the iterative scaling algorithm for approximately solving these problems. We then introduce MOT problems with graphical structure and provide several motivating examples for considering these types of problems.
The OT problem seeks an optimal transport plan that moves mass from a source distribution to a target one with minimum cost. In this work, we focus on the discrete OT problem (see, e.g., Villani (2003) for the continuous formulations). Thus, consider two non-negative vectors µ1 ∈ R d + and µ2 ∈ R d + with mass 1, i.e., P x1 µ1(x1) = P x2 µ2(x2) = 1, that describe discrete probability distributions. For simplicity of notation, we assume the marginals have equal dimension d throughout, though the extension to the setting with unequal dimensions is straightforward. Here µ1(x) and µ2(x) denote the amount of mass in the source distribution and target distribution, respectively, at location x. We also define a cost matrix C ∈ R d×d , where C(x1, x2) describes the transport cost of moving a unit mass from point x1 to x2. Similarly, we describe the transport by a non-negative matrix B ∈ R d×d + , where the element B(x1, x2) describes the amount of mass transported from point x1 to x2. This defines the optimal transport problem
$$\min{B \in \mathbb{R}{+}^{d \times d}} \operatorname{trace}(C^{\top}B) \quad \text{subject to} \quad B\mathbf{1} = \mu{1}, \quad B^{\top}\mathbf{1} = \mu{2}, \tag{1}$$
where 1 denotes a vector of ones of proper dimension. The constraints enforce that B is a feasible transport plan between µ1 and µ2, and the objective function trace(C ⊤B) = P x1,x2 C(x1, x2)B(x1, x2) describes the total cost of transportation for the given transport plan.
The MOT problem extends the OT problem (1) to finding a transport plan between multiple distributions µ1, . . . , µJ , where J > 2. Let µj ∈ R d for j = 1, . . . , J. Then the multi-marginal transport cost and transport plan are described by J-mode tensors C ∈ R d×···×d and M ∈ R d×···×d + . Here, C(x1, x2, . . . , xJ ) and M(x1, x2, . . . , xJ ) denote the cost and amount of transport associated with the locations x1, x2, . . . , xJ . Thus, analogously to (1), the MOT problem is defined as
$$\min{\mathbf{B} \in \mathbb{R}{+}^{d \times \dots \times d}} \langle \mathbf{C}, \mathbf{B} \rangle \qquad \text{subject to } P{j}(\mathbf{B}) = \mu{j}, \text{ for } j \in \Gamma,$$ (2)
where Γ ⊂ {1, 2, . . . , J} is an index set specifying which marginal distributions are given, and the projection on the j-th marginal of B is computed as
$$Pj(\mathbf{B}) = \sum{x1, \dots, x{j-1}, x_{j+1}, x_J} \mathbf{B}(x1, \dots, x{j-1}, xj, x{j+1}, \dots, x_J).$$ (3)
In the original MOT formulation (Gangbo and Swiech ´ 1998, Pass 2015), constraints are given on all marginal distributions, that is the index set in the constraints is Γ = {1, 2, . . . , J}. However,
Compute K = exp(−C/ϵ), initialize u1, u2, . . . , uJ to 1
while not converged do
for j ∈ Γ do
Update U ← u1 ⊗ u2 ⊗ · · · ⊗ uJ
Update uj ← uj ⊙ µj ./Pj (K ⊙ U ⊙ M)
as we will see in the next section, in many applications not all marginal distributions are given, thus in this work Γ is typically a strict subset of {1, 2, . . . , J}. Finally, we note that the standard
bi-marginal OT problem (1) is a special case of the MOT problem (2) with J = 2 and Γ = {1, 2}.
return B = K ⊙ U ⊙ M
We note that MOT is a standard linear program. However, as it has d J variables, solving it directly scales exponentially with the number of marginals J. The computational burden of solving (2) can be partly alleviated by regularizing the objective with an entropy term. This approach has first been proposed for the bi-marginal problem (1) in Cuturi (2013), Chen et al. (2021b), and has later been adapted to multi-marginal problems (Benamou et al. 2015, Haasler et al. 2021b,c). The entropy regularized MOT problem is the strictly convex optimization problem
$$\min_{\mathbf{B} \in \mathbb{R}^{d \times \dots \times d}} \langle \mathbf{C}, \mathbf{B} \rangle + \epsilon \mathcal{H}(\mathbf{B}|\mathbf{M}) \qquad \text{subject to } P_j(\mathbf{B}) = \mu_j, \text{ for } j \in \Gamma,$$ (4)
where ϵ > 0 is a small regularization parameter and
$$\mathcal{H}(\mathbf{B}|\mathbf{M}) = \sum_{x_1,\dots,x_J} \mathbf{B}(x_1,\dots,x_J) \left( \log \left( \mathbf{B}(x_1,\dots,x_J) \right) - \log \left( \mathbf{M}(x_1,\dots,x_J) \right) - 1 \right). \tag{5}$$
with M(x1, . . . , xJ ) = Q j∈Γ µj (xj ). The entropy regularized problem (4) can be solved by an iterative scaling algorithm, presented in Algorithm 1. This method is also known as Sinkhorn iterations, especially in the bi-marginal setting (1).
In the following, we briefly sketch how to derive Algorithm 1 as a block coordinate ascent method in the dual of (4), for details we refer the reader to Elvander et al. (2020), Haasler et al. (2021b,c). Namely, utilizing Lagrangian duality theory, one can show that the optimal solution to (4) is of the form
$$\mathbf{B} = \mathbf{K} \odot \mathbf{U} \odot \mathbf{M},$$ (6)
where ⊙ denotes element-wise multiplication and the tensors are given by
$$\mathbf{K} = \exp(-\mathbf{C}/\epsilon) \qquad \mathbf{U} = u_1 \otimes u_2 \otimes \cdots \otimes u_J,$$ (7)
where the vectors uj ∈ R d are given by
$$u{j} = \begin{cases} \exp\left(\frac{\lambda{j}}{\epsilon}\right), & \text{if } j \in \Gamma \ \mathbf{1}, & \text{otherwise,} \end{cases}$$ (8)
and λj ∈ R d is the dual variable corresponding to the constraint Pj (B) = µj on the j-th marginal. The corresponding dual problem of (4) is
$$\max_{{\lambdaj, j \in \Gamma}} -\epsilon \langle \mathbf{K}, \mathbf{U} \rangle + \sum{j \in \Gamma} \lambda_j^\top \mu_j, \tag{9}$$
where U depends on the dual variables as in (7)-(8). Algorithm 1 is then derived as a block coordinates ascent in this dual and is therefore guaranteed to converge globally.
In this work we consider multi-marginal optimal transport problems with a cost that is composed as a sum of cost terms, which each depend on a subset of the marginals. That is, we consider cost tensors of the form
$$\mathbf{C}(\mathbf{x}) = \sum{\alpha \in F} C{\alpha}(\mathbf{x}_{\alpha}), \tag{10}$$
where F is a set of subsets of the tensors indices, i.e., α ⊂ {1, . . . , J} for each element α ∈ F. If we associate the index set of the transport tensor with a set of vertices, V = {1, . . . , J}, then the set F can be understood as the set of factors in a factor graph (Haasler et al. 2021c, Fan et al. 2022a). In particular, all subsets α ∈ F that contain two elements can be illustrated as edges in a graph. Moreover, we allow for constraints that act on projections onto several marginals of the transport tensor B. For a set α ⊂ V , we define the projection Pα(B) = P xj , j∈V <sup>α
B(x1, . . . , xJ ). Let FΓ ⊂ F denote the set of factors which are constrained, and for α ∈ FΓ, let µα ∈ R d + be the given tensor, where the number of modes equals the cardinality |α| of α. This let us define the graph-structured MOT problem as$$\min{\mathbf{B} \in \mathbb{R}{+}^{d} \times \cdots \times d} \langle \mathbf{C}, \mathbf{B} \rangle \qquad \text{subject to } P{\alpha}(\mathbf{B}) = \mu{\alpha}, \text{ for } \alpha \in F_{\Gamma},$$ (11)
where the cost tensor decouples as in (10), and FΓ ⊂ F. Note that the graph structured MOT problem (11) is a special case of the standard MOT problem (2) in terms of the cost function. On the other hand, problem (11) allows for a more general class of constraints than the original problem (2).
Remark 1. We assume that FΓ ⊂ F for convenience of notation. However, in many cases there is not a cost associated with every factor in FΓ. In this case we can simply define Cα(xα) = 0 for the corresponding α ∈ FΓ. In the examples and applications in this paper, we do not explicitly define these vanishing parts of the cost tensor.
The most widely studied graph-structured MOT problem is the Wasserstein barycenter problem (Agueh and Carlier 2011). However, in a recent line of research, graph-structured MOT problems have been utilized in a wide range of applications. For instance, in Haasler et al. (2021b,a) different control and estimation problems for large systems of indistinguishable agents are modeled as graph-structured MOT problems. Even different classes of agents can be modeled with this framework, e.g., in network flow problems (Haasler et al. 2024) and mean field games (Ringh et al. 2021b, 2023). In the following we briefly review a few more exampled of optimal transport problems that fall into the category of graph-structured MOT. In particular, we note that to the best of our knowledge, the Euler problem and Wasserstein splines have not yet been classified as such. Throughout this paper, we illustrate constrained marginals as gray nodes, and white nodes are estimated in the respective problems.
One of the first applications of multi-marginal optimal transport, the generalized Euler flow (Benamou et al. 2015, Section 4.3), falls in fact into the category of a graph-structured problem. The Euler flow studies a relaxation of the incompressible Euler equation, where one seeks trajectories of particles between constrained initial and final data, that minimize the kinetic energy (Brenier 1989). Discretizing the problem into J time points, we describe the flow by a Jmode multi-marginal optimal transport plan B ∈ R d + , where each marginal describes the density
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 x1 and xJ in (a) may represent either a bi-marginal constraint, as in (12)-(13), or a cost interaction, as in (14)-(15).
of the flow at time instance j, discretized into d points. Note that the density of an incompressible flow, discretized uniformly over a given space, is described by a uniform distribution. Thus, we constrain each marginal of the transport plan by Pj (B) = µj = 1 d 1. In the Euler flow, it is enforced that each particle at time 0 is transported to a given location at time T, given by a permutation σ : {1, . . . , d} → {1, . . . , d}. This is modeled through a bi-marginal constraint P1,J (B) = µ{1,J} = Πσ, where Πσ ∈ R d×d is the permutation matrix for the permutation σ. The Euler flow problem is thus formulated as the MOT problem
$$\min{\mathbf{B} \in \mathbb{R}{+}^{d{X} \dots \times d}} \langle \mathbf{C}, \mathbf{B} \rangle$$ subject to $P{j}(\mathbf{B}) = \mu{j}$ for $j = 1, \dots, J$ ; $P{1,J}(\mathbf{B}) = \mu_{{1,J}}$ ,
where
$$\mathbf{C}(x_1, \dots, xJ) = \sum{j=1}^{J-1} C(xj, x{j+1}), \tag{13}$$
and the elements of the matrix C ∈ R d×d contain the squared Euclidean distances between the discretization points, namely, Cxj ,xj+1 = ∥xj − xj+1∥ .
Benamou et al. (2015) also propose an alternative relaxed formulation, where the bi-marginal constraint in (12) is replaced by an additional cost term in (13), and the problem becomes
$$\min{\mathbf{B} \in \mathbb{R}{+}^{d \times \cdots \times d}} \langle \mathbf{C}, \mathbf{B} \rangle \qquad \text{subject to } P{j}(\mathbf{B}) = \mu{j} \text{ for } j = 1, \dots, J,$$ (14)
where
$$\mathbf{C}(x_1, \dots, xJ) = \sum{j=1}^{J-1} C(xj, x{j+1}) + C^{\sigma}(x_1, x_J), \tag{15}$$
and the elements of the matrix C σ ∈ R d×d are given by C σ (x1, xJ ) = ∥xσ(1) − xJ ∥ 2
Both problems (12)-(13) and (14)-(15) can be described by a cyclic graph as illustrated in Fig. 1a. In particular, we note that although the authors in Benamou et al. (2015) present only the computational method for the relaxed formulation (14)-(15), our proposed framework encompasses both formulations.
Consider a least square fitting problem in the Wasserstein-2 space. Given a set of probability measures ρt1 , ρt2 , . . . , ρtJ with 0 ≤ t1 < t2 < · · · < tJ ≤ 1, we aim to find two distributions µ0 and µ1 such that the set of measures is closest to the displacement interpolation between µ0 and µ1 in terms of Wasserstein-2 distance. More specifically, we seek the solution to
$$\min_{\mu_0,\mu1} \sum{j=1}^{J} W2(\rho{tj}, \mu{t_j})^2 \tag{16}$$
where µt denotes the displacement interpolation connecting µ0 and µ1. This least square problem can be reformulated as
$$\min_{\pi,\mu_0,\mu_1} \int c(x_0, x_1, \dots, xJ, x{J+1}) d\pi(x_0, x_1, \dots, xJ, x{J+1}) Pj(\pi) = \rho{t_j}, \quad j = 1, 2, \dots, J. \quad P_0(\pi) = \mu0, \quad P{J+1}(\pi) = \mu_1.$$ (17)
By discretizing the integral in (17) to a summation, we can transform the problem into the form presented in (11). To enforce that the optimality of the coupling between µ0 and µ1, we can add a large penalty on it, leading to the modified cost
$$c(x_0, x_1, \dots, xJ, x{J+1}) = \sum_{j=1}^{J} |x_j - (1 - t_j)x_0 - ti x{J+1}|^2 + \alpha |x0 - x{J+1}|^2$$ (18)
with α > 0 being a sufficiently large number.
The Wasserstein spline developed in Chen et al. (2018) is a natural notion of cubic spline over the space of probability distributions. Given a set of points xt0 , xt2 , . . . , xtJ with 0 = t0 < t1 < · · · < tJ = 1, a cubic spline is a piece-wise cubic polynomial {x(t) | 0 ≤ t ≤ 1} that is smooth and matches the given data points in the sense that x(tj ) = xtj for all j = 0, 1, . . . , J. The cubic spline can also be characterized as the interpolating curve of the data points with minimum mean square acceleration. Based on this characterization, Chen et al. (2018) generalized the notion of cubic spline to interpolate a set of probability distributions ρt0 , ρt1 , . . . , ρtJ . The definition has a static reformulation in the phase space as
$$\min_{\pi} \int c(x_0, v_0, \cdots, x_J, v_J) d\pi(x_0, v_0, \cdots, x_J, vJ) \quad \text{subject to} \quad P{xj}(\pi) = \rho{t_j}, \quad j = 0, 1, \dots, J$$
$$\tag{19}$$
where c(x0, v0, · · · , xJ , vJ ) =
$$\sum{i=0}^{J-1} \frac{1}{t{j+1} - tj} \cdot \left[ 12 \left| \frac{x{j+1} - xj}{t{j+1} - t_j} - vj \right|^2 - 12 \left\langle \frac{x{j+1} - xj}{t{j+1} - t_j} - vj, v{j+1} - vj \right\rangle + 4 |v{j+1} - v_j|^2 \right]. \tag{20}$$
This problem is associated with the graphical model in Fig. 1c, where the constraints are on the position nodes x0, x1, . . . , xJ .
A line of recent works have drawn a connection between entropy-regularized multi-marginal optimal transport problems (4) with a graph-structured cost (10) and probabilistic graphical models (PGMs) (Haasler et al. 2021c, Fan et al. 2022a, Haasler et al. 2021b, Singh et al. 2022, Ringh et al. 2023, Zhou and Parno 2024). In this section we review and generalize these connections and the computational methods for solving these types of optimal transport problems.
A PGM provides a compact representation for the joint distribution of a collection of random vectors, which have dependencies that can be described by a graph (Koller and Friedman 2009). Define the tensors Kα(xα) = exp(−Cα(xα)/ϵ) for α ∈ F, where Cα are the cost tensors in (10), and ϵ is the regularization parameter in (4). Then, K in (7) has the form K(x) = Q α∈F Kα(xα). Moreover, to define the entropy regularization (5) for the graph-structured problem (11) we define the tensor M as M(x) = Q α∈FΓ µα(xα). The factorizations of the tensors K and M can be used to define the joint probability distribution of a PGM,
$$p(\mathbf{x}) = \frac{1}{Z} \left( \prod{\alpha \in F} K{\alpha}(\mathbf{x}{\alpha}) \right) \left( \prod{\alpha \in F{\Gamma}} \boldsymbol{\mu}{\alpha}(\mathbf{x}_{\alpha}) \right), \tag{21}$$
where Z is a normalization constant. Moreover, note that the solution tensor (6) can be written as K ⊙ U ⊙ M with U(x) = Q α∈FΓ uα(xα). Thus, the scaling factors uα, for α ∈ FΓ can be interpreted as local potentials of the modified graphical model K(x)U(x)M(x) The Iterative Scaling algorithm 1 finds potentials uα, for α ∈ FΓ, such that the graphical model defined by the tensor B = K ⊙ U ⊙ M satisfies all the constraints Pα(B) = µα, for α ∈ FΓ.
Each step in the Iterative Scaling algorithm involves calculating a projection Pj (K ⊙ U ⊙ M), which in terms of the PGM means inferring the distribution of the j-th variable. This Bayesian inference problem can be addressed by the Belief Propagation algorithm (Mezard and Montanari 2009, Part D), (Koller and Friedman 2009, Sec. 20.5.1). Belief propagation is known to converge when the underlying graph is a tree, i.e., does not have any loops, and reduces the complexity of computing the projections in the Iterative Scaling algorithm immensely. More precisely, whereas a brute force approach to compute the projections (3) scales exponentially, Belief Propagation scales only linearly in the number of marginals (Haasler et al. 2021c).
In Fan et al. (2022a), a similar approach has be utilized for general graphical models with possible loops, by partitioning the graph as a so-called junction tree. A junction tree (also called tree decomposition) describes a partitioning of a graph, where several nodes are clustered together, such that the interactions between the clusters can be described by a tree structure.
Definition 1 (Junction tree (Dawid 1992)). A junction tree T = (C, E) over a graph G = (V, F) is a tree whose nodes are cliques c ∈ C associated with subsets xc ⊂ V , and satisfying the properties:
For two adjoining cliques ci and cj , define the separation set sij = {v ∈ V : v ∈ ci ∩ cj}.
A graph typically has several possible junction trees. In particular, one trivial junction tree is always given by a single cluster containing all nodes. For our computational method, we set up the junction tree such that each constraint marginal or clique is represented by a leaf in the junction tree. The entropy-regularized graph structured optimal transport problem with junction tree decomposition T = (C, E) is thus of the form
$$\min{\mathbf{B} \in \mathbb{R}{\perp}^{d \times \dots \times d}} \langle \mathbf{C}, \mathbf{B} \rangle + \epsilon \mathcal{H}(\mathbf{B}|\mathbf{M}) \qquad \text{subject to } P_c(\mathbf{B}) = \boldsymbol{\mu}c, \text{ for } c \in \mathcal{C}{\Gamma},$$ (22)
Figure 2: Junction tree for the graph from the Euler flow problem in Fig. 1a, with leaf nodes corresponding to the constrained marginals.
Figure 3: Junction tree for the graph for the Wasserstein least square problem from Fig. 1b. The junction tree is constructed such that the leaves of the tree correspond to the constraints in the least square problem (17).
where CΓ ⊂ C is the set of constrained cliques, and the cost tensor decouples as
$$\mathbf{C}(\mathbf{x}) = \sum_{c \in \mathcal{C}} \mathbf{C}_c(\mathbf{x}_c). \tag{23}$$
Such a junction tree representation exists for any graph, however it may require adding artificial clusters to the junction tree decomposition.
Example 1. Some examples of valid junction trees for our method are shown in Fig. 2 (for the Euler flow example in Fig. 1a), Fig. 3 (for the Wasserstein least squares problem (17) in Fig. 1b) , and Fig. 4 (for the Wasserstein splines problem in Fig. 1c).
Similarly to the factorization (21), we can then write the probability distribution over the
Figure 4: Junction tree of Wasserstein splines in Fig. 1c.
Algorithm 2 Iterative Scaling Belief Propagation Algorithm for graphical MOT with global regularization
Input: Given a junction tree T = (C, E) with constraints on the leaf cliques CΓ ⊂ C Initialize messages mi→j (sij ), mj→i(sji), ∀(i, j) ∈ E
while not converged do
$$m{i \to j}(s{ij}) \leftarrow \sum{x{ci} \setminus s{ij}} \mathbf{K}_{ci}(\mathbf{x}{ci}) \prod{\ell \in N(i) \setminus j} m{\ell \to i}(s{\ell i}), \qquad ci \in \mathcal{C} \setminus \mathcal{C}{\Gamma}$$ (24)
$$m{i \to j}(s{ij}) \leftarrow \mu_{ci}(\mathbf{x}{ci})/m{j \to i}(\mathbf{x}_{c_i}), \qquad ci \in \mathcal{C}{\Gamma}.$$ (25)
return mi→j (sij ), mj→i(sji), ∀(i, j) ∈ E
PGM following the junction tree decomposition T = (C, E) as a product
$$p(\mathbf{x}) = \frac{1}{Z} \left( \prod_{c \in \mathcal{C}} \mathbf{K}_c(\mathbf{x}c) \right) \left( \prod{c \in \mathcal{C}_{\Gamma}} \boldsymbol{\mu}_c(\mathbf{x}_c) \right). \tag{26}$$
Let N(i) be the neighboring nodes of the node i. By extending the Iterative Scaling Belief Propagation methods in Haasler et al. (2021c), Fan et al. (2022a), we propose Algorithm 2 to solve (22). Upon convergence of Algorithm 2, the optimal transport plans on the clusters are given by
$$\mathbf{B}{c}(\mathbf{x}{c}) = \mathbf{K}{c}(\mathbf{x}{c}) \prod{\ell \in N(i)} m{\ell \to i}(s_{\ell i}). \tag{27}$$
The algorithm sends messages between neighboring cliques. The messages are thus of the same size as the separation sets. Recall that the leaves of the junction tree are associated with constraints in the optimal transport problem. Thus, the messages sent from these leaves correspond to an update step similar to the one in Algorithm 1. The messages sent within the graph encode the computation of the projections. More precisely, picking one leaf c and sending the messages from all other leaves across the separation sets towards our leaf amounts to one exact computation of the projection Pc(K ⊙ U ⊙ M). However, the algorithm converges with any message schedule (Haasler et al. 2021c, Fan et al. 2022a).
Note that in the case where all cost are pair-wise and constraints act on one-marginals of the transport tensors, Algorithm 2 boils down to the Iterative Scaling Belief Propagation algorithms in Haasler et al. (2021c), Fan et al. (2022a). Here, we have extend on these methods and provide a more general algorithm for the case where we may have constraints on joint distributions, and where cost terms can depend on more than two marginals.
The Iterative Scaling Belief Propagation algorithm utilizes the graphical structure in problem (22)-(23) to significantly improve on the complexity of the Sinkhorn iterations. Recall that the computational bottleneck of Algorithm 1 lies in the computation of the projections, which require O(d J ) operations. Algorithm 2 reduces this complexity by utilizing the graphical structure for computing the projections. The complexity of the problem depends then on the size of the messages, i.e., the size of the separation sets. This quantity can also be described by the so-called tree-width.
Definition 2 (Tree width (Bertele and Brioschi 1973, Halin 1976)). For a junction tree T = (C, E), we define its width as width(T ) = maxc∈C |c| − 1. For a graph G, we define its tree-width as w(G) = min{width(T ) | T is a junction tree for G}.
In order for Algorithm 2 to be efficient we thus would like to pick a junction tree factoriza-
tion with small tree-width. Given a graph, a junction tree decomposition that minimizes the tree-width is also called a minimal junction tree. In general, finding a junction tree with minimal width is a NP hard problem (Koller and Friedman 2009). However, for all applications in this paper it is easy to find minimal junction trees.
Example 2. The junction trees in Figures 2, 3, and 4 are minimal junction trees for the respective graphs. Since the cliques in the former two junction trees are of cardinality 3, the respective graphs' tree-widths are 2. Similarly, the last graph's tree width is 3.
For Algorithm 2, we can directly apply the complexity analysis in Fan et al. (2022a). We will discuss the complexity result in remark 2.
In the context where the MOT problem is tree structured, i.e., structured as a graph without cycles, the innate characteristics of tensor inner products facilitate a natural reformulation of the problem. This reformulation manifests as a sum of coupled bi-marginal OT problems. Considering a set of fixed marginals, denoted as µj for each j belonging to the vertex set Γ, the problem of coupled bi-marginal OT can be written as:
$$\min{\substack{B{(j,k)},(j,k)\in E\\mu{j},j\in V\setminus\Gamma}}\sum{(j,k)\in E}\langle C{(j,k)},B{(j,k)}\rangle \qquad \text{subject to} \quad B{(j,k)}\mathbf{1}=\mu{j}, \quad B{(j,k)}^{\top}\mathbf{1}=\mu{k} \quad \text{for } (j,k)\in E,$$
$$(28)$$
where E is a set of all the edges of the graph. The matrices C(j,k) ∈ R d×d and B(j,k) ∈ R d×d represent the transport cost and transport plan corresponding to the edge connecting nodes j and k, respectively. We implicitly assume the convention B(j,k) = B⊤ (k,j) and C(j,k) = C ⊤ (k,j) . In contrast to the global entropy regularization presented in Sec. 3 for a standard MOT problem (22), applying local regularization on each edge appears as a more direct and natural approach within this framework. Notably, in tree-structured optimal transport scenarios, such as barycenter problems, this form of regularization is frequently employed, as described in existing literature (Kroshnin et al. 2019).
Drawing inspiration from Kroshnin et al. (2019), we developed a parallel algorithm to address the coupled bi-marginal OT problem for general trees. Herein, we propose to use a modified entropy regularization similar to Fan et al. (2022a). The reference joint distribution between nodes j and k is denoted as:
$$M_{(j,k)} := \gamma_j \otimes \gamma_k$$ , (29)
where γj = µj if j ∈ Γ and 1 otherwise. As discussed before, we always assume that marginals may only be constrained if they are on the leaves of the tree. It is important to note that therefore the scenario where (j, k) ∈ E and both j, k ∈ Γ does not exist. In such cases, the OT problem simplifies into several independent subproblems. Consequently, this leads us to formulate the regularized problem with an enforced mass constraint.
$$\min{\substack{B{(j,k)},(j,k)\in E\\muj,j\in V\backslash\Gamma}} \sum{(j,k)\in E} \langle C{(j,k)},B{(j,k)}\rangle + \epsilon \mathcal{H}(B{(j,k)}|M{(j,k)})$$ subject to $$B_{(j,k)}\mathbf{1} = \muj, \ B{(j,k)}^{\top}\mathbf{1} = \mu_k \text{ for } (j,k)\in E; \quad \mathbf{1}^{\top}\mu_j = 1 \text{ for } j\in V\backslash\Gamma.$$ (30)
The regularization, as formulated, is applied locally to each edge of the tree, in contrast to the previous formulations (22), which regularizes the multi-marginal optimal transport tensor. Utilizing a standard Lagrangian dual theory enables the derivation of the dual problem. Denote
$$B{(j,k)} = \operatorname{diag}(u{(j,k)} \odot \gammaj) K{(j,k)} \operatorname{diag}(u_{(k,j)} \odot \gammak), \qquad K{(j,k)} = \exp(-C_{(j,k)}/\epsilon)$$ (31)
where u(j,k) ∈ R d . Furthermore, u = {u(j,k)}j̸=k and ρ = {ρj}j∈V \Γ. For a node j ∈ Γ, its unique neighbor is denoted by kj . The dual problem is then formulated as:
$$\min{u,\rho} \quad f(u,\rho) := \sum{(j,k)\in E} |B_{(j,k)}|1 - \sum{j\in\Gamma} \log(u_{(j,k_j)})^\top \muj - \sum{j\in V\setminus\Gamma} \rhoj$$ subject to $$\sum{k\in N(j)} \log(u_{(j,k)}) = \rho_j \mathbf{1}, \quad \text{for } j\in V\setminus\Gamma.$$ (32)
We then introduce another set of dual variables and apply the Lagrangian dual theory again to the dual problem (32), and obtain the update rules. The dual variables we introduce are reflected in vj in (34).
Theorem 1. The optimal dual variables can be found as the limit point of the following iterative scheme: For j ∈ Γ we update
$$u_{(j,kj)} = 1./(K{(j,kj)}u{(k_j,j)}).$$ (33)
For j ∈ V \ Γ we update the set of vectors u(j,k) , where k ∈ N(j), according to
$$u_{(j,k)} = e^{\rho_j/|N(j)|} vj./(K{(j,k)}(u_{(k,j)} \odot \gamma_k)), \qquad vj = \left( \bigodot{k \in N(j)} \left( K{(j,k)}(u{(k,j)} \odot \gamma_k) \right) \right)^{1/|N(j)|}.$$ (34)
and we update ρj = −|N(j)| log 1 ⊤vj .
Based on Theorem 1, we note that for any node j, the updates of {u(j,k)}k∈N(j) or ρj depend only on the variables u(k,j) from neighboring nodes k ∈ N(j). This observation is critical for enabling parallel updates in the algorithm, as it allows for independently updating dual variables associated with two distinct parts of the tree. The concept of a bipartite graph becomes instrumental here, providing the framework for such parallelism.
Definition 3. (Bipartite Graph (Metcalf and Casey 2016)). A bipartite graph is a graph where the vertices can be divided into two disjoint sets such that all edges connect a vertex in one set to a vertex in another set. There are no edges between vertices in the disjoint sets.
Figure 5: Bipartite graph with sets S1, S2.
To leverage this, the algorithm unfolds in the following manner: We begin by representing the tree as a bipartite graph, a natural representation since all trees are inherently bipartite according to Scheinerman (2012). The algorithm then iteratively updates until convergence, alternating between: i) Updating all dual variables from the first partition to the second. ii) Updating all dual variables from the second partition to the first. Denote the intermediate computational variables
$$q{(j,k)}^t = B{(j,k)}^t \mathbf{1}, \qquad qj^{t+1} = \left( \bigodot{k \in N(j)} q_{(j,k)}^t \right)^{\frac{1}{|N(j)|}}.$$ (35)
Our method is detailed in Algorithm 3, termed Bipartite Iterative Scaling. A supporting
Algorithm 3 Bipartite Iterative Scaling Algorithm for graphical coupled bi-marginal OT with local regularization
Input: Probability vectors {µj}j∈Γ. Stopping criteria δ ′ . Transport cost {C(j,k)}. Regularization param. ϵ.
Construct a bipartite graph where the nodes are divided into two disjoint groups {S1, S2}. Initialization of variables: t = 1, u0 (j,k) = 1 ∈ R d for j ̸= k.
while $$\sum{j \in \Gamma} |B{(j,k_j)}^t \mathbf{1} - \mu_j|1 + \sum{j \notin \Gamma} \sum{k \in N(j)} |q{(j,k)}^t - \bar{q}_j^t|_1 \ge \delta'$$ do
Let S = S1 if t is odd, otherwise S = S2
Calculate q t {j,k} , qt+1 j according to (35) for j ∈ S
$$u_{(j,kj)}^{t+1} = u{(j,k_j)}^t \odot \muj / q{(j,k_j)}^t \quad \text{for } j \in \Gamma \bigcap S$$ (36)
$$u{(j,k)}^{t+1} = u{(j,k)}^{t} \odot \frac{q_j^{t+1}}{|q_j^{t+1}|1} / q{(j,k)}^{t} \quad \text{for } j \in (V \backslash \Gamma) \bigcap S$$ (37)
t ← t + 1
Construct Bt (j,k) , (j, k) ∈ E according to (31).
return Bt (j,k) for all (j, k) ∈ E
proposition 1 underpins its theoretical foundation.
Proposition 1. The updates in Theorem 1 can be equivalently expressed as the updates (36)-(37) in Algorithm 3 together with the variables (35).
The updates of the dual variables, as specified in (36)-(37), are independent within the same group S, facilitating simultaneous updates in parallel. Provided there is sufficient memory to store all variables, concurrent updating is feasible.
In each iteration Algorithm 3 updates the variables corresponding to one of the partitions S1 or S2. This update guarantees that the constraints are matched exactly for the nodes in this partition. However, this means that after any finite number of iterations of Algorithm 3 the marginal constraints will not be matched for the nodes in the partition that has not been updated in the last iteration. To address this, we adopt a strategy akin to the ones in Altschuler et al. (2017), Fan et al. (2022a) and employ a Rounding Algorithm 4. This ensures the generation of transport plans that adhere to all stipulated constraints in (30). The complexity of the rounding process is O(|E|d 2 ). Algorithm 4 differs from the method in Fan et al. (2022a) in that it requires modification of the transport plans across all edges.
In this section, we provide a rigorous demonstration that our Algorithm 3, when used in conjunction with the Rounding Algorithm 4, requires T = Oe |E| 2/δ2 iterations to achieve a δ-approximation of the problem (28). First, we will illustrate that, through the specific mass normalization technique employed in (37), the transport plan maintains unit mass in each iteration.
Lemma 1. At each update t of Algorithm 3, we have $$||B_{(j,k)}^t||1 = ||q{(j,k)}^t||_1 = 1$$ , $\forall (j,k) \in E$ .
We now define the variable λ t (j,k) = ϵ log(u t (j,k) ) for (j, k) ∈ E, and u t (j,k) are the iterates in Algorithm 3. Lemma 2 provides bounds on the scaling variables.
/* Post-update of nodes in one partition, the projections on nodes in the opposite partition do not match the marginals. This partition is denoted as S. */
Input: The transportation plans B(j,k) for all (j, k) ∈ E, where P(j,k) has mismatch with marginal qj when j ∈ S. A sequence of probability vectors {µj}j∈Γ
for j ∈ S do
if j ∈ Γ then
Round B(j,kj ) by Algorithm 2 in Altschuler et al. (2017) such that Bb (j,kj ) ∈ Π(µj , B⊤ (j,kj ) 1) .
else
Round B(j,k) by Algorithm 2 in Altschuler et al. (2017) such that Bb (j,k) ∈ Π(q¯j (B), B⊤ (j,k) 1) for ∀k ∈ N(j), where ¯qj (B) = 1 |N(j)| P k∈N(j) B(j,k)1.
return Bb (j,k) for all (j, k) ∈ E
Lemma 2. Denote C∞ := max(j,k)∈E ∥C(j,k)∥∞. For all iterates λ t (j,k) , and their fixed points λ ∗ (j,k)
$$\max{i=1,\dots,d} (\lambda{(j,k)}^t)i - \min{i=1,\dots,d} (\lambda_{(j,k)}^t)i \le 2C{\infty}, \qquad \max{i=1,\dots,d} (\lambda{(j,k)}^)i - \min{i=1,\dots,d} (\lambda_{(j,k)}^)i \le 2C{\infty}.$$ (38)
The following lemma links the dual objective error in Theorem 1 to the stopping criteria in Algorithm 3.
Lemma 3. For all iterates (u t , ρt ) of the method in Algorithm 3 and their fixed points (u ∗ , ρ∗ it holds that
$$f(u^{t}, \rho^{t}) - f(u^{}, \rho^{}) \leq \frac{2C{\infty}}{\epsilon} \left( \sum{j \in \Gamma} |B{(j,k{j})}^{t} \mathbf{1} - \mu{j}|{1} + \sum{j \in V \setminus \Gamma} \sum{k \in N(j)} |B{(j,k)}^{t} \mathbf{1} - \bar{q}{j}^{t}|_{1} \right)$$ (39)
where $$\bar{q}j^t = \frac{1}{|N(j)|} \sum{k \in N(j)} q{(j,k)}^t = \frac{1}{|N(j)|} \sum{k \in N(j)} B_{(j,k)}^t \mathbf{1}$$ .
The increase in the dual objective between iterations is also closely related to the stopping criteria in Algorithm 3.
Lemma 4. For the iterates in Algorithm 3, it holds that
$$f(u^{t}, \rho^{t}) - f(u^{t+1}, \rho^{t+1}) \ge \frac{1}{22|E|} \left( \sum{j \in \Gamma} |B{(j,k{j})}^{t} \mathbf{1} - \mu{j}|{1} + \sum{j \in V \setminus \Gamma} \sum{k \in N(j)} |B{(j,k)}^{t} \mathbf{1} - \bar{q}{j}^{t}|{1} \right)^{2}.$$ (40)
Unlike Kroshnin et al. (2019), our Lemma 3 and Lemma 4 apply universally for any t, obviating the need to consider odd and even t distinctly.
Theorem 2. Algorithm 3 generates {Bt (j,k) } satisfying P j∈Γ ∥Bt (j,kj ) 1−µj∥1+ P j /∈Γ P k∈N(j) ∥q t (j,k)− q¯ t j ∥1 ≤ δ ′ within t = 2 + 88|E|C∞ ′ϵ iterations.
The subsequent two lemmas quantify the error introduced to the objective value by the rounding process.
Lemma 5. Let {B(j,k)} be a sequence of transportation plans on a bipartite tree. Algorithm 4 returns {Bb (j,k)} satisfying Bb⊤ (j,k) 1 = qj for j ∈ V , where qj is the marginal distri-
bution of node j. Moreover, it holds that $$\left| \sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, \widehat{B}{(j,k)} \rangle \right| \leq 2C{\infty} \left( \sum{j\in\Gamma} |\muj - B{(j,k_j)}\mathbf{1}|1 + \sum{j\notin\Gamma} \sum_{k\in N(j)} |\bar{q}j(B) - B{(j,k)}\mathbf{1}|_1 \right), \text{ where } \bar{q}j(B) = \frac{1}{|N(j)|} \sum{k\in N(j)} B_{(j,k)}\mathbf{1}.$$
Lemma 6. Let {Be (j,k)} be the output of Algorithm 3 and let {Bb (j,k)} be the output of Algorithm 4 with input {Be (j,k)}, and denote {B∗ (j,k) } the optimal solution to the unregularized MOT problem (28). Then it holds that
$$\sum{(j,k)\in E} \langle C{(j,k)}, \hat{B}{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)}^* \rangle$$ (41a)
$$\leq 2\epsilon |E| \log(d) + 4C{\infty} \left( \sum{j \in \Gamma} |\muj - \widetilde{B}{(j,k_j)} \mathbf{1})|1 + \sum{j \notin \Gamma} \sum_{k \in N(j)} |\bar{q}j(\widetilde{B}) - \widetilde{B}{(j,k)} \mathbf{1}|_1 \right). \tag{41b}$$
Given a predetermined final error threshold δ, we select parameters ϵ = 4|E| log(d) and δ ′ = 8C∞ for use in Algorithm 3. This results in an output of Be (j,k) , for each edge pair (j, k) ∈ E. Next, we employ Algorithm 4 to round Be (j,k) for each (j, k) ∈ E, thereby producing Bb (j,k) . With these procedures in place, we are equipped to compute the final complexity for the coupled bi-marginal OT problem as described in (28).
Theorem 3. Algorithm 3, in conjunction with Algorithm 4, yields a δ-approximate solution to the tree-structured MOT problem (28) in T arithmetic operations, where T = O |E| 3C 2∞d log(d) .
Remark 2 (Comparison of complexity with global regularization (Fan et al. 2022a)). We compare our complexity results with the ones in Fan et al. (2022a), for graph-structured MOT with global entropic regularization in terms of the total complexity and the number of iterations. i) Total Complexity: According to Theorem 2 in Fan et al. (2022a), global entropy regularization requires Oe(d(G)|E||Γ| 2C 2 ∞d 2/δ2 ) arithmetic operations with high probability, where d(G) denotes the maximum node distance in graph G. Under the assumptions |E| = O(|Γ|) and d(G) = O(1), akin to the Wasserstein barycenter problem, both the global and local regularization methods offer similar complexity bounds. However, a notable distinction is that Theorem 3 offers a deterministic bound. ii) Iterations: Algorithm 3 with local regularization, as per Theorem 2 and Theorem 3, necessitates O |E| 2C 2∞ log(d) iterations, less than the requirement for global regularization at O |E||Γ| 2C 2∞ log(d) . The reduced iteration count is a benefit of algorithm parallelization, albeit with a higher computational load per iteration.
Building on the concepts from Sec. 4, we extend our approach to OT problems with a general factor graph structure as indicated in (10). We employ a strategy that breaks down the MOT problem into interconnected smaller MOT problems, introducing local regularization for each sub-problem. To facilitate this, we define a specialized type of junction tree, elaborated in Definition 1.
Definition 4. A modified junction tree is a class of junction tree T = (C, E) with two disjoint sets of cliques QC ⊂ C and separators QS ⊂ C such that
Figure 6: Modified junction tree for the graph of Euler flow example from Fig. 1a.
Figure 7: Modified junction tree for the graph of Wasserstein least square example from Fig. 1b.
Note that each separator set in Definition 1 is connected to exactly two cliques. On the other hand, by Definition 4, our cliques in QS can be connected to an arbitrary number of cliques, and instead, the cost-cliques in QC can only be connected to two cliques. This latter assumption is necessary for efficient parallelization of the algorithm we introduce and can be relaxed as we will explain in remark 3.
Example 3. The junction trees in Figures 6, 7, and 8 are modified junction trees for the respective graphs. In contrast to junction trees in Figures 2, 3, and 4, each clique (shown as circles) can only be connected to two separators (shown as squares), while separator can connect to multiple cliques. Moreover, note that constraints in the respective MOT problems are indicated by gray boxes in Figures 6, 7, and 8. This is explained in more detail in the following section.
We structure a coupled MOT problem (42), which is equivalent to MOT problem (11), on the junction tree in Definition 4 by associating every cost term with a clique in QC . Moreover, we assume that each constraint is embedded in a separator in QS as described in the following. A constraint of the optimal transport problem can act on projections of the solution tensors Bc onto a set of marginals. Recall that Γ ⊂ {1, 2, . . . , J} denote the set of index-sets corresponding to the constraints (see (2)). Then we assume that there exists a partition of Γ such that for
Figure 8: Modified junction tree for the graph of Wasserstein spline example from Fig. 1c.
each γ in the partition there is a cγ ∈ QS such that γ ⊆ cγ, and any cγ can contain only one constrained set γ. Moreover, we denote the separators with constraints as QΓ = {cγ}⊆ QS. Note that γ may not only contain one node, for example, if cγ = {x1, x2, x3}, we allow γ = {x1, x3}. Thus, we define the coupled MOT problem
$$\begin{aligned} \min{\mathbf{B}{c},c \in \mathcal{Q}} \sum{c \in \mathcal{Q}{C}} \langle \mathbf{C}{c}, \mathbf{B}{c} \rangle \ \text{subject to } P{c{i} \cap c{j}}(\mathbf{B}{c{i}}) = P{c{i} \cap c{j}}(\mathbf{B}{c{j}}), \quad \text{for each } c{i} \in \mathcal{Q}{S}, \text{ and all } c{j} \in N(c{i}), \ P{\gamma}(\mathbf{B}{c{\gamma}}) = \mu{\gamma}, \quad \text{for } c{\gamma} \in \mathcal{Q}{\Gamma} \subseteq \mathcal{Q}{S}, \qquad \langle \mathbf{B}{c}, \mathbf{1} \rangle = 1, \quad \text{for } c \in \mathcal{Q}{S} \backslash \mathcal{Q}{\Gamma}. \end{aligned}$$
The separators in QS enforce that the cost cliques are consistent with each other, and thereby play a similar role as the separator sets in Definition 1. To develop an algorithm for problem (42) we add a local regularization term for each clique that can contain a cost, i.e., for each c ∈ QC . Therefore, we generalize the expression for the reference joint distribution in (29). For c ∈ QC , let cj , ck ∈ QS be the two neighbours, and define Mc =cj ⊗ck , where ci = µγ ⊗ 1 if γ ⊂ ci for some γ ⊂ Γ, and otherwise ci = 1. We use the notation ⊗1 to broadcast a tensor to a bigger tensor of appropriate size. In the following, the broadcasting dimensions will always be obvious from the context. For example, if ci = {x1, x2, x3}, γ = {x2}, then ci (x1, x2, x3) = µγ (x2).
Moreover, we add a mass constraint similar to (30). The regularized problem thus reads
$$\min{\mathbf{B}{c}, c \in \mathcal{Q}} \sum{c \in \mathcal{Q}{C}} \langle \mathbf{C}{c}, \mathbf{B}{c} \rangle + \epsilon \mathcal{H}(\mathbf{B}{c} | \mathbf{M}{c})$$ subject to $P{c{i} \cap c{j}}(\mathbf{B}{c{i}}) = P{c{i} \cap c{j}}(\mathbf{B}{c{j}})$ , for each $c{i} \in \mathcal{Q}{S}$ , and all $c{j} \in N(c{i})$
$$P{\gamma}(\mathbf{B}{c{\gamma}}) = \mu{\gamma}, \quad \text{for } c{\gamma} \in \mathcal{Q}{\Gamma} \subseteq \mathcal{Q}{S}, \quad \langle \mathbf{B}{c}, \mathbf{1} \rangle = 1, \quad \text{for } c \in \mathcal{Q}{S} \backslash \mathcal{Q}{\Gamma}.$$ (43)
The problem in (43) is similar to (30), and we will extend the algorithm developed in Section 4 to handle general graphs. More precisely, following the derivations in Sec. 4, we can derive a method for approximating the solution of (43) similar to Algorithm 3. Note that the second set of constraints in (43) correspond to the constraints on Γ in the tree-structured optimal transport problem (30). Special care needs to be applied to the first set of constraints. In particular, in order to derive an efficient algorithm we need the neighbour sets in Definition 4.4 to be nested as described in the following assumption.
Assumption 1. For all c ∈ QS with at least two neighbors, there is a inclusion sequence (c ∩ cj1 ) = (c ∩ cj2 ) ⊇ · · · ⊇ (c ∩ cjℓ ) that includes all ℓ neighbours cj1 , . . . , cjℓ of c. In case a constraint set γ ⊂ Γ is embedded in the clique c, then this can be added after the smallest set in the inclusion, i.e., (c ∩ cjℓ ) ⊇ γ.
Example 4. We give an example of the nested structure in Assumption 1 in Fig. 9. Here, c = {x1, x3, x4}, c1 = {x1, x2, x3, x4}, c2 = {x1, x3, x4}, c3 = {x0, x1, x3}, and finally γ = {x3}.
Figure 9: Example on nested structure.
Proposition 2. The solution to (43) is of the form Bc = Kc ⊙Uc ⊙Mc, for each c ∈ QC , where Kc = exp(−Cc/ϵ) and Uc = u cj ij ⊗ 1 u ck ik ⊗ 1 for two smaller tensors u cj ij and u ck ik . These tensors are dual variables associated with the two neighbors cj , ck ∈ QS of c, and ij , ik denote the position in the corresponding inclusion sequence as defined in Assumption 1. More precisely, denoting ℓc the length of the inclusion sequence for each c ∈ QS, the dual problem for (43) reads
$$\max{\mathbf{u}{i}^{c}, i=1, \dots, \ell{c} \text{ for } c \in \mathcal{Q}{S}} - \sum{c \in \mathcal{Q}{C}} |\mathbf{K}{c} \odot \mathbf{U}{c} \odot \mathbf{M}{c}|{1} + \sum{c{\gamma} \in \mathcal{Q}{\Gamma}} \langle \boldsymbol{\mu}{\gamma}, \log(\mathbf{u}{\gamma}^{c}) \rangle + \sum{c \in \mathcal{Q}{S} \setminus \mathcal{Q}{\Gamma}} \rho{c} \mathbf{u}{\gamma}^{c} \text{ for } c \in \mathcal{Q}{S} \setminus \mathcal{Q}{\Gamma}$$
$$subject \text{ to } \sum{i=1}^{\ell{c}} \log(\mathbf{u}{i}^{c} \otimes \mathbf{1}) = \log(\mathbf{u}{\gamma}^{c} \otimes \mathbf{1}), \text{ for } c \in \mathcal{Q}_{\Gamma}$$
$$\sum{i=1}^{\ell{c}} \log(\mathbf{u}{i}^{c} \otimes \mathbf{1}) = \rho{c} \mathbf{1}, \text{ for } c \in \mathcal{Q}{S} \setminus \mathcal{Q}{\Gamma}.$$
$$(44)$$
The proof of proposition 2 follows the standard Lagrangue dual theory. Recall that any separator set c ∈ QS has a set of neighbours c1, . . . , cℓ ∈ QC that define an inclusion sequence (c∩c1) = (c∩c2) ⊇ · · · ⊇ (c∩cℓ) as in Assumption 1. To find the optimal dual variables of the dual problem (44) we propose to iterate through all c ∈ QS, and update the dual variables u c 1 , . . . u c ℓc
Due to the fourth condition in Definition 4, the modified junction tree forms a bipartite graph, when we consider the separators in QS as nodes, and the cliques in QC as edges. For example, in Fig. 7, the square node {0, J + 1} alone is one set in a bipartite graph, and the bottom gray nodes belong to another set. The bipartite structure allows for a parallelized algorithm. In each update we can apply coordinate ascent to all the variables associated with separators in one set together. Moreover, because those variables do not depend on each other, we can consider a subproblem of the dual, as we will describe in the following. Note that the iterative Bregman projection (Benamou et al. 2015) applies a similar idea to the Wasserstein barycenter problem. To be precise, for the updates related to each c ∈ QS it suffices to consider a subproblem of the dual (44) , which for c ∈ QΓ can be formulated as
$$\max{\substack{\mathbf{u}{i}, i=1, \dots, \ell \ \mathbf{v}{\ell+1}}} \sum{i=1}^{\ell} -\langle \mathbf{k}{i}, \mathbf{u}{i} \rangle + \langle \boldsymbol{\mu}{\gamma}, \log(\mathbf{v}{\ell+1}) \rangle \quad \text{subject to } \sum{i=1}^{\ell} \log(\mathbf{u}{i} \otimes \mathbf{1}) = \log(\mathbf{v}_{\ell+1} \otimes \mathbf{1}). \quad (45)$$
Here we have omitted the superscript of the dual variables. Moreover, we denote the dual variable uγ = vℓ+1, the reason for which will become apparent in the algorithm. Also, we introduce the notation ki := Pc∩ci (Kci ⊙ Mci ⊙ (u¯i ⊗ 1)), where u¯i denotes the second dual variable defining the tensor Uci together with ui , as described in Proposition 2. Thus, we can write ∥Kci ⊙ Uci ⊙ Mci ∥1 = ⟨ki , ui⟩. Note that the case c ∈ QS\QΓ is a special case of (45), where vℓ+1 = exp(ρc) ∈ R is a scalar and µγ = 1.
The optimal dual variables of the dual problem (44) can now be found by iterating through all c ∈ QS, and solving a problem of the form (45). This can be done by performing the steps described in the following.
Proposition 3. The solution to (45), where ℓ > 1, can be found by performing the following steps.
(1) Calculate geometric mean variables from the largest intersection cluster c ∩ c1 to the smallest c ∩ cℓ as
$$\mathbf{q}_2 = (\mathbf{k}_1 \odot \mathbf{k}_2)^{1/2}; \qquad \mathbf{q}i = (P{c \cap ci}(\mathbf{q}{i-1})^{i-1} \odot \mathbf{k}_i)^{1/i}, \qquad i = 3, \dots, \ell$$ (46)
(2) Calculate intermediate variables from aℓ, . . . , a1 and vℓ, . . . , v3 from the smallest intersection
cluster c ∩ cℓ to the largest cluster c ∩ c1 as
$$\mathbf{a}{i} = (\mathbf{v}{i+1} \otimes \mathbf{1})^{1/i} \odot \mathbf{q}_{i}, \qquad \text{for } i = 2, \dots, \ell$$
$$\mathbf{v}{\ell+1} = (\boldsymbol{\mu}{\gamma} \cdot / P{\gamma}(\mathbf{q}{\ell}))^{\ell}, \qquad \mathbf{v}{i} = (\mathbf{a}{i} \cdot / P{c \cap c{i}}(\mathbf{q}_{i-1}))^{i-1}, \quad \text{for } i = 3, \dots, \ell$$
$$(47)$$
(3) Update u1, . . . , uℓ as
$$\mathbf{u}_1 = \mathbf{a}_2 \cdot / \mathbf{k}_1, \quad \mathbf{u}_i = \mathbf{a}_i \cdot / \mathbf{k}_i \quad \text{for } i = 2, \dots, \ell,$$ (48)
Proposition 3 only discusses the case of ℓ > 1. In the case ℓ = 1, the constraint in (45) asserts that u1 = v2. Thus, the objective in (45) is optimized by u1 = µγ ./k1.
Remark 3. The method can be generalized to the case where the cost cliques have a larger number of neighbors, i.e., point 3. in Definition 4 is generalized to: For all c ∈ CC it holds |N(c)| ≤ k for some k ≥ 2. A parallelized algorithm for optimal transport problems on this type of junction tree can be set up by considering a k-partite partition, instead of the bipartite partition in the algorithm above.
We present numerical results for two types of MOT problems. We examine the Wasserstein barycenter problem as discussed in Fan et al. (2022a, Example 1) and structured according to the graph in Fan et al. (2022a, Figure 1), and the Wasserstein Least Squares (WLS) example detailed in Sec. 2.2.2. The original graph and the modified junction tree for the WLS example are depicted in Fig. 1b and Fig. 7, respectively. For the sake of completeness, we also show modified junction tree for the Wasserstein spline example in Fig. 8. Notably, the barycenter problem is tree-structured, while the WLS problem forms a graph with a tree-width of two.
For the barycenter problem, the cost matrices C(j,k) in (28) are set as the squared Euclidean distances. In the WLS problem, the cost matrices follow (18) with coefficient α = 10. The constrained marginal distributions {µk}k∈Γ, supported on a uniform grid of d points between 0 and 1, are generated from the log-normal distribution and normalized to sum to one. We set the accuracy threshold δ = 0.2 for the unregularized MOT problem.
Our comparison involves Algorithm 2, which solves a standard MOT problem with global regularization, and Algorithm 3, which addresses coupled bi-marginal OT with local regularization. We employ a random update rule for the global regularization method, consistent with the algorithm complexity analysis in Fan et al. (2022a). Our focus is on the number of iterations, anticipating that the parallel mechanism introduced by the local regularization method could reduce them. Each experiment is repeated five times with different random seeds, and the average number of iterations is reported in Fig. 10. Theoretical complexity bounds for the barycenter problem are also illustrated as dashed lines, with bounds for global and local regularization derived from Fan et al. (2022a, Theorem 2) and Theorem 3, respectively.
From the observations in Fig. 10a, it is evident that the experimental results are consistent with the theoretical bounds. Notably, local regularization demonstrates the potential to reduce the number of iterations by a factor of |E|, leading to a lower iteration count beyond a certain threshold |E|. Furthermore, both global and local regularization exhibit a logarithmic dependency on the dimension, with a log(d) relationship. In contrast, for the results presented in Fig. 10b, we have not established theoretical bounds for local regularization within this paper; this constitutes an area for future investigation.
Figure 10: Number of iterations for converging to the accuracy δ = 0.2. The left column shows the number of iterations as a function of d when |E| is fixed (|E| = 3 in (a) and |E| = 12 in (b)), and the right column vice versa with d = 10.
In this paper, we study the multi-marginal optimal transport (MOT) problem with graphstructured costs. We propose a novel approach to solve the MOT problem using local entropy regularization. Specifically, we reformulate the entire MOT problem as a sum of smaller MOT problems. When the underlying graph is a tree, these smaller problems involve dual variables that can be divided into two groups, forming a bipartite graph. We develop an efficient parallel algorithm that alternately updates these two groups of dual variables. A key advantage of this approach is its parallelizability, making it well-suited for GPU implementation. We also provide a rigorous complexity analysis of this approach and extend it to handle MOT problems associated with general graphs.
Fan J, Haasler I, Karlsson J, Chen Y (2022a) On the complexity of the optimal transport problem with graph-structured cost. International conference on artificial intelligence and statistics, 9147–9165 (PMLR).
Fan J, Liu S, Ma S, Zhou H, Chen Y (2021) Neural monge map estimation and its applications. arXiv preprint arXiv:2106.03812 .
Fan J, Taghvaei A, Chen Y (2020) Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462 .
Fan J, Zhang Q, Taghvaei A, Chen Y (2022b) Variational wasserstein gradient flow. International Conference on Machine Learning, 6185–6215 (PMLR).
Gangbo W, Swiech A (1998) Optimal maps for the multidimensional Monge-Kantorovich ´ problem. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences 51(1):23–45.
Haasler I, Karlsson J, Ringh A (2021a) Control and estimation of ensembles via structured optimal transport. IEEE Control Systems Magazine 41(4):50–69.
Haasler I, Ringh A, Chen Y, Karlsson J (2021b) Multimarginal optimal transport with a tree-structured cost and the Schr¨odinger bridge problem. SIAM Journal on Control and Optimization 59(4):2428–2453.
Haasler I, Ringh A, Chen Y, Karlsson J (2024) Scalable computation of dynamic flow problems via multimarginal graph-structured optimal transport. Mathematics of Operations Research 49(2):986–1011.
Haasler I, Singh R, Zhang Q, Karlsson J, Chen Y (2021c) Multi-marginal optimal transport and probabilistic graphical models. IEEE Transactions on Information Theory 67(7):4647–4668.
Halin R (1976) S-functions for graphs. Journal of geometry 8(1-2):171–186.
Koller D, Friedman N (2009) Probabilistic graphical models: principles and techniques (MIT press).
Korotin A, Selikhanovych D, Burnaev E (2022) Neural optimal transport. arXiv preprint arXiv:2201.12220 .
Kroshnin A, Tupitsa N, Dvinskikh D, Dvurechensky P, Gasnikov A, Uribe C (2019) On the complexity of approximating wasserstein barycenters. International conference on machine learning, 3530–3540 (PMLR).
Metcalf L, Casey W (2016) Cybersecurity and applied mathematics (Syngress).
Mezard M, Montanari A (2009) Information, physics, and computation (Oxford University Press).
Nenna L (2016) Numerical methods for multi-marginal optimal transportation. Ph.D. thesis.
Noble M, De Bortoli V, Doucet A, Durmus A (2023) Tree-based diffusion schr¨odinger bridge with applications to wasserstein barycenters. arXiv preprint arXiv:2305.16557 .
Pass B (2012) On the local structure of optimal measures in the multi-marginal optimal transportation problem. Calculus of Variations and Partial Differential Equations 43(3-4):529–536.
Pass B (2015) Multi-marginal optimal transport: theory and applications. ESAIM: Mathematical Modelling and Numerical Analysis 49(6):1771–1790.
Perrot M, Courty N, Flamary R, Habrard A (2016) Mapping estimation for discrete optimal transport. Advances in Neural Information Processing Systems 29.
Peyr´e G, Cuturi M, et al. (2017) Computational optimal transport. Center for Research in Economics and Statistics Working Papers (2017-86).
Rabin J, Peyr´e G, Delon J, Bernot M (2012) Wasserstein barycenter and its application to texture mixing. Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29–June 2, 2011, Revised Selected Papers 3, 435–446 (Springer).
Ringh A, Haasler I, Chen Y, Karlsson J (2021a) Efficient computations of multi-species mean field games via graph-structured optimal transport. 2021 60th IEEE Conference on Decision and Control (CDC), 5261–5268 (IEEE).
Ringh A, Haasler I, Chen Y, Karlsson J (2021b) Graph-structured tensor optimization for nonlinear density control and mean field games. arXiv preprint arXiv:2112.05645 .
Ringh A, Haasler I, Chen Y, Karlsson J (2023) Mean field type control with species dependent dynamics via structured tensor optimization. IEEE Control Systems Letters .
Scheinerman EA (2012) Mathematics: a discrete introduction (Cengage Learning).
Seguy V, Damodaran BB, Flamary R, Courty N, Rolet A, Blondel M (2017) Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283 .
Singh R, Haasler I, Zhang Q, Karlsson J, Chen Y (2022) Inference with aggregate data in probabilistic graphical models: An optimal transport approach. IEEE Transactions on Automatic Control 67(9):4483–4497.
Srivastava S, Cevher V, Dinh Q, Dunson D (2015) Wasp: Scalable bayes via barycenters of subset posteriors. Artificial Intelligence and Statistics, 912–920 (PMLR).
Trillos NG, Jacobs M, Kim J (2023) The multimarginal optimal transport formulation of adversarial multiclass classification. Journal of Machine Learning Research 24(45):1–56.
Villani C (2003) Topics in optimal transportation. Number 58 (American Mathematical Soc.).
Yang S, Zha H (2023) Estimating latent population flows from aggregated data via inversing multi-marginal optimal transport. Proceedings of the 2023 SIAM International Conference on Data Mining (SDM), 181–189 (SIAM).
Zanetti F, Gondzio J (2023) An interior point–inspired algorithm for linear programs arising in discrete optimal transport. INFORMS Journal on Computing 35(5):1061–1078.
Zhou B, Parno M (2024) Efficient and exact multimarginal optimal transport with pairwise costs. Journal of Scientific Computing 100(1):1–27.
Proof. Proof of Theorem 1 We perform a block-coordinate descent in (32). First, minimizing (32) with respect to u(j,kj ) for j ∈ Γ, while keeping the other variables fixed, yields the first set of updates. To minimize the objective with respect to ρj and u(j,k) for j /∈ Γ, relax the constraint with Lagrangian multiplier β, which gives the Lagrangian
$$L({\rhoj}, {u{(j,k)}}, {\betaj}) = \sum{(j,k)\in E} ||B_{(j,k)}||1 - \sum{j\in\Gamma} \log(u_{(j,k_j)})^\top \mu_j$$ (49)
$$-\sum_{j \in V \setminus \Gamma} \left( \rho_j + \betaj^{\top} \left( \sum{k \in N(j)} \log(u_{(j,k)}) - \rho_j \mathbf{1} \right) \right)$$ (50)
The Lagrangian is minimized with respect to u(j,k) for j /∈ Γ if
$$\betaj = u{(j,k)} \odot (K{(j,k)}(u{(k,j)} \odot \gamma_k)), \quad \text{for } k \in N(j).$$ (51)
Multiplying these relations for all k ∈ N(j) yields
$$\bigodot_{k \in N(j)} \beta_j = e^{\rhoj} \bigodot{k \in N(j)} \left( K{(j,k)} (u{(k,j)} \odot \gamma_k) \right).$$
Denote vj = J k∈N(j) (K(j,k)(u(k,j) ⊙ γk))1/|N(j)| , it follows that βj = e ρj /|N(j)|vj , which together with (51) leads to the updates for u(j,k) , j /∈ Γ. Note also that the Lagrangian (49) is only bounded in ρj if it holds β ⊤ j 1 = 1. Thus, e ρj /|N(j)|v ⊤ j 1 = 1 and the update for ρj follows.
Proof. Proof of Theorem 2 Define ∆t := f(u , ρt ) − f(u ∗ , ρ∗ ). Then according to the Lemma 3, Lemma 4, there is ∆t − ∆t+1 ≥ 1 22|E| max n ϵ∆t 2C∞ , δ′ o2 . On the one hand, we have ϵ 2∆t+1 88|E|C2∞ ≤ ϵ 2∆t 88|E|C2∞ − ϵ 2∆t 88|E|C2∞ 2 . Note that 0 ≤ at+1 ≤ at − a t implies 0 ≤ at ≤ 1. This implies whenever ∆t+1 ≥ 0, we have 0 ≤ 2∆t 88|E|C2∞ ≤ 1 for any t > 1. Also note that 0 ≤ at+1 ≤ at − a 2 t implies 1/at+1 − 1/at ≥ 1 and further implies 1/at − 1/a1 ≥ t − 1. Thus,
$$t \le 1 + \frac{88|E|C_{\infty}^2}{\epsilon^2 \Deltat} - \frac{88|E|C{\infty}^2}{\epsilon^2 \Delta_1}. \tag{52}$$
On the other hand,
$$\Delta_{t+\tau} \le \Delta_t - \frac{\delta'^2 \tau}{22|E|}, \quad t, \tau \ge 0 \tag{53}$$
Fixing an number ϕ > 0, and considering using (52) for ∆t ≥ ϕ and using (53) for ∆t < ϕ. We minimize the sum of two estimates over the parameter ϕ ∈ (0, ∆1]. Combining the above two inequalities, we get the total number of iterations satisfies
$$t \le \min_{0 \le \phi \le \Delta1} \left( 1 + \frac{88|E|C\infty^2}{\epsilon^2 \phi} - \frac{88|E|C_\infty^2}{\epsilon^2 \Delta1} + 1 + \frac{22|E|\phi}{\delta'^2} \right) \le 2 + \frac{88|E|C\infty}{\epsilon \delta'}. \tag{54}$$
Proof. Proof of Theorem 3 With the specific choices ϵ = 4|E| log(d) and δ ′ = 8C∞ , Lemma 6 gives
$$\sum{(j,k)\in E} \langle C{(j,k)}, \widehat{B}{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, B^*{(j,k)} \rangle \le \delta.$$ (55)
By Theorem 2, the stopping time t satisfies t ≤ 2 + 88|E|C∞ δ ′ϵ = O |E| 2C 2∞ log(d) δ . Since each iteration of Algorithm 3 requires O(|E|d 2 ) operations, a solution of Algorithm 3 is achieved in O |E| 3C 2∞d log(d) δ . On the other hand, Algorithm 4 takes O(|E|d ) (see Lemma 7 in Altschuler et al. (2017)). Therefore, the bound on total complexity follows.
Proof. Proof of Proposition 1 Firstly, when j ∈ Γ, we have γj = µj , and
$$u_{(j,kj)}^{t+1} \stackrel{(33)}{=} \mathbf{1}./(K{(j,kj)}u{(kj,j)}^t) \stackrel{(31)}{=} \mathbf{1}./((B{(jkj)}^t\mathbf{1})./(u{(j,k_j)}^t\odot \muj)) \stackrel{(35)}{=} u{(j,k_j)}^t\odot \muj./q{(j,k_j)}^t. \tag{56}$$
When j /∈ Γ, we have γj = 1. It follows that
$$vj^{t+1} \stackrel{(34)}{=} \left( \bigodot{k \in N(j)} \left( K{(j,k)}(u{(k,j)}^t \odot \gammak) \right) \right)^{\frac{1}{|N(j)|}} \stackrel{(31)}{=} \left( \bigodot{k \in N(j)} \left( B{(j,k)}^t \mathbf{1} . / u{(j,k)}^t \right) \right)^{\frac{1}{|N(j)|}}$$ (57)
$$= \left( \bigodot{k \in N(j)} q{(j,k)}^t \right)^{\frac{1}{|N(j)|}} / e^{\rho_j^t/|N(j)|} = q_j^{t+1} / e^{\rho_j^t/|N(j)|}, \tag{58}$$
where we use the constraint P k∈N(j) log(u(j,k)) = ρj1 of the dual problem (32) in the third equality, and the definition of q t+1 in the last equality. This implies that
$$\rho_j^{t+1} = -|N(j)|\log\left(\mathbf{1}^\top v_j^{t+1}\right) = \rho_j^t - |N(j)|\log(\mathbf{1}^\top q_j^{t+1}) = \rho_j^t - |N(j)|\log(|q_j^{t+1}|_1). \tag{59}$$
Furthermore,
$$u_{(j,k)}^{t+1} \stackrel{(34)}{=} e^{\rho_j^{t+1}/|N(j)|} vj^{t+1} \cdot / (K{(j,k)}(u_{(k,j)}^t \odot \gamma_k))$$
$$\tag{60}$$
$$\stackrel{\text{(31)}}{=} e^{(\rho_j^{t+1} - \rho_j^t)/|N(j)|} qj^{t+1} \cdot / (B{(j,k)}^t \mathbf{1} \cdot / u{(j,k)}^t) \stackrel{\text{(59)}}{=} u{(j,k)}^t \odot \frac{q_j^{t+1}}{|q_j^{t+1}|1} \cdot / q{(j,k)}^t. \tag{61}$$
Proof. Proof of Proposition 3 We introduce auxiliary variables vi , i = 1, . . . , ℓ, and rewrite the constraints (45), resulting in the following reformulation of the problem
$$\max{\substack{\mathbf{u}{i}, i=1,\dots,\ell\\mathbf{v}{i}, i=3,\dots,\ell+1}} -\sum{i=1}^{\ell} \langle \mathbf{k}{i}, \mathbf{u}{i} \rangle + \langle \boldsymbol{\mu}{\boldsymbol{\gamma}}, \log(\mathbf{v}{\ell+1}) \rangle$$ (62)
subject to $$log(\mathbf{u}_1) + log(\mathbf{u}_2) = log(\mathbf{v}_3 \otimes \mathbf{1}),$$ (63)
$$log(\mathbf{v}_i) + log(\mathbf{u}i) = log(\mathbf{v}{i+1} \otimes \mathbf{1}), \quad \text{for } i = 3, ..., \ell - 1$$ (64)
$$\log(\mathbf{v}{\ell}) + \log(\mathbf{u}{\ell}) = \log(\mathbf{v}_{\ell+1} \otimes \mathbf{1}). \tag{65}$$
Note here that vi and ui have the same dimension. The Lagrange relaxed problem of (62) is
$$\max{\substack{\mathbf{u}{i}, i=1, \dots, \ell \ \mathbf{v}{i}, i=3, \dots, \ell+1}} - \sum{i=1}^{\ell} \langle \mathbf{k}{i}, \mathbf{u}{i} \rangle + \langle \boldsymbol{\mu}{\gamma}, \log(\mathbf{u}{\ell+1}) \rangle + \langle \mathbf{a}{2}, \log(\mathbf{u}{1}) + \log(\mathbf{u}{2}) - \log(\mathbf{v}{3} \otimes \mathbf{1}) \rangle \ + \sum{i=3}^{\ell-1} \langle \mathbf{a}{i}, \log(\mathbf{v}{i}) + \log(\mathbf{u}{i}) - \log(\mathbf{v}{i+1} \otimes \mathbf{1}) \rangle + \langle \mathbf{a}{\ell}, \log(\mathbf{v}{\ell}) + \log(\mathbf{u}{\ell}) - \log(\mathbf{v}_{\ell+1} \otimes \mathbf{1}) \rangle,$$ (66)
where ai are newly introduced Lagrange variables. Maximizing with respect to vi we get ai = Pc∩ci (ai−1) for i = 3, . . . , ℓ. Maximizing with respect to vℓ+1, we get that Pγ(aℓ) = µγ . Further, maximizing with respect to ui , for i = 1, . . . , ℓ, we get (48). To satisfy the constraint (63) we need u1 ⊙ u2 = v3 ⊗ 1. Plugging in the expressions for u1 and u2 in (48), we get the expression for a2 in (47) together with (46). The update of v3 in (47) follows from a3 = Pc∩ci (a2). Similarly, the remaining updates for qi , ai , and vi , for i = 3, . . . , ℓ, follow from the constraints vi⊙ui = uℓ+1⊗1 and ai = Pc∩ci (ai−1). Finally, the expression for vℓ+1 is a consequence of µγ = Pγ(aℓ).
Proof. Proof of Lemma 1 When j ∈ Γ S, the updates of u(j,k) would keep ∥B t+1 (j,kj ) ∥1 = 1 since it matches the projection B t+1 (j,kj ) 1 and the true marginal µj . When j ∈ (V \Γ) T S, by proposition 1, the updates of u(j,k) are equivalent to (34) in Theorem 1, thus they satisfy (51), which lead to ∥B t+1 (j,k) ∥ = 1 ⊤β t+1 . Since 1 ⊤β t+1 j = 1 holds, we get ∥B t+1 (j,k) ∥ = 1.
Proof. Proof of Lemma 2 By (33), for j ∈ Γ, following the logic in Dvurechensky et al. (2018, Lemma 1), there is
$$\max{i=1,\dots,d} (\lambda{(j,k)}^t)i - \min{i=1,\dots,d} (\lambda_{(j,k)}^t)i \le ||C{(j,k)}||_{\infty}.$$ (67)
By (34), for j ∈ V \Γ, following Kroshnin et al. (2019, Lemma 4), there is
$$\max{i=1,\dots,d} (\lambda{(j,k)}^t)i - \min{i=1,\dots,d} (\lambda_{(j,k)}^t)i \le ||C{(j,k)}||{\infty} + \frac{1}{|N(j)|} \sum{i \in N(j)} ||C{(j,i)}||{\infty}.$$ (68)
Combining these results gives the bound.
Proof. Proof of Lemma 3 By the gradient inequality for convex functions it holds that
$$f(u^{t}, \rho^{t}) - f(u^{}, \rho^{}) = f(\lambda^{t}, \rho^{t}) - f(\lambda^{}, \rho^{})$$
$$\leq \langle \nabla f(\lambda^{t}, \rho^{t}), ((\lambda^{t}, \rho^{t}) - (\lambda^{}, \rho^{}))^{\top} \rangle$$
$$= \sum{j \in \Gamma} \langle \lambda^{t}{(j,k{j})} - \lambda^{*}{(j,k{j})}, \frac{1}{\epsilon} B^{t}{(j,k{j})} \mathbf{1} - \frac{1}{\epsilon} \mu{j} \rangle$$
$$+ \sum{j \notin \Gamma} \sum{k \in N(j)} \langle \lambda^{t}{(j,k)} - \lambda^{*}{(j,k)}, \frac{1}{\epsilon} q^{t}{(j,k)} \rangle - \sum{j \notin \Gamma} (\rho^{t}{j} - \rho^{*}{j})$$
$$= \frac{1}{\epsilon} \sum{j \in \Gamma} \langle \lambda^{t}{(j,k{j})} - \lambda^{*}{(j,k{j})}, B^{t}{(j,k{j})} \mathbf{1} - \mu{j} \rangle$$
$$+ \sum{j \notin \Gamma} \sum{k \in N(j)} \left\langle \frac{1}{\epsilon} \lambda^{t}{(j,k)} - \frac{1}{|N(j)|} \rho^{t}{j} \mathbf{1} - \frac{1}{\epsilon} \lambda^{}_{(j,k)} - \frac{1}{|N(j)|} \rho^{}{j} \mathbf{1}, q^{t}{(j,k)} \right\rangle.$$ (69)
Here, we have used that ⟨1, qt (j,k) ⟩ = 1, and thus −ρ t j + ρ ∗ j = ⟨−ρ t j1 + ρ ∗ j1, qt (j,k) ⟩. We bound the first sum as in Fan et al. (2022a, proof of Lemma 2) (compare also Dvurechensky et al. (2018, proof of Lemma 2)), which leads to
$$\sum{j \in \Gamma} \langle \lambda{(j,kj)}^t - \lambda{(j,kj)}^*, B{(j,k_j)}^t \mathbf{1} - \muj \rangle \le 2C\infty \sum{j \in \Gamma} |B{(j,k_j)}^t \mathbf{1} - \mu_j|_1.$$ (70)
To bound the second sum in the final expression of (69) note that for any j /∈ Γ,
$$\sum{k \in N(j)} \left\langle \frac{1}{\epsilon} \lambda{(j,k)}^t - \frac{1}{|N(j)|} \rho_j^t \mathbf{1} , \ \bar{q}j^t \right\rangle = \left\langle \sum{k \in N(j)} \frac{1}{\epsilon} \lambda_{(j,k)}^t - \rho_j^t \mathbf{1} , \ \bar{q}_j^t \right\rangle = 0,$$
since P k∈N(j) λ t (j,k) /ϵ = ρ t j1. The same equality holds with λ ∗ (j,k) and ρ ∗ . Moreover, note that it holds
$$\sum{k \in N(j)} \langle \mathbf{1}, q{(j,k)}^t - \bar{q}j^t \rangle = \left\langle \mathbf{1}, \sum{k \in N(j)} q_{(j,k)}^t - |N(j)| \bar{q}_j^t \right\rangle = 0.$$
Using these relations it follows
$$\sum{k \in N(j)} \left\langle \frac{1}{\epsilon} \lambda{(j,k)}^t - \frac{1}{|N(j)|} \rhoj^t \mathbf{1} - \frac{1}{\epsilon} \lambda{(j,k)}^* - \frac{1}{|N(j)|} \rhoj^* \mathbf{1} , q{(j,k)}^t \right\rangle$$ (71)
$$= \frac{1}{\epsilon} \sum{k \in N(j)} \left\langle \lambda{(j,k)}^t - \lambda{(j,k)}^*, \ q{(j,k)}^t - \bar{q}j^t \right\rangle \le \frac{2C\infty}{\epsilon} \sum{k \in N(j)} |B{(j,k)}^t \mathbf{1} - \bar{q}_j^t|_1. \tag{72}$$
Here, the last step follows with the same approach as for (70) (see also Kroshnin et al. (2019, proof of Lemma 5)). Finally, combining (69), (70) and (72) completes the proof.
Proof. Proof of Lemma 4 Using the relations in Algorithm 3 and Lemma 1, we get
$$f(u^t, \rho^t) - f(u^{t+1}, \rho^{t+1}) = \sum{j \in \Gamma} (\log u^{t+1}{(j, kj)} - \log u^t{(j, k_j)})^\top \muj + \sum{j \notin \Gamma} (\rho^{t+1}_j - \rho^t_j)$$ (73)
$$= \sum_{j \in \Gamma} \left( \mathcal{H} \left( \muj | B{(j,kj)}^t \mathbf{1} \right) + 1 \right) - \sum{j \notin \Gamma} |N(j)| \log \left( ||q_j^{t+1}||_1 \right). \tag{74}$$
By Pinsker's inequality the first sum can be bounded as
$$\sum_{j \in \Gamma} \left( \mathcal{H} \left( \muj | B{(j,kj)}^t \mathbf{1} \right) + 1 \right) \ge \frac{1}{2} \sum{j \in \Gamma} | B_{(j,k_j)}^t \mathbf{1} - \mu_j |_1^2$$
Moreover, since ∥q t+1 j ∥1 ≤ 1 and ∥q t (j,k) ∥1 = 1, we can bound
$$-\log\left(|q_j^{t+1}|_1\right) \ge 1 - |q_j^{t+1}|1 = \frac{1}{|N(j)|} \sum{k \in N(j)} |q_{(j,k)}^t|_1 - |q_j^{t+1}|_1 \tag{75}$$
$$= \mathbf{1}^{\top} \left( \frac{1}{|N(j)|} \sum{k \in N(j)} q{(j,k)}^t - qj^{t+1} \right) \ge \frac{1}{11|N(j)|} \sum{k \in N(j)} |q_{(j,k)}^t - \bar{q}_j^t|_1^2, \tag{76}$$
where the last inequality is proved in (Kroshnin et al. 2019, proof of Lemma 6). Therefore,
$$f(u^{t}, \rho^{t}) - f(u^{t+1}, \rho^{t+1}) \ge \frac{1}{2} \sum{j \in \Gamma} |B{(j,k{j})}^{t} \mathbf{1} - \mu{j}|{1}^{2} + \frac{1}{11} \sum{j \notin \Gamma} \sum{k \in N(j)} |q{(j,k)}^{t} - \bar{q}{j}^{t}|{1}^{2}$$
$$(77)$$
$$\geq \frac{1}{11} \left( \sum{j \in \Gamma} |B{(j,k_j)}^t \mathbf{1} - \mu_j|1^2 + \sum{j \notin \Gamma} \sum{k \in N(j)} |q{(j,k)}^t - \bar{q}_j^t|_1^2 \right)$$ (78)
$$\geq \frac{1}{22|E|} \left( \sum{j \in \Gamma} |B{(j,k_j)}^t \mathbf{1} - \mu_j|1 + \sum{j \notin \Gamma} \sum{k \in N(j)} |q{(j,k)}^t - \bar{q}_j^t|_1 \right)^2, \quad (79)$$
where we use Cauchy-Schwarz in the final inequality.
Proof. Proof of Lemma 5 Denote the partition where the nodes have mismatch as S. Without loss of generality, we assume j ∈ S and k ∈ V \S.
$$\sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, \widehat{B}{(j,k)} \rangle = \sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)} - \widehat{B}{(j,k)} \rangle$$ (80)
$$= \sum{j \in \Gamma} \langle C{(j,kj)}, B{(j,kj)} - \widehat{B}{(j,kj)} \rangle + \sum{j \notin \Gamma} \sum{k \in N(j)} \langle C{(j,k)}, B{(j,k)} - \widehat{B}{(j,k)} \rangle. \tag{81}$$
By H¨older's inequality and Lemma 7 in Altschuler et al. (2017), the first part can be bounded and
$$\sum{j \in \Gamma} \langle C{(j,kj)}, B{(j,kj)} - \widehat{B}{(j,kj)} \rangle \le \sum{j \in \Gamma} |C_{(j,kj)}|{\infty} |B_{(j,kj)} - \widehat{B}{(j,k_j)}|1 \le 2C{\infty} \sum_{j \in \Gamma} |\muj - B{(j,k)}\mathbf{1}|_1.$$ (82)
Similarly, for the second part,
$$\sum{j \notin \Gamma} \sum{k \in N(j)} \langle C{(j,k)}, B{(j,k)} - \widehat{B}{(j,k)} \rangle \le \sum{j \notin \Gamma} \sum{k \in N(j)} |C{(j,k)}|{\infty} |B{(j,k)} - \widehat{B}{(j,k)}|{1}$$ (83)
$$\leq 2 \sum{j \notin \Gamma} \sum{k \in N(j)} |C{(j,k)}|{\infty} |\bar{q}j - B{(j,k)} \mathbf{1}|1 \leq 2C{\infty} \sum{j \notin \Gamma} \sum{k \in N(j)} |\bar{q}j - B{(j,k)} \mathbf{1}|_1.$$ (84)
Summing them up gives
$$\sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, \widehat{B}{(j,k)} \rangle$$ (85)
$$\leq 2C{\infty} \left( \sum{j \in \Gamma} |\muj - B{(j,k_j)} \mathbf{1})|1 + \sum{j \notin \Gamma} \sum_{k \in N(j)} |\bar{q}j(B) - B{(j,k)} \mathbf{1}|_1 \right). \tag{86}$$
It follows the exact same way to obtain the same bound for P (j,k)∈E ⟨C(j,k) , Bb (j,k)⟩−P (j,k)∈E ⟨C(j,k) , B(j,k)⟩.
Proof. Proof of Lemma 6 Let {Yb (j,k)} denote the output from Algorithm 4 with inputs {B∗ (j,k) and {Be (j,kj )1}j∈Γ. Note that Be (j,k) is the optimal solution to
$$\min{B \in \Pi(\tilde{B}{(j,k)}\mathbf{1}, \; \tilde{B}{(j,k)}^{\top}\mathbf{1})} \langle C{(j,k)}, B \rangle + \epsilon \mathcal{H}(B|M_{(j,k)}), \quad \text{for } \forall (j,k) \in E$$ (87)
which can be verified by checking KKT conditions. Thus it holds that
$$\langle C{(j,k)}, \widetilde{B}{(j,k)} \rangle + \epsilon \mathcal{H}(\widetilde{B}{(j,k)}|M{(j,k)}) \le \langle C{(j,k)}, \widehat{Y}{(j,k)} \rangle + \epsilon \mathcal{H}(\widehat{Y}{(j,k)}|M{(j,k)}). \tag{88}$$
Since ⟨Be (j,k) , log(Be (j,k))⟩ ≥ −2 log(d) and ⟨Yb (j,k) , log(Yb (j,k))⟩ ≤ 0 it follows that
$$\langle C{(j,k)}, \widetilde{B}{(j,k)} \rangle - \langle C{(j,k)}, \widehat{Y}{(j,k)} \rangle \le \epsilon \mathcal{H}(\widehat{Y}{(j,k)}|M{(j,k)}) - \epsilon \mathcal{H}(\widetilde{B}{(j,k)}|M{(j,k)})$$ (89)
$$\leq -\epsilon \langle \widetilde{B}{(j,k)}, \log(\widetilde{B}{(j,k)}) \rangle + \epsilon \langle \widetilde{B}{(j,k)} - \widehat{Y}{(j,k)}, \log M_{(j,k)} \rangle \tag{90}$$
$$\leq 2\epsilon \log(d) + \epsilon \langle \log \gammaj, \widetilde{B}{(j,k)} \mathbf{1} - \widehat{Y}_{(j,k)} \mathbf{1} \rangle + \epsilon \langle \log \gammak, \widetilde{B}{(j,k)}^{\top} \mathbf{1} - \widehat{Y}_{(j,k)}^{\top} \mathbf{1} \rangle = 2\epsilon \log(d), \quad (91)$$
where the last equality holds because Be (j,k)1 = Yb (j,k)1 whenever j ∈ Γ, and γj = 1 whenever j /∈ Γ.
The plans {B∗ (j,k) } already satisfy the condition q¯j (B∗ ) = B∗ (j,k) 1 for ∀j ∈ Γ, k ∈ N(j). Thus Lemma 5 gives
$$\sum{(j,k)\in E} \langle C{(j,k)}, \hat{Y}{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, B{(j,k)}^* \rangle$$ (92)
$$\leq 2C{\infty} \left( \sum{j \in \Gamma} |\muj - \widetilde{B}{(j,k_j)} \mathbf{1})|_1 \right), \tag{93}$$
$$\sum{(j,k)\in E} \langle C{(j,k)}, \widehat{B}{(j,k)} \rangle - \sum{(j,k)\in E} \langle C{(j,k)}, \widetilde{B}{(j,k)} \rangle$$ (94)
$$\leq 2C{\infty} \left( \sum{j \in \Gamma} |\muj - \widetilde{B}{(j,k_j)} \mathbf{1})|1 + \sum{j \notin \Gamma} \sum_{k \in N(j)} |\bar{q}j(\widetilde{B}) - \widetilde{B}{(j,k)} \mathbf{1}|_1 \right). \tag{95}$$