SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Joint Feature and Differentiable  $k$ -NN Graph Learning using Dirichlet Energy
School of Computer Science & School of Artificial Intelligence, OPtics and ElectroNics (iOPEN) Northwestern Polytechnical University Xi'an 710072, P.R. China solerxl1998@gmail.com
Lei Chen
School of Computer Science  
Nanjing University of Posts and Telecommunications  
Nanjing 210003, P.R. China  
chenlei@njupt.edu.cn
Rong Wang
School of Artificial Intelligence, OPtics and ElectroNics (iOPEN) Northwestern Polytechnical University Xi'an 710072, P.R. China wangrong07@tsinghua.org.cn
Feiping Nie*
School of Artificial Intelligence, OPtics and ElectroNics (iOPEN) & School of Computer Science Northwestern Polytechnical University Xi'an 710072, P.R. China feipingnie@gmail.com
Xuelong Li
School of Artificial Intelligence, OPtics and ElectroNics (iOPEN) Northwestern Polytechnical University Xi'an 710072, P.R. China li@nwpu.edu.cn
Feature selection (FS) plays an important role in machine learning, which extracts important features and accelerates the learning process. In this paper, we propose a deep FS method that simultaneously conducts feature selection and differentiable  $k$ -NN graph learning based on the Dirichlet Energy. The Dirichlet Energy identifies important features by measuring their smoothness on the graph structure, and facilitates the learning of a new graph that reflects the inherent structure in new feature subspace. We employ Optimal Transport theory to address the non-differentiability issue of learning  $k$ -NN graphs in neural networks, which theoretically makes our method applicable to other graph neural networks for dynamic graph learning. Furthermore, the proposed framework is interpretable, since all modules are designed algorithmically. We validate the effectiveness of our model with extensive experiments on both synthetic and real-world datasets.
Feature selection (FS) is a critical technique in machine learning that identifies informative features within the original high-dimensional data. By removing irrelevant features, FS speeds up the learning process and enhances computational efficiency. In many real-world applications such as
image processing, bioinformatics, and text mining [1-3], FS techniques are widely used to identify important features, thereby providing some explanations about the results and boosting the learning performance [4-7].
While numerous FS methods have been proposed in both supervised and unsupervised settings, as several studies highlighted [8, 9], the nature of FS is more unsupervised due to the unavailability of task-specific labels in advance. The selected features should be versatile to arbitrary downstream tasks, which motivates us to focus on the unsupervised FS in this study. Related works have recently resorted to neural networks to exploit the nonlinear information within feature space. For example, AEFS [9] uses the group Lasso to regularize the parameters in the first layer of the autoencoder, so as to reconstruct original features based on the restricted use of original features. Another well-known method is CAE [8], which selects features by learning a concrete distribution over the input features. However, most unsupervised deep methods rely on the reconstruction performance to select useful features. On the one hand, if there exists noise in the dataset, the reconstruction performance will be terrible even if useful features are selected, since the noise cannot be reconstructed with these informative features (see our reconstruction experiments on Madelon in Section 4.3). On the other hand, it is difficult to explain why selected features reconstruct the original data well. These issues prompt us to seek a new target for unsupervised deep FS.
As the saying goes, "birds of a feather flock together", the homophily principle [10] suggests that similar samples tend to be connected in a natural graph structure within real-world data. This graph structure is useful for describing the intrinsic structure of the feature space, and is commonly used in machine learning studies [11, 12]. Building upon this graph structure, He et al. [13] introduce the Dirichlet Energy, which they call "locality preserving power", as a powerful tool for unsupervised FS that is able to identify informative features reflecting the intrinsic structure of the feature space.
In many practical applications, the graph structure is not naturally defined and needs to be constructed manually based on the input features using some similarity measurements. The quality of features affects the quality of the constructed graph. As highlighted in [14], the useless features increase the amount of unstructured information, which hinders the exploration of inherent manifold structure within data points and deteriorates the quality of constructed graph. Therefore, the reference [14] proposes the UDFS method to discard such nuisance features and constructs a  $k$ -nearest-neighbor (NN) graph on the selected features using the heat kernel. Despite the good performance of UDFS, constructing graphs using the heat kernel may not reflect the intrinsic structure of the feature space. Besides, the sorting algorithms in learning the  $k$ -NN graph in UDFS is non-differentiable in neural networks, which restricts its application in downstream networks.
In this paper: (1) We propose a deep unsupervised FS network that performs simultaneous feature selection and graph learning by minimizing the Dirichlet Energy, thereby revealing and harnessing the intrinsic structure in the dataset. (2) Within the network, a Unique Feature Selector (UFS) is devised to approximate discrete and distinct feature selection using the Gumbel Softmax technique combined with decomposition algorithms. (3) Moreover, a Differentiable Graph Learner (DGL) is devised based on the Optimal Transport theory, which is capable of obtaining a differentiable  $k$ -NN graph that more accurately reflects the intrinsic structure of the data than traditional graph constructing methods. Due to the differentiability, DGL is also theoretically capable of serving as a learnable graph module for other graph-based networks. (4) The entire framework is developed algorithmically. Unlike most deep learning networks with complex components that are tough to decipher, each core module in our framework has an algorithmic and physically interpretable design, which greatly facilitates observing and understanding the network's internal operations during the learning process. (5) Experimental results on both synthetic datasets and real-world datasets demonstrate the effectiveness of our method.
Notations. For an arbitrary matrix  $M \in \mathcal{R}^{a \times b}$ ,  $m^i$ ,  $m_i$ , and  $m_{i,j}$  denote the  $i$ -th row, the  $i$ -th column, and the  $(i,j)$ -th entry of  $M$ , respectively. Given a vector  $m \in \mathcal{R}^b$ , its  $\ell_2$ -norm is defined as  $\| m \|_2 = \sqrt{\sum_{i=1}^b m_i^2}$ . Based on this, the Frobenius norm of  $M$  is defined as  $\| M \|_F = \sqrt{\sum_{i=1}^a \| m^i \|_2^2}$ . When  $a = b$ , the trace of  $M$  is defined as  $\operatorname{tr}(M) = \sum_{i=1}^a m_{i,i}$ . Given two matrices  $M, N \in \mathcal{R}^{a \times b}$ , we define their inner product as  $\langle M, N \rangle = \sum_{i=1}^a \sum_{j=1}^b m_{i,j} n_{i,j}$ .  $1_b$  denotes a  $b$ -dimensional column vector with all entries being 1, and  $I_b$  denotes a  $b$ -order identity matrix.  $\text{Bool}(\text{cond})$  is a boolean operator that equals 1 if  $\text{cond}$  is true, otherwise it equals 0. Moreover, given a vector  $m \in \mathcal{R}^b$ , we define its sorting permutation in ascending order as  $\sigma \in \mathcal{R}^b$ , namely,  $m_{\sigma_1} \leq m_{\sigma_2} \leq \dots \leq m_{\sigma_b}$ .
(b)
Figure 1: Illustration of the Dirichlet Energy on various graph structures and graph signals. Blue points, black edges, and red bars represent nodes, connections, and signal values on nodes, respectively. Upside bars represent positive values, and downside bars represent negative values.
2 Dirichlet Energy
Let  $\mathbf{X} \in \mathcal{R}^{n \times d}$  be the data matrix with the  $n$  samples and  $d$ -dimensional features. In this paper, we assume that the features have zero means and normalized variances, $^2$  namely,  $\mathbf{1}_n^\top \mathbf{x}_i = 0$  and  $\mathbf{x}_i^\top \mathbf{x}_i = 1$  for  $i \in \{1, \dots, d\}$ . According to the homophily assumption [10], we assume that data  $\mathbf{X}$  forms an inherent graph structure  $\mathcal{G}$  with nodes standing for samples and edges standing for their correlations. In  $\mathcal{G}$ , similar samples are more likely to connect to each other than dissimilar ones. Graph  $\mathcal{G}$  can be represented with a similarity matrix  $S \in \mathcal{R}_+^{n \times n}$ , where  $s_{i,j}$  denotes the similarity between  $\mathbf{x}^i$  and  $\mathbf{x}^j$ . If  $s_{i,j} = 0$ , it means that there is no connection between  $\mathbf{x}^i$  and  $\mathbf{x}^j$  in  $\mathcal{G}$ , which is common in  $k$ -NN graphs since we only consider the local structure of the data space. Given an adjacency matrix  $S$ , $^3$  we define the Laplacian matrix [15] as  $L_S = D - S$ , where  $D$  is a diagonal matrix whose diagonal entries represent the degrees of data points, namely,  $d_{i,i} = \sum_{j=1}^{n} s_{i,j}$ .
Based on the Laplacian matrix  $\pmb{L}_S$ , we introduce the Dirichlet Energy [16] as a powerful tool to identify important features. Specifically, given the Laplacian matrix  $\pmb{L}_S$ , the Dirichlet Energy of a graph signal  $\pmb{v}$  is defined as
$$
\mathcal {L} _ {d i r} (\boldsymbol {v}) = \frac {1}{2} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} s _ {i, j} \left(v _ {i} - v _ {j}\right) ^ {2} = \boldsymbol {v} ^ {\top} \boldsymbol {L} _ {S} \boldsymbol {v}. \tag {1}
$$
In graph theory, each dimensional feature  $\boldsymbol{x}_i \in \mathcal{R}^n$  can be seen as a graph signal on  $\mathcal{G}$ . The Dirichlet Energy in Eq. (1) provides a measure of the local smoothness [11] of each feature on graph  $\mathcal{G}$ , which is small when the nodes that are close to each other on  $\mathcal{G}$  have similar feature values. Hence, the Dirichlet Energy can be used to identify informative features by evaluating the consistency of the distribution of feature values with the inherent data structure. To demonstrate this, we provide an example in Fig. 1, where we generate a 2-NN graph  $\mathcal{G}$  including two bubbles, and compose the data  $\boldsymbol{X}$  using the two-dimensional coordinates of the graph nodes. Then we set the graph signal  $\boldsymbol{v}$  as the first coordinate  $\boldsymbol{x}_1$  and visualize it on  $\mathcal{G}$  in Fig. 1(a). In Fig. 1(b) we change  $\boldsymbol{v}$  to a random noise vector. While in Fig. 1(c), we change the graph structure to a random 2-NN graph. We compute the Laplacian matrix  $L_S$  of each figure and present the corresponding Dirichlet Energy in the figures. We can see that Fig. 1(a) achieves the best smoothness, whereas both Fig. 1(b) and Fig. 1(c) have poor smoothness due to a mismatch between the graph signal and the graph structure.
Based on the Dirichlet Energy, a well-known FS method called Laplacian Score (LS) is proposed in [13]. However, the Laplacian matrix in LS is precomputed and fixed. If  $X$  contains too many irrelevant features, the quality of the graph  $\mathcal{G}$  will be poor and not reflect the underlying structure. As illustrated in Fig. 1(c), a poor-quality graph will lead to the poor smoothness even if the right feature is selected, this insight motivates us to learn graph and features jointly during the learning process.
In this paper, we devise a collaborative neural network driven by the Dirichlet Energy for joint feature and graph learning, as illustrated in Fig. 2. Generally, the proposed framework consists of two
2Note that constant features (if any) will be removed during the feature preprocessing stage.
Note that the similarity matrix  $S$  of a  $k$ -NN graph is not symmetric. When calculating the objective in Eq. (3), we obtain the Laplacian matrix  $L_{S}$  using the symmetrized similarity matrix  $\hat{S} = (S + S^{\top}) / 2$ .
<sup>4</sup>LS considers both the smoothness and the variance of each feature. However, as we assume that each feature has unit variance according to  $\mathbf{1}_n^\top \mathbf{x}_i = 0$  and  $\mathbf{x}_i^\top \mathbf{x}_i = 1$ , minimizing LS is equivalent to minimizing Eq. (1).
Figure 2: (1) Top Panel: Overview of the proposed framework, where smiley faces denote the value 1 representing that the feature is selected, while sad faces denote the value 0 representing that the feature is unused. (2) Bottom Left Panel: Illustration of the Unique Feature Selector (UFS), where green bars denote the value distributions of different vectors. (3) Bottom Right Panel: Illustration of the Differentiable  $k$ -NN Graph Learner (DGL), where the "Differentiable  $k$ -NN Selector" in deep blue shows how to learn  $k$  nearest neighbors with the Optimal Transport theory.
modules: the Unique Feature Selector (UFS) and the Differentiable  $k$ -NN Graph Learner (DGL). At the beginning, the input features  $X$  are selected with the learnable feature mask  $F$  generated by UFS, which is carefully designed to avoid the duplicate feature selection. Based on the selected data  $\hat{X}$ , we measure the distances between different samples, and feed the resulting distance vectors of each sample into DGL to learn their  $k$  nearest neighbors. The adaptive graph structure and informative features are learned jointly under the Dirichlet Energy, so as to identify the optimal feature subset that effectively captures the underlying data structure.
3.1 Unique Feature Selector
Based on the original data  $\mathbf{X}$ , the goal of FS is to identify a feature subset  $\hat{\mathbf{X}} \in \mathcal{R}^{n \times m}$  from the original features by minimizing a prespecified target  $\mathcal{L}_{obj}(\hat{\mathbf{X}})$ :
$$
\min  _ {\boldsymbol {F}} \mathcal {L} _ {o b j} (\hat {\boldsymbol {X}}) \quad \text {s . t .} \hat {\boldsymbol {X}} = \boldsymbol {X} \boldsymbol {F}, \boldsymbol {F} \in \{0, 1 \} ^ {d \times m}, \boldsymbol {F} ^ {\top} \boldsymbol {F} = \boldsymbol {I} _ {m}, \tag {2}
$$
where  $m \leq d$  denotes the number of selected features, and  $\pmb{F} \in \mathcal{R}^{d \times m}$  denotes the selection matrix selecting  $m$  features from  $\pmb{X}$ . Different from existing methods that use the reconstruction error as  $\mathcal{L}_{obj}(\hat{\pmb{X}})$ , in this paper, we utilize the Dirichlet Energy in Eq. (1) for FS as follows:
$$
\mathcal {L} _ {o b j} (\hat {\boldsymbol {X}}) = \sum_ {i = 1} ^ {m} \mathcal {L} _ {d i r} (\hat {\boldsymbol {x}} _ {i}) = \operatorname {t r} \left(\hat {\boldsymbol {X}} ^ {\top} \boldsymbol {L} _ {S} \hat {\boldsymbol {X}}\right). \tag {3}
$$
Given the selection number  $m$ ,  $\mathcal{L}_{obj}$  updates the network parameters by minimizing the Dirichlet Energy, thereby selecting  $m$  features that best reflect the intrinsic structure.
The constraints in problem (2) indicate that an ideal result  $\pmb{F}$  should be exact and unique. Exact means the result should exactly be the original features, instead of their linear combinations. Unique means each feature should be selected only once under a given number  $m$ . These two properties require  $\pmb{F}$  to be a binary and column-full-rank matrix including  $m$  orthogonal one-hot column vectors.
3.1.1 Approximating Discrete Feature Selection
It is difficult to learn a discrete  $\pmb{F}$  in neural networks due to its non-differentiable property. Inspired by [8], we propose to learn the discrete distribution using the Gumbel Softmax [17, 18] technique:
$$
\hat {\boldsymbol {f}} _ {i} = \operatorname {s o f t m a x} \left(\left(\log \boldsymbol {w} _ {i} + \boldsymbol {g} _ {i}\right) / T\right) \quad \text {w i t h} g _ {i, j} = - \log (- \log u _ {i, j}), u _ {i, j} \sim \operatorname {U n i f o r m} (0, 1), \tag {4}
$$
where  $\pmb{W} = [\pmb{w}_1, \pmb{w}_2, \dots, \pmb{w}_m]$  denotes a learnable parameter. The random vector  $\pmb{g}_i$  consists of  $d$  Gumbel-distributed variables  $g_{i,j}$ , which is generated with  $u_{i,j}$  sampled from Uniform distribution. Based on  $\pmb{w}_i$  and  $\pmb{g}_i$ , we obtain the approximated FS vector  $\hat{\pmb{f}}_i$  that represents the  $i$ -th selected feature. The distribution of  $\hat{\pmb{f}}_i$  is controlled by a non-negative temperature parameter  $T$ . A smaller value of parameter  $T$  will generate a better approximation of the one-hot vector, but will be more likely to be stuck in a poor local minimum. As suggested in [8], we employ the annealing schedule on  $T$  by initializing it with a high value and then gradually decreasing it during the learning process.
3.1.2 Selecting Unique Features
Despite having obtained the approximated FS vectors in neural networks, Eq. (4) does not consider the uniqueness requirement of FS. This is because Eq. (4) learns each selected feature separately, and does not consider the orthogonal constraint between columns in  $\hat{\pmb{F}}$ , which is prone to result in the repeated selection of the same features. To address this issue, we develop a unique feature selector (UFS) in Algorithm 1, where  $\hat{\mathbf{0}} \in \mathcal{R}^{(d - m)\times m}$  denotes the zero matrix. First, we add a small enough perturbation  $\epsilon I_{m}(\epsilon > 0)$  on  $\hat{\pmb{F}}^{\top}\hat{\pmb{F}}$ . Next, we perform
Algorithm 1 UFS
1: procedure UFS(ˆF, ε)
2: PΛP^T =ˆF^TˆF + εIm
3: LL^T =ˆF^TˆF + εIm
4: F = [ Λ^{1/2}P^T] (L^{-1})^T
5: return F
6: end procedure
the eigendecomposition (line 2) and the Cholesky decomposition (line 3) on the perturbed result respectively, and correspondingly obtain the diagonal matrix  $\Lambda \in \mathcal{R}^{m\times m}$ , the orthogonal matrix  $P\in \mathcal{R}^{m\times m}$ , and the lower triangle matrix  $L\in \mathcal{R}^{m\times m}$ . Based on  $\Lambda$ ,  $P$ , and  $L$ , we obtain the selection matrix  $F$  in line 4 and have the following conclusion:
Proposition 3.1. Given any real matrix  $\hat{\mathbf{F}}\in \mathcal{R}^{d\times m}$ , one can always generate a column-orthogonal matrix  $\mathbf{F}$  through Algorithm 1.
The proof of Proposition 3.1 is provided in Appendix S1. On the one hand, the small perturbation  $\epsilon I_{m}$  guarantees the column-full-rank property of  $\pmb{F}$ , thereby avoiding the duplicate selection results. On the other hand, the orthogonality property in Proposition 3.1 facilitates the approximation of discrete FS based on the matrix  $\hat{\pmb{F}}$ . We verify the efficacy of UFS in Section 4.2.
3.2 Differentiable  $k$ -NN Graph Learner
The existence of noise and irrelevant features may negatively affect the quality of the constructed graph. As depicted in Fig. 1, a low-quality graph structure can significantly perturb the smoothness of features and undermine the performance of feature selection. Hence, we propose to learn an adaptive graph during the learning process using the selected features.
3.2.1 Learning an Adaptive  $k$ -NN Graph Using Dirichlet Energy
Considering the objective function in Eq. (3), a natural way is to learn the similarity matrix  $S$  based on the Dirichlet Energy in  $\mathcal{L}_{obj}$ . However, this may yield a trivial solution where, for sample  $x^i$ , only the nearest data point can serve as its neighbour with probability 1, while all the other data points will not be its neighbours. To avoid this trivial solution, we propose to learn an adaptive graph by incorporating the Tikhonov regularization [19] of  $S$  into the Dirichlet Energy:
$$
\min  _ {\boldsymbol {S}} \operatorname {t r} \left(\hat {\boldsymbol {X}} ^ {\top} \boldsymbol {L} _ {S} \hat {\boldsymbol {X}}\right) + \frac {\alpha}{2} \| \boldsymbol {S} \| _ {F} ^ {2} \quad \text {s . t .} \boldsymbol {S} \mathbf {1} _ {n} = \mathbf {1} _ {n}, s _ {i, j} \geq 0, s _ {i, i} = 0, \tag {5}
$$
where  $\alpha$  denotes the trade-off parameter between the Dirichlet Energy and the Tikhonov regularization. Note that each row  $s^i$  in  $S$  can be solved separately, instead of tuning  $\alpha$  manually, we model  $\alpha$  as
In practice, we calculate  $\pmb{F}$  with  $\pmb {F} = \hat{\pmb{F}} (\pmb{L}^{-1})^{\top}$  without eigendecomposition to achieve discrete FS vectors. We provide a discussion in Appendix S2 to explain this treatment.
a sample-specific parameter  $\alpha_{i}$  and determine it algorithmically, which plays an important role in learning  $k$  nearest neighbors for each sample. Based on problem (5), we define the distance matrix  $\pmb{E}$  with its entries being  $e_{i,j} = \|( \hat{\pmb{x}}^i -\hat{\pmb{x}}^j)\| _2^2$ , then we solve each row  $\boldsymbol{s}^i$  in problem (5) separately as
$$
\min  _ {\boldsymbol {s} ^ {i}} \frac {1}{2} \| \boldsymbol {s} ^ {i} + \frac {\boldsymbol {e} ^ {i}}{2 \alpha_ {i}} \| _ {2} ^ {2} \quad \text {s . t .} \boldsymbol {s} ^ {i} \mathbf {1} _ {n} = 1, s _ {i, j} \geq 0, s _ {i, i} = 0. \tag {6}
$$
Problem (6) can be solved easily by constructing the Lagrangian function and then using the Karush-Kuhn-Tucker(KKT) conditions [20]. By doing so, we obtain the solution of  $s_{i,j}$  as
$$
s _ {i, j} = \left(\frac {1}{k} + \frac {1}{k} \frac {\boldsymbol {e} ^ {i} \boldsymbol {\delta} _ {i} ^ {(k)}}{2 \alpha_ {i}} - \frac {e _ {i , j}}{2 \alpha_ {i}}\right) _ {+} \quad \text {w i t h} \quad \delta_ {i, j} ^ {(k)} = \operatorname {B o o l} \left(e _ {i, j} \leq e _ {i, \sigma_ {k}}\right), \tag {7}
$$
where  $\pmb{\sigma} = [\sigma_1, \dots, \sigma_n]$  denotes the sorting permutation over  $e^i$ , i.e.  $e_{i,\sigma_1} \leq \dots \leq e_{i,\sigma_n}$  and  $\pmb{\delta}_i^{(k)}$  denotes the selection vector identifying the  $k$  minimal values in  $e^i$ .
Recall that we aim to learn  $k$  nearest neighbors for each sample, which implies that there are only  $k$  nonzero elements in  $s^i$  corresponding to the nearest neighbors. To this end, we determine the trade-off parameters  $\alpha_i$  such that  $s_{i,\sigma_k} > 0$  and  $s_{i,\sigma_{k + 1}}\leq 0$ . Then we have:
$$
\frac {1}{2} \left(k e ^ {i} \boldsymbol {\xi} _ {i} ^ {(k)} - e ^ {i} \boldsymbol {\delta} _ {i} ^ {(k)}\right) <   \alpha_ {i} \leq \frac {1}{2} \left(k e ^ {i} \boldsymbol {\xi} _ {i} ^ {(k + 1)} - e ^ {i} \boldsymbol {\delta} _ {i} ^ {(k)}\right) \quad \text {w i t h} \quad \xi_ {i, j} ^ {(k)} = \operatorname {B o o l} \left(e _ {i, j} = e _ {i, \sigma_ {k}}\right), \tag {8}
$$
where  $\pmb{\xi}_i^{(k)}$  denotes an indicator vector identifying the  $k$ -th minimal value in  $\pmb{e}^i$ . Setting  $\alpha_{i}$  as the maximum and substituting it into Eq. (7), we obtain the final solution as:
$$
s _ {i, \sigma_ {j}} = \frac {\mathbf {e} ^ {i} \boldsymbol {\xi} _ {i} ^ {(k + 1)} - e _ {i , \sigma_ {j}}}{k \mathbf {e} ^ {i} \boldsymbol {\xi} _ {i} ^ {(k + 1)} - \mathbf {e} ^ {i} \boldsymbol {\delta} _ {i} ^ {(k)}} \cdot \operatorname {B o o l} (1 \leq j \leq k). \tag {9}
$$
The detailed derivation of solution (9) can be found in Appendix S3. We note that the formulation in problem (6) bears similarity to CLR proposed in [21]. In Appendix S4, we discuss the connection between our method and CLR, and highlight the differences between the two w.r.t. the feature utilization and the sorting operation. Remarkably, the  $k$ -NN can be obtained easily in CLR using off-the-shelf sorting algorithms, which is not the case for neural networks due to the non-differentiability of sorting algorithms. To address this issue, we propose to transform the  $k$ -NN selection into a differentiable operator utilizing the Optimal Transport (OT) [22] technique as follows.
3.2.2 Differentiable  $k$ -NN Selector
Let  $\pmb{\mu} = [\mu_1, \mu_2, \dots, \mu_{n_1}]^\top$  and  $\pmb{\nu} = [\nu_1, \nu_2, \dots, \nu_{n_2}]^\top$  be two discrete probability distributions defined on the supports  $\mathcal{A} = \{\mathfrak{a}_i\}_{i=1}^{n_1}$  and  $\mathcal{B} = \{\mathfrak{b}_j\}_{j=1}^{n_2}$  respectively. The goal of OT is to find an optimal transport plan  $\Gamma \in \mathcal{R}^{n_1 \times n_2}$  between  $\mathcal{A}$  and  $\mathcal{B}$  by minimizing the following transport cost: