Learning transport cost from subset correspondence
Authors: Ruishan Liu
Akshay Balsubramani
James Zou
Venue:  (2020)

Abstract: 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 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. 

Published as a conference paper at ICLR 2020
LEARNING TRANSPORT COST FROM SUBSET CORRE- SPONDENCE.
Ruishan Liu Department of Electrical Engineering Stanford University
Akshay Balsubramani Department of Genetics Stanford University
James Zou
Department of Biomedical Data Science Stanford University
ABSTRACT
Learning to align multiple datasets an important problem with many applications, and it is especially useful when we need to integrate multiple experiments correct for confounding. Optimal transport (OT) is a principled approach to align datasets, but key challenge in applying OT is that need to specify transport cost function that accurately captures how the two datasets are related. Reliable cost functions 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 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 common annotated cell type single-cell data from two batches. We develop 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-matchinn and single-cell RNA-seq, our method substantially outperform state-of-the-art benchmarks.
INTRODUCTION
In many applications, we have multiple related datasets from different sources or domains, and learning efficient computational mappings between these datasets is an important problem (Long let al.] 2017 Zamir et al. 2018). For example, 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 map cells between the datasets merge them, so that 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 first lab, there 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 principled analytical framework align heterogeneous datasets (San- Tambrogio 2015). It has been increasingly applied to problems in domain adaptation and transfer learning (Seuuy al. 2017] Genevay al. [017 Courty et al.) 2017b; Li et al.) (019). Optimal transport is an approach for taking two datasets, and computing mapping between them in the form of a "transport plan" y. The mapping optimal in the sense that among all reasonable mappings (precisely defined in Section 2), it minimizes the cost of aligning the two datasets. The transport cost 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 specified, then there are powerful methods for finding 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 Fuclidean other hand-crafted cost functions, which could give misleading mappings.

---

Published as a conference paper at ICLR 2020
Our contributions. We propose 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 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 be crudely annotated in both datasets, which is reasonable; we don't need to know that particular T cell should be mapped to another specific T cell. This gives rises to the side information that we can leverage in our algorithm mapped to another subset of points in dataset Y.
certain subset of points in dataset X should be
We present the first 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 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 efficiently 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-an methods for mapping datasets.
Related Work Optimal transport been well studied in the mathematics, statistics and optimization literature (Villami 2008) Courty et al.] 2017a Li et al. (2019 [Courty et al.] 2017b). OT can be used to define 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 significant work developing fast algorithms for efficient computation of the optimal transport plan (Cuturi 2013; Altschuler et al. [2117] Staib et al. (017). 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.) [0033] Bilenko et al.) (2044) and alignment methods (Ham et al. 2005 Wang & Mahadevan 2008; 2009) have realized benefits 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 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 AAvarez et al. 2019). GaN-based approaches have been used to align single-cell RNA-seq data when pairs of cells are known 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 flexibility to handle nonlinear artifacts and the side information cannot be utilized. In contrast, major benefit of our approach is its graceful adaptation to partial subset correspondence information, where we frame the problem as semi-supervised.
2 LEARNING COST METRICS
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 0 and adaptively learns using side information about the training data. The side information we consider is subset correspondence 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 prove generalizability.
common situation when some subsets of
2

---

Published as a conference paper at ICLR 2020
2.1 OPTIMAL TRANSPORT
y(ur)) }. Here we use Rix and ny to denote the number of datapoints; each sample x(()) or yot could be a vector well. We briefly recall the optimal transport framework in this setting. Given probability vectors and py, the transport polytope is defined as
Consider learning mapping between two datasets X = , x(1x)) and Y (,,,),,,,
UV == {> E R X h = -xx Iny py
where Lnx (Iay) ) is the nx (nz dimensional vector of ones. Here the probability vector μ. (ny) is in the simplex En == {p E R2 1, }} for n nx (ny). For two random variables with distribution μa and jay, the transport polytope U fyy represents the set of all possible joint probabilities of the two variables. In this paper, we consider μAx and jy to represent the empirical distributions of the samples X and Y, respectively, and set μr = (1/n)1n.
Given nx X wy cost matrix C. the classical optimal transport plan between jax and jy is defined as 7" = arg min (eU(xx by ) (7; C), where (.,)) denotes the Frobenius inner product. y" is also called coupling, Despite its intuitive formulation, the computation of this linear program quickly becomes prohibitive especially in the common situation when x and ny, the sizes of the datasets, exceed few hundred. For computational efficiency, Sinkhorn-Knop iteration is widely used to compute the optimal transport (Cuturi 2013). Sinkhorn-Knnn Knopp a fast iterative algorithm for approximately solving the optimization problem with entropy regularization Santambrogio (2015):
arg min (7,, C) - h(7).
(1)
TEU .Y)
7ij denotes the entropy. The regularized solution 7 converges to the classical one 7 when the regularization diminishes, i.e., X oc, with exponential convergence rate (Cominetti & San Martín 1994). The transport > treats X and Y symmetrically.
where A 0 is a regularization parameter and h(7) - ΣnN Ej= Yij Ty
The cost matrix is retrieved from the cost function C.j = c(x(0) ((())) ()) A good choice of the cost function is the key to influencing the learned mapping 7> 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.2 SIDE INFORMATION
Subset correspondence describes a common situation when certain subsets points are known be related. For example, images with the same objects should always be mapped together in domain adaptation tasks (Courty al. 2017b), while cells in a single-cell dataset need be aligned to those with the same cell type, where the cell type annotation is available.
Given m corresponding subsets, we write se()) and s((),
L,... m. denote the sets of data indices in the corresponding subsets, i.c., (t) ii E s and {y(() l s }. Note that s and
(k) could have different probability mass. If
|
my we take this side information to be that SX should be mapped into s(*) In other words, all the other entries of the transport matrix y that Mathematically, this side information corresponds to the constraint that Ta 0 for (i,)) E So. where
nx
maps s((k) to outside of s should be O. Everything is swapped if (g(
(s())
X
{ nx
S
S
So = (((.)) { X & S
< ny
PU((1)) & ( .J E S
x
nly
n1
1=1
(2)
Note that pair matching is extreme case of subset correspondence with subset size 1 that is, the exact pairwise relation is known. The pair matching problem has been addressed 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 significantly aid cost function learning.
3

---

Published as a conference paper at ICLR 2020
We investigate how to learn the cost function using subset correspondence information between training datasets Xrain and Yrain- The learned cost function is evaluated via the mapping quality on the test datasets Xest and Y fest. Note that the training and test sets 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.3 THE OT-SI ALGORITHM
such that the computed optimal transport Twain Our ultimate goal is to learn cost function such satisfies 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 ma- trix computed as Cij = ,(x(), = Eh k)) where i the data dimension. To learn the cost function systematically, we parametrize it as Cit(0) (0) = c(xx(*) y(²) 0) with weight 0. Here the function form e(-) can be chosen by users. To illustrate the improvement over the commonly-used (2) Euclidean cost, we parametrize c(x!"), ,);;;; 0) as polynomial in (a1 ..... Ld 3/1 ..... Vd, with co- efficients 0 and degree 2 for www-dimessional data. The Euclidean cost is equivalent to specific choice of 0g which is set as the initialization; see Appendix for more diseussions. For high-dimensional data, the memory required to store the second order polynomials becomes too large and we use fully connected neural network to parametrize c(x(1) ,)();; 0) with imput (z .....
(6)
...,,z/ od vf (()) and weights 0. Throughout this paper, the polynomial parametrization is used not specified.
The optimal transport solution is characterized by 0 as
(0)
arg min (7, C(0)) h(7).
(3)
-EU(xx
Then the problem can be formulated optimizing 9 to make the transport (0) approximately satisfy the conditions defined in Eq. @ penalizing deviation of the solution from these constraints with the loss
(0) = Σ
(4)
(o.
Theorem 1. For any X > O, the optimal transport plan (0): Rd - U(px, is O in the interior of its domain.
The infinite diffcrentiability of the Sinkhorn distance is mreviously-nnowe (Luise et al. 1 proves that the Sinkhorm transport plan also has this desirable property. Because Thm. guaarantces that th (0)) is infinitely differentiable, practice, we iterate Sinkhorn's 's update a sufficient number of times to converge to >> Each iteration matrix operation involving the cost matrix C(0). and when the number of iterations is fixed, we can propagate the gradient Va through all of the iterations using the chain rule. Updating 0 by one forward and backward pass has complexity of Onn)) 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 al. 2019: Alvarez 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 Algorithm and details are in the Appendix. The proof of Thm. [ mimic learning algorithm, details, and discussions about convergence are also in the Appendix.
2018); Thm.
are able optimize T. (,) in Eq. 44) by gradient descent. In
2.4 EXPERIMENTS SETUP
The OT-SI algorithm is carried out in Pytorch (Paszke al. 2017) and trained with GPU. The model is fitted 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 optimal transport with Euclidean cost function baseline for comparison, referred as "OT-baseline". We also compare our result with oftt--of-the-ar GAN-basea data alignment methods, MAGAN (Amodio & Krishnaswamy 2018) and CycleGAN (Zhu et al.
4

---

Published as a conference paper at ICLR 2020
Algorithm 1 OT-SI
Require: training datasets Xtrain and Yusin, corresponding subsets index SX S and YYruan Exain
s(k)
(A: I, ..., Marain), step size a, total training steps T, weights 0, from initialization proce- dure, Sinkhorn regularization parameter À and number of Sinkhorn iterations ninnkhorn-
1: nx = length(Xurain) ny = length(Yraiil) μx = 1 Py = (1/ny) lay
2: Initialize 0 0n
3: for 1= to do
4: Compute cost matrix C(0) with entries (1) = cxx(*) y0);
5: Solve Ttrain (0) arg min (eU( (n, C(0)) - 11((7) with Sinkhorn's update with
1lSinkhorn iterations.
6: Derive Venimin (0) by backpropagating the gradient through all Sinkhorn Sinkhorn-Knopp iterations.
7: Update weights 0 = 0 a (Li)))SSo tra (0) (0) V 7 trattin)))
8:
end for
2017). as well as the OT-based methods, RIOT (Li al. 2019) which developed for specific pair matching applications and the Procrustes-based OT (Grave 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 five comparison methods, OT-baseline, OT 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 the test (not shown to the al- gorithm), we evaluate a transport plan y by how much it satisfies the correspondence. Mathematically, define subset matching accuracy:
k= EiES(M EsEsyn)
Vightest
Accuracy
Neem
Yiest
E k= min min(1s Xincs (onool////n 1/nxxwss
(s(A) Yoses 1/NYica
From the definition, 0 & Accuracy E 1 gives the probability of mapping the correct corresponding subsets. When all the test datapoints are mapped into the correct subsets, the accuracy is 1; when all the data 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.
3 BENCHMARK ON SYNTHETIC DATASETS
We first 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 al. 2013 Courty et al. 2017b). The dataset is simulated with two domains, source and target. As shown in Fig. La 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, of each moon. We set the parameter À = 10 and the number of Sinkhorn-Knon 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).
generate the training, test and validation datasets with 100, 100, and 50 samples 200.
iterations
Baseline performance. The optimal transport plan under Euclidean cost function is depicted in Fig. 1b Datapoints learned be matched are connected by solid curves. The red curves indicate wrong
5

---

Published as a conference paper at ICLR 2020
Table 1: Subset matching and pair matching accuracy on data for two moon datasets. Here the subset (pair) matching accuracy corresponds to the proportion of the data points that are aligned to the correct moon (data points) on the test set. Higher better. We generated 10 independent datasets and the standard deviation is shown.
Side
OT:
OT-SI
OT-Procrustes
MAGAN
CycleGAN RIOT
Information
Baseline
Subsets
72%
100%
59% + 14%
N/A
48%
N/A
Subset Matching
1 Pair 10 Pairs
72%
93% = 1%
59% + 14%
46% = 2% 71% = 2%
48%
N/A
72%
99.7% + 0.1%
59% + 14%
48%
46.9%
Subsets
2%
87% = 1%
46% + 14%
N/A 0.94% 0.08%
0%
N/A
Pair Matching
1 Pair 10 Pairs
2%
53% = 5%
46% + 14%
0%
N/A
2%
86% = 2%
46% + 14%
1.2% + 0.2%
0%
0%
-19
-18
-14
-30
(a) Data
(b) OT-baseline.
(c) OT-SI
(d) OT-Procristes.
datasets and the optimal transport result. The target domain (blue) is built by adding noise to the source (orange) and rotating by 60 degrees. Corresponding subsets denoted by circles and crossings. (b-e) Optimal transport plan under (b) Euclidean cost (OT-baseline) (c) learned cost function by OT-SI and (d) Procrustes-basee OT. Points learned to be matched are connected by solid curves. When a datapoint matched to the wrong subset, i.e, the other moon, the connection curve is colored by red.
Figure 1: Illustration of the
transports which map the data into the wrong subset, i.c., the other moon. As shown by Fig. lb 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 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 and only 2% the data are matched to their corresponding pairs, as given Table
Subset correspondence. We first evaluate our method OT-SI when the side information only that
the data on the corresponding moons are known to be related i.e. s and s are the two moons in dataset one, and s(()) and s(2) are the two moons in dataset two. With the learned cost function, almost all the datapoints are mapped into the corresponding moon, i.c., the subset matching accuracy on the test data achieves 100% for both methods, as shown in Table ! The results are averaged for 10 independent
Interestingly, although there is pair matching information provided during training, the learned cost function significantly 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. Ic The rotation property of the datasets is correctly captured. In contrast, OT-Procrustes sometimes learns as good as OT-SI, similar Fig. Ic but sometimes mistakenly learns rotation as flipping, indicated by Fig. Id 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 and 10 pairs are known out of the total 100 training pairs. The matching pairs are randomly selected from the training data. OT-SI significantly outperforms all four bench methods particularly when the number of known pairs is very limited. Even when only matching pair is provided, the learned cost function greatly improves the OT performance, as given in Table [ The improvement here is largely attributable to the unlabeled data, i.e., the datapoints without any pair

---

Published as a conference paper at ICLR 2020
Table 2: Subset matching and pair matching test accuracy for the alignment of protein and RNA expression data in CITE-seq CBMCs experiment. Here subset (pair) matching accuracy denotes the proportion of the cells that are aligned to the correct cluster (cells). Higher is better. For alignment in the original space, the expression of 5001 mRNAs are mapped to the expression of 13 proteins. For alignment the embedding space, the first 10 principal components are used for both RNA and protein expression datasets.
Embedding Space (10:10)
Original Space (5001:13)
OT-Baseline
OT-SI OT-Procrss
OT-Baseline
OT-ŞI
OT-Procrustes
Subset Matching 9.9%
56.1%
44.3%
43.9%
N/A
N/A
Pair Matching
0 %
3.2%
1.0%
0.8%
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 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.
4 BIOLOGICAL MANIFOLD ALIGNMENT
In this section, implement our method OT-SI to learn 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.1 ALIGNMENT OF PROTEIN AND RNA SEQUENCING DATA
How to align datasets with different data types has been major topic in many fields. 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 al. 2017). The dataset is subset 8,005 human cells with the expression of 5,001 highly variable genes and 13 proteins. Fifteen clusters are identified 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 learned in the original data space with the expression of 5001 mRNA and 13 proteins for each cell, we use fully connected neural network to parametrize the cost function, with two hidden layers of 100 and neurons.
We align the RNA and protein expression datasets in two scenarios: i) the embedding space where we use the first 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 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.
the effect and relation between
7

---

Published as a conference paper at ICLR 2020
Table 3: Subset matching accuracy on the two held-out cell types, T cell and immature T cell, for aligning FACS and droplet data. The algorithms OT-SI and OT-PPrccrustes are trained on 2, 5 and 8 other cell types. Higher is better.
#Training Cell Types
2
(OT-Baseline)
OT-SI
70.0%
75.0%
83.8%
70.0%
ProPrucss
36.9%
56.8%
61.5%
Table 4: Subset matching accuracy on the two held-out digits 3 and 5 when trained our model on the rest eight digits for MNIST dataset. OT is used to align the original images with the perturbed ones.
original
watering
swirl
sphere
flip
9
1
P
I
2
3
3
S
EE
OT-baseline OT-SI
70.0%
62.5%
57.5%
35.0%
72.5%
75.0%
65.0%
57.5%
OT-Procrustes
66.3%
48.6%
39.4%
83.9%
4.2 BATCH ALIGNMENT
Another fascinating biological application of optimal transport is to align data from different batches. In biological studies, the samples processed or measured different batches usually result in non- biological variations, known as batch effect (Chen et al.) 201 Haghverdi et al. (0118). Here we use OT to align two batches of data data collected with fluorescence activated cell sorting (FACS) and droplet methods. In this case, the cell types are used as the subset correspondence information. For illustration purposes, we subsample the top 10 celltypes with 400 samples of each. There are 1,682 genes after filtering and the first 10 principal components
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 rypes that were used to train the cost. This is 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 all of the data. To do this, we choose held-out and train on the rest. Among all the ground truth expert annotations, T cell and immature T cell are most difficult to aligned. With OT OT-baseline, only 70% cells are mapped into the correct cell types. Substantial improvement is achieved, as shown in Table 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.
cell types T cell and immature T cell as
EXPERIMENTS ON IMAGES AND TABULAR DATA ALIGNMENT
5.1 IMAGE ALIGNMENT FOR NEW DIGITS
To illustrate the use of OT-SI in image alignment, we use it to learn cost metrics in aligning images with partial annotation the MNIST dataset which contains 28 x 28 images of handwritten digits. images from each digit class and split them into training, validation, and test sets with ratio 50%, 20% and 30%. For illustration purposes, the alignment analysis. We generate four different types of perturbations to the original images, as plotted in Table Then we align the original images with the perturbed ones using OT, respectively.
We subsample 200
use the first ten principal components for
8

---

Published as a conference paper at ICLR 2020
Table 5: Root mean square error (RMSE) and mean absolute error (MAE) of pair matching algorithms marriage-matching
for marriage- matching dataset. Lower is better.
Random
PMF
SVD
itemKNN
RIOT
FM
OT-SI
RMSE MAE
54.7
77.8
109.0
2.4
2.4
9.5
2.4
36.5
36.1
62.0
1.6
1.5
7.5
1.5
The experiment setting is the same in Sec. B Similar Sec. 4
test how well algorithm generalizes to new classes that were not used in learning the cost function. In the experiment, we hold digits 3 and 5, and demonstrate the metric learned the other eight digits can help the alignment of digits and 5. We achieve consistent improvement over the baseline on all four distribution types, indicated by Table 4
5.2 MARRIAGE DATASET FOR PAIR MATCHING
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 finally benchmark 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 & Salakbutdinov 2008). item- based collaborative filtering model (itemKNN) (Cremonesi 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. (2019). We follow the same experimental protocol as in [i 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 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 When used for pair matching purposes, OT-SI report comparable performance dataset was a primary motivating dataset used to design RIOT (Li et al. (2019). and therefore we expect RIOT to perform very well for this task.
state-of-the-an matching algorithms. Note that this marriage
DISCUSSION
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, the best of our knowledge. Previous works rely on more restricted information such as that specific pairs of points should be aligned. In settings such as genomics and images, it is often difficult to say that 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 flexible and principled method to learn the transport cost with side information. Experiments demonstrate that they work significantly better than state-of-the-art methods when indeiinformation is very limited, which is often the case. We compare against state-of-the-ar 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 efficiently using all of the unannotated datapoints that are in any pairs or subsets. These unannotated data act as regularization (similar to in semi-supervised leerringg), which enables the model avoid overfitting 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.
REFERENCES
Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. In Advances in Neural Information Processing Systems, pp. 1964-1974, 2017.

---

Published as a conference paper at ICLR 2020
David Alvarez-Melis, Tommi Jaakkola, and Stefanie Jegelka. Structured optimal transport. In International Conference on Artificial Intelligence and Statistics, pp. 1771-1780, 2018.
David Alvarez-Melis, Stefanie Jegelka, and Tommi S Jaakkola. Towards optimal transport with global invariances. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1870-1879, 2019.
Matthew Amodio and Smita Krishnaswamy. Magan: Aligning biological manifolds. In International Conference on Machine Learning, pp. 215-223 2018.
Mikhail Bilenko, Sugato Basu, and Raymond J Mooney. Integrating constraints and metric learning in semi-supervised clustering. Proceedings of the awenty-first international conference on Machine learning, pp. 11. ACM, 2004.
Andrew Butler, Paul Hoffman, Peter Smibert, Efthymia Papalexi, and Rahul Satija. Integrating single-cell transcriptomic data across different conditions, technologies, and species. Nature biotechnology, 36(5):411, 2018.
Chao Chen, Kay Grennan, Judith Badner, Dandan Zhang, Elliot Gershon, Li Jin, and Chunyu Liu. Removing batch effects analysis of expression microarray data: an evaluation of six batch adjustment methods. PloS one, 6(2):e17238, 2011.
Yunjey Choi, Minje Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, and Jaegul Choo. Star- gan: Unified generative adversarial networks for multi-domain image-to-image translation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 8789-8797, 2018.
Roberto Cominetti and Jaime San Martín. Asymptotic analysis of the exponential penalty trajectory in linear programming, Mathematical Programming, 67(1-3):199-187, 1994.
Nicolas Courty, Rémi Flamary, and Mélanie Ducoffe. Learning wasserstein embeddings. arXiv preprint arXiv:1710.0744 2017a
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853 1865, 2017b.
Paolo Cremonesi, Yehuda Koren, and Roberto Turrin. Performance of recommender algorithms top-n recommendation tasks. In Proceedings of the fourth ACM conference on Recommender systems, pp. 39-46. ACM, 2010.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-232220 2300, 2013.
Marco Cuturi and David Avis. Ground metric learning. The Journal of Machine Learning Research, 15(1):533-564, 2014.
Joel Franklin and Jens Lorenz. On the scaling of multidimensional matrices. Linear Algebra and its applications, 14:7177773,, 1989.
Alfred Galichon and Bernard Salanié. Matching with trade-offs: Revealed preferences over competing characteristics. 2010.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. arXiv preprint arXiv: 1706. 00292, 2017.
Pascal Germain, Amaury Habrard, François Laviolette, and Emilie Morvant. A pac-bayesian approach for domain adaptation with specialization to linear classifiers. In International conference on machine learning, pp. 738-746. 2013.
Edouard Grave, Armand Joulin, and Quentin Berthet. Unsupervised alignment of embeddings with wasserstein procrustes. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1880-1890, 2019.
10

---

Published as a conference paper at ICLR 2020
Laleh Haghverdi, Aaron TL Lun, Michael D Morgan, and John C Marioni. Batch effects in single-cell massequenncng data are corrected by matching mutual nearest neighbors. Nature biotechnology. 36(5):421, 2018.
Jihun Ham, Daniel D Lee, and Lawrence Saul. Semisupervised alignment of manifolds. In AISTATS, pp. 120-127, 2005.
Yehuda Koren, Robert Bell, and Chris Volinsky, Matrix factorization techniques for recommender systems. Computer, (8):30-37, 2009.
Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha. Learning to match via inverse optimal transport. Journal of Machine Learning Research, 20(80):1-37, 2019.
Mingsheng Long, Han Zhu, Jianmin Wang, and Michael Jordan. Deep transfer learning with joint adaptation networks. In Proceedings of the 34th International Conference on Machine Learning-Volume Volume 70, pp. 2208 2017.221..00 JMLR. org, 2017.
Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. In Advances in Neural Information Processing Systems, pp. 5859-5870, 2018.
Andriy Mnih and Ruslan R Salakhutdinov. Probabilistic matrix factorization. Advances in neural information processing systems, pp. 1257-1264, 2008.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vito,o Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Anand Rangarajan, Haili Chui, and Fred L Bookstein. The softassign procrustes matching algorithm. In Biennial International Conference on Information Processing in Medical Imaging, pp. 29-42. Springer, 1997.
Steffen Rendle. Factorization machines with libfm. ACM Transactions on Intelligent Systems and Technology (TIST), 3(3):57,2013
Filippo Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling, volume 87. Birkhäuser, 2015.
Nicholas Schaum, Jim Karkanias, Norma F Neff, Andrew P May, Stephen R Quake, Tony Wyss- Coray, Spyros Darmanis, Joshua Batson, Olga Botvinník, Michelle B Chen, et al. Single-cell transcriptomics of 20 mouse organs creates a tabula muris: The tabula muris consortium. Nature, 562(7727):367, 2018.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint 17X0021711.02283, 2017.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-344 1967.
Matthew Staib, Sebastian Claici, Justin M Solomon, and Stefanie Jegelka. Parallel streaming wasserstein barycenters. In Advances in Neural Information Processing Systems, pp. 2647-2658,
2017.
Marlon Stoeckius. Christoph Hafemeister, William Stephenson, Brian Houck-Loomis, Pratip Chattopadhyay, Harold Swerdlow, Rahul Satija, and Peter Smibert. Simultaneous epitope and transcriptome measurement in single cells. Nature methods, 14(9):865, 2017.
Vayer Titouan, Nicolas Courty, Romain Tavenard, Chapel Laetitia, and Rémi Flamary. Optimal transport for structured data with application graphs. In International Conference Machine Learning, pp. 6275-6284, 2019.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
11

---

Published as a conference paper at ICLR 2020
Chang Wang and Sridhar Mahadevan. Manifold alignment using procrustes analysis. In Proceedings of the 25th international conference on Machine learning ICML '08, pp. 1120-1127, New York, New York, USA. jul 2008. ACM Press. ISBN 9781605582054, doi: 10.1145/1390156.1909 13902297 httl:ttp::/poot
acm. org/ httgg/iita ctt.t on c cfm.mmmmm..mmmmmmmmmmmm www.fmoooiiiioocccoccoooccooo 30 391166 13902997
Chang Wang and Sridhar Mahadevan. Manifold alignment without correspondence. In IJCAI, volume 2, pp. 3, 2009.
Eric Xing, Michael Jordan, Stuart Russell, and Andrew Ng. Distance metric learning with application to clustering with side-information. In Advances in neural information processing systems, pp. 521-528, 2003.
Amir R Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese.
Taskonomy: Disentangling task transfer learning. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 3712-3722, 2018.
Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. Proceedings of the IEEE international conference on computer vision, pp. 2223-222222222 2223,2232, 2017.
12

---

Published as a conference paper at ICLR 2020
A COST FUNCTION
In Sec. 2.3 the cost function parametrized as Cir(0) ,(0); with weight 0. For low-
dimensional data, choose c(x!1) ,y0; polynomial in (at Ld v1 Ddy with coefficients 9 and degree
(0) (5)
de dyy
(O) = Σ
On.
III )a) ( k=1/=1
(5)
a, 0.330
1g1
Here do not require the dimension of the two datasets de and dy to be the same. When the two datasets X and 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(A) as a small fully connected neural network and achieved very similar performance.
The Euclidean cost is equivalent to a specific choice of 0, when the two datasets have the same feature space du dy = d. The entry of the cost matrix computed as Ci = , = Ei (2 + (y) ) (()) )
B OT-SI ALGORITHM
The Lagrangian dual of Eq. 3 is
max u My
A(C;,(0)-a4-2)
(6)
By Sinkhorn's 's scaling theorem (innkhorn Knopp 1967 the optimal transport plan (^ is computed as
(0) diag(e" (aaag(""""u (0))e xC(0) diag(e (o)))))
where u² (0) and o0 (0) are the solutions to the dual problem in Eq. (6)
B. PROOF OF THEOREM
The analysis for Theorem follows the strategy of the proofs of Theorem 2 in (Luise et al. 2018
Theorem 1, For any 0, the optimal transport plan > (0): Re - ((xx;Hy) is y in the interior of its domain.
Proof. Based on Eq. 7 the optimal transport plan (0) is smooth function when u (0), o² (0)
and C(0) are smooth. In the meantime, the cost C(0) is linear function of 0, as indicated by Eq. (5).
Thus to prove the smoothness of (0). we only need to demonstrate u (() and v (0) are both smooth in 0.
Here we define
nx By
u) uTD
ΣΣ
(0
The dual problem in Eq. (6) becomes min,, a (0: u, v). From the definition, (0; u, v) is smooth
and strictly convex in (u, v). Note that C(0) is linear in B. Then for any fixed A in the interior of Rdo there exits (u* (0), (0)) such that (0; u² (0), oz (0)) = min o u, v). The function V (u, (0; u, v) E 0000 due to the smoothness of a (0; u, v). Now we fix (wo, Vo, vo, 20) such that
(u, (0o; uo, vo) 0. The strict convexity of (0; v) ensures that V Pa. (000) uo. vo) is invertible.
From implicit function theorem, we can always find a function f and subset Un, E R such that f(0)) (wo, vo): ii) (0, (()) 0 for any 0 E Ung iii) f E (Co(U))) That is, f (() is
13

---

Published as a conference paper at ICLR 2020
stationary point of the function for any 0 in Ud.- Together with the strict convexity of a, derive
f() = (u" (0), t (f)). Recalling f E C (Up,). : we prove that (u" (0), o" (0)) is O in the interior of its domain.
B.2 CONVERGENCE PROPERTIES
The gradient of (0) in Eq. 7 is computed
V Vo (0) >n (0) (V) + Vo0 - VoCi)(0))
The convergence of scaling factors ( and evv()) is linear (i.e. exponential in NiSinkhhrn). with the bounded rate given by (Franklin & Lorenz 1989): experiments use linear C(0), and we do not find 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 r, a relevant open question.
B.3 COMPUTING GRADIENT
In practice, we iterate Sinkhorn's 's update N times to converge to >>. More specifically, the dual solutions u² (0) and v² (0) for Eq. (6) are computed by iterating for N times
u - Mx/E 2
v uy/ (ey/(e-c(0)yy
For the optimal plan (0) given in Eq. 77. its gradient Vo) is obtained by propagating through all the iterations using the chain rule and implemented by Pytorch. Thus updating 0 by one forward and backward pass has complexity of (n²²) up to logarithmic terms. OT-SI has the same complexity as the Procrustes based OT methods which alternatively optimize over coupling matrix and lincar transformation matrix. Taking the image alignment for watering MNIST dataset as example, the running time for OT-SI and OT-Procrustes 229.5s and 196.7s, respectively.
C MIMIC LEARNING FOR INITIALIZATION
In this section, we derive a mimic learning an fast initialization method to boost the performance and accelerate the learning. While Algorithm requires 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 7" arg min (E .y:) C') as example.
As discussed in Sec. 2.3 our ultimate goal is to learn a cost function, such that the optimal transport 7 satisfies the side information defined in Eq. (2) as faithfully possible. From another perspective, we force an additional set constraints on the transport plan y to fulfill the condition in Eq. (2):
((()
S
Uh = n (ring = 0, (()) $
S(K)
nx
1=1
(8)
|
((k)
n (717 0,i ((k) ,j E s
To quantify how much the learned y" Eq. I follows the side information, we compare it with
arg min vEU by jnt
(7, C).
Here is interpreted as the optimal transport plan when the side information is completely satisfied, and (f: C*) is the smallest transport distance under the constraint. With the cost function parametrized
as (0) = c(x(", c(x(9.00);8) ()); The optimal transport solution in Eq. (9) is characterized by as (0) arg min yEU(u))) (n; C(0)).
14

---

Published as a conference paper at ICLR 2020
Algorithm 2 Mimic Learning for Initialization
Require: training datasets Xtrain and Yuain. corresponding subsets index Xxsain
s(k) and ((k)
Yerain (k
1, M.rain) step size a, total steps Tiiit optimal transport solver OTSolver.
1: nx length(Xtraiu). ny length(Yuraill μxx = (1/nx) I . My (1/ny) Ayy
2: Define transport polytopes U = ('xxx,y and constraints U from Eq. (
3: Initialize 0 such that the cost function is equivalent the Euclidean cost.
4: for = 1 to Tint do
5: Compute cost matrix C as Ctj = c(z(9) y())
6: ? = O'SSolvveeCC.. (C, U*), y OTSoolver (y
7: Limit (0) (A,C(0) (x".C(0))
(oC(0))((".c0 Update weights = 8 VV (8)
8:
9:
end for
Then we also expect a good cost function to make the distance under constraint (9/ C) to be small as the lowest one (n", C?) as possible optimize A to minimize the loss
Limit (0) = (F(O),((0) - (> C())
(10)
We refer to this method as mimic learning, because its objective to make the 7 mimic the cost performance of y
Note that (0) the optimal solution for any transport matrix in ()xx,,y) That is, the optimal distance (0), C(O)) < (7, C(O)) for any 7 E ((xx,Wy) The equality holds true only when y = (0) for the convex transport problem. In the meantime, we have (0) E U(Hx, py) Thus the loss Liuit (0) is always larger or equal to O. When zero loss is achieved, we have (0) (().
coinciding with the optimal solution for [() in Eq. 3
Equation 10) describes the absolute difference between the two transport distances, but relative difference is more desirable in practice to adjust for the scale of the objective function around (7", C). For example, scaling the cost matrix C(0) by a constant does change the solutions 7 and 5, but does scale the loss defined in Eq. (10) by the same constant. We modify the loss to be invariant to such scaling:
Linie (0) (()) (x))
(11)
Here 7y is a uniform nx X ny matrix used to stand for the averaged performance of random transport plans. Eq. (IT) captures how close the distance under constraint (y*, C) 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 y" and f while fixing #; (ii) take one gradient step with respect Linit (0) with fixed y* and f. 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 y and y in the first step, requiring no gradient propagation. Given such estimates of transport mappings 7 and f, the second step can be interpreted as learning a cost function which equates their transport costs, i.e., makes the behavior of f mimic that of 7 In the experiments, we set a 1 and Tinit 10 for initialization purpose.
15


Reviews:
[{"review_id": "r1lYrm46FS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors consider the alignment problem for multiple datasets with side information via entropic optimal transport (Sinkhorn). The authors formulate it as a transport cost learning in optimal transport framework with constraints giving by side information. Empirically, the authors illustrated the effectiveness of the proposed approach over state-of-the-art on several applications (e.g. single-cell RNA-seq, marriage-matching, MNIST with its perturbed one. \n\nThe paper show good results of the proposed method for dataset alignment on several applications. However, the motivation of the dataset alignment seems quite weak (e.g. it is unclear the advantage for jointly analysis when one merge datasets by alignment). The side information is not clearly easy to obtain in practice, since it simply enforces the constraints on the transportation polytope (i.e. how one knows that which subset of a given class C_A on dataset A is matched with a certain subset of another class C_B on dataset B, it seems easy to obtain the side information about matching over the whole class C_A with C_B, and it seems nontrivial to obtain the subset matching). Additionally, some important information about the proposed method is not in the main text but in the appendix which makes the paper hard to read without checking appendix constantly (especially the main proposed method). Overall, I lean on the negative side.\n\nBelow are some of my concerns:\n\n1) As stated above, I am curious how one can obtain the side information about the subset matching? For example, in Figure 1, it seems trivial to have a side information about the whole class matching and quite nontrivial to obtain a matching for subsets of any given class.\n\n2) The side information enforces the constraints on the feasible set of the transportation polytope. It seems to me that one can solve the variant of optimal transport (or Sinkhorn) with those extra constraints directly, it is unclear to me how this kind of side information can navigate the transport cost learning. Moreover, given a new pairs of input, the proposed method needs to optimize the transport cost again which make the benefits to learn to transport cost quite limited. Overall, a trivial baseline for the proposed method is to solve directly optimal transport with those extra constrains (from side information) directly (maybe with the Euclidean cost metric), but not the optimal transport without extra constraints as in the current setting. \n\n3) The authors may need to consider to reorganize the presentation. Much important information of the proposed method is not in the main text, but in the appendix which make the paper hard to follow without checking the appendix constantly. The derivation of gradient in line 6 of Algorithm 1 is not trivial. It is better if the authors give more detail information and discussion. I am quite confused about the idea to back propagate the \"gradient\" through Sinkhorn-Knopp iterations?\n\n4) It seems that the complexity of gradient update of Algorithm 1 is quite high. The authors should give a comparison about time consumption in the experiments (besides the accuracy matching)\n\n5) It seems that there is nothing in the code sharing folder?\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SklG_8hatS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary\nThis paper proposes a new way to learn the optimal transport (OT) cost between two datasets that can utilize subset correspondence information. The main idea is to use a neural network to define a transport cost between two examples instead of using the popular Euclidean distance and have a L2 regularizer that encourages the resulting transport plan to be small if two examples belong to the same subset. The results on synthetic datasets (2D, MNIST image) and real-world datasets (RNA sequence) show that the learned transport cost outperforms the Euclidean distance and Procrustes-based OT. \n\n# Originality\n- This paper considers a new problem which is to learn OT given \"subset correspondence\" information. The prior work does not have a mechanism to utilize such information or requires more side information such as pair-wise matching. \n\n# Quality\n- The idea of parameterizing the transport cost using a neural network is novel to my knowledge. The idea of regularizing the transport cost based on the subset correspondence information sounds sensible.\n- The empirical result is quite comprehensive and convincing in that it includes from illustrative examples to real-world datasets (RNAs) and includes many relevant baselines.\n\n# Clarity\n- The paper is well-organized and well-written.\n- Since this paper considers a new problem (learning OT with subset correspondence data), it would be great to motivate and explain the problem better in the introduction. RNA-seq examples in the introduction were not easy to parse as I do not have any background on biology. It would be better if the paper presented more illustrative/intuitive examples ideally with a figure that describes the problem. \n- \"an principled\" -> \"a principled\"\n\n# Significance\n- This paper seems to present a neat and sensible idea for learning OT with neural networks given subset correspondence information. The problem and method are new, and the empirical results look comprehensive. One minor issue is that it is unclear how significant this new problem setup is (learning OT cost with subset correspondence). Although the paper presents real-world examples (RNAs), the paper could be stronger if it presented more real-world examples/results and emphasized that this is an important problem to solve. ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "BJx8MZsRFS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper presents a gradient-based method for learning the cost function for optimal transport, applied to dataset alignment. The algorithm is based on the Sinkhorn-Knopp iterative procedure for approximating the optimal transport plan. The current paper proves that the Sinkhorn transport plan with a parameterized cost function is infinitely differentiable. The cost function can thus be optimized by unrolling the Sinkhorn iteration for a fixed number of steps and then backpropagating through the iterative updates. The cost function is optimized via a loss function derived from side information, specifically subsets of the two datasets to be aligned that contain elements that should be matched in the optimal transport plan. Experiments on a variety of synthetic and real datasets show that the proposed method is often better than, and almost always competitive with, other modern approaches.\n\nI recommend that the paper be accepted. The paper presents a well-justified and seemingly novel approach to solving an important problem, and demonstrates empirically that it compares favorably to other approaches. The proposed technique seems to be a fairly small extension to existing work, but together with its analysis and experiments it is sufficiently novel.\n\nDetails / Questions:\n* An interesting point of comparison would be methods for semi-supervised metric learning that are not necessarily tailored to the OT setting. E.g. [1] (picked fairly arbitrarily)\n\n* Is it possible to adapt algorithms that need pairwise matches to your subset matching setting by e.g. choosing pairs in the matching subsets at random and calling them matching pairs?\n\n* Do you have any intuitions for the pattern of results in Table 4? Why might flips be a difficult transforation for OT-SI to learn?\n\nTrivia:\n* Page 6 is blank when I print the paper. Not sure if anyone else has this problem.\n\nReferences:\n[1] Bilenko, M., Basu, S., & Mooney, R. J. (2004, July). Integrating constraints and metric learning in semi-supervised clustering. In Proceedings of the twenty-first international conference on Machine learning (p. 11). ACM.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes an algorithm for learning a transport cost function that accurately captures how two datasets are related by leveraging side information such as a subset of correctly labeled points. The reviewers believe that this is an interesting and novel idea. There were several questions and comments, which the authors adequately addressed. I recommend that the paper be accepted.
Decision: Accept (Poster)