Learning transport cost from subset correspondence

1 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 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 first 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 defined 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 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 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 specific 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 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 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 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-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 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 on developing fast algorithms for efficient 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 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 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 flexibility to handle nonlinear artifacts and the side information can not be utilized . In contrast , a 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 θ 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.1 OPTIMAL TRANSPORT . Consider learning a mapping between two datasets X = { x ( 1 ) , ... , x ( nX ) } and Y = { y ( 1 ) , ... , y ( nY ) } . Here we use nX and nY to denote the number of datapoints ; each sample x ( i ) or y ( j ) could be a vector as well . We briefly recall the optimal transport framework in this setting . Given probability vectors µX and µY , the transport polytope is defined as U ( µX , µY ) : = { γ ∈ RnX×nY+ |γ1nX = µX , γT1nY = µY } , where 1nX ( 1nY ) is the nX ( nY ) dimensional vector of ones . Here the probability vector µX ( µY ) is in the simplex ∑ n : = { p ∈ Rn+ : pT1n = 1 } for n = nX ( nY ) . 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 ) 1n . Given a nX×nY cost matrix C , the classical optimal transport plan between µX and µY is defined as γ∗ = arg min γ∈U ( µX , µY ) 〈γ , C〉 , where 〈· , ·〉 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 nX and nY , the sizes of the datasets , exceed a few hundred . For computational efficiency , 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 ) : γλ = arg min γ∈U ( µX , µY ) 〈γ , C〉 − 1 λ h ( γ ) . ( 1 ) where λ > 0 is a regularization parameter and h ( γ ) = − ∑nX i=1 ∑nY j=1 γ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 Cij = c ( x ( i ) , y ( j ) ) . A good choice of the cost function is the key to influencing 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.2 SIDE INFORMATION . 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 | nX ≤ |S ( k ) Y | nY , 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 | nY ≤ |S ( k ) X | nX . Mathematically , this side information corresponds to the constraint that γij = 0 for ( i , j ) ∈ S0 , where S0 = m⋃ k=1 { ( i , j ) |i ∈ S ( k ) X , j /∈ S ( k ) Y , |S ( k ) X | nX ≤ |S ( k ) Y | nY } ∪ { ( i , j ) |i /∈ S ( k ) X , j ∈ S ( k ) Y , |S ( k ) Y | nY ≤ |S ( k ) X | nX } ( 2 ) 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 significantly aid cost function learning . We investigate how to learn the cost function using subset correspondence information between training datasets Xtrain and Ytrain . The learned cost function is evaluated via the mapping quality on the test datasets Xtest and Ytest . 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 .

Summary: The 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.