On Robust Optimal Transport: Computational Complexity and Barycenter Computation

2 ε ) time , in which n is the number of supports of the probability distributions and ε is the desired error . Furthermore , we investigate a fixed-support robust barycenter problem between m discrete probability distributions with at most n number of supports and develop an approximating algorithm based on iterative Bregman projections ( IBP ) . For the specific case m = 2 , we show that this algorithm can approximate the optimal barycenter value in Õ ( mn 2 ε ) time , thus being better than the previous complexity Õ ( mn 2 ε2 ) of the IBP algorithm for approximating the Wasserstein barycenter . 1 Introduction . The recent advance in computation with optimal transport ( OT ) problem [ 12 , 3 , 13 , 7 , 22 , 26 , 20 ] has led to a surge of interest in using that tool in various domains of machine learning and statistics . The range of its applications is broad , including deep generative models [ 4 , 16 , 36 ] , scalable Bayes [ 33 , 34 ] , mixture and hierarchical models [ 24 ] , and other applications [ 32 , 29 , 10 , 17 , 37 , 35 , 8 ] . The goal of optimal transport is to find a minimal cost of moving masses between ( supports of ) probability distributions . It is known that the estimation of transport cost is not robust when there are outliers . To deal with this issue , [ 38 ] proposed a trimmed version of optimal transport . In particular , they search for truncated probability distributions such that the transport cost between them is minimized . However , their trimmed optimal transport is non-trivial to compute , which hinders its usage in practical applications . Another line of works proposed using unbalanced optimal transport ( UOT ) to solve the sensitivity of optimal transport to outliers [ 5 , 31 ] . More specifically , their idea is to assign as small as possible masses to outliers by relaxing the marginal constraints of OT through a penalty function such as the Kullback-Leibler ( KL ) divergence . This direction of robust optimal transport has been shown to have good performance in generative models and domain adaptation [ 5 ] . Although this approach achieved considerable success , the full picture of its computational complexity has remained missing. ? Khang Le and Huy Nguyen contributed equally to this work . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our Contribution : In the paper , we provide a comprehensive study of the computational complexity of robust optimal transport and its corresponding barycenter problem when the probability distributions are discrete and have at most n components . Our contribution is twofold and can be summarized as follows : ( 1 ) On robust optimal transport , we consider two versions corresponding to two ways of relaxing marginal constraints in the standard optimal transport problem via the KL divergence . We show that two scaling algorithms computing these robust formulations have the complexities Õ ( n2/ε ) , where ε denotes the desired error for the computed cost . These complexities are lower than the complexity of the Sinkhorn algorithm for solving the optimal transport problem , which is Õ ( n2/ε2 ) [ 13 ] , and match the complexity of the Sinkhorn algorithm that solves the UOT problem [ 27 ] . Furthermore , we show how the above complexity can be improved by utilizing the low-rank approximation method to speed up the matrix-vector computations in the loop similar to [ 2 ] , and obtain the improved computing time of Õ ( nr2 + nrε ) , where r is the approximated rank . ( 2 ) On robust barycenter problem , where the goal is to determine a probability measure that minimizes its robust optimal cost to a given set of m ≥ 2 probability measures , we propose ROBUSTIBP algorithm for solving the robust barycenter problem , which is inspired by the iterative Bregman projection ( IBP ) algorithm for solving the traditional barycenter problem [ 6 ] . We show that whenm = 2 , the complexity of ROBUSTIBP algorithm is at the order of Õ ( mn2/ε ) , better than that of the IBP algorithm for solving the traditional barycenter problem [ 19 ] , which is Õ ( mn2/ε2 ) . To the best of our knowledge , the ROBUSTIBP is also the first practical algorithm obtaining the near-optimal complexity Õ ( mn2/ε ) for solving the barycenter problem even under only the setting m = 2 . Organization : The paper is organized as follows . In Section 2 , we provide the background on the optimal transport problem and some of its variants that have robust effects . In Section 3 , we discuss in-depth the variant where only one marginal constraint is relaxed , study the computational complexity of a Sinkhorn-based algorithm that solves it , and then briefly introduce the fully-relaxed formulation . We also establish the complexities of these algorithms after applying Nyström method . Subsequently , we present our study of the robust barycenter problem in Section 4 . In Section 5 , we carry out empirical studies to illustrate the theories before concluding with a few discussions in Section 6 . The proofs of our theoretical results are in the supplementary material . Notation : We let [ n ] stand for the set { 1 , 2 , . . . , n } while Rn+ indicates the set of all vectors with nonnegative entries . For a vector x ∈ Rn and p ∈ [ 1 , ∞ ) , we denote ‖x‖p as its ` p-norm and diag ( x ) as the diagonal matrix with x on the diagonal . The natural logarithm of a vector a = ( a1 , ... , an ) ∈ Rn+ is denoted by log a = ( log a1 , ... , log an ) , 1n stands for a vector of length n that all of its entries equal to 1 , and ∂xf refers to the partial differentiation of function f with respect to x . For any given space X ⊂ Rd , we denote by P ( X ) the space of all probability measures on X . Given an integer n > 0 and a real number ε > 0 , the notation a = O ( b ( n , ε ) ) means that a ≤ C · b ( n , ε ) where C is independent of n and ε . Meanwhile , the notation a = Õ ( b ( n , ε ) ) indicates the previous inequality may depend on a logarithmic function of n and ε . For any two probability measures x = ( x1 , . . . , xn ) and y = ( y1 , . . . , yn ) with the same supports , the generalized Kullback-Leibler divergence is defined as KL ( x‖y ) = ∑n i=1 [ xi log ( xi yi ) − xi + yi ] . Finally , the entropy of a matrix X is given by H ( X ) = ∑n i , j=1−Xij ( logXij − 1 ) . 2 Background on Optimal Transport . In this section , we review optimal transport and its unbalanced formulation , then from that deriving formulations for robust optimal transport . For any P and Q in P ( X ) for a space X , the OT distance between P and Q takes the following form OT ( P , Q ) : = min π∈Π ( P , Q ) ∫ c ( x , y ) dπ ( x , y ) , ( 1 ) where Π ( P , Q ) is the set of joint probability distributions in X × X such that their marginal distributions are P and Q , and c : X × X → [ 0 , ∞ ) is a cost function . Unbalanced Optimal Transport : When P or Q is not a probability distribution , the OT formulation between P and Q in equation ( 1 ) is no longer valid . One solution to this issue is using the unbalanced optimal transport ( UOT ) [ 9 ] , which is given by : UOT ( P , Q ) : = min π∈M+ ( X×X ) ∫ c ( x , y ) dπ ( x , y ) + τ1KL ( π1‖P ) + τ2KL ( π2‖Q ) , ( 2 ) whereM+ ( X × X ) denotes the set of joint non-negative measures on the space X × X ; π1 , π2 are the marginal distributions of π and respectively correspond to P and Q ; τ1 , τ2 are regularized positive parameters . Note that , we can replace the KL divergence in equation ( 2 ) by any Csiszár-divergence [ 11 ] . However , we only consider the case of KL divergence in this work . Robust Optimal Transport : Optimal transport is well-known for not being robust in the present of outliers . A way to deal with this issue is using the approach of unbalanced optimal transport ( UOT ) , which has demonstrated favorable practical performance in generative models and domain adaptation [ 5 ] . More specifically , when P and Q are probability distributions in X , the Robust Unconstrained Optimal Transport ( ROT ) admits the following form ROT ( P , Q ) : = inf P1 , Q1∈P ( X ) min π∈Π ( P1 , Q1 ) ∫ c ( x , y ) dπ ( x , y ) + τ1KL ( P1‖P ) + τ2KL ( Q1‖Q ) , ( 3 ) where τ1 , τ2 > 0 are some given regularized parameters . The reason to name it robust unconstrained optimal transport is that instead of looking for an optimal transport plan moving masses from P to Q , we seek another plan that optimally transports masses between their approximations , which are probability measures P1 and Q1 , under the KL divergence . This formulation is closely related to the ones studied in [ 5 ] and [ 23 ] : the former used χ2-divergence for the relaxation and the latter used total variation distance ( note that those three divergences all together belong to the family of f -divergence ) . By relaxing only one marginal constraint regarding ( presumably ) on P , we have another version of ROT , named Robust Semi-constrained Optimal Transport ( RSOT ) , which is given by RSOT ( P , Q ) : = inf P1∈P ( X ) min π∈Π ( P1 , Q ) ∫ c ( x , y ) dπ ( x , y ) + τKL ( P1‖P ) , ( 4 ) where τ > 0 is a regularized parameter . We could also define RSOT ( Q , P ) similarly with a remark that although RSOT ( P , Q ) can be different from RSOT ( Q , P ) , the techniques for obtaining the computational complexity of both are similar . UOT vs ROT/RSOT : Though the formulations ROT/RSOT and UOT seem to be similar , they serve different purposes . The goal of UOT is to deal with unbalanced measures , thus there is no condition on the “ transport plan ” . Hence , the meaning of the optimal plan π of UOT problem is dependent on the interpreter . For example , in applications such as [ 30 ] , the UOT is used to figure out the developmental trajectory of cells . Meanwhile , ROT/RSOT aim to seek an accurate transport plan between two possibly corrupted probability distributions . The toy example in Figure 1 illustrates this difference . In particular , the marginals of the “ transport plan '' obtained by the latter ( see plots ( b ) , ( d ) ) are very different from the two original probability measures a , b . On the other hand , the solution of the former leads to good approximations of a and b ( see plots ( a ) , ( c ) ) while removing some bumps in both tails which are presumably outliers .

Summary: This paper proposes two robust formulation of the optimal transport problem (namely RSOT and ROT) based on the relaxation of the marginal constraints of the optimal transport plan using Kullback-Leibler divergence. The authors introduces algorithms based on Sinkhorn iterations in order to solve RSOT and ROT problems, and analyse their computational complexities. The fixed-support Wasserstein barycenter problem is also tackled using the proposed RSOT as a divergence, and the algorithm shows a better complexity than the one based on entropy regularised OT. Numerical experiments are also provided to support the theoretical results and applied task are performed (color transfer and generative modeling).