L EARNING TRANSPORT COST FROM SUBSET CORRESPONDENCE .
A BSTRACT
Learning to align multiple datasets is an important problem with many applications, and it is especially useful when we need to integrate multiple experiments or correct for confounding.
Optimal transport (OT) is a principled approach to align datasets, but a key challenge in applying OT is that we need to specify a transport cost function that accurately captures how the two datasets are related.
Reliable cost functions are typically not available and practitioners often resort to using hand-crafted or Euclidean cost even if it may not be appropriate.
In this work, we investigate how to learn the cost function using a small amount of side information which is often available.
The side information we consider captures subset correspondence—i.e. certain subsets of points in the two data sets are known to be related.
For example, we may have some images labeled as cars in both datasets; or we may have a common annotated cell type in single-cell data from two batches.
We develop an end-to-end optimizer (OT-SI) that differentiates through the Sinkhorn algorithm and effectively learns the suitable cost function from side information.
On systematic experiments in images, marriage-matching and single-cell RNA-seq, our method substantially outperform state-of-the-art benchmarks.
I NTRODUCTION
In many applications, we have multiple related datasets from different sources or domains, and learning efﬁcient computational mappings between these datasets is an important problem (Long et al., 2017; Zamir et al., 2018).
For example, we might have single-cell RNA-Seq datasets generated for the same tissue type from two different labs.
Since data come from the same type of tissue, we would like to map cells between the two datasets to merge them, so that we could analyze them jointly.
However, there are often complex nonlinear batch artifacts generated by the different labs.
Moreover the cells are not paired—for each cell measured in the ﬁrst lab, there is not an identical clone in the second lab.
How to integrate or align these two datasets is therefore a challenging problem.
Optimal transport (OT) is a principled analytical framework to align heterogeneous datasets (Santambrogio, 2015).
It has been increasingly applied to problems in domain adaptation and transfer learning (Seguy et al., 2017; Genevay et al., 2017; Courty et al., 2017b;
Li et al., 2019).
Optimal transport is an approach for taking two datasets, and computing a mapping between them in the form of a "transport plan" γ .
The mapping is optimal in the sense that among all reasonable mappings (precisely deﬁned in Section 2), it minimizes the cost of aligning the two datasets.
The transport cost is given by the user and encodes expert knowledge about how datasets relate to each other.
For example, if the expert believes that one data Y is essentially data X with added Gaussian noise, then
Euclidean cost could be natural.
If the cost is correctly speciﬁed, then there are powerful methods for ﬁnding the global optimal transport (Villani, 2008).
A major challenge in practice, e.g. for single-cell
RNA-seq, is that we and experts do not know what cost is appropriate.
Users often resort to using
Euclidean or other hand-crafted cost functions, which could give misleading mappings.
Our contributions.
We propose a novel approach to automatically learn good transport costs by leveraging side information we may have about the data.
In particular, aligning datasets across different conditions is a very important problem in biology, especially in single cell analysis (Butler et al., 2018), where we often have annotations that certain cluster of cells in a dataset corresponds to one cell type (e.g. B cells and T cells) based on known marker genes (Schaum et al., 2018).
Then we can deduce that T cells from dataset X should be at least mapped to T cells from dataset Y .
We only need T cells to be crudely annotated in both datasets, which is reasonable; we don’t need to know that a particular T cell should be mapped to another speciﬁc T cell.
This gives rises to the side information that we can leverage in our algorithm — a certain subset of points in dataset X should be mapped to another subset of points in dataset Y .
We present the ﬁrst algorithm, OT-SI, to leverage subset correspondence as a general form of side information.
In contrast, previous works mainly focus on pair matching problems (Li et al., 2019;
Galichon & Salanié, 2010) — the extreme case of subset correspondence when the subset sizes are 1.
In practice, exact one-to-one matching labels are often expensive to obtain or even intractable.
OT-SI is an end-to-end framework that learns the transport cost.
The intuition is to optimize over a parametrized family of transport costs, to identify the cost under which the annotated subsets are naturally mapped to each other via optimal transport.
OT-SI efﬁciently leverages even a small amount of side information and it generalizes well to new, unannotated data.
The learned transport cost is also interpretable.
We demonstrate in extensive experiments across image, single-cell, marriage and synthetic datasets that our method OT-SI substantially outperform state-of-the-art methods for mapping datasets.
Related Work Optimal transport been well studied in the mathematics, statistics and optimization literature (Villani, 2008; Courty et al., 2017a;
Li et al., 2019; Courty et al., 2017b).
OT can be used to deﬁne a distance metric between distributions (e.g. Wasserstein distance) or to produce an explict mapping across two datasets.
The latter is the focus of our paper.
In machine learning, there has been signiﬁcant work on developing fast algorithms for efﬁcient computation of the optimal transport plan (Cuturi, 2013; Altschuler et al., 2017; Staib et al., 2017), and analyzing the properties of the transport plan under various structures and constraints on the optimization problem (Alvarez-Melis et al., 2018;
Titouan et al., 2019).
The previous work on learning the transport cost is done on a very different setting from ours – learning feature histogram distances between many pairs of datapoints (Cuturi & Avis, 2014).
Some classical clustering (Xing et al., 2003; Bilenko et al., 2004) and alignment methods (Ham et al., 2005; Wang & Mahadevan, 2008; 2009) have realized beneﬁts by including side information, but these semi-supervised approaches differ from our explicit parametrization and optimization of the transport cost function.
Separately, there have been recent efforts to directly map between datasets, without learning a transport cost.
The standard alignment methods can be divided into two categories: GANs-based (Zhu et al., 2017; Choi et al., 2018) and OT-based (Grave et al., 2019; Alvarez-Melis et al., 2019).
GAN-based approaches have been used to align single-cell RNA-seq data when pairs of cells are known to be related (Amodio & Krishnaswamy, 2018).
However the exact pairing of individual cells is always not readily available or even intractable.
To address this issue, our method OT-SI allows for more general correspondence between subsets, i.e., clusters, cell types and also individual cells.
In the meantime, the OT-based methods always rely on Procrustes analysis (Rangarajan et al., 1997) — a linear transformation between the datasets is assumed, which lacks the ﬂexibility to handle nonlinear artifacts and the side information cannot be utilized.
In contrast, a major beneﬁt of our approach is its graceful adaptation to partial subset correspondence information, where we frame the problem as semi-supervised.
L EARNING C OST M ETRICS
A good choice of the cost function for optimal transport is the key to a successful mapping between two datasets.
In this section, we present the algorithm OT-SI, which parametrizes the cost function with weight θ and adaptively learns θ using side information about the training data.
The side information we consider is subset correspondence — a common situation when some subsets of training points are known to be related; pair matching is included as an extreme case.
The learned cost function is further evaluated on the unseen test data to prove generalizability.
2. O PTIMAL T RANSPORT
Consider learning a mapping between two datasets X
= { x (1) , ..., x ( n X ) } and Y = { y (1) , ..., y ( n Y ) } .
Here we use n X and n Y to denote the number of datapoints; each sample x ( i ) or y ( j ) could be a vector as well.
We brieﬂy recall the optimal transport framework in this setting.
Given probability vectors µ X and µ Y , the transport polytope is deﬁned as
[Equation]
where 1 n X ( 1 n Y ) is the n X ( n Y ) dimensional vector of ones.
Here the probability vector µ X ( µ Y ) is in the simplex (cid:80) n := { p ∈ R n + : p T 1 n = 1 } for n = n X ( n Y ).
For two random variables with distribution µ X and µ Y , the transport polytope U ( µ X , µ Y ) represents the set of all possible joint probabilities of the two variables.
In this paper, we consider µ X and µ Y to represent the empirical distributions of the samples X and Y, respectively, and set µ = (1 /n ) 1 n .
Given a n X × n Y cost matrix C , the classical optimal transport plan between µ X and µ Y is deﬁned as γ ∗ =
arg min γ ∈ U ( µ X ,µ Y ) (cid:104) γ, C (cid:105) , where (cid:104)· , ·(cid:105) denotes the Frobenius inner product.
γ ∗ is also called a coupling.
Despite its intuitive formulation, the computation of this linear program quickly becomes prohibitive especially in the common situation when n X and n Y , the sizes of the datasets, exceed a few hundred.
For computational efﬁciency, Sinkhorn-Knopp iteration is widely used to compute the optimal transport (Cuturi, 2013).
Sinkhorn-Knopp is a fast iterative algorithm for approximately solving the optimization problem with entropy regularization Santambrogio (2015):
[Equation]
where λ > 0 is a regularization parameter and h ( γ )
= − (cid:80) n X i =1 (cid:80) n Y j =
γ
ij log γ ij denotes the entropy.
The regularized solution γ λ converges to the classical one γ ∗ when the regularization diminishes, i.e., λ → ∞ , with exponential convergence rate (Cominetti & San Martín, 1994).
The transport γ λ treats
X and Y symmetrically.
The cost matrix is retrieved from the cost function C ij =
c ( x ( i ) , y ( j ) ) .
A good choice of the cost function is the key to inﬂuencing the learned mapping γ λ .
However, reliable cost functions are typically not available and Euclidean cost is mostly used.
In this paper, the representation of the cost function is adaptively learned using side information about the data.
2. S IDE I NFORMATION
Subset correspondence describes a common situation when certain subsets of points are known to be related.
For example, images with the same objects should always be mapped together in domain adaptation tasks (Courty et al., 2017b), while cells in a single-cell dataset need to be aligned to those with the same cell type, where the cell type annotation is available.
Given m corresponding subsets, we write S ( k ) X and S ( k ) Y , k = 1 , ..., m , to denote the sets of data indices in the corresponding subsets, i.e., { x ( i ) | i ∈ S ( k ) X } and { y ( i ) | i ∈ S ( k ) Y } .
Note that S ( k ) X and
S ( k ) Y could have different probability mass.
If | S ( k ) X | n X ≤ | S ( k ) Y | n Y , we take this side information to be that
S ( k ) X should be mapped into S ( k ) Y .
In other words, all the other entries of the transport matrix γ that maps S ( k ) X to outside of S ( k ) Y should be 0.
Everything is swapped if | S ( k ) Y | n Y [Equation]
Mathematically, this side information corresponds to the constraint that γ ij = 0 for ( i, j ) ∈ S 0 , where
[Equation]
Note that pair matching is an extreme case of subset correspondence with subset size 1 — that is, the exact pairwise relation is known.
The pair matching problem has been addressed in literature (Li et al., 2019; Galichon & Salanié, 2010).
However, in practice, exact one-to-one matching labels are often expensive to obtain or even intractable.
In this paper, we show that subset correspondence, as a small amount of side information, can signiﬁcantly aid cost function learning.
We investigate how to learn the cost function using subset correspondence information between training datasets X train and Y train .
The learned cost function is evaluated via the mapping quality on the test datasets X test and Y test .
Note that the training and test sets are not necessarily under the same distribution.
We demonstrate the power of OT-SI in generalizing to new subsets that were not seen in the training process.
2. T HE OT-SI A LGORITHM
Our ultimate goal is to learn a cost function c ( · ) , such that the computed optimal transport γ λ train satisﬁes the side information given in Eq.
(2) as faithfully as possible.
Cost function parametrization.
When the cost function c ( · ) is Euclidean, the entry of the cost matrix is computed as C ij =
c
( x ( i ) , y ( j ) )
= (cid:80) dk
=1 ( x ( i ) k − y ( j ) k ) 2 , where d is the data dimension.
To learn the cost function systematically, we parametrize it as C ij ( θ ) =
c ( x ( i ) , y ( j ) ; θ ) with weight θ .
Here the function form c ( · ) can be chosen by users.
To illustrate the improvement over the commonly-used
Euclidean cost, we parametrize c ( x ( i ) , y ( j ) ; θ ) as a polynomial in ( x ( i ) 1 , ..., x ( i ) d x , y ( j ) 1 , ..., y ( j ) d y ) with coefﬁcients θ and degree 2 for low-dimensional data.
The Euclidean cost is equivalent to a speciﬁc choice of θ 0 which is set as the initialization; see Appendix for more discussions.
For high-dimensional data, the memory required to store the second order polynomials becomes too large and we use a fully connected neural network to parametrize c ( x ( i ) , y ( j ) ; θ ) with input ( x ( i ) 1 , ..., x ( i ) d x , y ( j ) 1 , ..., y ( j ) d y ) and weights θ .
Throughout this paper, the polynomial parametrization is used if not speciﬁed.
The optimal transport solution is characterized by θ as
[Equation]
Then the problem can be formulated as optimizing θ to make the transport γ λ ( θ ) approximately satisfy the conditions deﬁned in Eq.
(2), penalizing deviation of the solution from these constraints with the loss
[Equation]
Theorem 1.
For any λ > 0 , the optimal transport plan γ λ ( θ ) :
R d θ → U ( µ X , µ Y ) is C ∞ in the interior of its domain.
The inﬁnite differentiability of the Sinkhorn distance is previously-known (Luise et al., 2018);
Thm.
1 proves that the Sinkhorn transport plan also has this desirable property.
Because Thm. 1 guarantees that γ λ ( θ ) is inﬁnitely differentiable, we are able to optimize L λ ( θ ) in Eq.
(4) by gradient descent.
In practice, we iterate Sinkhorn’s update a sufﬁcient number of times to converge to γ λ .
Each iteration is a matrix operation involving the cost matrix C ( θ ) , and when the number of iterations is ﬁxed, we can propagate the gradient ∇ θ through all of the iterations using the chain rule.
Updating θ by one forward and backward pass has complexity of O ( n 2 ) up to logarithmic terms.
Hence OT-SI has the same complexity as the Procrustes-based OT methods which alternatively optimize over coupling matrix and linear transformation matrix (Grave et al., 2019; Alvarez-Melis et al., 2019).
To further boost the performance, we propose to use a mimic learning method for initialization, which does not need to propagate the gradient.
The pseudocode for OT-SI is in Algorithm 1 and details are in the Appendix.
The proof of Thm. 1, mimic learning algorithm, details, and discussions about convergence are also in the Appendix.
2. E XPERIMENTS S ETUP
The OT-SI algorithm is carried out in Pytorch (Paszke et al., 2017) and trained with GPU.
The model is ﬁtted on training set, and evaluated on test set.
We use validation set for hyperparameter selection and early stopping.
We evaluate OT-SI with different types of data and correspondence information.
Comparison methods.
We use optimal transport with Euclidean cost function as a baseline for comparison, referred as "OT-baseline".
We also compare our result with state-of-the-art GAN-based data alignment methods, MAGAN (Amodio & Krishnaswamy, 2018) and CycleGAN (Zhu et al.,
Require: training datasets X train and Y train , corresponding subsets index S ( k ) X train and S ( k ) Y train ( k = 1 , ..., m train ), step size α , total training steps T , weights θ 0 from initialization procedure, Sinkhorn regularization parameter λ and number of Sinkhorn iterations n Sinkhorn .
1: n X = length
( X train ) ,
n Y = length
( Y train ) , µ X =
(1 /n X ) 1 n X , µ Y =
(1 /n Y ) 1 n Y 2:
Initialize θ = θ3: for t = 1 to T do 4:
Compute cost matrix C ( θ ) with entries C ij ( θ )
=
c ( x ( i ) , y ( j ) ; θ ) .
Solve γ λ train ( θ ) =
arg min γ ∈ U ( µ X ,µ Y ) (cid:104) γ, C ( θ ) (cid:105) − 1 λ h ( γ ) with Sinkhorn’s update with n
Sinkhorn iterations.
Derive ∇ θ γ train by backpropagating the gradient through all Sinkhorn-Knopp iterations.
Update weights θ := θ − α (cid:80) ( i,j ) ∈ S 0 , train γ λij, train ( θ ) ∇ θ γ λij, train ( θ ) .
8: end for 2017), as well as the OT-based methods, RIOT (Li et al., 2019) which is developed for speciﬁc pair matching applications and the Procrustes-based OT (Grave et al., 2019), referred as "OT-Procrustes".
For MAGAN and CycleGAN, the matching point for a source sample is set as its nearest neighbor in the target after mapping. Among the ﬁve comparison methods, OT-baseline, OT-Procrustes and CycleGAN do not use any side information; MAGAN makes use of matching pairs; RIOT requires the one-to-one matching labels for all the datapoints.
Because MAGAN and RIOT requires pairwise correspondence, they are not applied in some experiments and these are marked as N/A.
We use the same settings and hyperparameters for the comparison methods as in their original implementations.
Evaluation metrics.
When the subset correspondence is known on the test set (not shown to the algorithm), we evaluate a transport plan γ by how much it satisﬁes the correspondence.
Mathematically, we deﬁne subset matching accuracy :
[Equation]
From the deﬁnition, 0 (cid:54) Accuracy (cid:54) 1 gives the probability of mapping to the correct corresponding subsets.
When all the test datapoints are mapped into the correct subsets, the accuracy is 1; when all the data are matched to the wrong subsets, accuracy is 0.
As an extreme example, pair matching is equivalent to subset correspondence with subset sizes 1, referred to as pair matching accuracy .
In the next few sections, we thoroughly evaluate OT-SI and several state-of-the-art methods in extensive and diverse experiments—aligning single-cell RNA-seq data to correct for batch effects, aligning single-cell gene expression and protein abundances, a marriage data, an image dataset, and the synthetic twin-moon data for illustration.
B ENCHMARK ON S YNTHETIC D ATASETS
We ﬁrst experiment with the benchmark toy example for domain adaptation — two moon datasets—to illustrate the challenges of data alignment, before we move onto complex real-world data (Germain et al., 2013; Courty et al., 2017b).
The dataset is simulated with two domains, source and target.
As shown in Fig.
1a, each domain contains two standard entangled moons.
The two moons are associated with two different classes, denoted by circle and crossing respectively.
The target (colored in orange), is built by adding noise to the source (colored in orange) and rotating by 60 degrees.
In the experiment, we generate the training, test and validation datasets with 100, 100, and 50 samples of each moon.
We set the parameter λ = 10 3 and the number of Sinkhorn-Knopp iterations N = 200 .
The algorithm is run for 100 epochs with step size 1.
There are two types of side information available for OT tasks: (i) subset correspondence — datapoints are known to be mapped into the corresponding moon class; (ii) pair matching — known matched datapoints after rotation.
The result is averaged over 10 (50) independent runs when the side information is subset correspondence (pair matching).
Baseline performance.
The optimal transport plan under Euclidean cost function is depicted in Fig.
1b.
Datapoints learned to be matched are connected by solid curves.
The red curves indicate wrong [Table] [Figure] transports which map the data into the wrong subset,
i
[Table]
e., the other moon.
As shown by Fig.
1b, most wrong transports are between the points at the edges of the moons.
In the euclidean space, the edge of one moon becomes “closer" to the other moon after rotation, which leads to a small cost between datapoints in different moon classes.
A new cost function which captures the rotation property is expected.
Quantitatively, only 72% of the data are mapped into the correct moon it belongs to and only 2% of the data are matched to their corresponding pairs, as given in Table 1.
Subset correspondence.
We ﬁrst evaluate our method OT-SI when the side information is only that the data on the corresponding moons are known to be related—i.e. S (1) X and S (2) X are the two moons in dataset one, and S (1) Y and S (2) Y are the two moons in dataset two.
With the learned cost function, almost all the datapoints are mapped into the corresponding moon, i.e., the subset matching accuracy on the test data achieves 100% for both methods, as shown in Table 1.
The results are averaged for 10 independent runs.
Interestingly, although there is no pair matching information provided during training, the learned cost function signiﬁcantly improves the matching performance.
As shown in
Table 1, 87% datapoints are transported into their exact matching points after rotation.
The learned mapping of OT-SI is depicted in Fig. 1c.
The rotation property of the datasets is correctly captured.
In contrast, OT-Procrustes sometimes learns as good as OT-SI, similar to Fig. 1c, but sometimes mistakenly learns rotation as ﬂipping, indicated by Fig.
1d.
This results in overall worse accuracy and large variance for OT-Procrustes.
Pair matching.
OT-SI demonstrates substantial improvement when the pair matching information is provided—only 1 and 10 pairs are known out of the total 100 training pairs.
The matching pairs are randomly selected from the training data.
OT-SI signiﬁcantly outperforms all four bench methods particularly when the number of known pairs is very limited.
Even when only 1 matching pair is provided, the learned cost function greatly improves the OT performance, as given in Table 1.
The improvement here is largely attributable to the unlabeled data, i.e., the datapoints without any pair [Table] [Table] matching information.
For comparison, we carry out another experiment with only 3 unlabeled data, with all other settings unchanged.
The algorithms are not able to learn the right cost function anymore — the test pair matching accuracy is only 0 .
5% after learning, even worse than the Euclidean baseline.
With 199 unlabeled data, the accuracy achieves 53% .
In contrast, the competing methods
MAGAN and RIOT learn barely any patterns, because too few labeled datapoints are available and the unlabeled ones are wasted.
B IOLOGICAL M ANIFOLD A LIGNMENT
In this section, we implement our method OT-SI to learn a cost function that aligns biological manifolds with partial supervision — annotations of some cell types or clusters, which is the common situation in biological studies.
The pair matching methods, MAGAN and RIOT, are not compared here because the the cell-to-cell matching information is not available.
Similar to Sec.
3, the
CycleGAN does not learn correctly for these data types and is not presented here.
We show that
OT-SI has substantial improvement for aligning datasets with different data types and aligning data from different batches.
4. A LIGNMENT OF P ROTEIN AND RNA SEQUENCING DATA
How to align datasets with different data types has been a major topic in many ﬁelds.
For example, in single-cell studies, RNA and protein sequencing can both be done at cellular resolution.
How to map between those two types of data, i.e., map cells with certain mRNA level to cells with certain protein level, becomes critical for downstream studies such as RNA and protein co-expression analysis.
We demonstrate the power of OT-SI in learning cost function that aligns two different data types —
RNA and protein expression in the CITE-seq cord blood mononuclear cells (CBMCs) experiments (Stoeckius et al., 2017).
The dataset is subset to 8,005 human cells with the expression of 5, highly variable genes and 13 proteins.
Fifteen clusters are identiﬁed using the Louvain modularity clustering.
The CITE-seq technology has enabled the simultaneous measurement of RNA and protein expression at single-cell level, hence the ground truth about the cell pairing is available.
To emulate the common situation, we only use the information of cluster correspondence in the training and report the performance of both subset (cluster) matching and pair matching for test.
We randomly sampled 500 and 500 cells for validation and test purpose.
When OT-SI is learned in the original data space with the expression of 5001 mRNA and 13 proteins for each cell, we use a fully connected neural network to parametrize the cost function, with two hidden layers of 100 and 5 neurons.
We align the RNA and protein expression datasets in two scenarios: i) the embedding space where we use the ﬁrst 10 PCs for both datasets; ii) the original expression space.
Table 2 shows that
OT-SI substantially outperforms other methods.
OT-SI is able to learn good cost function when the dimensions of the two datasets are highly unbalanced as 5001:13.
The learned cost metrics in the original expression space can be used for future biological analysis on the effect and relation between
RNA and protein expressions.
Although the single-cell sequencing data are noisy and the algorithm is not informed of any matching pair during training, the accuracy of test pair matching is improved. [Table] [Table] [Table]
B ATCH A LIGNMENT
Another fascinating biological application of optimal transport is to align data from different batches.
In biological studies, the samples processed or measured in different batches usually result in nonbiological variations, known as batch effect (Chen et al., 2011; Haghverdi et al., 2018).
Here we use
OT to align two batches of data 1 — data collected with ﬂuorescence activated cell sorting (FACS) and droplet methods.
In this case, the cell types are used as the subset correspondence information.
For illustration purposes
[Table]
we subsample the top 10 celltypes with 400 samples of each.
There are 1,682 genes after ﬁltering and the ﬁrst 10 principal components are used for analysis.
The dataset is split into training, validation and test sets with ratio 50%, 20% and 30%.
The experiment setting is the same as in Sec. 3.
We demonstrate the power of our learned metric in generalizing to entirely new cell types that were not used to train the cost.
This is a hard task (a zero-shot learning task), and is more realistic because in most settings we only have partial annotations for cell types and we would like the mapping to generalize to all of the data.
To do this, we choose two cell types — T cell and immature T cell — as held-out and train on the rest.
Among all the ground truth expert annotations, T cell and immature T cell are most difﬁcult to be aligned.
With OT-baseline, only 70% cells are mapped into the correct cell types.
Substantial improvement is achieved, as shown in Table 3.
From a small number of annotated cells types, OT-SI is able to learn a transport cost that captures the batch artifacts between FACS and droplet which generalized to mapping these two new cell types.
We anticipate future uses of our formulation to further investigate the cost function, particularly in biological discovery applications like isolating genes that mark single-cell heterogeneity.
E XPERIMENTS ON I MAGES AND T ABULAR D ATA A LIGNMENT 5. I MAGE A LIGNMENT FOR N EW D IGITS
To illustrate the use of OT-SI in image alignment, we use it to learn cost metrics in aligning images with partial annotation on the MNIST dataset which contains 28 × 28 images of handwritten digits.
We subsample 200 images from each digit class and split them into training, validation, and test sets with ratio 50%, 20% and 30%.
For illustration purposes, we use the ﬁrst ten principal components for the alignment analysis.
We generate four different types of perturbations to the original images, as plotted in Table 4.
Then we align the original images with the perturbed ones using OT, respectively.
[Table]
The experiment setting is the same as in Sec. 3.
Similar to Sec.
4, we test how well our algorithm generalizes to new classes that were not used in learning the cost function.
In the experiment, we hold out digits 3 and 5, and demonstrate the metric learned on the other eight digits can help the alignment of digits 3 and 5.
We achieve consistent improvement over the baseline on all four distribution types, as indicated by Table 4.
5. M ARRIAGE D ATASET FOR P AIR M ATCHING
While OT-SI is designed for the more general form of side information — subset correspondence, it can also be used for pair matching purpose.
We ﬁnally benchmark it from the comparison with other state-of-the-art pair matching methods, including RIOT, factorization machine model (FM) (Rendle, 2012), probabilistic matrix factorization model (PMF) (Mnih & Salakhutdinov, 2008), item-based collaborative ﬁltering model (itemKNN) (Cremonesi et al., 2010), classical SVD model (Koren et al., 2009) and baseline random predictor model.
These methods were also used as comparisons in Li et al.
We follow the same experimental protocol as in Li et al.
(2019) for the Dutch Household
Survey (DHS) dataset.
The exact matching matrix between 50 datapoints with 11 features are known.
We note that the coupling matrix is treated as continuous and OT-Procrustes is not applicable.
The performance is evaluated by the root mean square error (RMSE) and the mean absolute error (MAE) for the predicted matching matrix, as given in Table 5.
When used for pair matching purposes,
OT-SI report comparable performance to state-of-the-art matching algorithms.
Note that this marriage dataset was a primary motivating dataset used to design RIOT (Li et al.
, and therefore we expect RIOT to perform very well for this task.
D ISCUSSION
In this paper, we study the problem of learning the transport cost using side information in the form of a small number of corresponding subsets between the two datasets.
This is a new problem formulation, to the best of our knowledge.
Previous works rely on more restricted information such as that speciﬁc pairs of points should be aligned.
In settings such as genomics and images, it is often difﬁcult to say that a single point in dataset one should be mapped onto a particular point in dataset two.
It is more common to have partial annotation of subsets of points—e.g. T cells are annotated in two single-cell RNA-seq datasets—which motivates our generalization.
We propose a ﬂexible and principled method to learn the transport cost with side information.
Experiments demonstrate that they work signiﬁcantly better than state-of-the-art methods when the side-information is very limited, which is often the case.
We compare against state-of-the-art methods for the special case when the side information consists of matching pairs, since we are not aware of other published OT methods that deal with the more general subset correspondence.
One interesting reason for the improved performance is that by learning the transport cost directly, our algorithms are more efﬁciently using all of the unannotated datapoints that are not in any pairs or subsets.
These unannotated data act as regularization (similar to in semi-supervised 1learning), which enables the model to avoid overﬁtting to the limited side information.
An interesting direction of future work is to interpret the learned cost function for insights on how the datasets differ.
A C OST FUNCTION
In Sec. 2.3, the cost function is parametrized as C ij ( θ ) =
c ( x ( i ) , y ( j ) ; θ ) with weight θ .
For lowdimensional data, we choose c ( x ( i ) , y ( j ) ; θ ) as a polynomial in ( x ( i ) 1 , ..., x ( i ) d x , y ( j ) 1 , ..., y ( j ) d y ) with coefﬁcients θ and degree 2:
[Equation]
Here we do not require the dimension of the two datasets d x and d y to be the same.
When the two datasets X and Y have different features, the learned weights indicate the coupling between different features in the data mapping.
In general, the function form c ( · ) can be chosen by users.
We have also investigated parametrizing C ( θ ) as a small fully connected neural network and achieved very similar performance.
The Euclidean cost is equivalent to a speciﬁc choice of θ , when the two datasets have the same feature space d x = d y = d .
The entry of the cost matrix is computed as C ij =
c
( x ( i ) , y ( j ) ; θ ) = (cid:80) dk
=1 (( x ( i ) k )
+ ( y ( j ) k ) 2 − 2 x ( i ) k
y ( j ) k ) .
B OT-SI A LGORITHM
The Lagrangian dual of Eq.
(3) is
[Equation]
By Sinkhorn’s scaling theorem (Sinkhorn & Knopp, 1967), the optimal transport plan γ λ ( θ ) is computed as
[Equation]
where u ∗ ( θ ) and v ∗ ( θ ) are the solutions to the dual problem in Eq.
B. P ROOF OF T HEOREM
The analysis for Theorem 1 follows the strategy of the proofs of Theorem 2 in (Luise et al., 2018).
Theorem 1.
For any λ > 0 , the optimal transport plan γ λ ( θ ) :
R d θ → U ( µ X , µ Y ) is C ∞ in the interior of its domain.
Proof.
Based on Eq.
(7), the optimal transport plan γ λ ( θ ) is a smooth function
when u ∗ ( θ ) , v ∗ ( θ ) and C ( θ ) are smooth.
In the meantime, the cost C ( θ ) is a linear function of θ , as indicated by Eq.
Thus to prove the smoothness of γ λ ( θ ) , we only need to demonstrate u ∗ ( θ ) and v ∗ ( θ ) are both smooth in θ .
Here we deﬁne
[Equation]
The dual problem in Eq.
(6) becomes min u,v σ ( θ ; u, v ) .
From the deﬁnition, σ ( θ ; u, v ) is smooth and strictly convex in ( u, v ) .
Note that C ( θ ) is linear in θ .
Then for any ﬁxed θ in the interior of R d θ , there exits ( u ∗ ( θ ) , v ∗ ( θ ))
such that σ ( θ ; u ∗ ( θ ) , v ∗ ( θ )) =
min u,v σ ( θ ; u, v ) .
The function ∇ ( u,v ) σ ( θ ; u, v ) ∈
C
∞
due to the smoothness of σ ( θ ; u, v ) .
Now we ﬁx ( u 0 , v 0 , x 0 ) such that ∇ ( u,v ) σ ( θ 0 ; u 0 , v 0 ) = 0 .
The strict convexity of σ ( θ ; u, v ) ensures that ∇ 2( u,v ) σ ( θ 0 ; u 0 , v 0 ) is invertible.
From implicit function theorem, we can always ﬁnd a function f and a subset U θ 0 ∈ R d θ such that i) f
( θ 0 ) =
( u 0 , v 0 ) ; ii) ∇ ( u,v ) σ ( θ ; f ( θ ))
=
for any θ ∈ U θ 0 ; iii) f ∈ C ∞ ( U θ 0 ) .
That is, f ( θ ) is a stationary point of the function σ for any θ in U θ 0 .
Together with the strict convexity of σ , we derive f ( θ ) =
( u ∗ ( θ ) , v ∗ ( θ )) .
Recalling f ∈ C ∞ ( U θ 0 ) , we prove that ( u ∗ ( θ ) , v ∗ ( θ )) is C ∞ in the interior of its domain.
B. C ONVERGENCE PROPERTIES
The gradient of γ λ ( θ ) in Eq.
(7) is computed as
[Equation]
The convergence of scaling factors e λu ( θ ) and e λv ( θ ) is linear (i.e. exponential in n Sinkhorn ), with the bounded rate given by (Franklin & Lorenz, 1989); our experiments use a linear C ( θ ) , and we do not ﬁnd Sinkhorn to bottleneck convergence.
Besides this point, convergence is determined by the loss landscape of C .
Analyzing convergence to the global optimum, and the role of primal regularization of γ , is a relevant open question.
B. C OMPUTING G RADIENT
In practice, we iterate Sinkhorn’s update N times to converge to γ λ .
More speciﬁcally, the dual solutions u ∗ ( θ ) and v ∗ ( θ ) for Eq.
(6) are computed by iterating for N times
[Equation]
For the optimal plan γ λ ( θ ) given in Eq.
(7), its gradient ∇ θ γ λ ( θ ) is obtained by propagating through all the iterations using the chain rule and implemented by Pytorch.
Thus updating θ by one forward and backward pass has complexity of O ( n 2 ) up to logarithmic terms.
OT-SI has the same complexity as the Procrustes-based OT methods which alternatively optimize over coupling matrix and linear transformation matrix.
Taking the image alignment for watering MNIST dataset as an example, the running time for OT-SI and OT-Procrustes is 229.5s and 196.7s, respectively.
C M IMIC L EARNING FOR I NITIALIZATION
In this section, we derive a mimic learning as an fast initialization method to boost the performance and accelerate the learning.
While Algorithm 1 requires to differentiate through Sinkhorn updates, the mimic learning approach does not need to propagate the gradient through all the iterations and is applicable for any kind of OT algorithm.
Here we take the classical optimal transport plan γ ∗ = arg min γ ∈ U ( µ X ,µ Y ) (cid:104) γ, C (cid:105) as an example.
As discussed in Sec. 2.3, our ultimate goal is to learn a cost function, such that the optimal transport γ ∗ satisﬁes the side information deﬁned in Eq.
(2) as faithfully as possible.
From another perspective, we force an additional set of constraints on the transport plan γ to fulﬁll the condition in Eq.
[Equation]
To quantify how much the learned γ ∗ in Eq.
(1) follows the side information, we compare it with
[Equation]
[Equation]
Here ˆ γ is interpreted as the optimal transport plan when the side information is completely satisﬁed, and (cid:104)
ˆ γ, C (cid:105) is the smallest transport distance under the constraint.
With the cost function parametrized as C ij ( θ )
=
c ( x ( i ) , y ( j ) ; θ ) , The optimal transport solution in Eq.
(9) is characterized by θ as
[Table]
[Table]
Update weights
9: end
Then we also expect a good cost function to make the distance under constraint (cid:104) ˆ γ, C (cid:105) to be as small as the lowest one (cid:104) γ ∗ , C (cid:105) as possible — optimize θ to minimize the loss
[Equation]
We refer to this method as mimic learning, because its objective is to make the ˆ γ mimic the cost performance of γ ∗ .
Note that γ ∗ ( θ ) is the optimal solution for any transport matrix in U ( µ X , µ Y ) .
That is, the optimal distance (cid:104)
γ ∗ ( θ ) , C ( θ ) (cid:105) ≤ (cid:104) γ, C ( θ ) (cid:105) for any γ ∈ U ( µ X , µ Y ) .
The equality holds true only when γ
= γ ∗ ( θ ) for the convex transport problem.
In the meantime, we have ˆ γ ( θ ) ∈
U ( µ X , µ Y ) ∩ U c .
Thus the loss L init ( θ ) is always larger or equal to 0.
When zero loss is achieved, we have γ ∗ ( θ ) =
ˆ γ ( θ ) , coinciding with the optimal solution for L ( θ ) in Eq.
Equation (10) describes the absolute difference between the two transport distances, but a relative difference is more desirable in practice to adjust for the scale of the objective function around (cid:104) γ ∗ , C (cid:105) .
For example, scaling the cost matrix C ( θ ) by a constant does not change the solutions γ ∗ and ˆ γ , but does scale the loss deﬁned in Eq.
(10) by the same constant.
We modify the loss to be invariant to such scaling:
[Equation]
Here ¯ γ is a uniform n X × n Y matrix used to stand for the averaged performance of random transport plans.
Eq.
(11) captures how close the distance under constraint (cid:104) γ ∗ , C (cid:105) is to the best one, compared to other random transports.
The mimic learning is approximately solved by alternating minimization.
As described in Algorithm 2, we iterate over two steps: (i) compute the value of γ ∗ and ˆ γ while ﬁxing θ ; (ii) take one gradient step with respect to L init ( θ ) with ﬁxed γ ∗ and ˆ γ .
The computation for optimal transport plans and the optimization of θ are carried out in alternating fashion.
The OT solver is used only to estimate the value of γ ∗ and ˆ γ in the ﬁrst step, requiring no gradient propagation.
Given such estimates of transport mappings γ ∗ and ˆ γ , the second step can be interpreted as learning a cost function which equates their transport costs, i.e., makes the behavior of ˆ γ mimic that of γ ∗ .
In the experiments, we set α = 1 and T init = 10 for initialization purpose.