On the Convergence of Projected Alternating Maximization for Equitable and Optimal Transport

This paper studies the equitable and optimal transport ( EOT ) problem , which has many applications such as fair division problems and optimal transport with multiple agents etc . In the discrete distributions case , the EOT problem can be formulated as a linear program ( LP ) . Since this LP is prohibitively large for general LP solvers , Scetbon et al . Scetbon et al . ( 2021 ) suggests to perturb the problem by adding an entropy regularization . They proposed a projected alternating maximization algorithm ( PAM ) to solve the dual of the entropy regularized EOT . In this paper , we provide the first convergence analysis of PAM . A novel rounding procedure is proposed to help construct the primal solution for the original EOT problem . We also propose a variant of PAM by incorporating the extrapolation technique that can numerically improve the performance of PAM . Results in this paper may shed lights on block coordinate ( gradient ) descent methods for general optimization problems . 1 INTRODUCTION . Optimal transport ( OT ) is a classical problem that recently finds many emerging applications in machine learning and artificial intelligence , including generative models Arjovsky et al . ( 2017 ) , representation learning Ozair et al . ( 2019 ) , reinforcement learning Bellemare et al . ( 2017 ) and word embeddings Alvarez-Melis et al . ( 2019 ) etc . More recently , Scetbon et al . Scetbon et al . ( 2021 ) proposed an equitable and optimal transport ( EOT ) problem that targets to fairly distribute the workload of OT when there are multiple agents . In this problem formulation , there are multiple agents working together to move mass from measures µ to ν and each agent has its unique cost function . A very important issue that needs to be considered here is the fairness , which aims at finding transportation plans such that the workloads among all the agents are equal to each other . This can be achieved by minimizing the largest transportation cost among all agents , which leads to a convex-concave saddle point problem . The EOT problem has wide applications in economics and machine learning , such as fair division or the cake-cutting problem Moulin ( 2003 ) ; Brandt et al . ( 2016 ) , multi-type resource allocation Mackin & Xia ( 2015 ) , internet minimal transportation time and sequential optimal transport Scetbon et al . ( 2021 ) . We now describe the EOT problem formally . Given two discrete probability measures µn =∑n i=1 aiδxi and νn = ∑n i=1 biδyi , the EOT studies the problem of transporting mass from µ to ν by N agents . Here , { x1 , x2 , ... , xn } ⊂ Rd and { y1 , y2 , ... , yn } ⊂ Rd are the support points of each measure and a = [ a1 , a2 , ... , an ] ⊤ ∈ ∆n , b = [ b1 , b2 , ... , bn ] ⊤ ∈ ∆n are corresponding weights for each measure , where ∆n denotes the probability simplex in Rn . Moreover , throughout this paper , we assume vector b > 0 . For each agent k , we denote its unique cost function as ck ( x , y ) , k ∈ [ N ] = { 1 , . . . , N } and its cost matrix as Ck , where Cki , j = ck ( xi , yj ) . Moreover , we define the following coupling decomposition set ΠNa , b : = { π = ( πk ) k∈ [ N ] ∣∣∣∣∣ r ( ∑ k πk ) = a , c ( ∑ k πk ) = b , πkij ≥ 0 , ∀i , j ∈ [ n ] } , where r ( π ) = π1 , c ( π ) = π⊤1 are the row sum and column sum of matrix π respectively . Mathematically , the EOT problem can be formulated as min π∈ΠNa , b max 1≤k≤N ⟨πk , Ck⟩ . ( 1 ) When N = 1 , ( 1 ) reduces to the standard OT problem . Note that ( 1 ) minimizes the point-wise maximum of a finite collection of functions . It is easy to see that ( 1 ) is equivalent to the following constrained problem : min π∈ΠNa , b max λ∈∆N+ ℓ ( π , λ ) : = N∑ k=1 λk⟨πk , Ck⟩ . ( 2 ) The following proposition shows an important property of EOT : at the optimum of the minimax EOT formulation ( 2 ) , the transportation costs of the agents are equal to each other . Proposition 1 ( Scetbon et al. , 2021 , Proposition 1 ) Assume that all cost matrices Ck , k ∈ [ N ] have the same sign . Let π∗ ∈ ΠNa , b be the optimal solution of ( 2 ) . It holds that ⟨ ( π∗ ) i , Ci⟩ = ⟨ ( π∗ ) j , Cj⟩ , ∀i , j ∈ [ N ] . ( 3 ) Note that Proposition 1 requires all cost matrices to have the same sign . When the cost matrices are all non-negative , ( 2 ) solves the transportation problem with multiple agents . When the cost matrices are all non-positive , the cost matrices are interpreted as the utility functions and ( 2 ) solves the fair division problem Moulin ( 2003 ) . The discrete OT is a linear programming ( LP ) problem ( in fact , an assignment problem ) with a complexity of O ( n3 log n ) Tarjan ( 1997 ) . Due to this cubic dependence on the dimension n , it is challenging to solve large-scale OT in practice . A widely adopted compromise is to add an entropy regularizer to the OT problem Cuturi ( 2013 ) . The resulting problem is strongly convex and smooth , and its dual problem can be efficiently solved by the celebrated Sinkhorn ’ s algorithm Sinkhorn & Knopp ( 1967 ) ; Cuturi ( 2013 ) . This strategy is now widely used in the OT community due to its computational advantages as well as improved sample complexity Genevay et al . ( 2019 ) . Similar ideas were also used for computing the Wasserstein barycenter Benamou et al . ( 2015 ) , projection robust Wasserstein distance Paty & Cuturi ( 2019 ) ; Lin et al . ( 2020 ) ; Huang et al . ( 2021a ) , projection robust Wasserstein barycenter Huang et al . ( 2021b ) . Motivated by these previous works , Scetbon et al . Scetbon et al . ( 2021 ) proposed to add an entropy regularizer to ( 2 ) , and designed a projected alternating maximization algorithm ( PAM ) to solve its dual problem . However , the convergence of PAM has not been studied . Scetbon et al . Scetbon et al . ( 2021 ) also proposed an accelerated projected gradient ascent algorithm ( APGA ) for solving a different form of the dual problem of the entropy regularized EOT . Since the objective function of this new dual form has Lipschitz continuous gradient , APGA is essentially the Nesterov ’ s accelerated gradient method and thus its convergence rate is known . However , numerical experiments conducted in Scetbon et al . ( 2021 ) indicate that APGA performs worse than PAM . We will discuss the reasons in details later . Our Contributions . There are mainly three issues with the PAM and APGA algorithms in Scetbon et al . ( 2021 ) , and we will address all of them in this paper . Our results may shed lights on designing new block coordinate descent algorithms . Our main contributions are given below . ( i ) The PAM algorithm in Scetbon et al . ( 2021 ) only returns the dual variables . How to find the primal solution of ( 2 ) , i.e. , the optimal transport plans π , was not discussed in Scetbon et al . ( 2021 ) . In this paper , we propose a novel rounding procedure to find the primal solution . Our rounding procedure is different from the one widely used in the literature Altschuler et al . ( 2017 ) . ( ii ) We provide the first convergence analysis of the PAM algorithm , and analyze its iteration complexity for finding an ϵ-optimal solution to the EOT problem ( 2 ) . In particular , we show that it takes at most O ( Nn2ϵ−2 ) arithmetic operations to find an ϵ-optimal solution to ( 2 ) . This matches the rate of the Sinkhorn ’ s algorithm for computing the Wasserstein distance Dvurechensky et al . ( 2018 ) . ( iii ) We propose a variant of PAM that incorporates the extrapolation technique as used in Nesterov ’ s accelerated gradient method . We name this variant as PAM with Extrapolation ( PAME ) . The iteration complexity of PAME is also analyzed . Though we are not able to prove a better complexity over PAM at this moment , we find that PAME performs much better than PAM numerically . Notation . For vectors a and b with the same dimension , a./b denotes their entry-wise division . We denote c∞ : = maxk ∥Ck∥∞ . Throughout this paper , we assume vector b > 0 , and we denote ι : = minj log ( bj ) . We use 1n to denote the n-dimensional vector whose entries are all equal to one . We use IX ( x ) to denote the indicator function of set X , i.e. , IX ( x ) = 0 if x ∈ X , and IX ( x ) = ∞ otherwise . We denote ct = c ( ∑N k=1 π k ( f t+1 , gt , λt ) ) . For integer N > 0 , we denote [ N ] : = { 1 , . . . , N } . We also denote π ( f , g , λ ) = [ πk ( f , g , λ ) ] k∈ [ N ] . 2 PROJECTED ALTERNATING MAXIMIZATION ALGORITHM . The PAM algorithm proposed in Scetbon et al . ( 2021 ) aims to solve the entropy regularized EOT problem , which is given by min π∈ΠNa , b max λ∈∆N+ ℓη ( π , λ ) : = N∑ k=1 pkη ( π k , λ ) ( 4 ) where η > 0 is a regularization parameter , pkη ( π k , λ ) : = λk⟨πk , Ck⟩ − ηH ( πk ) , and the entropy function H is defined as H ( π ) = − ∑ i , j πi , j ( log πi , j − 1 ) . The entropy regularization was first introduced into the OT problem by Cuturi ( 2013 ) and is now widely used in the OT community . By adding an entropy regularizer , the primal problem becomes strongly convex and the dual problem is unconstrained and is suitable for alternating maximization . This leads to the Sinkhorn ’ s algorithm which has low per-iteration complexity and thus is scalable . The PAM algorithm proposed by Scetbon et al . ( 2021 ) used the same idea for the EOT problem . Note that ( 4 ) is a strongly-convexconcave minimax problem whose constraint sets are convex and bounded , and thus the Sion ’ s minimax theorem Sion ( 1958 ) guarantees that min π∈ΠNa , b max λ∈∆N+ ℓη ( π , λ ) = max λ∈∆N+ min π∈ΠNa , b ℓη ( π , λ ) . ( 5 ) Now we consider the dual problem of minπ∈ΠNa , b ℓη ( π , λ ) . First , we add a redundant constraint∑ k , i , j π k i , j = 1 and consider the dual of min π∈ΠNa , b , ∑ k , i , j π k i , j=1 ℓη ( π , λ ) . ( 6 ) The reason for adding this redundant constraint is to guarantee that the dual objective function is Lipschitz smooth . It is easy to verify that the dual problem of ( 6 ) is given by max f , g min∑ k , i , j π k i , j=1 , π∈ ( Rn×n+ ) N N∑ k=1 λk⟨πk , Ck⟩−ηH ( π ) +f⊤ ( a− r ( ∑ k πk ) ) +g⊤ ( b− c ( ∑ k ( πk ) ⊤ ) ) , ( 7 ) where f and g are the dual variables and H ( π ) = ∑ k H ( π k ) . It is noted that problem ( 7 ) admits the following solution : πk ( f , g , λ ) = ζk ( f , g , λ ) ∑ k ∥ζk ( f , g , λ ) ∥1 , ∀k ∈ [ N ] , ( 8 ) where ζk ( f , g , λ ) = exp ( f1⊤n + 1ng ⊤ − λkCk η ) , ∀k ∈ [ N ] . ( 9 ) By plugging ( 8 ) into ( 7 ) , we obtain the following dual problem of ( 6 ) : max f∈Rn , g∈Rn ⟨f , a⟩+ ⟨g , b⟩ − η log ( N∑ k=1 ∥ζk ( f , g , λ ) ∥1 ) − η . ( 10 ) Plugging ( 10 ) into ( 5 ) , we know that the entropy regularized EOT problem ( 4 ) is equavalent to a pure maximization problem : max f∈Rn , g∈Rn , λ∈∆N F ( f , g , λ ) : = ⟨f , a⟩+ ⟨g , b⟩ − η log ( N∑ k=1 ∥ζk ( f , g , λ ) ∥1 ) − η . ( 11 ) Function F ( f , g , λ ) is a smooth concave function with three block variables ( f , g , λ ) . We use ( f∗ , g∗ , λ∗ ) to denote an optimal solution of ( 11 ) , and we denote F ∗ = F ( f∗ , g∗ , λ∗ ) . The PAM algorithm proposed in Scetbon et al . ( 2021 ) is essentially a block coordinate descent ( BCD ) algorithm for solving ( 11 ) . More specifically , the PAM updates the three block variables by the following scheme : f t+1 ∈ argmax f F ( f , gt , λt ) , ( 12a ) gt+1 ∈ argmax g F ( f t+1 , g , λt ) , ( 12b ) λt+1 : = Proj∆N ( λt + τ∇λF ( f t+1 , gt+1 , λt ) ) . ( 12c ) Each iteration of PAM consists of two exact maximization steps followed by one projected gradient step . Importantly , the two exact maximization problems ( 12a ) - ( 12b ) have numerous optimal solutions , and we choose to use the following ones : f t+1 = f t + η log a r ( ∑N k=1 ζ k ( f t , gt , λt ) ) , ( 13 ) gt+1 = gt + η log b c ( ∑N k=1 ζ k ( f t+1 , gt , λt ) ) . ( 14 ) Furthermore , the optimiality conditions of ( 12a ) - ( 12b ) imply that a− r ( ∑N k=1 ζ k ( f t+1 , gt , λt ) ) ∑ k ∥ζk ( f t+1 , gt , λt ) ∥1 = 0 , b− c ( ∑N k=1 ζ k ( f t+1 , gt+1 , λt ) ) ∑ k ∥ζk ( f t+1 , gt+1 , λt ) ∥1 = 0 , ∀t . ( 15 ) However , we need to point out that the PAM ( 12 ) only returns the dual variables ( f t , gt , λt ) . One can compute the primal variable π using ( 8 ) , but it is not necessarily a feasible solution . That is , π computed from ( 8 ) does not satisfy π ∈ ΠNa , b . How to obtain an optimal primal solution from the dual variables was not discussed in Scetbon et al . ( 2021 ) . For the OT problem , i.e. , N = 1 , a rounding procedure for returning a feasible primal solution has been proposed in Altschuler et al . ( 2017 ) . However , this rounding procedure can not be applied to the EOT problem directly . In the next section , we propose a new rounding procedure for returning a primal solution based on the dual solution ( f t , gt , λt ) . This new rounding procedure involves a dedicated way to compute the margins .

Summary: The paper addresses the equitable and optimal transport (EOT) problem. One approach to this problem involves perturbation via the addition of an entropy regularization. For this formulation, researchers have proposed a projected alternating maximization algorithm (PAM). The authors of this paper provide a convergence analysis of the PAM algorithm. They also introduce a novel rounding procedure to construct the primal solution of the EOT problem. The paper also contributes a variant of PAM that leads to numerical performance improvements. Numerical experiments are provided for a synthetic dataset.