{'fraction_non_alphanumeric': 0.04622155730290168, 'fraction_numerical': 0.01567392665785262, 'mean_word_length': 4.426184632811025, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Multi-marginal optimal transport and probabilistic graphical models

We study multi-marginal optimal transport problems from a probabilistic graphical model perspective. We point out an elegant connection between the two when the underlying cost for optimal transport allows a graph structure. In particular, an entropy regularized multi-marginal optimal transport is equivalent to a Bayesian marginal inference problem for probabilistic graphical models with the additional requirement that some of the marginal distributions are specified. This relation on the one hand extends the optimal transport as well as the probabilistic graphical model theories, and on the other hand leads to fast algorithms for multi-marginal optimal transport by leveraging the well-developed algorithms in Bayesian inference. Several numerical examples are provided to highlight the results.

From a numerical point of view, this connection allows for adapting existing PGM algorithms to this class of MOT problems with graphical structured cost. In this work, we focus on belief propagation (BP) [35], [30] and an extension of it known as norm-product algorithm [36]. These belong to the so called message passing algorithms which exploit the underlying graphical structure, and rely on exchanging information between the nodes. Thus, only local updates are needed, which greatly reduces the computational complexity of the inference problem. If the graphical model is a tree, these algorithms converge globally to the exact conditional marginal distributions in a finite number of iterations. For general PGMs with cycles, there is no convergence guarantee, but both methods usually work well in practice and provide relatively accurate approximations of the marginals. In this work, we develop algorithms for solving entropy regularized MOT problems, or equivalently constrained inference problems, by combining these message passing algorithms with the iterative scaling method. Moreover, we build on earlier results and establish global convergence of our algorithms for tree graphs. Similar constrained inference problems have previously been studied in [37]. Interestingly, the algorithm presented therein is comparable to our proposed extension of the BP algorithm. With the connections to MOT, we provide a new motivation for studying this problem, which also leads to a more complete picture of the algorithms.
A promising application of our framework are in the inference problems for collective dynamics, for instance, in the estimation of the behavior of large groups from only aggregate measurements. Such types of filtering methods are crucial for collective dynamics since it is usually impossible to track the trajectories of each single agent in a large population, due to exploding computational complexity, lack of sensor data or for privacy considerations. Related problems have been studied under the name collective graphical models (CGMs) [38], [39], [40]. These works consider a large collection of identical graphical models, which are observed simultaneously, and aim to infer aggregate distribution over the nodes. Several heuristic algorithms [40] have been proposed to solve the resulting inference problems. Our MOT framework suggest a different observation model [41], [42], which is reasonable in many scenarios. More importantly, our algorithms, which enjoy global convergence guarantee, provide a reliable machine to estimate collective dynamics in these models.
The rest of this paper is structured as follows. In Section II we review some background knowledge in optimal transport and probabilistic graphical models. In Section III we provide the main theoretical result in this paper, which is the equivalence between entropy regularized MOT and the inference problem for PGMs. We also modify the belief propagation algorithm to solve MOT problems. Another algorithm based on the norm-product algorithm is introduced in Section IV. We test and verify our results through several numerical examples in Section V. This is followed by a brief concluding remark in Section VI.
Notation: The notation used throughout is mostly standard. However, with exp(·), ln(·), , and ./ we denote the element-wise exponential, logarithm, multiplication, and division of vectors, matrices, and tensors, respectively. Moreover, ⊗ denotes the outer product. By 1 we denote a vector of ones, the size of which will be clear from the context. Throughout, we use bold symbols to represent vectors, e.g., bj, µ j , and regular symbol for the corresponding entries, e.g., bj(xj), µj(xj).

II. PRELIMINARIES
In this section, we provide a quick overview of optimal transport theory, probabilistic graphical models and belief propagation algorithm. We only cover material that is most relevant to this work. The reader is referred to [2], [25], [26] for more details.

A. Optimal transport
In optimal transport (OT) problems, one seeks an optimal plan that transports mass from a source distribution to a target distribution with minimum cost. In its original formulation [43], OT was studied over Euclidean space. However, in general, OT problems can be formulated in both continuous space and discrete space. In this work, we focus on optimal transport over discrete space.
Let µ 1 ∈ R d 1 + , µ 2 ∈ R d 2 + be two discrete distributions, viz., nonnegative vectors, with equal mass, that is i 1 µ1(i1) = i 2 µ2(i2). Here µ1(i) denotes the amount of mass in the source distribution at location i and µ2(i) denotes the amount of mass in the target distribution at location i. Without loss of generality, we assume that both µ 1 and µ 2 are probability vectors, that is, the total mass is The transport cost of moving a unit mass from point i1 to i2 is denoted by C(i1, i2), and collected in the matrix C = [C(i1, i2)] ∈ R d 1 ×d 2 . In the Kantorovich formulation [44] of OT, the goal is to find a transport plan between the two marginal distributions µ 1 and µ 2 that minimizes the total transport cost. A transport plan is encoded in a joint probability matrix B = [B(i1, i2)] ∈ R d 1 ×d 2 + of µ 1 , µ 2 . Then the total transport cost is i 1 ,i 2 C(i1, i2)B(i1, i2) = trace(C T B) and therefore the OT problem reads min where 1 denotes a vector of ones of proper dimension. The constraints are to enforce that B is a joint distribution between µ 1 and µ 2 .
Even though the above OT problem (1) is a linear program, in many practical applications it is too difficult to be solved directly using standard solvers due to the large number of variables [23], [17], especially in the case where the marginal distributions µ 1 , µ 2 come from discretizations of continuous measures. Recently, a regularization of OT [10] was proposed that greatly reduces the computational complexity of (approximately) solving OT problems over discrete space. In this method, an entropy term is added to regularize the problem, leading to min where > 0 is a regularization parameter. The entropy regularized OT problem (3) is strictly convex and thus the solution is unique. More importantly, it can be solved efficiently via the Sinkhorn algorithm [45], [10], also known as iterative scaling [11], [46].
, then the iterative scaling updates alternate between the two steps where ./ denotes element-wise division. The algorithm converges linearly to a unique pair of vectors u1 ∈ R d 1 , u2 ∈ R d 2 up to a normalization [46]. Given the limit point of the iteration, the solution to (3) has the form

B. Probabilistic graphical models
A probabilistic graphical model (PGM) is a graph-based representation of a collection of random vectors that captures the conditional dependencies between them. It provides a compact representation of their joint distributions through factorization: a graphical model consists of a collection of distributions that factorize according to an underlying graph structure. In this work we focus on undirected graphs, which represent Markov random fields [25]. Note that directed graphs represent Bayesian networks, and can always be transformed into a Markov random field [26].
Among the many representations for Markov random fields, the factor graph representation has been widely used due to its elegance and flexibility [25], and is also used in this paper. Consider a graphical model with underlying factor graph G = (V, F, E) where V denotes the set of variable nodes, F denotes the set and factor nodes, and E stands for the edges connecting them (see Figure 1 for an example). In such a factor graph G, the neighbors of a node j ∈ V consists of factor nodes, N (j) ⊂ F , and the neighbors of a factor node α ∈ F are variable nodes N (α) ⊂ V . Therefore, G is a bipartite graph [47]. Each variable node j ∈ V is associated with a random variable xj which can be either discrete or continuous. Here we consider only the discrete cases and assume that the random variable xj can take dj possible values. Each factor node α ∈ F corresponds to the dependence between the variable nodes connected to α, which are compactly denoted by xα := {xj ; j ∈ N (α)}. In Markov random fields with underlying factor graph G, the joint probability is assumed to be of the form where φ j is the node/local potential corresponding to node j, ψ α is the factor node potential corresponding to factor node α, and Z is a normalization constant. A factor node potential ψα describes the dependence between random variables in {xj ; j ∈ N (α)}. The node potentials normally come from two sources: prior belief and evidence from measurements. In the latter, φj(xj) is short for φj(xj, yj) [33] with yj being the measurement. In cases where all the local potentials are induced by evidence, a more precise formula for the model is Since the measurement is usually specified in inference problems, yj is often neglected to simplify the notation. In principle, the node potentials can be fully absorbed into the factor potentials, that is, the joint distribution becomes p(x) ∝ α∈F ψα(xα). For the ease of presentation, we adopt the formulation (6).
Apart from factor graphs, another popular representation of PGMs is the standard graph where the nodes are all variables. In the standard graph representation, the dependence between the variables are fully captured by the edges of the graphs. The two representations are equivalent and one can transform one to another easily as the following example illustrates. Example 1. The factor graph in Figure 1a models the joint distribution To convert this into a standard graph representation, the dependence among variables induced by the three factors have to be translated to edges. This is straightforward for factors ψ α 2 and ψ α 3 . The factor ψ α 1 involves 4 variables and is more complicated. Without further assumptions on the structure of this factor, it may induce dependence among all of the 4 variables, and thus a complete graph connecting them is required (see Figure 1b).
The two fundamental problems in PGMs are inference and learning. Learning refers to estimating the underlying graphical models (often the parametrized factors) using available data sampled from the models. In inference problems, the parameters of the models are given. Instead, the goal is to infer the statistics of the node variables. The two main approaches to inference problems are maximum a posteriori estimation and Bayesian marginal inference [25]. Given a set of measurements {y1, y2, . . . , yJ }, the aim of maximum a posteriori estimation x 1 x 2 x 5 x 4 x 3 x 1 x 2 x 5 x 4 x 3 x 6 (b) Fig. 1: Factor graph representation vs standard graph representation.
is to find the most likely variable value {x1, x2, . . . , xJ } given the model and measurements. Instead, Bayesian inference calculates the posterior marginal distributions of each node. The focus of this paper is most relevant to Bayesian/marginal inference.
Formally, given a graphical model (6), the objective of Bayesian inference is to calculate marginal distributions p(xj) for j ∈ V . In cases where the nodes variables are discrete, the marginal p(xj) is defined to be The Bayesian inference problem can be reformulated as the optimization problem of minimizing with and over the space of all the probability distributions on x. By convention, F, U, H are known as free energy, average energy and entropy respectively due to their similarity to thermodynamics [48]. To see this, we note that the Kullback-Leibler (KL) divergence [49] between b(x) and p(x) is Since the KL divergence is nonnegative and equals 0 only if b = p, the unique minimizer of F(b) is p with the associated minimum being − ln Z. The optimization formulation (9) of Bayesian inference is the basis for variational inference [25], one of the most popular approximate inference techniques. In the variational inference approach, the approximate distribution b is usually assumed to have some simple structure to ease the optimization, e.g., the mean field approximation b(x) = b1(x1) · · · bJ (xJ ) [50]. This work is not concerned with variational inference; (9) simply serves as a link to connect Bayesian inference with optimization. In the PGM literature [51], [36], it is common to introduce a temperature coefficient > 0 into (9), which leads to a slightly more general optimization problem It corresponds to the Bayesian inference for the model From an optimization point of view, (12) is a regularized version of the linear program Interestingly, this linear program in fact corresponds to the maximum a posteriori problem [36] for the model (6).

C. Belief Propagation
In principle, Bayesian inference is achievable through the definition (8) by calculating the marginals using brute force summation. The complexity of this summation however scales exponentially as the number of variable nodes J goes up [33]. Also the normalization factor Z is extremely difficult to calculate when J is large due to the same reason.
During the last two decades, many methods have been developed to solve or approximately solve Bayesian marginal inference problems. One of the most widely used methods is a message-passing algorithm called Belief Propagation [35]. It updates the marginal distribution of each node through communications of beliefs/messages between them. In the factor graph representation, it reads where mα→j(xj) denotes the message from factor node α to variable node j, and nj→α(xj) represents the message from variable node j to factor node α. The symbol ∝ means "proportional to" and indicates that often a normalization is applied in the Belief Propagation algorithm. The messages in (13) are updated iteratively over the factor graph.
The Belief Propagation algorithm was first invented to solve Bayesian inference program over trees, in which case global convergence is guaranteed [35], [33]. This method was later generalized to deal with inference problems involving general graphs under the name Loopy Belief Propagation [30]. Even though there is no convergence proof and the algorithm does diverge in some occasions, it works well in practice and is widely adopted. When the algorithm converges, one can calculate the beliefs on the variables and factors by In cases where the factor graph has no cycles (i.e., it is a tree), the beliefs in (14) coincide with the true posterior marginals, that is, For general graphs with cycles, convergence is not guaranteed and even if it does converge, the beliefs in (14) are only approximations of the true marginals p(xj), p(xα). A remarkable discovery [30], [31] related to (Loopy) Belief Propagation is that if the updates (13) converge, then the beliefs in (14) form a fixed point of the Bethe free energy [30], [31] where U Bethe (b) is the Bethe average energy and H Bethe (b) is the Bethe entropy with Nj denoting the degree of the variable node j, i.e., Nj = |N (j)|. In (16), we define b = {bj, bα : j ∈ V, α ∈ F }. This is different to b in (9), which is a J-mode tensor. For the sake of conciseness, by abuse of notation, we use b in both settings. For a factor tree, the two are connected through the relation b(x) ∼ ( α∈F bα(xα))( j∈V bj(xj) 1−N j ) [25]. In terms of Bethe free energy, the Bayesian inference problem reads The constraint (19b) is to ensure that bα, bj are compatible and (19c) is to guarantee that they are in the probability simplex. It is easy to check that when the factor graph has no cycles, the Bethe free energy (16) is strictly convex in the feasible set defined by the constraints (19b)-(19c), and is equal to the free energy (9), i.e., F Bethe = F. Thus, (19) is again a convex optimization problem. For general graphs, the Bethe free energy serves as a good approximation of the free energy [31], but is no longer convex.

III. MULTIMARGINAL OPTIMAL TRANSPORT AS BAYESIAN INFERENCE
Multimarginal optimal transport (MOT) extends the OT framework (1) to the setting involving multiple distributions. In particular, in MOT, one aims to find a transport plan among a set of marginals µ 1 , . . . , µ J with J ≥ 2. In this setting, the transport cost is encoded in a tensor C = [C(i1, i2, . . . , iJ )] ∈ R d 1 ×d 2 ×···×d J where C(i1, i2, . . . , iJ ) denotes the unit transporting cost corresponding to the locations i1, i2, . . . , iJ , and the transport plan is in the same way described by a J-mode tensor B ∈ R d 1 ×d 2 ×···×d J + . For a given transport plan B, the total cost of transportation is Thus, similar to (1), MOT has a linear programming formulation 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 Note that the standard bi-marginal OT problem (1) is a special case of the MOT problem (20) with J = 2 and Γ = {1, 2}.

A. Entropy regularized MOT
Although MOT (20) is a standard linear program, its complexity grows exponentially as J increases. This computational burden can be partly alleviated in an analogous manner as for the classical bi-marginal problem (1), which again yields an iterative scaling algorithm. In particular, similarly to (3), one can add an entropy term to (20) to regularize the problem, resulting in the strictly convex optimization problem with > 0 being a regularization parameter.
For the bi-marginal case, (23) reduces to problem (3). The iterative scaling algorithm (4) can be generalized to the multi-marginal setting [46] in order to solve (23). From an optimization perspective, the iterative scaling algorithm amounts to a coordinate ascent method [54] in the dual problem of (23). The introduction of the entropy term in (23) allows for closed-form expressions for the updates of the dual variables [17]. Utilizing Lagrangian duality theory, one can show that the optimal solution to (23) is of the form where denotes element-wise multiplication and the tensors are given by and where the vectors uj ∈ R d j are given by and λj ∈ R d j is the dual variable corresponding to the constraint Pj(B) = µ j on the j-th marginal. Moreover, the dual of (23) is We emphasis that in (28), U is a function of the multipliers {λj, j ∈ Γ} as defined in (26)- (27).
The iterative scaling algorithm iteratively updates the vectors uj, for j ∈ Γ, in (27) according to for all j ∈ Γ. For future reference, we summarize the steps in Algorithm 1. The Iterative Scaling algorithm (Algorithm 1) is a special case of the iterative Bregman projection algorithm [55], [23], which itself is a special case of a dual block coordinate ascent method [56], [57], and thus enjoys a global convergence guarantee [55], [57].
Algorithm 1 Iterative Scaling Algorithm for MOT Note that the standard Sinkhorn iterations (4) for the two-marginal case (3) is a special case of Algorithm 1 when J = 2 and Γ = {1, 2}. Indeed, in this case, Algorithm 1 boils down to iterating which coincide with (4).
Although Algorithm 1 is easy to implement and considerably faster than general linear programming solvers, its complexity still scales exponentially as J grows since the number of elements in B are d1d2 . . . dJ . The computational bottleneck of it lies in the calculation of the projections Pj(B), for j ∈ Γ, in (21). Generally, this computational burden is inevitable. However, in some cases it is possible to utilize structures in the cost tensor C to make the computation of the projections more accessible [23], [17], [52]. In Section III-B we consider MOT problems with cost tensors that can be decomposed according to a graph. This graphical structure allows us to leverage the Bayesian inference tools [26] in PGMs to compute the projections efficiently. Other than providing a workhorse for solving MOT problems with graphical structured cost, this connection between MOT and PGMs also presents new elements and perspective to Bayesian inference in PGMs, which is discussed in details in Section III-C.

B. MOT with graphical structures
Consider the cases where the cost tensor C can be decomposed according to a factor graph. More specifically, the cost tensor C has the form where F denotes the set of factors of a graph. Here, to be consistent with the notations in PGMs, we write the cost of associating i1, i2, . . . , iJ by C(x) = C(x1, x2, . . . , xJ ) instead of C(i1, i2, . . . , iJ ), but the two have exactly the same meaning; both xj and ij take values in a set with dj elements. Thus, by abuse of notation, we use C(x) and C(i1, i2, . . . , iJ ) interchangeably.
A graph structured cost tensor (30) occurs in various applications of the OT framework [23], [17]. For instance, in Barycenter problems [53], the cost C can be decomposed into the sum of pairwise costs between the target distribution and each given marginal distribution. For general cost functions, it might be possible to approximate them using the structured cost (30). Thus the framework we establish can also be viewed as an efficient method to approximate the solution to any MOT problem.
Denote the factor graph associated with the cost (30) by G = (V, F, E). Then the j-th mode of C corresponds to node j ∈ V and the marginal distribution of the j-th mode is the same as the marginal distribution of xj at node j. In this paper, we only consider the cases where the factor graph G is connected but does not have any loop, that is, G is a factor tree. We associate the cost C with a probabilistic graphical model where Clearly, K in (25) has the form and From a PGM point of view, the (transformed) Lagrangian multipliers uj, for j ∈ Γ, introduced by the Iterative Scaling algorithm are local potentials of the modified graphical model K(x)U (x). The Lagrangian approach of solving the constrained optimization problem (23) seeks multipliers uj, for j ∈ Γ, such that the tensor B = K U satisfies all the constraints Pj(B) = µ j , for j ∈ Γ. Thus, in the language of PGMs, to solve the MOT problem (23), one can search for a proper set of artificial local potentials uj, for j ∈ Γ, such that the modified graphical model K(x)U (x) in (32) has the specified marginal distribution µ j on the j-th variable node for each j ∈ Γ. Note that uj = exp(−1/J)1 is a uniform potential for all j / ∈ Γ and thus does not affect the graphical model K U.
For fixed multipliers u1, u2, . . . , uJ , calculating (with proper normalization) the projection Pj(K U) is exactly a Bayesian inference problem of inferring the j-th variable node over the modified graphical model K(x)U (x). When G does not have any loops, a condition we assume throughout, Bayesian inference can be achieved efficiently using the Belief Propagation algorithm. Generally, the marginal constraints Pj(B) = µ j can be imposed on any variable node j ∈ V . However, a marginal constraint on a non-leaf node will decompose the MOT problem (23) into several independent MOT problems with constraints only on leaf nodes, see [52]. Thus, without loss of generality, we assume marginal constraints on leaf nodes only, that is, Γ ⊂ L where L ⊂ V denotes the set of leaf nodes of G. Example 2. Leveraging the graphical structure (30) of the cost, based on the discussions above, we obtain a simple strategy to solve the MOT problem (23): We apply the Iterative Scaling algorithm and utilize the Belief Propagation algorithm to carry out the computation of Pj(K U) with the current multiplier U. The acceleration is tremendous for MOT problems with a large number of marginals; the Belief Propagation algorithm scales well for large problem while the complexity of the brute force projection using definition (21) grows exponentially as the dimension increases. It turns out that some more tricks can be adopted to further improve the speed of the projection step Pj(K U). The full algorithm, which we call Iterative Scaling Belief Propagation (ISBP) algorithm will be presented and discussed in details in Section III-D.

C. MOT and Bayesian inference
In the previous section, we have seen that in cases where the cost tensor C in MOT problem (23) has a graphical structure, one can take advantage of PGM methods, in particular the Belief Propagation algorithm, to accelerate the Iterative Scaling algorithm. In this section, we establish further connections between MOT and PGMs. These links add novel components to both the MOT theory and PGM theory. These connections also bring new insight and interpretation of Iterative Scaling Belief Propagation.
Clearly, the objective function of the entropy regularized MOT problem (23) is exactly the free energy F in (12) with Thus, the MOT problem (23) can be written as Therefore, the entropic regularized MOT problem (23) with cost function that decouples according to a graph structure as in (30) is equivalent to a Bayesian inference problem in a PGM with additional constraints on the marginal distributions of a set of variable nodes. In other words, (34) is a constrained version of a Bayesian inference problem.
On the other hand, any Bayesian inference problem in a PGM can be rewritten in the constrained form (34). More specifically, consider the problem of inferring the posterior distribution of where yj is the observation associated with the variable node potential φj. When the observations are fixed, say yj =ŷj, then the standard Bayesian inference method replaces φj(xj, yj) by a local potential φj(xj) and infers the marginal distributions of the resulting graphical model with only the nodes {xj, xα}. Alternatively, the measurement yj =ŷj can be viewed as constraints on the node yj of an augmented The graph is augmented by the nodes y 1 , y 2 , y 3 and factors ψ α4 = φ 1 (x 1 , y 1 ), ψ α5 = φ 2 (x 2 , y 2 ), and ψ α6 = φ 3 (x 3 , y 3 ). The measurements become marginal constraints graphical model which includes also the observation variable node yj. In particular, the constraint is of the form p(yj) = δ(yj −ŷj), where δ denotes Dirac distribution. Thus, the posterior distribution can also be obtained by solving the constrained Bayesian inference problem (34) over the augmented graphical models under the constraints that p(yj) = δ(yj −ŷj), for j ∈ Γ. This equivalence is illustrated in Figure  3. Therefore, from this point of view, the constrained Bayesian inference problem (34) can also be viewed as a generalization of standard Bayesian inference.
When the underlying factor graph associated with the cost function (30) is in fact a factor tree, the free energy F(B) is equal to the Bethe free energy (see Section II-C) where Bα is the marginal distribution on factor node α and Bj denotes the marginal on variable node j, namely, Bα(xα) = x\xα B(x), and Bj(xj) = x\x j B(x). In (35), B is the collection of marginal distributions, that is, B = {Bj, Bα : j ∈ V, α ∈ F }. This is different from the J-mode tensor B in (34). However, with slight abuse of notation, we use the same symbol B in both contexts. Again, due to tree structure, the two have a one-to-one correspondence with each other.
The marginals Bα and Bj capture only local information around a factor variable or node variable and they have to satisfy certain conditions in order to be feasible marginal distributions of some joint distribution. In particular, they have to be compatible in the sense xα\x j Bα(xα) = Bj(xj), ∀j ∈ V, ∀xj.
In terms of local marginals, the constraints in (34) read Therefore, the MOT problem (34) can be reformulated as where the last constraint (36d) is to ensure that the optimization variables {Bj, Bα : j ∈ V, α ∈ F } are in the probability simplex.
Since the Bethe free energy is convex for factor trees, and the constraints are linear, Problem (36) is a convex optimization problem. One advantage of (36) over (34) is that the size of optimization variables in (36) is considerably smaller than that in (34). More specifically, the optimization variables of (36) are local marginals which are either vectors Bj or low-dimensional tensors Bα, which is in contrast to the high-dimensional J-mode tensor B in (34).

D. Iterative Scaling Belief Propagation algorithm
In this section, we present the full Iterative Scaling Belief Propagation algorithm for the entropy regularized MOT problem (23) (or equivalently (34) and (36)). We start with a characterization of the solution to (36). Theorem 3. The solution to the MOT problem (36) is given by where mα→j, nj→α are fixed points of the following iterations Here ∝ indicates that a normalization step is needed.
Proof. In order to solve the constrained optimization problem (36), we introduce Lagrange multipliers ηα for the simplex constraints (36d), λj,α for the marginalization compatibility constraints (36c), and νj for the fixed-marginal constraints (36b), yielding the Lagrangian (39) Note that we have used a scaled version 1 F Bethe of the objective function. In view of (33), (35) and (31), Setting the derivatives of the Lagrangian with respect to the local marginals Bα and Bj to zero, we get that the minimizer satisfies The relation (37a) follows immediately from (40a) and (41a). This together with the constraint (36c) and (41b) leads to By (40c), we obtain nj→α(xj) = 1 if Nj = 1, j / ∈ Γ.

It follows that
x 1 x 2 x 5 x 4 x 3

Fig. 4: Messages in ISBP
which is (37b) for leaf nodes. We next show (37b) when Nj > 1. To this end, we plug (40a) and (40b) into the constraint (36c) and arrive at For fixed j, the above holds for all γ ∈ N (j). Multiplying the above equation for all γ ∈ N (j)\α yields which is (38b) after canceling out equal terms. Thus, in view of (42), Finally, (37c) is clearly true due to constraints. This together with (42) leads to (38c), which completes the proof.
The updates in (38) resemble the standard Belief Propagation algorithm (13). In particular, the updates (38a) and (38b) are exactly the same as (13). The update (38c) is new and is due to the constraints (36b) on the marginal distributions. Pictorially the message mα→j sent to a constrained node j from node α bounces back to α, in form of nj→α. This is illustrated in Figure 4. The update (38c) in fact corresponds to the scaling step (29) of the Iterative Scaling algorithm (Algorithm 1). In particular, the multipliers {uj : j ∈ Γ} in (27) relate to the messages as uj = nj→α, for j ∈ Γ. To see this, we note that the projection Pj(K U) requires solving a Bayesian inference problem with respect to the modified graphical model Upon convergence of the Belief Propagation algorithm (13), it holds Pj(K U) = ujmα→j, where α is the only factor node in N (j) since j ∈ Γ is a leaf node. Thus, the projection step (29) reads uj µ j ./Pj(K U) = µ j ./mα→j = nj→α.
Therefore, the updates (38) contain all the components of our ISBP algorithm with (38a)-(38b) being the Belief Propagation part and (38c) being the Iterative Scaling part. ISBP is a scheduling of these updates in a certain order. As discussed in Section III-B, the key idea of ISBP is to implement the projection Pj(K U) in the iterative scaling step (29) using Belief Propagation. In the contexts of the updates (38), it is equivalent to run (38a)-(38b) sufficiently many iterations to obtain the precise projection Pj(K U) and then run (38c), which is essentially (29). How many iterations of (38a)-(38b) are enough? One option is to run Belief Propagation over the whole graph G with the most recent modified model K U to compute Pj(K U) for all j ∈ V . This is clearly sufficient but it is not necessary. Let j1, j2, . . . be a sequence taking values in Γ in arbitrary order and suppose the Iterative Scaling algorithm is carried out in this order. Then after the k-th step, uj k is updated, and the only projection required in the next step is Pj k+1 (K U). It turns out that to evaluate Pj k+1 (K U), it suffices to update all the messages on the path from j k to j k+1 . Compared to the naive Belief Propagation over the whole graph, this local updating strategy is considerably faster. The steps of the ISBP algorithm are summarized in Algorithm 2.

Algorithm 2 Iterative Scaling Belief Propagation (ISBP) Algorithm for MOT
Initialize the messages m α→j (x j ) and n j→α (x j ) Update m α→j (x j ) and n j→α (x j ) using (38a)-(38b) until convergence while not converged do Update n j k →α (x j k ) using (38c) Update all the messages on the path from node j k to node j k+1 according to (38a) and (38b) end while Upon convergence of Algorithm 2, the solution to Problem (36) can be obtained through (37). The whole belief tensor B can also be obtained through B = K U with U = u1 ⊗ u2 ⊗ . . . ⊗ uJ , where uj = nj→α for j ∈ Γ, and uj = exp(− 1 J )1 otherwise.

IV. CONSTRAINED NORM-PRODUCT ALGORITHM
One potential drawback of the ISBP algorithm lies in the fact that it is a two-loop algorithm with the outer loop being iterative scaling and inner loop being belief propagation. Such a two-loop structure might slow down the convergence rate, especially when the underlying graph is large. Moreover, the two loops have to coordinate closely to guarantee convergence. Such coordination is even more difficult, or impossible if a distributed implementation is needed. Thus, we seek to develop a single loop algorithm for the entropy regularized MOT problems. A natural question to ask is whether we can borrow ideas from the Bayesian inference literature. After all, the Belief Propagation algorithm is not the only algorithm for Bayesian inference.
The answer is affirmative. In this section, we examine the Norm-product algorithm [36], another powerful Bayesian inference method, and extend it to a single loop algorithm for our MOT problems. Below we first review Norm-product algorithm for standard Bayesian inference problems in Section IV-A. The extensions to entropy regularized MOT, or equivalently, constrained Bayesian inference problems are presented in Section IV-B.

A. The Norm-product algorithm
Consider the Bayesian inference problem (9). The Norm-product algorithm [36] for Bayesian inference is based on the so called fractional entropy approximation of the entropy H(b) in (11). The coefficientscα, for α ∈ F , andcj, for j ∈ V are defined as for a set of real numbers cα, cj, and cjα, for j ∈ V and α ∈ F , which are known as counting numbers [58]. Clearly, an equivalent formulation of the fractional entropy (44) is The fractional entropy resembles the Bethe entropy (18). In fact, for the choice of counting numbers cj = 1 − Nj, cjα = 0, and cα = 1, the fractional entropy (46) reduces to the Bethe entropy. Moreover, just like the Bethe entropy, the fractional entropy approximation can be made exact when the underlying graph is a tree (see Section IV-D).
With the fractional entropy representation (46), the total free energy (12) is modified to the fractional free energy where is the average energy defined as in (10). Thus, in terms of fractional free energy, the Bayesian inference problem (9) reads The constraints (48b)-(48c) are to ensure that bα, bj are indeed marginal distributions of some certain joint distribution. For a given graphical model, there are infinitely many different fractional free energy approximations determined by the counting numbers cj, cα, cjα; some of them are convex and some of them are not. A sufficient condition for the convexity of the fractional free energy is as follows. Lemma 4 ( [36]). When the counting numbers satisfy cj ≥ 0, cjα ≥ 0, and cα > 0, for j ∈ V and α ∈ F , then the fractional free energy F frac is strictly convex over the set defined by the constraints (48b)-(48c).
Denote the two sets corresponding to the constraints (48b) and (48c) by and respectively, and definef andĥ where δ is the indicator function. Then the Bayesian inference problem (48) can be reformulated as hj(b).
In cases where cj ≥ 0, cjα ≥ 0, and cα > 0, for j ∈ V and α ∈ F , by Lemma (4),f is strictly convex andĥj is convex for each 1 ≤ j ≤ J for b ∈ M ∩ P. The Norm-product algorithm relies on the reformulation (53). In particular, it leverages a powerful primal-dual ascent algorithm (stated below in Lemma 5) that is well studied in the convex optimization community to solve problem with the special structure of (53). The primal-dual ascent in Lemma 5 is derived from a more general algorithm known as dual block ascent [57] and thus inherit the nice convergence property of the latter. We refer the read to [36] for more details on these algorithms.

Lemma 5 ([36]). Consider the convex optimization problem min
The primal-dual ascent algorithm initializes λ1 = 0, . . . , λJ = 0 and repeatedly iterates the following steps for j = 1, . . . , J until convergence: Supposef is strictly convex and smooth, and hj, j = 1, . . . , J are convex and continuous over their domains, then b * in the above iteration converges to the unique global minimizer of f + J j=1 hj.
The Norm-product algorithm [36] (Algorithm 3) is a direct application of the primal-dual ascent algorithm to the formulation (53) of the Bayesian inference problem (48). It can be seen as a message-passing type algorithm for problem (48), where the dual variables λj in the primal-dual ascent algorithm work as "messages" between neighboring nodes. To see this, note that hj(b) depends only on bα, where α ∈ N (j), and thus the corresponding dual variable λj depends only on xα, where α ∈ N (j). This sparsity is encoded by the representation λj = {λj,α(xα)}. The relation between the dual variables λj in Lemma 5 and messages in Algorithm 3 is given by nj→α(xα) = exp(−λj,α(xα)).
For more details on the derivation of the Norm-product algorithm as a primal-dual ascent method, see [36]. Moreover, our development of the constrained Norm-product algorithm (Algorithm 4) is similar to this, and is provided in the appendix. Upon convergence of Algorithm 3, the solution to (48) has the form where again ∝ indicates that a normalization step might be needed. Finally, note that for the special choice of counting numbers cj = 1−Nj, cjα = 0, and cα = 1, the Norm-product algorithm reduces to the Belief Propagation algorithm (13). However, note that this choice of counting number does not satisfy the conditions in Lemma 4 for convergence of the algorithm, although it is well known that the Belief Propagation algorithm is guaranteed to converge for trees [50]. Therefore, even though formally the Norm-product algorithm can be viewed as a unifying framework for many message-passing algorithms, its convergence proof is restricted in some sense due to the strong requirement on the counting numbers.

B. Constrained Norm-product algorithm
In this section, we develop a Norm-product type algorithm for the entropy regularized MOT problem (23), or equivalently the constrained Bayesian inference problem (34). Consider a modification of problem (48) with constrained marginal distributions, which reads Problem (56) can be seen in the light of the entropy regularized MOT problem formulated as a free energy minimization problem (34). In particular, if the free energy F is approximated by the fractional free energy F frac , then (34) becomes (56). Recall that in the MOT problem the factor and node potentials are ψ α = exp(−Cα) and φ j ≡ 1 (cf. (33)). However, the constrained Norm-product algorithm, which we develop in the following solves the Bayesian inference problem (56) for any potentials ψ α and φ j .
Note that compared to (48), the modified problem (56) is only augmented by one linear constraint (56b). Thus, problem (56) can be formulated as in (53) instead of (49), and defining all other components of (53) as in (50)- (52). The primal-dual ascent algorithm in Lemma 5 can then be applied to (56). The resulting Constrained Norm-product (CNP) algorithm is presented in Algorithm 4. For a detailed derivation of the method see Appendix A.
Upon convergence of Algorithm 4, the solution to (56) is of the form (55), as in the standard Norm-product algorithm. Moreover, the optimal marginal calculated through (55) satisfies the constraint bj = µ j for all j ∈ Γ. Algorithm 4 is presented for general constrained Bayesian inference problems (56). Recall that the entropy regularized MOT problem (34) is recovered as the special case, where the potentials are given by ψ α = exp(−Cα) and φ j ≡ 1, as in (33).
Compared with the standard Norm-product algorithm, the messages from variable nodes with marginal constraint to the neighboring factor nodes, i.e., nj→α, for j ∈ Γ, α ∈ N (j), depend not only on the incoming messages to j and α, but also the given marginal µ j . Moreover, in the case when the marginal constraint (56b) is absent, namely, Γ = ∅, Algorithm 4 reduces to the standard Norm-product belief algorithm 3. Remark 6. The message updates in Algorithm 4 can be problematic when the denominators become zero. This scenario can occur when either the factor or node potentials ψα(xα) or φj(xj) contain zero elements. Note that zero entries in the potential let the average energy (10) be unbounded if bj(xj) or bα(xα) are nonzero on the corresponding entries. In implementations, this can be avoided by ignoring the updates involving zero denominators. See [36, Appendix F] for a more detailed discussions of this issue.

C. Relations to Iterative Scaling Belief Propagation algorithm
Compared to the ISBP algorithm, the CNP algorithm is a single loop algorithm. Each iteration of Algorithm 4 requires visiting every variable node only once. In contrast, since Algorithm 2 has a double-loop structure and each inner-loop iteration requires updating throughout Algorithm 4 Constrained Norm-product (CNP) algorithm Set n j→α (x α ) = 1 for all j = 1, · · · , J, α ∈ N (j) and x α while not converged do for j = 1, 2, . . . , J do , ∀α ∈ N (j), ∀x α end if end for end while an entire path between two leaf nodes, the messages associated with most variable nodes will be updated multiple times in one iteration of the algorithm. Thus, the iteration complexity of the ISBP algorithm is higher than that of the CNP algorithm. This difference becomes more significant as the diameter/size of the underlying graph increases; for larger graphs, the inner-loop iteration of ISBP algorithm takes more updates. Apart from the iteration complexity, another potential advantage of the CNP algorithm is that its single loop structure allows for more flexible scheduling of the message passing/updating. In particular, it does not require any communication between inner and outer loop updates. Thus, it is easier to parallelize the Constrained Norm-product algorithm or develop a distributed version of it.
Recall from Section IV-A that the standard Norm-product method with counting numbers chosen as cα = 1, cj = 1 − Nj and cjα = 0 reduces to the standard Belief propagation method as given in (13). It turns out that similar results can be established to relate the Iterative Scaling Belief Propagation algorithm and the Constrained Norm-product algorithm. In particular, with this set of counting numbers, the constrained Norm-Product algorithm reads Note that in general the messages nj→α in the Constrained Norm-product algorithm depend on xα, but for this special choice of counting numbers, they depend only on xj. The messages (58) are exactly the same as the messages (38) in the ISBP algorithm. If the messages in (58) are scheduled in a specific way, then this becomes the Iterative Scaling Belief Propagation Algorithm 2. In particular, this is achieved by cycling through the nodes in Γ, where for two successive nodes j1, j2 ∈ Γ, one schedules the messages (58a) and (58b) on the path from j1 to j2, and finally the message nj 2 →α as in (58c). In this light, Algorithm 2 may not only be understood as Iterative scaling Belief propagation, but also as constrained Belief propagation, i.e., an extension of the standard Belief propagation method, where the marginals on some nodes are fixed.
What if we update the messages (58) following the scheduling of Algorithm 4? In fact, this is a single-loop version of the ISBP algorithm, and we have empirically observed good convergence properties of it. However, the choice of counting numbers cα = 1, cj = 1 − Nj and cjα = 0 does not yield a strictly convex objective function decomposition in the associated fractional variational inference problem (56) as discussed in Lemma 4. Thus, the convergence proof of Algorithm 4 does not apply to this setting, and a global convergence proof remains an open problem.

D. Counting numbers of fractional entropy
One way to guarantee the convergence of the Constrained Norm-product algorithm is to choose the counting numbers for the fractional entropy H frac such that they satisfy the convexity conditions in Lemma 4. Thus, a crucial question is whether, given a graphical model, such a choice of counting numbers exists, and how to find them. This question has been discussed in [36,Appendix E] where several optimization based methods have been proposed. In this section, we present a structured method to construct a feasible set of counting numbers that satisfy the assumptions in Lemma 4, viz., cj ≥ 0, cjα ≥ 0, cα > 0, for factor graphs, which are trees. In particular, we provide a closed form expression for the choice of counting numbers, which makes parameter tuning for the Constrained-norm product algorithm simple and intuitive.
The fractional entropy decomposition requires the fractional entropy H frac (b) to be equal to the entropy H(b), that is

cjα(H(bα) − H(bj)).
On the other hand, for a factor tree, the entropy equals the Bethe entropy, namely, Hence, by identifying the coefficients, we see that finding a set of feasible convex counting numbers is achieved by finding cα > 0, cj ≥ 0, cjα ≥ 0 that satisfy the following equations A direct consequence of (59) is To see this, sum up (59a) and (59b) over all variable nodes and factor nodes. The left hand side becomes j∈V cj + α∈F cα as all the terms cjα get canceled. The right hand side becomes The last equality is due to the fact that the factor graph is a tree (acyclic).
The property (60) can be generalized to subgraphs of G. Let (j * , α * ) be any edge of G. If we cut this edge, then the tree G is split into two trees G1 and G2, where G1 contains the variable node j * and G2 contains the factor node α * . This is illustrated in Figure 5. Let G1 = (V1, F1, E1), then cj * α * = This relation (61) can be established similarly to (60). It determines values for cjα given cj, j ∈ V and cα, α ∈ F . Moreover, it guarantees that cjα is non-negative, as long as cj and cα are non-negative. Hence, based on (60) and (61), we obtain a remarkably simple strategy to get a set of convex counting numbers cα > 0, cj ≥ 0, cjα ≥ 0. i) Choose cα > 0, cj ≥ 0 for j ∈ V, α ∈ F such that (60) is satisfied; ii) Iterate over each edge in the graph, split the graph along the edge and calculate the corresponding cjα through (61).
Proposition 1 makes constructing a feasible set of counting numbers that induces convex fractional free energy (see Lemma 4) as easy as finding cα > 0, cj ≥ 0, j ∈ V, α ∈ F that satisfy (60). One choice we found effective is for all j ∈ V, α ∈ F . For the specific example in Figure 5, this choice leads to cj = cα = 1 11 . The value of c2α 2 is 7 11 by (61).

V. NUMERICAL EXAMPLES
In this section we present two sets of numerical experiments based on our framework. The first set of experiments is to validate the correctness of Constrained Norm-Product and Iterative Scaling Belief Propagation algorithm, and compare their performance. The second experiment is to illustrate potential applications of our framework in nonlinear filtering for collective dynamics.

A. Performance evaluation
We implement three algorithms: Constrained Norm-Product (CNP) (Algorithm 4), Iterative Scaling Belief Propagation (ISBP) (Algorithm 2) and Vanilla Iterative Scaling (Vanilla IS) (Algorithm 1) on four different type of graphs: line graphs, hidden Markov models (HMMs), and two star shape graphs (see Figure 6). For the line graph (Figure 6a), the constraints of marginal distribution are on the head and tail nodes. This corresponds to a standard OT problem with two marginals. HMMs (Figure 6b) are widely used in many real applications. In the standard HMM framework, the measurements are deterministic values, which can be equivalently viewed as Dirac distributions on the observation/measurement nodes. In our MOT framework, these observation nodes are associated with marginal distribution constraints, which can be viewed as a relaxation of the standard HMM where the deterministic measurement are replaced with "soft" stochastic measurements. The star graph (Figure 6c) structure has marginal constraints on all the leaf nodes. This corresponds to the Barycenter problem over the Wasserstein space [53], which has found applications in information fusion [17].
We test the algorithms with several different configurations. In particular, we vary the number of discrete states at each variable node d1 = d2 = · · · = dJ = d as well as the number of nodes J in the tests. Throughout, we let = 1. The factor potentials and the counting numbers are set consistently for all the experiments. In particular, the factor potentials are chosen in a way such that the variable nodes connecting to a common factor node are strongly correlated. In our examples, all the factors are connected to only two variable nodes. This choice amounts to taking diagonally dominant matrices as potentials. The counting numbers are selected using the strategy in Proposition 1 by setting cj = 0, ∀j ∈ V and cα = c β , ∀α, β ∈ F . In all our experiments, we observe that the three algorithms converge to the same solutions. To fairly compare the computation complexity of the three algorithms, we use a unified stopping criteria; the algorithms stop when the relative error with respect to the "ground truth" b * in terms of the 1-norm is less than 10 −4 . The "ground truth" b * is obtained by running one of the algorithms (e.g., Vanilla IS) for sufficiently many iterations so that the duality gap is less than 10 −8 . Figure 7 depicts the evaluation results of the three algorithms under different configurations. The y-axis shows the total time consumption of the algorithms before they stop. The x-axis represents the size of the graphs, more specifically, the number of nodes J of the graphs being used. Thus, each subplot showcases the relation between computational efficiency and the number of nodes of the graphs. The dependence of the computational complexity on the number d of discrete states at each node can be understood by comparing the subplots along each rows. Each row of subplots corresponds to a type of graph. Thus, the effect of the graph topology on the computational complexity can be captured by comparing the subplots in the same column. From the results it can be seen that, for all types of graphs, and all values of d, the complexity of the Vanilla Iterative Scaling grows exponentially as the number of nodes J increases. In contrast, both CNP and ISBP scale much better than Vanilla IS when J increases. Moreover, CNP and ISBP seem to be less sensitive to the number of discrete states d at each node, compared with Vanilla IS.
To comprehensively compare the performances of ISBP and CNP, we conduct several more experiments on graphs of larger sizes where the Vanilla Iterative Scaling is no longer applicable. Besides the three graphs used in the previous experiment, we study an additional star shape graph with more nodes on each branch (Figure 6d). The stopping criteria is the same as before; the algorithms stop when the relative error with respect to a "ground truth" b * in terms of 1-norm is less than 10 −4 . Since Vanilla IS is computationally forbidden for large J, we run the ISBP algorithm for sufficiently many iterations to obtain b * . The experiments results are summarized in Figure 8. The presentation of the results in Figure 8 is similar to that in Figure 7, so that we can understand the dependence of the computational complexity over the number of discrete states d, number of nodes J and graph topology. From the figures we can see that the two algorithms CNP and ISBP, have comparable performances. Both of them scale well when J and d increase. ISBP behaves better on line graphs and HMMs, while CNP is faster on star shaped graphs.

B. Filtering for collective dynamics
To demonstrate the performance of the methods developed in this paper, we present an example to demonstrate the effectiveness of our MOT framework in filtering problems for collective dynamics. Consider a synthetic particle ensemble with 10000 agents moving over a 20 × 20 grid, aiming from bottom-left corner to top-right corner. The dynamics of the agents follow a log-linear distribution characterized by four factors: the distance between two positions, the angle between the movements direction and an external force, the angle between the direction of movement and the direction to the goal, and the preference to stay in the original cell. The weights for the log-linear model associated with these four factors are set to be (3,5,5,10). This model has been used to model the migration of birds [40], [42].
There are 16 sensors placed over the grid as shown in Figure 9a. These sensors can not measure the exact locations of the agents. Instead, the measurement of each sensor is a count of agents it currently observes. The probability of an observation decreases exponentially as the distance between the sensor and the agent increases.This type of sensors show up in many real applications. For instance, the sensors can be Wi-Fi hotspots, or cell phone based stations, which can measure the number of phones connected to them. Our goal is to estimate the movement of the whole population using this limited sensor information.
This filtering problem for collective dynamics can be modeled as a MOT problem, or equivalently a constrained marginal inference problem in our framework. In particular, the agents form a HMM and the sensor measurements correspond to constraints on marginal distributions over the observation nodes. The number of discrete states at each node is d = 20 × 20 = 400 and the number of nodes depends on the number of time steps. We simulate the model for 15 time steps and run both ISBP and CNP to infer the marginal distributions of the free nodes in order to estimate the group behavior of the 10000 agents. The agents start in two clusters: one in the left-bottom and one in the center-bottom; both aim to reach the right-top corner of the grid in 15 time steps. The results are depicted in Figure 9b. Both CNP and ISBP give the same estimation result and thus we only display one of them in the figures. As can be seen from the plots, even though the sensor data (center column) is hard to interpret visually, our constrained marginal inference framework can still infer the population movements to a satisfying accuracy.

VI. CONCLUSION
We studied multi-marginal optimal transport problems and pointed out an unexpected connection to probabilistic graphical models. This relation between MOT on graphs and constrained PGMs provides a completely new perspective of both MOT and PGMs, which may have far reaching impact in the future development of both subjects, in both theory and applications. This connection also enables us to adapt the rich class of algorithms in PGMs to tackle difficult MOT problems. In this work, to highlight the key idea of this line of research, we focused on MOT on trees with discrete states. The next step is to generalize the results to more general graphs with cycles as well as continuous state spaces. These are more challenging problems in PGMs, for which exact Bayesian inference is usually too expensive, and one needs to turn to variational inference or sampling based methods [26].
We next derive a closed form solution b * j , b * α to (63). The constraint (63b) implies that bj is a marginal distribution of bα, thus, bα can be rewritten in terms of conditional distribution b α|j as bα(xα) = bj(xj)b α|j (xα | xj).
The entropy H(bα) can be rewritten as [59] H(bα) = H(bj) + together with the extra constraint bj = µ j if j ∈ Γ. One advantage of this reformulation is that the problem now can be optimized over b α|j first and then over bj.
Minimizing (65) over b α|j is a standard exercise, and the minimizer is Next, we move to step (54c) of the primal-dual ascent algorithm (Lemma 5), which reads λj ← −νj − ∇f (b * ) + A T σ with σ being an arbitrary vector. The vector A T σ spans the orthogonal space to the domain B (see Lemma 5) of f . In our problem, B = P is the probability simplex, thus A T σ is in alignment with 1, the vector with all 1 entries. It follows that λj,α(xα) = −νj,α(xα) − ∇f (b * α (xα)) + σα1.