Feature-Robust Optimal Transport for High-Dimensional Data

1 INTRODUCTION . Optimal transport ( OT ) is a machine learning problem with several applications in the computer vision and natural language processing communities . The applications include Wasserstein distance estimation ( Peyré et al. , 2019 ) , domain adaptation ( Yan et al. , 2018 ) , multitask learning ( Janati et al. , 2019 ) , barycenter estimation ( Cuturi & Doucet , 2014 ) , semantic correspondence ( Liu et al. , 2020 ) , feature matching ( Sarlin et al. , 2019 ) , and photo album summarization ( Liu et al. , 2019 ) . The OT problem is extensively studied in the computer vision community as the earth mover ’ s distance ( EMD ) ( Rubner et al. , 2000 ) . However , the computational cost of EMD is cubic and highly expensive . Recently , the entropic regularized EMD problem was proposed ; this problem can be solved using the Sinkhorn algorithm with a quadratic cost ( Cuturi , 2013 ) . Owing to the development of the Sinkhorn algorithm , researchers have replaced the EMD computation with its regularized counterparts . However , the optimal transport problem for high-dimensional data has remained unsolved for many years . Recently , a robust variant of the OT was proposed for high-dimensional OT problems and used for divergence estimation ( Paty & Cuturi , 2019 ; 2020 ) . In the robust OT framework , the transport plan is computed with the discriminative subspace of the two data matrices X ∈ Rd×n and Y ∈ Rd×m . The subspace can be obtained using dimensionality reduction . An advantage of the subspace robust approach is that it does not require prior information about the subspace . However , given prior information such as feature groups , we can consider a computationally efficient formulation . The computation of the subspace can be expensive if the dimensionality of data is high , for example , 104 . One of the most common prior information items is a feature group . The use of group features is popular in feature selection problems in the biomedical domain and has been extensively studied in Group Lasso ( Yuan & Lin , 2006 ) . The key idea of Group Lasso is to prespecify the group variables and select the set of group variables using the group norm ( also known as the sum of ` 2 norms ) . For example , if we use a pretrained neural network as a feature extractor and compute OT using the features , then we require careful selection of important layers to compute OT . Specifically , each layer output is regarded as a grouped input . Therefore , using a feature group as prior information is a natural setup and is important for considering OT for deep neural networks ( DNNs ) . In this paper , we propose a high-dimensional optimal transport method by utilizing prior information in the form of grouped features . Specifically , we propose a feature-robust optimal transport ( FROT ) problem , for which we select distinct group feature sets to estimate a transport plan instead of determining its distinct subsets , as proposed in ( Paty & Cuturi , 2019 ; 2020 ) . To this end , we formulate the FROT problem as a min–max optimization problem and transform it into a convex optimization problem , which can be accurately solved using the Frank–Wolfe algorithm ( Frank & Wolfe , 1956 ; Jaggi , 2013 ) . The FROT ’ s subproblem can be efficiently solved using the Sinkhorn algorithm ( Cuturi , 2013 ) . An advantage of FROT is that it can yield a transport plan from high-dimensional data using feature selection , using which the significance of the features is obtained without any additional cost . Therefore , the FROT formulation is highly suited for high-dimensional OT problems . Through synthetic experiments , we initially demonstrate that the proposed FROT is robust to noise dimensions ( See Figure 1 ) . Furthermore , we apply FROT to a semantic correspondence problem ( Liu et al. , 2020 ) and show that the proposed algorithm achieves SOTA performance . Contribution : . • We propose a feature robust optimal transport ( FROT ) problem and derive a simple and efficient Frank–Wolfe based algorithm . Furthermore , we propose a feature-robust Wasserstein distance ( FRWD ) . • We apply FROT to a high-dimensional feature selection problem and show that FROT is consistent with the Wasserstein distance-based feature selection algorithm with less computational cost than the original algorithm . • We used FROT for the layer selection problem in a semantic correspondence problem and showed that the proposed algorithm outperforms existing baseline algorithms . 2 BACKGROUND . In this section , we briefly introduce the OT problem . Optimal transport ( OT ) : The following are given : independent and identically distributed ( i.i.d . ) samples X = { xi } ni=1 ∈ Rd×n from a d-dimensional distribution p , and i.i.d . samples Y = { yj } mj=1 ∈ Rd×m from the d-dimensional distribution q . In the Kantorovich relaxation of OT , admissible couplings are defined by the set of the transport plan : U ( µ , ν ) = { Π ∈ Rn×m+ : Π1m = a , Π > 1n = b } , where Π ∈ Rn×m+ is called the transport plan , 1n is the n-dimensional vector whose elements are ones , and a = ( a1 , a2 , . . . , an ) > ∈ Rn+ and b = ( b1 , b2 , . . . , bm ) > ∈ Rm+ are the weights . The OT problem between two discrete measures µ = ∑n i=1 aiδxi and ν = ∑m j=1 bjδyj determines the optimal transport plan of the following problem : min Π∈U ( µ , ν ) n∑ i=1 m∑ j=1 πijc ( xi , yj ) , ( 1 ) where c ( x , y ) is a cost function . For example , the squared Euclidean distance is used , that is , c ( x , y ) = ‖x − y‖22 . To solve the OT problem , Eq . ( 1 ) ( also known as the earth mover ’ s distance ) using linear programming requires O ( n3 ) , ( n = m ) computation , which is computationally expensive . To address this , an entropic-regularized optimal transport is used ( Cuturi , 2013 ) . min Π∈U ( µ , ν ) n∑ i=1 m∑ j=1 πijc ( xi , yj ) + H ( Π ) , where ≥ 0 is the regularization parameter , and H ( Π ) = ∑ni=1 ∑m j=1 πij ( log ( πij ) − 1 ) is the entropic regularization . If = 0 , then the regularized OT problem reduces to the EMD problem . Owing to entropic regularization , the entropic regularized OT problem can be accurately solved using Sinkhorn iteration ( Cuturi , 2013 ) with a O ( nm ) computational cost ( See Algorithm 1 ) . Wasserstein distance : If the cost function is defined as c ( x , y ) = d ( x , y ) with d ( x , y ) as a distance function and p ≥ 1 , then we define the p-Wasserstein distance of two discrete measures µ = ∑n i=1 aiδxi and ν = ∑m j=1 bjδyj as Wp ( µ , ν ) = min Π∈U ( µ , ν ) n∑ i=1 m∑ j=1 πijd ( xi , yj ) p 1/p . Recently , a robust variant of the Wasserstein distance , called the subspace robust Wasserstein distance ( SRW ) , was proposed ( Paty & Cuturi , 2019 ) . The SRW computes the OT problem in the discriminative subspace . This can be determined by solving dimensionality-reduction problems . Owing to the robustness , it can compute the Wasserstein from noisy data . The SRW is given as SRW ( µ , ν ) = min Π∈U ( µ , ν ) max U∈Rd×k , U > U=Ik n∑ i=1 m∑ j=1 πij‖U > xi −U > yj‖22 1 2 , ( 2 ) where U is the projection matrix with k ≤ d , and Ik ∈ Rk×k is the identity matrix . The SRW or its relaxed problem can be efficiently estimated using either eigenvalue decomposition or the Frank–Wolfe algorithm . 3 PROPOSED METHOD . This paper proposes FROT . We assume that the vectors are grouped as x = ( x ( 1 ) > , . . . , x ( L ) > ) > and y = ( y ( 1 ) > , . . . , y ( L ) > ) > . Here , x ( ` ) ∈ Rd ` and y ( ` ) ∈ Rd ` are the d ` dimensional vectors , where ∑L ` =1 d ` = d. This setting is useful if we know the explicit group structure for the feature vectors a priori . In an application in L-layer neural networks , we consider x ( ` ) and y ( ` ) as outputs of the ` th layer of the network . If we do not have a priori information , we can consider each feature independently ( i.e. , d1 = d2 = . . . = dL = 1 and L = d ) . All proofs in this section are provided in the Appendix . 3.1 FEATURE-ROBUST OPTIMAL TRANSPORT ( FROT ) . The FROT formulation is given by min Π∈U ( µ , ν ) max α∈ΣL n∑ i=1 m∑ j=1 πij L∑ ` =1 α ` c ( x ( ` ) i , y ( ` ) j ) , ( 3 ) where ΣL = { α ∈ RL+ : α > 1L = 1 } is the probability simplex . The underlying concept of FROT is to estimate the transport plan Π using distinct groups with large distances between { x ( ` ) i } ni=1 and { y ( ` ) j } mj=1 . We note that determining the transport plan in nondistinct groups is difficult because the data samples in { x ( ` ) i } ni=1 and { y ( ` ) j } mj=1 overlap . By contrast , in distinct groups , { x ( ` ) i } ni=1 and { y ( ` ) j } mj=1 are different , and this aids in determining an optimal transport plan . This is an intrinsically similar idea to the subspace robust Wasserstein distance ( Paty & Cuturi , 2019 ) , which estimates the transport plan in the discriminative subspace , while our approach selects important groups . Therefore , FROT can be regarded as a feature selection variant of the vanilla OT problem in Eq . ( 1 ) , whereas the subspace robust version uses dimensionality-reduction counterparts . Algorithm 1 Sinkhorn algorithm . 1 : Input : a , b , C , , tmax 2 : Initialize K = e−C/ , u = 1n , v = 1m , t = 0 3 : while t ≤ tmax and not converge do 4 : u = a/ ( Kv ) 5 : v = b/ ( K > u ) 6 : t = t+ 1 7 : end while 8 : return Π = diag ( u ) Kdiag ( v ) Algorithm 2 FROT with the Frank–Wolfe . 1 : Input : { xi } ni=1 , { yj } mj=1 , η , and . 2 : Initialize Π , compute { C ` } L ` =1 . 3 : for t = 0 . . . T do 4 : Π̂ = argminΠ∈U ( µ , ν ) 〈Π , MΠ ( t ) 〉 + H ( Π ) 5 : Π ( t+1 ) = ( 1− γ ) Π ( t ) + γΠ̂ 6 : with γ = 22+t . 7 : end for 8 : return Π ( T ) Using FROT , we can define a p-feature robust Wasserstein distance ( p-FRWD ) . Proposition 1 For the distance function d ( x , y ) , FRWDp ( µ , ν ) = min Π∈U ( µ , ν ) max α∈ΣL n∑ i=1 m∑ j=1 πij L∑ ` =1 α ` d ( x ( ` ) i , y ( ` ) j ) p 1/p , ( 4 ) is a distance for p ≥ 1 . Note that we can show that 2-FRWD is a special case of SRW with d ( x , y ) = ‖x − y‖2 ( See Appendix ) . The key difference between SRW and FRWD is that FRWD can use any distance , while SRW can only use d ( x , y ) = ‖x− y‖2 .

Summary: This work proposes variants of robust OT/p-wasserstein-dist (3)/(4), where the ground cost is in some sense the maximum over costs with (prefixed) groups of features. The motivation is similar to that for feature selection: where perhaps only few of these groups of features are critical/sufficient for OT purposes. So it can also be understood as joint feature-group selection with OT. The resulting convex problem is proposed to be solved using FW, whose details are presented (including convergence).