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: This paper considers the problem of approximate optimal transport, where population probability measures are replaced by discrete PMFs at collections of samples, and there is an additional "equitable" constraint that seeks to ``fairly" distribute the probability mass. This admits a formulation as a linear program and may be solved efficiently through its dual. However, recovering the primal solution is nontrivial, and there exists no rate analysis to-date. This work develops a stochastic dual ascent method to solve an entropy-regularized version of this LP, where entropy regularization is introduced to ensure strong concavity, as well as avoid a notion of over-fitting. The crux of this method is an alternating block coordinate ascent algorithm, which is then combined with an extrapolation/Nesterov acceleration scheme to accelerate convergence. Numerical experiments demonstrate the practical utility of the proposed approach.