F EATURE -R OBUST O PTIMAL T RANSPORT FOR H IGH -D IMENSIONAL D ATA
Anonymous authors
Paper under double-blind review
A BSTRACT
Optimal transport is a machine learning problem with applications including distribution comparison, feature selection, and generative adversarial networks.
In this paper, we propose feature-robust optimal transport (FROT) for highdimensional data, which solves high-dimensional OT problems using feature selection to avoid the curse of dimensionality.
Speciﬁcally, we ﬁnd a transport plan with discriminative features.
To this end, we formulate the FROT problem as a min–max optimization problem.
We then propose a convex formulation of the
FROT problem and solve it using a Frank–Wolfe-based optimization algorithm, whereby the subproblem can be efﬁciently solved using the Sinkhorn algorithm.
Since FROT ﬁnds the transport plan from selected features, it is robust to noise features.
To show the effectiveness of FROT, we propose using the FROT algorithm for the layer selection problem in deep neural networks for semantic correspondence.
By conducting synthetic and benchmark experiments, we demonstrate that the proposed method can ﬁnd a strong correspondence by determining important layers.
We show that the FROT algorithm achieves state-of-the-art performance in real-world semantic correspondence datasets.
I NTRODUCTION
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 (Peyr ´ e 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 ∈ R d × n and Y ∈ R d × 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 efﬁcient formulation.
The computation of the subspace can be expensive if the dimensionality of data is high, for example, 10 4 .
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 (cid:96) 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.
Speciﬁcally, each
[Figure]
(cid:101)
(cid:101) y where true features and noisy are grouped, 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.
Speciﬁcally, 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 efﬁciently 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 signiﬁcance 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 efﬁcient 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.
B ACKGROUND
In this section, we brieﬂy introduce the OT problem.
Optimal transport (OT): The following are given: independent and identically distributed (i.i.d.)
samples X = { x i } ni =1 ∈ R d × n from a d -dimensional distribution p , and i.i.d.
samples Y = { y j } mj =1 ∈ R d
× m from the d -dimensional distribution q .
In the Kantorovich relaxation of OT, admissible couplings are deﬁned by the set of the transport plan:
[Equation]
where Π ∈ R n × m + is called the transport plan, 1 n is the n -dimensional vector whose elements are ones, and a = ( a 1 , a 2 , . .
. , a n ) (cid:62) ∈ R n + and b = ( b 1 , b 2 , . . .
, b m ) (cid:62) ∈ R m + are the weights.
The
OT problem between two discrete measures µ = (cid:80) ni
=1 a i δ x i and ν = (cid:80) mj
=1 b j δ y j determines the optimal transport plan of the following problem:
[Equation]
where c ( x , y ) is a cost function.
For example, the squared Euclidean distance is used, that is, c ( x , y )
= (cid:107) x − y (cid:107) 22 .
To solve the OT problem, Eq.
(1) (also known as the earth mover’s distance) using linear programming requires O ( n 3 ) , ( n = m ) computation, which is computationally expensive.
To address this, an entropic-regularized optimal transport is used (Cuturi, 2013).
[Equation]
where (cid:15) ≥ 0 is the regularization parameter, and H ( Π ) =
[Equation]
mj
[Equation]
log [Equation] π
[Equation]
) − 1) is the entropic regularization.
If (cid:15) = 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 deﬁned as c ( x , y )
=
d ( x , y ) with d ( x , y ) as a distance function and p ≥ 1 , then we deﬁne the p -Wasserstein distance of two discrete measures µ = (cid:80)
[Equation]
i
[Equation] i and ν = [Equation] 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
[Equation]
where U is the projection matrix with k ≤ d , and I
k ∈ R
k × k is the identity matrix.
The SRW or its relaxed problem can be efﬁciently estimated using either eigenvalue decomposition or the
Frank–Wolfe algorithm.
P ROPOSED M ETHOD
This paper proposes FROT.
We assume that the vectors are grouped as x = ( x (1) (cid:62) , . . .
, x ( L ) (cid:62) )
(cid:62) and y =
( y (1) (cid:62) , . . .
, y ( L ) (cid:62) ) (cid:62) .
Here, x ( (cid:96) ) ∈ R d (cid:96) and y ( (cid:96) )
∈ R d (cid:96) are the d (cid:96) dimensional vectors, where
[Equation]
L(cid:96)
=1 d (cid:96) =
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 ( (cid:96) ) and y ( (cid:96) ) as outputs of the (cid:96) th layer of the network.
If we do not have a priori information, we can consider each feature independently (i.e., d 1 = d 2 = . . .
=
d L = 1 and L = d ).
All proofs in this section are provided in the Appendix.
3. F EATURE -R OBUST O PTIMAL T RANSPORT (FROT)
The FROT formulation is given by
[Equation]
where Σ L = { α ∈ R L + : α (cid:62) 1 L = 1 } is the probability simplex.
The underlying concept of FROT is to estimate the transport plan Π using distinct groups with large distances between { x ( (cid:96) )
i } ni =1 and { y ( (cid:96) )
j } mj =1 .
We note that determining the transport plan in nondistinct groups is difﬁcult because the data samples in { x ( (cid:96) ) i } ni =1 and { y ( (cid:96) )
j } mj =1 overlap.
By contrast, in distinct groups, { x ( (cid:96) )
i } ni =
and { y ( (cid:96) )
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.
Π ∈ U ( µ,ν ) (cid:104) Π , M Π ( t ) (cid:105)
[Equation]
(cid:15)H ( Π ) 5: Π ( t +1) = (1 − γ ) Π ( t ) + γ (cid:98) Π 6: with γ = 22+ t . 7: end for 8: return Π ( T )
Using FROT, we can deﬁne a p -feature robust Wasserstein distance ( p -FRWD).
Proposition 1 For the distance function d ( x , y ) ,
[Equation]
is a distance for p ≥ 1 .
Note that we can show that 2-FRWD is a special case of SRW with d ( x , y )
= (cid:107) x − y (cid:107) 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 )
= (cid:107) x − y (cid:107) 2 .
3. FROT O PTIMIZATION
Here, we propose two FROT algorithms based on the Frank–Wolfe algorithm and linear programming.
Frank–Wolfe: We propose a continuous variant of the FROT algorithm using the Frank–Wolfe algorithm, which can be fully differentiable.
To this end, we introduce entropic regularization for α and rewrite the FROT as a function of Π .
Therefore, we solve the following problem for α :
[Equation] [Equation]
where η ≥ 0 is the regularization parameter, and H ( α ) =
(cid:80)
L(cid:96)
=
1 α (cid:96) (log( α (cid:96) ) − 1) is the entropic regularization for α .
An advantage of entropic regularization is that the nonnegative constraint is naturally satisﬁed, and the entropic regularizer is a strong convex function.
Lemma
The optimal solution of the optimization problem
[Equation]
J η ( Π , α ) , with J η ( Π , α ) =
[Equation]
with a ﬁxed admissible transport plan Π ∈ U ( µ, ν ) , is given by
[Equation]
Using Lemma 2 (or Lemma 4 in Nesterov (2005)) together with the setting φ (cid:96)
[Equation]
[Equation]
the global problem is equivalent to
[Equation]
Note that this is known as a smoothed max-operator (Nesterov, 2005; Blondel et al., 2018).
Speciﬁcally, regularization parameter η controls the “smoothness” of the maximum.
Proposition 3 G η ( Π ) is a convex function relative to Π .
The derived optimization problem of FROT is convex.
Therefore, we can determine globally optimal solutions.
Note that the SRW optimization problem is not jointly convex (Paty & Cuturi, 2019) for the projection matrix and the transport plan.
In this study, we employ the Frank–Wolfe algorithm (Frank & Wolfe, 1956;
Jaggi, 2013), using which we approximate G η ( Π ) with linear functions at Π ( t ) and move Π toward the optimal solution in the convex set (See Algorithm 2).
The derivative of the loss function G η ( Π ) at Π ( t ) is given by
[Equation]
Then, we update the transport plan by solving the EMD problem:
[Equation]
where γ = 2 / (2 + k ) .
Note that M Π ( t ) is given by the weighted sum of the cost matrices.
Thus, we can utilize multiple features to estimate the transport plan Π for the relaxed problem in Eq.
Using the Frank–Wolfe algorithm, we can obtain the optimal solution.
However, solving the EMD problem requires a cubic computational cost that can be expensive if n and m are large.
To address this, we can solve the regularized OT problem, which requires O ( nm ) .
We denote the Frank–Wolfe algorithm with EMD as FW-EMD and the Frank–Wolfe algorithm with Sinkhorn as FW-Sinkhorn.
Computational complexity: The proposed method depends on the Sinkhorn algorithm, which requires an O ( nm ) operation.
The computation of the cost matrix in each subproblem needs an O ( Lnm ) operation, where L is the number of groups.
Therefore, the entire complexity is
O ( TLnm ) , where T is the number of Frank–Wolfe iterations (in general, T = 10 is sufﬁcient).
Proposition 4 For each t ≥ 1 , the iteration Π ( t ) of Algorithm 2 satisﬁes
[Equation]
where σ max ( Φ (cid:62) Φ ) is the largest eigenvalue of the matrix Φ (cid:62) Φ and Φ = ( vec ( C 1 ) , vec ( C 2 ) , .
, vec ( C L )) (cid:62) ; and δ ≥ 0 is the accuracy to which internal linear subproblems are solved.
Based on Proposition 4, the number of iterations depends on η , (cid:15) , and the number of groups.
If we set a small η , convergence requires more time.
In addition, if we use entropic regularization with a large (cid:15) , the δ in Proposition 4 can be large.
Finally, if we use more groups, the largest eigenvalue of the matrix Φ (cid:62) Φ can be larger.
Note that the constant term of the upper bound is large; however, the
Frank–Wolfe algorithm converges quickly in practice.
Linear Programming: Because lim
[Equation]
the FROT problem can also be written as
[Equation]
Because the objective is the max of linear functions, it is convex with respect to Π .
We can solve the problem via linear programming:
[Equation]
This optimization can be easily solved using an off-the-shelf LP package.
However, the computational cost of this LP problem is high in general (i.e., O ( n 3 ) , n = m ).
3. A PPLICATION : S EMANTIC C ORRESPONDENCE
We applied our proposed FROT algorithm to semantic correspondence.
The semantic correspondence is a problem that determines the matching of objects in two images.
That is, given input image pairs ( A, B ) , with common objects, we formulated the semantic correspondence problem to estimate the transport plan from the key points in A to those in B ; this framework was proposed in (Liu et al., 2020).
In Figure 2, we show an overview of our proposed framework.
Cost matrix computation C (cid:96) :
In our framework, we employed a pretrained convolutional neural network to extract dense feature maps for each convolutional layer.
The dense feature map of the (cid:96) th layer output of the s th image is given by
[Equation]
where w s and h s are the width and height of the s th image, respectively, and d (cid:96) is the dimension of the (cid:96) th layer’s feature map.
Note that because the dimension of the dense feature map is different for each layer, we sample feature maps to the size of the 1 st layer’s feature map size (i.e., h s × w s ).
The (cid:96) th layer’s cost matrix for images s and s (cid:48) is given by
[Equation]
[Figure] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] /B3mu7SefmbDXQZjmR5/LpS0nT1aWNZH1KnxiRteKOjuuwuuqG2lXeMsNS9Hlv1tiXOqzD8QJ3EpWtYQwVZ7yErYq/f4e//+I/VxYVidNhLLviZRw1fO9hPSayNUHp7fvLD6oiq+K+cXBwweLv3jwcOeXdz5aZvLrDfYjzltgi+x99hHb5K16n+fk37d+cKty691XW69evhq9mkjRr9yKdN5hqa9Xv/0/d1uOqg==</latexit>
A potential problem with FROT is that the estimation depends signiﬁcantly on the magnitude of the cost of each layer (also known as a group).
Hence, normalizing each cost matrix is important.
Therefore, we normalized each feature vector by f ( (cid:96),s ) i ← f ( (cid:96),s ) i / (cid:107) f
( (cid:96),s ) i (cid:107) 2 .
Consequently, the cost matrix is given by [ C (cid:96) ] ij = [Equation] 2 f ( (cid:96),s ) i (cid:62) f
( (cid:96),s (cid:48) )
j .
We can use distances such as the L 1 distance.
Computation of a and b with staircase re-weighting: For semantic correspondence, setting a ∈ R h s w s and b ∈ R
h s (cid:48)
w s (cid:48) is important because semantic correspondence can be affected by background clutter.
Therefore, we generated the class activation maps (Zhou et al., 2016) for the source and target images and used them as a and b , respectively.
For CAM, we chose the class with the highest classiﬁcation probability and normalized it to the range [0 , 1] .
R ELATED W ORK
OT algorithms: The Wasserstein distance can be determined by solving the OT problem.
An advantage of the Wasserstein distance is its robustness to noise; moreover, we can obtain the transport [Figure] plan, which is useful for many machine learning applications.
To reduce the computation cost for the Wasserstein distance, the sliced Wasserstein distance is useful (Kolouri et al., 2016).
Recently, a tree variant of the Wasserstein distance was proposed (Evans & Matsen, 2012; Le et al., 2019;
Sato et al., 2020); the sliced Wasserstein distance is a special case of this algorithm.
The approach most closely related to FROT is a robust variant of the Wasserstein distance, called the subspace robust Wasserstein distance (SRW) (Paty & Cuturi, 2019).
SRW computes the OT problem in a discriminative subspace; this is possible by solving dimensionality-reduction problems.
Owing to the robustness, SRW can successfully compute the Wasserstein distance from noisy data.
The max–sliced Wasserstein distance (Deshpande et al., 2019) and its generalized counterpart (Kolouri et al., 2019) can also be regarded as subspace-robust Wasserstein methods.
Note that SRW (Paty &
Cuturi, 2019) is a min–max based approach, while the max–sliced Wasserstein distances (Deshpande et al., 2019; Kolouri et al., 2019) are max–min approaches.
The FROT is a feature selection variant of the Wasserstein distance, whereas the subspace approaches are used for dimensionality reduction.
As a parallel work, a general minimax optimal transport problem called the robust Kantorovich problem (RKP) was recently proposed (Dhouib et al., 2020).
RKP involves using a cutting-set method for a general minmax optimal transport problem that includes the FROT problem as a special case.
The approaches are technically similar; however, our problem and that of Dhouib et al.
(2020) are intrinsically different.
Speciﬁcally, we aim to solve a high-dimensional OT problem using feature selection and apply it to semantic correspondence problems, while the RKP approach focuses on providing a general framework and uses it for color transformation problems.
As a technical difference, the cutting-set method may not converge to an optimal solution if we use the regularized
OT (Dhouib et al., 2020).
By contrast, because we use a Frank–Wolfe algorithm, our algorithm converges to a true objective function with regularized OT solvers.
The multiobjective optimal transport (MOT) is an approach (Scetbon et al., 2020) parallel to ours.
The key difference between FROT and
MOT is that MOT tries to use the weighted sum of cost functions, while FROT considers the worst case.
Moreover, as applications, we focus on the cost matrices computed from subsets of features, while MOT considers cost matrices with different distance functions.
E XPERIMENTS 5. S YNTHETIC D ATA
We compare FROT with a standard OT using synthetic datasets.
In these experiments, we initially generate two-dimensional vectors x ∼ N ( µ x , Σ x ) and y ∼ N ( µ y , Σ y ) .
Here, we set µ x = ( − 5 , 0) (cid:62) , µ y =
(5 , 0) (cid:62) , Σ x = Σ y
=
((5 , 1) (cid:62) , (4 , 1) (cid:62) ) .
Then, we concatenate z x ∼ N ( 0 8 , I 8 ) and z y ∼ N ( 0 8 , I 8 ) to x and y , respectively, to give
[Equation]
For FROT, we set η = 1 . 0 and the number of iterations of the Frank–Wolfe algorithm as T = 10 .
The regularization parameter is set to (cid:15) = 0 .
02 for all methods.
To show the proof-of-concept, we set the true features as a group and the remaining noise features as another group.
Fig.
1a shows the correspondence from x and y with the vanilla
OT algorithm.
Figs.
1b and 1c show the correspondence of FROT and OT with (cid:101) x and (cid:101) y , respectively.
Although FROT can identify [Table] a suitable matching , the OT fails to obtain a signiﬁcant correspondence.
We observed that the α parameter corresponding to a true group is α 1 = 0 . 9999 .
Moreover, we compared the objective scores of the FROT with LP, FW-EMD, and FW-Sinkhorn ( (cid:15) = 0 . 1 ).
Figure 3a shows the objective scores of FROTs with the different solvers, and both FW-EMD and FW-Sinkhorn can achieve almost the same objective score with a relatively small η .
Moreover, Figure 3b shows the mean squared error between the LP method and the FW counterparts.
Similar to the objective score cases, it can yield a similar transport plan with a relatively small η .
Finally, we evaluated the FW-Sinkhorn by changing the regularization parameter η .
In this experiment, we set η = 1 and varied the (cid:15) values.
The result shows that we can obtain an accurate transport plan with a relatively small (cid:15) .
5. S EMANTIC CORRESPONDENCE
We evaluated our FROT algorithm for semantic correspondence.
In this study, we used the SPair71k (Min et al., 2019b).
The SPair-71k dataset consists of 70 , 958 image pairs with variations in viewpoint and scale.
For evaluation, we employed a percentage of accurate key points (PCK), which counts the number of accurately predicted key points given a ﬁxed threshold (Min et al., 2019b).
All semantic correspondence experiments were run on a Linux server with NVIDIA P100.
For the optimal transport based frameworks, we employed ResNet
(He et al., 2016) pretrained on ImageNet (Deng et al., 2009) for feature and activation map extraction.
The ResNet101 consists of 34 convolutional layers and the entire number of features is d = 32 , 576 .
Note that we did not ﬁne-tune the network.
We compared the proposed method with several baselines (Min et al., 2019b) and the SRW 1 .
Owing to the computational cost and the required memory size for SRW, we used the ﬁrst and the last few convolutional layers of ResNet101 as the input of SRW.
In our experiments, we empirically set T = 3 and (cid:15)
= 0 .
1 for FROT and SRW, respectively.
For SRW, we set the number of latent dimension as k = 50 for all experiments.
HPF (Min et al., 2019a) and OT-HPF (Liu et al., 2020) are state-of-the-art methods for semantic correspondence.
HPF and OT-HPF required the validation dataset to select important layers, whereas SRW and FROT did not require the validation dataset.
OT is a simple optimal transport-based method that does not select layers.
Table 1 lists the per-class PCK results obtained using the SPair-71k dataset.
FROT ( η = 0 . 3) outperforms most existing baselines, including HPF and OT.
Moreover, FROT ( η = 0 . 3) is consistent with OT-HPF (Liu et al., 2020), which requires the validation dataset to select important layers.
In this experiment, setting η < 1 results in favorable performance (See Table 3 in the Appendix).
The computational costs of FROT is 0.29, while SRWs are 8.73, 11.73, 15.76, respectively.
Surprisingly,
FROT outperformed SRWs.
However, this is mainly due to the used input layers.
Therefore, scaling up SRW would be an interesting future work.
We further evaluated FROT by tuning hyperparameters η and (cid:15) using validation sets, where the maximum search ranges for η and (cid:15) are set to 0.2 to 2.0 and 0.1 to 0.6 with intervals of 0.1, respectively.
Figure 6 in Appendix shows the average PCK scores for ( η , (cid:15) ) pairs on the validation split of SPair71k.
By using hyperparameter search, we selected ( η = 0 .
2 , (cid:15) = 0 . 4) as an optimal parameter.
The
FROT with optimal parameters outperforms the state-of-the-art method (Liu et al., 2020).
[Figure] 5. F EATURE S ELECTION E XPERIMENTS Since FROT ﬁnds the transport plan and discriminative features between X and Y , we can use FROT as a feature-selection method.
We considered X ∈ R d × n and Y ∈
R
d × m as sets of samples from classes 1 and 2 , respectively.
The optimal important feature is given by
[Equation]
, with
[Equation]
where [ C (cid:96) ] ij = ( x ( (cid:96) )
i − y
( (cid:96) )
j ) 2 .
Finally, we selected the top K features by the ranking (cid:98) α .
Hence, α changes to a one-hot vector for a small η and to α k ≈ 1 L for a large η .
Here, we compared FROT with several baseline algorithms in terms of solving feature-selection problems.
In this study, we employed a high-dimensional and a few sample datasets with two class classiﬁcation tasks (see Table 2).
All feature selection experiments were run on a Linux server with an Intel Xeon CPU E7-8890 v4 with 2.20 GHz and 2 TB RAM.
In our experiments, we initially randomly split the data into two sets ( 75% for training and 25% for testing) and used the training set for feature selection and building a classiﬁer.
Note that we standardized each feature using the training set.
Then, we used the remaining set for the test.
The trial was repeated 50 times, and we considered the averaged classiﬁcation accuracy for all trials.
Considered as baseline methods, we computed the Wasserstein distance, maximum mean discrepancy (MMD) (Gretton et al., 2007), and linear correlation 2 for each dimension and sorted them in descending order.
Note that the Wasserstein distance is computed via sorting, which is computationally more efﬁcient than the Sinkhorn algorithm when d = 1 .
Then, we selected the top K features as important features.
For FROT, we computed the feature importance and selected the features that had signiﬁcant importance scores.
In our experiments, we set η = 1 . 0 and T = 10 .
Then, we trained a two-class SVM 3 with the selected features.
Fig.
4 shows the average classiﬁcation accuracy relative to the number of selected features.
From
Figure 4, FROT is consistent with the Wasserstein distance-based feature selection and outperforms the linear correlation method and the MMD for two datasets.
Table 2 shows the computational time(s) of the methods.
FROT is about two orders of magnitude faster than the Wasserstein distance and is also faster than MMD.
Note that although MMD is as fast as the proposed method, it cannot determine the correspondence between samples.
C ONCLUSION
In this paper, we proposed FROT for high-dimensional data.
This approach jointly solves feature selection and OT problems.
An advantage of FROT is that it is a convex optimization problem and can determine an accurate globally optimal solution using the Frank–Wolfe algorithm.
We used FROT for high-dimensional feature selection and semantic correspondence problems.
Through extensive experiments, we demonstrated that the proposed algorithm is consistent with state-of-theart algorithms in both feature selection and semantic correspondence.
A PPENDIX R ELATED WORK
In addition to accelerating the computation, structured optimal transport incorporates structural information directly into OT problems (Alvarez-Melis et al., 2018).
Speciﬁcally, they formulate the submodular optimal transport problem and solve the problem using a saddle-point mirror prox algorithm.
Recently, more complex structured information was introduced in the OT problem, including the hierarchical structure (Alvarez-Melis et al., 2020; Yurochkin et al., 2019).
These approaches successfully incorporate structured information into OT problems with respect to data samples.
By contrast, FROT incorporates the structured information into features.
OT applications: OT has received signiﬁcant attention for use in several computer vision tasks.
Applications include Wasserstein distance estimation (Peyr´e 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), photo album summarization (Liu et al., 2019), generative model (Arjovsky et al., 2017; Bunne et al., 2019), graph matching (Xu et al., 2019a;b), and the semantic correspondence (Liu et al., 2020).
P ROOF OF P ROPOSITION
For the distance function d ( x , y ) , we prove that
[Equation]

[Equation]
is a distance for p ≥ 1 .
The symmetry can be read directly from the deﬁnition as we used distances that are symmetric.
For the identity of indiscernibles, when FRWD p ( µ, ν ) = 0 with the optimal α and Π , there exists (cid:96) such that α (cid:96) > 0 (as α is in the simplex set).
As there is a max in the deﬁnition and (cid:80)
[Equation]
this
[Equation]
that
[Equation]
Therefore, we have µ = ν when FRWD p ( µ, ν ) = 0 .
When µ = ν , this means that x i
=
y i , a i =
b i , ∀ i , and n = m , and we have d ( x i , y j )
=
for i = j .
Thus, for any α (cid:96) ≥ 0 , the optimal transport plan is π ii > 0 for d ( x i , y j )
= 0 and π ij =
for d ( x i , y j ) > 0 .
Therefore, when µ = ν , we have FRWD p ( µ, ν ) = 0 .
T RIANGLE I NEQUALITY
Let µ = (cid:80) ni
=1 a i δ x i , ν = (cid:80) mj
[Equation]
[Equation]
k and α ∈ Σ L , we prove that
[Equation]
To simplify the notations in this proof, we deﬁne D (cid:96) as the distance ”matrix” such that [ D (cid:96) ] ij = d ( x ( (cid:96) )
i , y ( (cid:96) )
j ) is the i th-row and j th-column element of the matrix D (cid:96) , [ D (cid:96) ]
jk =
d ( y ( (cid:96) )
j , z ( (cid:96) ) k ) , and [ D (cid:96) ]
ik = d ( x ( (cid:96) )
i , z ( (cid:96) ) k ) .
Moreover, note that D p(cid:96) is the ”matrix,” where each element is an element of D (cid:96) raised to the power p .
Consider that P ∈ U ( µ, ν ) is the optimal transport plan of FRWD p ( µ, ν ) , and
Q ∈ U ( ν, γ ) is the optimal transport plan of FRWD p ( ν, γ ) , where γ = (cid:80) rk =1 c i δ
z i is a discrete measure.
Similar to the proof for the Wasserstein distance in (Peyr´e et al., 2019), let S =
P diag (1 / (cid:101) b ) Q with (cid:101) b be a vector such that
[Equation]
b j = b j if b j > 0 , and b j = 1 otherwise.
We can show that S ∈ U ( µ, γ ) .
[Table]
By letting g ijk(cid:96)
=
[ D (cid:96) ] ij ( α (cid:96) p ij q jk /
[Table]
b j ) 1 /p and h ijk(cid:96) =
[ D (cid:96) ] ij ( α (cid:96) p ij q jk /
[Table]
j
[Table]
p
[Table]
the right-hand side of this inequality can be rewritten as  
[Equation]
by the Minkovski inequality
[Equation]
This inequality is valid for all α .
Therefore [Equation] we have
[Equation]
FROT WITH L INEAR P ROGRAMMING
Linear Programming:
The FROT is a convex piecewise-linear minimization because the objective is the max of linear functions.
Thus, we can solve the FROT problem via linear programming:
[Equation]
This optimization can be easily solved using an off-the-shelf LP package.
However, the computational cost of this LP problem is high in general (i.e., O ( n 3 ) , n = m ).
The FROT problem can be written as
[Equation]
This problem can be transformed to an equivalent linear program by ﬁrst forming an epigraph problem:
[Equation]
Thus, the linear programming for FROT is given as
[Equation]
Next, we transform this linear programming problem into the canonical form.
For matrix Π =
( π 1 π 2 . . .
π n ) (cid:62) ∈ R n × m and π i ∈ R n , we can vectorize the matrix using the following linewise operator:
[Equation]
Using this vectorization operator [Equation] we can write (cid:104) Π , C (cid:96) (cid:105) ≤ t as
[Equation]
where 0 L ∈ R L is a vector whose elements are zero.
For the constraints Π1 m = a and Π (cid:62) 1 n = b , we can deﬁne vectors q 1 , . . .
, q n ∈ R nm and r 1 , . . .
, r m ∈ R nm such that q (cid:62) i vec ( Π ) =
a i and r (cid:62) j vec ( Π ) =
b j in this way:
[Equation]
and
[Equation]
We can collect these vectors to obtain the vectorized constraints: [Equation] [Equation] 

[Equation]
Thus, we can rewrite the linear programming as
[Equation]
where u = ( vec ( Π ) (cid:62) t ) (cid:62) ∈
R
nm +1 , e = ( 0 (cid:62)
nm 1)
(cid:62) ∈ R
nm +1 is the unit vector whose nm + 1th element is 1,
A = ( vec ( C 1 ) , . . .
, vec ( C L ))
∈
R nm × L , Q = ( q 1 , . . . , q n ) ∈ R nm × n , and
R = (
[Equation] [Equation]
P ROOF OF L EMMA
We optimize the function with respect to α :
[Equation]
where
[Equation]
Because the entropic regularization is a strong convex function and its negative counterpart is a strong concave function, the maximization problem is a concave optimization problem.
We consider the following objective function with the Lagrange multiplier (cid:15) :
[Equation]
Note that owing to the entropic regularization, the nonnegative constraint is automatically satisﬁed.
Taking the derivative with respect to α (cid:96) , we have
[Equation]
Thus, the optimal α (cid:96) has the form
[Equation]
α
[Equation]
satisﬁes the sum to one constraint.
[Equation]
Hence, the optimal α (cid:96) is given by
[Equation]
Substituting this into Eq.(8), we have
J ( α ∗ ) =
L
[Equation]
exp (cid:16)
[Equation]
φ (cid:96) (cid:17) (cid:80)
L(cid:96)
[Equation] 1 exp (cid:16) [Equation] (cid:96) (cid:48) (cid:17) φ
[Equation]
η
L (cid:88)
[Equation]
[Equation]
(cid:16) 1 η φ (cid:96) (cid:17) (cid:80)
L(cid:96) (cid:48)
=1 exp (cid:16) 1 η φ (cid:96) (cid:48) (cid:17)   log   exp (cid:16) 1 η φ (cid:96) (cid:17) (cid:80)
L(cid:96) (cid:48)
=1 exp (cid:16) 1 η φ (cid:96) (cid:48) (cid:17) 
 − 1 
 =
η log (cid:32)
L (cid:88) (cid:96) = exp (cid:18) η φ (cid:96) (cid:19)(cid:33) + η
Therefore, the ﬁnal objective function is given by
[Equation]
P ROOF OF P ROPOSITION
Proof: For 0 ≤ θ ≤ 1 and η > 0 , we have
L
[Equation]
Here, we use H¨older’s inequality with p = 1 /θ , q = 1 / (1 − θ ) , and 1 /p + 1 /q
= 1 .
Applying a logarithm on both sides of the equation and then premultiplying η , we have
[Equation]
P ROOF OF P ROPOSITION
Theorem 5 (Jaggi, 2013)
For each t ≥ 1 , the iterates Π ( t ) of Algorithms 1, 2, 3, and 4 in (Jaggi, 2013) satisfy
[Equation]
where Π ∗ ∈
D is an optimal solution to problem
[Equation]
C f is the curvature constant deﬁned as
[Equation]
δ ≥ 0 is the accuracy with which internal linear subproblems are solved.
Lemma 6 (Jaggi, 2013) Let f be a convex and differentiable function with its gradient ∇ f being
Lipschitz-continuous w.r.t.
some norm (cid:107) · (cid:107) over the domain D with Lipschitz-constant L > 0 .
Then,
[Equation]
Deﬁnition
The softmax function is given by
[Equation]
where λ > 0 is referred to as the inverse temperature constant.
Lemma 8 (Gao & Pavel, 2017)
The softmax function σ ( · ) is L -Lipschitz with respect to (cid:107) · (cid:107) 2 with
L = λ , that is for all z , z (cid:48) [Equation] R n ,
[Equation]
where λ is the inverse temperature constant.
The derivative of G η ( Π ) is given as
[Equation]
Thus, we have
[Equation]
, where
[Equation]
Π , C
[Equation]
exp
[Equation] (cid:17)   (cid:62) ∈ R L . Here, p Π is the softmax function σ ( z ) with z (cid:96) =
(cid:104) Π , C (cid:96) (cid:105) .
We have
[Equation]
Lemma 8 with λ = 1 η )
[Equation]
where (cid:107) · (cid:107) op is the operator norm.
We have (cid:107) Φ (cid:107) op
= (cid:107) Φ (cid:62) (cid:107) op , (cid:107) Φ (cid:62)
Φ (cid:107) op
= (cid:107) Φ (cid:107) 2 op , and (cid:107) vec ( Π ) − vec ( Π (cid:48) )
(cid:107) 2 ≤ √ 2 .
Therefore, the Lipschitz constant for the gradient is L = 1 η (cid:107) Φ (cid:107) 2 op =
η σ max ( Φ (cid:62) Φ ) , and the curvature constant is bounded above by C f ≤ 2 L , where σ max ( Φ (cid:62) Φ ) is the largest eigenvalue of the matrix Φ (cid:62) Φ .
By plugging C f in Theorem 5, we have
[Equation]
M AX /M IN FORMULATION
We deﬁne the max–min formulation of the FROT as
[Equation]
where Σ L = { α ∈ R L + : α (cid:62) 1 L = 1 } is the probability simplex, the set of probability vectors in R L .
This problem can be solved by computing the group that maximizes the optimal transport distance k ∗ =
argmax k W 1 ( µ ( (cid:96) )
, ν ( (cid:96) ) ) and then by considering α ∗ =
δ k ∗ as a one-hot vector.
The result of this formulation provides an intuitive idea (the same as for the robust Wasserstein method).
Hence, we maximize the group (instead of the subspace) that provides the optimal result.
However, the formulation requires solving the OT problem L times.
This approach may not be suitable if we have a large L .
Moreover, the argmax function is generally not differentiable.
Relation to the max-sliced Wasserstein distance: The max-sliced Wasserstein-2 distance can be deﬁned as (Deshpande et al., 2019)
[Equation]
where Ω ⊂ R d is the set of all possible directions on the unit sphere.
The max–sliced Wasserstein is a max–min approach.
That is, for each w , it requires solving the
OT problem.
The max–min approach is suited for simply measuring the divergence between two distributions.
However, it is difﬁcult to interpret features using the max–sliced Wasserstein, where it is the key motivation of FROT.
Relation to Subspace Robust Wasserstein (Paty & Cuturi, 2019):
Here, we show that 2-FRWD with d ( x , y )
= (cid:107) x − y (cid:107) 2 is a special case of SRW.
Let us deﬁne U = ( √ α 1 e 1 , √ α 2 e 2 , . .
. , √ α d e d ) (cid:62) ∈ R d × d , where e (cid:96) ∈
R
d is the one-hot vector whose (cid:96) th element is 1 and α (cid:62) 1 = 1 , α (cid:96) ≥ 0 .
Then, the objective function of SRW can be written as
[Equation]
Therefore, SRW and 2-FRWD are equivalent if we set U =
( √ α 1 e 1 , √ α 2 e 2 , . .
. , √ α d e d ) (cid:62) and d ( x , y )
=
(cid:107) x − y (cid:107) 2 .
[Table]
[Figure]
[Table]
[Figure]
A DDITIONAL S EMANTIC C ORRESPONDENCE E XPERIMENTS
Figure 5a shows an example of key points matched using the FROT algorithm.
Fig.5b shows the corresponding feature importance.
The lower the η value, the smaller the number of layers used.
The interesting ﬁnding here is that the selected important layer in this case is the third layer from the last.
More qualitative results are presented in the Figure
[Table]
[Figure]
[Table]
[Figure]