---
abstract: 'We propose a new algorithm for the distributed computation of Wasserstein Barycenters over networks. Assuming that each node in a graph has a probability distribution, we prove that every node can reach the barycenter of all distributions held in the network by using local interactions compliant with the topology of the graph. We provide an estimate for the minimum number of communication rounds required for the proposed method to achieve arbitrary relative precision both in the optimality of the solution and the consensus among all agents for undirected fixed networks.'
author:
- 'César A. Uribe, Darina Dvinskikh, Pavel Dvurechensky, Alexander Gasnikov and Angelia Nedić [^1]'
bibliography:
- 'IEEEfull.bib'
- 'wass.bib'
- 'opt\_dec2.bib'
title: '**Distributed Computation of Wasserstein Barycenters over Networks**'
---

Introduction
============

Optimal Transport (OT) distances (also known as *earth mover’s distances* or *Wasserstein distances*) design an optimal plan to move “mass" from one probability distribution to another. This problem can be traced back to the early work of Monge [@Monge1781] and Kantorovich [@Kantorovich1942] and has been of constant interest for allowing natural formulations to the problems of comparing, interpolating, and measuring distances of functions [@Levy2017]. On the other hand, computational OT has gained popularity for its applications in learning theory [@Frogner2015], computer vision [@Rabin2011], computer graphics [@Solomon2015], statistical inference [@Srivastava2015a], information fusion [@Bishop2014a]; and its relative complexity advantages with respect to classical methods [@Dvurechensky2018]. . Comprehensive accounts of the OT problem and its computational aspects can be found in [@Villani2008; @Solomon2017; @Peyre2017; @Levy2017].

One of the common uses of the Wasserstein distance is the aggregation of distributions by considering their barycenter [@Agueh2011], which itself is another distribution [@Cuturi2014]. Wasserstein Barycenters (WB) have been shown superior to traditional Euclidean means in a range of application such as image processing [@Agueh2011], economics and finance [@Beiglbock2013], and condensed matter physics [@Buttazzo2012]. Figure \[fig:sevens\] shows a sample of $30$ images of the digit $7$ from the MNIST dataset [@LeCun1998], and their respective Euclidean mean and Wasserstein mean. The WB better captures the structural features of the input images.

For discrete and finite distributions, the WB can be efficiently computed by solving a large linear program [@Anderes2016] or using regularization to approximate a solution efficiently and exploit its convenient algebraic properties [@Agueh2011; @Cuturi2014; @Cuturi2016]. Here, we consider the problem of computation of WB over a network. The flexibilities induced by the distributed setup make it suitable for problems involving large quantities of data with no centralized storage [@boy11; @ned16w; @ned17e; @Nedic2017a]. Mainly, we assume a group of agents is connected over a network, and each agent locally holds a probability distribution with finite support. The group seeks to compute the WB of all distributions in the network cooperatively. Figure \[fig:erdos\_sevens\] shows an Erdős-Rényi random graph with $160$ agents where each agent holds a sample of the digit $7$ from the MNIST dataset.

Distributed consensus with the Wasserstein metric was introduced in [@Bishop2014a; @Bishop2014]. In [@Bishop2014], the authors showed asymptotic converge to the WB of the initial distributions given some weak connectivity assumptions on the graph over which agents exchange information. Nevertheless, the proposed algorithm requires that each agent computes an exact WB of local distributions at each iteration. Although one can have closed-form solutions for some families of continuous distributions, in general, the problem can be intractable. On the other hand, a recent approach [@Staib2017] explores the computational advantages of a dual formulation of the WB and exploits the parallelizable structure of the problem to propose a scalable, and communication-efficient algorithm for its computation on arbitrary continuous distributions. Nevertheless, it requires a central fusion center that coordinates the actions of the parallel machines.

In contrast with existing literature [@Bishop2014; @Staib2017], propose a first-order algorithm that can be executed distributedly over a network with unknown topology. We derive an explicit convergence rate of the order $O(1/k^2)$ with an additional cost that depends on the condition number of the graph over which the agents interact. Additionally, we present two numerical examples to illustrate and validate our results. First, we show some basic properties of the algorithm for the problem of computing WB of univariate, discrete and truncated Gaussian distributions. Then, we show the result of applying our algorithm to a subset of the MNIST digit database on a large-scale network of $1000$ agents.

This paper is organized as follows. Section \[sec:problem\] presents basic definitions for the problem of computation of WB over networks. Section \[sec:results\] states auxiliary results, introduces the proposed algorithm and states our main results on its convergence rate and dependency on the problem parameters. Section \[sec:numerical\] shows two numerical examples that experimentally verify the theoretical properties of the algorithm. Conclusions and future work are presented in Section \[sec:conclusions\].

**Notation:** We assume that the agents are indexed from $1$ through $m$. The enumeration is not needed in the execution of the proposed algorithm. It is only used in the analysis. Superscripts $i$ or $j$ denote agent indices and subscript $k$ iteration indices. $[A]_{ij}$ denotes the $i$-th row and $j$-th column entry of a matrix $A$. $I_{n}$ is the identity matrix of size $n$. .

Problem Statement {#sec:problem}
=================

In this section, we recall some basic definitions of the optimal transport problem. We describe the Wasserstein distance and the WB of a set of discrete probability distributions with finite support. Finally, we introduce the problem of distributed computation of WB over networks.

Entropy Regularized Optimal Transport
-------------------------------------

Consider two probability distributions $p,q \in S_1(n)$ with support on a finite set of points $\{x_i \in \mathbb{R}^d \}_{i=1}^n$ such that $p(x_i) = p_i$ and $q(x_i) = q_i$, where $S_1(n)  = \{ p \in \mathbb{R}_+^n  \mid p^T \boldsymbol1 =1 \}$. Moreover, consider a non-negative symmetric matrix , where $[M]_{ij}\in \mathbb{R}_+$ accounts for the cost of moving mass from $p_i$ to bin $q_j$. Without loss of generality, in the numerical example we will consider the Euclidean costs where $[M]_{ij} = \|x_i - x_j\|_2^2$. Additionally, define the set of couplings or *transportation polytope* $U(p,q)$ as

The entropy-regularized OT problem [@Cuturi2013] seeks to minimize the transportation costs while maximizing the entropy (maximum-entropy principle) and is defined as

where $\gamma>0$, and

and $\forall x>0, h(x) \triangleq x\log x$ and $h(0) \triangleq 0$. A solution $\mathcal{W}_0(p,q)$ is called the Wasserstein distance between $p$ and $q$ , problem  is convex and admits a unique solution .

$$\begin{aligned}
    \mathcal{W}_{\gamma,q}(p) \triangleq  \mathcal{W}_{\gamma}(p,q).
    \end{aligned}$$

One particular advantage of entropy-regularizing the Wasserstein distance is that there exists closed-form representations for the dual problem and its gradients [@Agueh2011; @Cuturi2016] where the Fenchel-Legendre transform of  is defined as

\[thm:cuturi\] For $\gamma >0$, the Fenchel-Legendre dual function $\mathcal{W}^*_{\gamma,q}(y)$ is differentiable and its gradient $\nabla \mathcal{W}^*_{\gamma,q}(y)$ is $1/\gamma$-Lipschitz with

where $y \in \mathbb{R}^n$, $\alpha = \exp( {y}/{\gamma}) $ and .

We will use the result Theorem \[thm:cuturi\] to design an algorithm for the computation of the WB on graphs based recent ideas of dual approaches for convex optimization problems with affine constrains [@all17; @Gasnikov2016] and optimal algorithms for distributed optimization [@Uribe2017; @uribe2018dual].

Computation of a Wasserstein Barycenter over a Network
------------------------------------------------------

The uniform WB [@Agueh2011; @Cuturi2014] of a family of discrete distributions $q_i\in S_1(n)$, for $i=1,\ldots,m$; is defined as the solution to the following optimization problem

The WB is an extension of the Euclidean barycenter to nonlinear metric spaces to the empirical Fréchet mean [@Frechet1948]. The existence and uniqueness of WB has been studied in the literature [@Bigot2016]. Problem  is strictly convex and admits a unique solution, denoted by $p^*$ [@Cuturi2016].

${{{\mathtt{p}} = [p_1^T,\cdots,p_m^T]^T}}$ and ${{{\mathtt{q}} = [q_1^T,\cdots,q_m^T]^T}}$, where , and rewrite the problem in an equivalent form

We denote the unique solution of  by .

We seek to solve problem  in a distributed manner over a network, where each distribution $q_i$ is held by an agent $i$ on a network. We model such a network as a fixed *connected undirected graph* , where $V$ is the set of $m$ nodes, and $E$ is a set of edges. We assume that the graph $\mathcal{G}$ does not have self-loops. The network structure imposes information constraints; specifically, each node $i$ has access to $q_i$ only and a node can exchange information only with its immediate neighbors, i.e., a node $i$ can communicate with node $j$ if and only if $(i,j)\in E$.

We can represent the communication constraints imposed by the network by introducing a set equivalent to the constraints in . To do so, we define the Laplacian matrix $\bar W{\in \mathbb{R}^{m\times m}}$ of the graph $\mathcal{G}$ by [$$\begin{aligned}
        [\bar W]_{ij} = \begin{cases}
        -1,  & \text{if } (i,j) \in E,\\
        \text{deg}(i), &\text{if } i= j, \\
        0,  & \text{otherwise,}
        \end{cases}
        \end{aligned}$$]{}

where $\text{deg}(i)$ is the degree of the node $i$, i.e., the number of neighbors of the node. Finally, define the communication matrix (also referred to as an interaction matrix) by , where $\otimes$ indicates the Kronecker product.

Throughout the paper, [*we assume that graph $\mathcal{G} = (V,E)$ is undirected and connected*]{}. Under this assumption, the Laplacian matrix $\bar W$ is symmetric and positive . Furthermore, the vector $\boldsymbol{1}$ is the unique (up to a scaling factor) eigenvector associated with the eigenvalue $\lambda=0$. $W$ inherits the properties of $\bar W$, including symmetry and positive semidefiniteness. Moreover, $W{x} = 0$ if and only if [$x_1 = \cdots = x_m$]{}, and $\sqrt{W}{x} = 0$ if and only if [$x_1 = \cdots = x_m$]{}. Therefore, one can equivalently rewrite problem  as

Note that the constraint set is the same as the set , since due to the connectivity of the graph $\mathcal{G}$.

In the next section, we state the proposed algorithm for solving  and analyze its convergence rate.

Algorithm and Results {#sec:results}
=====================

In this section, we build on recent results on dual approaches for optimal distributed optimization [@Uribe2017; @ani17] to construct an algorithm that solves  over a network. Moreover, we analyze its convergence rates and provide explicit dependencies on the problem parameters and the network topology.

Dual Approach for Strongly Convex Functions and Affine Constraints {#S:DualApprMain}
------------------------------------------------------------------

In [@ani17], the authors proposed a novel analysis for the minimization of strongly convex functions with affine constraints of the form

$$\begin{aligned}
\label{eq:linear}
    \min_{Ax=0}f(x),
    \end{aligned}$$ where $f(x)$ is $1$-strongly convex with respect to the with the corresponding dual problem defined as

$$\begin{aligned}
    \label{eq:dual_to_linear}
    \min_y g(y)  \ \ \ \text{where}  \ \ g(y) = \max_x \{\left\langle A^Ty,x \right\rangle -f(x)\}.
    \end{aligned}$$

The dual function $g(y)$ is $L$-smooth with , where $A_i$ is the $i$-th column of $A$. Thus, one can use accelerated first order methods such as Nesterov’s Fast Gradient [@nes83] or one of its recent reformulations [@all14] to obtain an approximate solution. The novelty in [@ani17] lies in the statement of the convergence rate of the accelerated methods in terms of the duality gap and the constraint violation. Additionally, it was shown that for the linear coupling accelerated algorithm [@all17] one can guarantee that the solutions will remain in a closed ball around the optimal solution, with a radius proportional to the distance between the initial point of the algorithm and the optimal solution. Next, we state a technical result, based on [@ani17], that will help us in the design and analysis of our proposed algorithm for the distributed computation of the WB.

\[thm:gasnikov\] The fast gradient method based on linear coupling proposed in [@all14] with the change $y_k$ to $w_k$ and $x_k$ to $y_k$ and applied to , has the following properties: $ \dd{\forall ~} k \geq N$ and $\varepsilon >0$, it holds that $$\begin{aligned}
        g(\pd{w_k}) + f(\breve x_k) \leq \varepsilon \qquad \text{and} \qquad \|A\breve x_k\|_{\pd{2}}\leq  \varepsilon / R,
        \end{aligned}$$ where , $N \triangleq \sqrt{\pd{16} \cu{L} R^2/\varepsilon}$, and $y^*$ is the optimal point of $g(\cdot)$ with minimal norm.

\[S:DualWasser\]

Additionally, from Theorem \[thm:cuturi\] the gradient can be expressed in closed form as

where $p^*_j(\tilde y_j)  = \alpha(\tilde y_j) \circ ( K \cdot { q_i}/{(K \alpha(\tilde y_j) )} )$. Moreover, it holds that one can recover the solution $\mathtt{p}^*$ to the primal problem  from a solution $\mathtt{y}^*$ to the dual problem  as $
    \mathtt{p}^* = \mathtt{p}^*(\sqrt{W}\mathtt{y}^* )$.

The optimality relation between the dual and the primal problem follows from Theorem $3.1$ in [@Cuturi2016]. In general, the dual problem  can have multiple solutions of the form $\mathtt{y}^* + \ker(\sqrt{W})$ when the matrix $\sqrt{W}$ does not have a full row rank. When the solution is not unique, we [*will use $\mathtt{y}^*$ to denote the smallest norm solution*]{}, and we let $R$ be its norm, i.e. $R = \|\mathtt{y}^*\|_2$.

. As a consequence, the computation of the WB of a set of discrete probability distributions $\{q_i\}_{i=1}^m$ is equivalent to solving the dual decomposable $L$-smooth (with respect to the $2$-norm) optimization problem  with  [@Kakade2009]. Specifically, in this setup it holds that

\[S:AlgAndMain\] We can explicitly write the Nesterov’s Accelerated Gradient Method (FGM) [@nes13] for smooth functions. Particularly, we use follow the linear coupling approach recently proposed in [@all14]. Setting $\hat{\mathtt{w}}_{k} = \hat{\mathtt{z}}_{k}=\hat{\mathtt{y}}_{k} = \boldsymbol{0}$, the FGM generates iterates according to:

where $\alpha_{k+1} = (k+2)/(2L)$ and $\tau_k = 2/(k+2)$.

Unfortunately, algorithm  cannot be executed in a distributed manner. Although the entries of local gradient vectors can be computed independently by each node, the sparsity pattern of the matrix $\sqrt{W}$ need not be the same as the communication constraints induced by the graph $\mathcal{G}$. Thus, the variables $\hat {\mathtt{w}}_{k}$ and $\hat {\mathtt{z}}_{k}$ cannot be computed on the network. This problem is solved by a change of variables such that $\tilde{ \mathtt{y}} = \sqrt{W}\hat{\mathtt{y}}$, $\tilde{ \mathtt{w}} = \sqrt{W}\hat{\mathtt{w}}$ and $\tilde{ \mathtt{z}} = \sqrt{W}\hat{\mathtt{z}}$.

Algorithm \[alg:main\] presents the resulting distributed accelerated gradient method for the dual problem of the WB problem.

. Now, we are ready to state our main result that provides a convergence rate for Algorithm \[alg:main\] with explicit dependencies on the problem parameters and the topology of the network.

\[thm:main\] Let $\varepsilon>0$ and assume that $\|\nabla \mathcal{W}^*_{\gamma,\mathtt{q}}(\tilde{\mathtt{y}})\|_2 \leq \dd{G} $ on a ball $B_R(0)$. Then, it holds that that after

iterations, the outputs of Algorithm \[alg:main\], i.e. $\mathtt{p}^*_N = [(p_N^*)_1^T,\cdots,(p_N^*)_m^T]^T$ and $\mathtt{y}^*_N = [(y_N^*)_1^T,\cdots,(y_N^*)_m^T]^T$ have the following properties: $$\begin{aligned}
        \mathcal{W}_{\gamma,\mathtt{q}}(\mathtt{p}^*_N)+ \mathcal{W}^*_{\gamma,\mathtt{q}}(\mathtt{y}^*_N)  
        \leq  \varepsilon \ \ \text{and} \ \ \|\sqrt{W}\mathtt{p}^*_N\|_2 \leq {\varepsilon}/{R}.
        \end{aligned}$$

The dual function $\mathcal{W}^*_{\gamma,\mathtt{q}}(\mathtt{y})$ is $(\cu{d_{\max}}/\gamma)$-smooth. Thus, from Theorem \[thm:gasnikov\] it follows that $$\begin{aligned}
        \mathcal{W}_{\gamma,\mathtt{q}}(\mathtt{p}^*_N)+ \mathcal{W}^*_{\gamma,\mathtt{q}}(\mathtt{y}^*_N)  
        \leq  \varepsilon \ \ \text{and} \ \ \|\sqrt{W}\mathtt{p}^*_N\|_2 \leq {\varepsilon}/{R}.
        \end{aligned}$$ holds for $k \geq \sqrt{16 \cu{d_{\max} }R^2/(\gamma \varepsilon)}$.\
Moreover, Thus, we require $ k \geq {\scriptscriptstyle \sqrt{\frac{16 \cu{G}^2}{\gamma \cdot \varepsilon}\cu{\frac{d_{\max}}{d_{\min}} }}}$ and the desired result follows.

Theorem \[thm:main\] provides an estimate of the minimum number of iterations required for the proposed algorithm to reach some arbitrary relative accuracy in the solution of the distributed WB problem. The convergence rate is shown to be of the order $O(1/k^2)$ which has been established to be optimal for smooth convex optimization problems [@nes13] with an additional cost proportional to the .

In general, one might be interested in finding a WB for the original Wasserstein distance with no regularization term. That is, to solve problem  with $\gamma=0$. The next theorem explains a choice of $\gamma$ that provides a convergence rate results with respect to the non-regularized optimal transport based on the iterates of Algorithm .

Let $\varepsilon>0$, and assume that $\|\nabla \mathcal{W}^*_{\gamma,\mathtt{q}}(\tilde{\mathtt{y}})\|_2 \leq \dd{G} $ on a ball $B_R(0)$. Moreover, set $\gamma = \varepsilon/(4\pd{ m} \log n)$.

iterations, the outputs of Algorithm \[alg:main\], i.e. $\mathtt{p}^*_N = [(p_N^*)_1^T,\cdots,(p_N^*)_m^T]^T$ and $\mathtt{y}^*_N = [(y_N^*)_1^T,\cdots,(y_N^*)_m^T]^T$ have the following properties [ $$\begin{aligned}
            \mathcal{W}_{0,\mathtt{q}}(\mathtt{p}^*_N)- \pd{\mathcal{W}_{0,\mathtt{q}}(\mathtt{p}^*)}
            \leq  \varepsilon \ \ \text{and} \ \ \|\sqrt{W}\mathtt{p}^*_N\|_2 \leq {\varepsilon}/{(\pd{2}R)}.
            \end{aligned}$$ ]{}

Numerical Experiments {#sec:numerical}
=====================

In this section, we show two numerical experiments to validate the results of Theorem \[thm:main\]. We explore the problem of computing WB of univariate, discrete and truncated Gaussian densities and the computation of the WB of a subsample of $1000$ digit images from the MNIST dataset.

Barycenter of Gaussian Distributions
------------------------------------

Initially, we explore the computation of WB for sets of univariate, discretized and truncated Gaussian densities [@Agueh2011]. We consider a network of agents where each agent $i$ holds a univariate, discretized and truncated Gaussian distribution, with mean $\mu_i \in [-5,5]$, standard deviation $\sigma \in [0.1,2]$ and equally spaced support of $100$ points in $[-5,5]$. The entropy regularization parameter is set to $\gamma = 0.1$. Figure \[fig:results\] shows the distance to optimality and the distance to consensus of Algorithm \[alg:main\] for the star, cycle, complete, and Erdős-Rényi random graph graphs with a fixed size of $50$ nodes. Also, Figure \[fig:results\] shows the scalability of the algorithm, i.e., the number of iterations required to reach an $\varepsilon$ accuracy in the distance to optimality and consensus for networks of increasing size.

MNIST Dataset
-------------

We randomly sample $1000$ images for each digit of the MNIST dataset [@LeCun1998; @LeCun1998a]. Each image has $28\times28$ pixels and is scaled uniformly at random between $0.5$ and $2$ of its size and randomly located on a larger $56\times56$ blank image. The pixel values of the image are normalized to add up to $1$. We assign one sample from each digit to each agent on a group of $1000$ agents, and the objective is to jointly compute the WB for each digit of the $1000$ samples present in the network. The agents are connected over an Erdős-Rényi random graph with $1000$ nodes and connectivity parameter $4/1000$. The entropy regularization parameter is set to $\gamma = 0.01$. Figure \[fig:mnist\] shows the local barycenter of the $9$ digits for a subset of $3$ agents in the network for $0$, $60$, and $300$ iterations. As the number of iteration increases, all agents converge to a common image, namely, the WB of the images held by the agents in the network.

at (0,0) [ $\ \ \ \ N=0$]{};

\

at (0,0) [ $\ \ \  N=60$]{};

\

at (0,0) [ $\ \  N=300$]{};

Discussion and Conclusions {#sec:conclusions}
==========================

We developed a novel algorithm for the distributed computation of WB over networks where a group of agents connected over a network and each agent holds some local probability distribution with finite support. Our results provably guarantee that all agents in the network will converge to the WB of all distribution held by the agents in the network. We provide an explicit and non-asymptotic convergence rate of the order $O(1/k^2)$ with an additional cost proportional to the ratio between the maximum and minimum degree among the nodes in the graph over which the agents exchange information.

The case where spectral information of the network is not available or when the graphs are directed or change with time require further study. The use of second-order information can also be exploited to get better performance. Also, recently proposed stochastic approaches can provide more efficient algorithms [@Claici2018; @Staib2017; @dvurechensky2018decentralize; @rogozin2018optimal].

[^1]: C.A. Uribe (*cauribe@mit.edu*) is with the Laboratory for Information and Decision Systems (LIDS), Massachusetts Institute of Technology. A. Gasnikov (*gasnikov@yandex.ru*) is with the Moscow Institute of Physics and Technology, and the Institute for Information Transmission Problems RAS. A. Nedić (*angelia.nedich@asu.edu*) is with the ECEE Department, Arizona State University, and Moscow Institute of Physics and Technology. The work of A. Gasnikov in Section III-C was supported by the grant of the president of Russian Federation no. MD-1320.2018.1. The work of A. Nedić and C.A. Uribe is supported by the National Science Foundation under grant no. CPS 15-44953.
