Unbalanced Optimal Transport through Non-negative Penalized Linear Regression

1 Introduction . Optimal Transport ( OT ) theory provides powerful tools for comparing probability distributions and has been successfully employed in a wide range of machine learning applications such as supervised learning ( Frogner et al. , 2015 ) , clustering ( Ho et al. , 2017 ) , generative modelling ( Arjovsky et al. , 2017 ) , domain adaptation ( Courty et al. , 2017 ) , learning of structured data ( Maretic et al. , 2019 ; Vayer et al. , 2019 ) or natural language processing ( Kusner et al. , 2015 ) , among many others . One reason for those recent successes is the introduction of entropy-regularized OT that can be solved with the efficient Sinkhorn-Knopp matrix scaling algorithm ( Cuturi , 2013 ) . However , the classical OT problem seeks the optimal cost to transport all the mass from a source distribution to a target one ∗First two authors have equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ( Villani , 2009 ) , greatly limiting its use in scenarii where the measures have different masses or when they contain noisy observations or outliers . Unbalanced Optimal Transport ( UOT ) ( Benamou , 2003 ) has been introduced to tackle this shortcoming , allowing some mass variation in the transportation problem . It is expressed as a relaxation of the Kantorovich formulation ( Kantorovich , 1942 ) by penalizing the divergence between the marginals of the transportation plan and the given distributions . Several divergences can be considered , such as the Kullback-Leiber ( KL ) divergence ( Frogner et al. , 2015 ; Liero et al. , 2018 ) , the ` 1 norm corresponding to the partial optimal transport problem ( Caffarelli and McCann , 2010 ; Figalli , 2010 ) , or the squared ` 2 norm ( Benamou , 2003 ) . Regarding numerical solutions , Chizat et al . ( 2018 ) considered an entropicregularized version of UOT leading to a class of scaling algorithms in the vein of the Sinkhorn-Knopp approach ( Sinkhorn and Knopp , 1967 ) . The introduction of this entropic regularization improves the scalability of OT , but involves a spreading of the mass and a loss of sparsity in the OT plan . When a sparse transport plan is sought , the convergence is slowed down , necessitating the use of acceleration strategies ( Thibault et al. , 2021 ) . Regarding UOT with the ( squared ) ` 2 norm , Blondel et al . ( 2018 ) showed that the resulting OT plan is sparse and proposed to use an efficient L-BFGS-B algorithm ( Byrd et al. , 1995 ) to address this case . Note that the L-BFGS-B method can be used to solve UOT with differentiable divergences even without the entropic-regularization on the OT plan that induces the Sinkhorn-like iterations . Also note that , as for balanced OT , UOT can be solved more efficiently when the data has a specific structure , such as unidimensional distributions ( Bonneel and Coeurjolly , 2019 ) or distributions supported on trees ( Sato et al. , 2020 ) . Finally , recent work investigated UOT between Gaussians and provided closed form solutions for the regularized Janati et al . ( 2020 ) and unregularized ( Janati , 2021 , Eq . 2.72 ) versions of UOT associated with a KL divergence . Contributions . In this paper , we show after some preliminaries that UOT can be recast as a convex penalized linear regression problem with non-negativity constraints ( Section 2.2 ) . The main interest of this reformulation resides in the fact that non-negative linear regression has been extensively studied in inverse problems and machine learning , offering a large panel of tools for devising new numerical algorithms . Our reformulation involves a design/dictionary matrix that is structured and sparse . Leveraging this structure , we propose two new families of algorithms for solving the exact ( i.e. , without regularization of the plan ) UOT problem in Section 3 . We first derive in Section 3.1 a new Majorization-Minimization ( MM ) algorithm for solving UOT with Bregman divergences , and more specifically KL and ` 2-penalized UOT . The MM approach results in multiplicative updates that have appealing features : i ) they are easy to implement , ii ) have low complexity per iteration and can be instantiated on GPU , iii ) ensure monotonicity of the objective function and inherit existing convergence results . Our methodology is inspired by well-known algorithms in image restoration ( Richardson , 1972 ; De Pierro , 1993 ) and non-negative matrix factorization ( NMF ) ( Lee and Seung , 2001 ; Dhillon and Sra , 2005 ; Févotte and Idier , 2011 ) . Interestingly , the resulting multiplicative updates bear a similarity with the celebrated Sinkhorn scaling algorithm , with some key differences that are discussed . Next , we derive in Section 3.2 an efficient algorithm to compute the regularization path in ` 2- penalized UOT . To do so , we build on our proposed reformulation and more precisely on the fact that ` 2-penalized UOT can be reformulated as a weighted Lasso problem . We propose a new methodology inspired by LARS ( Efron et al. , 2004 ; Hastie et al. , 2004 ) , which , to the best of our knowledge , is the first regularization path algorithm for OT problems . It brings a novel understanding of the properties of the evolution of the support of OT plans , besides the practical interest of computing the complete regularization path when hyperparameter validation is necessary . Our new families of algorithms ( MM for general UOT , LARS for ` 2-penalized UOT ) are showcased in the numerical experiments of Section 4 . Python implementation of the algorithms , provided in supplementary , will be released with MIT license on GitHub . The connection between UOT and linear regression that we reveal in the paper opens the door to further fruitful developments and in particular to more efficient algorithms , thanks to the large literature dealing with non-negative penalized linear regression . We discuss those possible research directions in Section 5 , before concluding the paper . Notations . Vectors such as m are written with lower case and bold font , with coefficients mi or [ m ] i , according to context . The |A|-dimensional sub-vector with indexes in set A is written mA . Matrices such asM are written with upper case and bold font , with coefficients Mi , j . We introduce a vectorization operator defined by m = vec ( M ) = [ M1,1 , M1,2 , . . . , Mn , m−1 , Mn , m ] > , i.e. , the concatenation of the rows of the matrix , following the Numpy/C memory convention . 1n is a vector of n ones and M ≥ 0 denotes entry-wise non-negativity . Finally , Dϕ is the Bregman divergence generated by the strictly convex and differentiable function ϕ , i.e. , Dϕ ( u , v ) = ∑ i dϕ ( ui , vi ) =∑ i [ ϕ ( ui ) − ϕ ( vi ) − ϕ′ ( vi ) ( ui − vi ) ] . 2 Reformulation of UOT as non-negative penalized linear regression . 2.1 Background on Optimal Transport . Let us consider two clouds of pointsX = { xi } ni=1 and Y = { yj } mj=1 . Let a ∈ R+n and b ∈ R+m be two discrete distributions of mass onX and Y , such that ai ( resp . bj ) is the mass at xi ( resp . yj ) . The balanced OT problem , as defined by Kantorovich ( 1942 ) , is a linear problem that computes the minimum cost of moving a to b : OT ( a , b ) = min T≥0 〈C , T 〉 such that ( s.t . ) T1m = a , T > 1n = b ( 1 ) where 〈· , ·〉 is the Frobenius inner product , T ∈ R+n×m is the transport plan and C ∈ R+n×m is the cost matrix . The entry Ci , j of C represents the cost of moving point xi to yj . The Wasserstein 1- distance ( also known as the earth mover ’ s distance ) is obtained for Ci , j = ‖xi−yj‖ . The constraints on the transport plan T require that ‖a‖1 = ‖b‖1 and that all the mass from a is transported to b . These constraints can be alleviated through relaxation , leading to UOT ( Benamou , 2003 ) : UOTλ ( a , b ) = min T≥0 〈C , T 〉+ λ1Dϕ ( T1m , a ) + λ2Dϕ ( T > 1n , b ) . ( 2 ) The deviations from the true marginals are penalized by means of a given Bregman divergence Dϕ , as introduced in Chizat et al . ( 2018 ) , where λ1 and λ2 are hyperparameters that represent the strengths of penalization . Note that , in the case of ‖a‖1 = ‖b‖1 , balanced OT ( Eq . 1 ) is recovered when λ1 = λ2 → ∞ . Furthermore , when λ1 or λ2 → ∞ , we recover semi-relaxed OT ( Rabin et al. , 2014 ) . In practice , authors often set λ1 = λ2 = λ for UOT in order to reduce the necessity of hyperparameter tuning . Various divergences have been considered in the literature . The ` 1 norm gives rise to so-called partial optimal transport ( Caffarelli and McCann , 2010 ) . The squared ` 2 norm provides a sparse and smooth transport plan ( Blondel et al. , 2018 ) when introducing a strongly convex term in Eq . ( 2 ) . Chizat et al . ( 2018 ) derive efficient algorithms to solve Eq . ( 2 ) for several divergences by adding an additional regularization term λregDϕ ( T , ab > ) . In particular , entropic regularization is obtained when the KL divergence is used , promoting a dense transport plan unlike exact UOT . 2.2 Reformulation of UOT . UOT cast as regression . Let t = vec ( T ) , c = vec ( C ) and y > = [ a > , b > ] . Problem ( 2 ) can be re-written as min t≥0 Fλ ( t ) def = 1 λ c > t+Dϕ ( Ht , y ) ( 3 ) and as such be expressed as a non-negative penalized linear regression problem , where the design matrixH = [ H > r , H > c ] > is the concatenation of the matricesHr andHc that compute sums of the rows and columns of T , respectively ( see expressions in Section ? ? of the supplementary material ) . Note that , for the sake of simplicity , we consider here λ1 = λ2 = λ but this hypothesis could be easily alleviated for a given family of divergences ( see Sec . 5 for a discussion ) . Important features of Eq . ( 3 ) should be discussed . First , Fλ ( t ) is convex thanks to the convexity of Bregman divergences w.r.t . their first argument . Second , H is very structured and sparse ( with a ratio of only 1m+n non-zero coefficients ) which will allow for more efficient computations and updates than with a dense H . Finally , since t ≥ 0 and c ≥ 0 , the linear term can be expressed as 1λc > t = 1λ ∑ i citi = 1 λ ∑ i ci|ti| . This corresponds to a weighted ` 1 regularization , promoting sparsity in t and hence in the transport plans . Note that the “ sparse ” regularization is here controlled by 1λ ( instead of λ in classical penalized linear regression ) , meaning that the sparsity promoting term will be more aggressive for small λ . Solving problem ( 3 ) . Problems of the form of Eq . ( 3 ) are well-known in inverse problems and NMF . In inverse problems , t typically acts as a clean image degraded by operatorH ( e.g. , a convolution ) and noise . The data fitting term Dϕ ( Ht , y ) captures assumptions about the noise corrupting the observed image y. Sparsity is a common regularizer of t. In NMF , given a set of nonnegative samples { yl } one wants to learn a non-negative dictionary H and non-negative lower-dimensional embeddings { tl } such that yl ≈Htl ( Lee and Seung , 1999 ) . Updating the latter involves optimization problems of form ( 3 ) ( with or without sparse regularization ) . In contrast to problem ( 3 ) , the data fitting term is more commonly Dϕ ( y , Ht ) instead of Dϕ ( Ht , y ) in inverse problems and NMF . This is because the former is a log-likelihood in disguise for the mean-parametrized exponential family , and takes important noise models as special cases , such as Poisson , additive Gaussian or multiplicative Gamma noise ( Févotte and Idier , 2011 ) . Using such penalizations with reversed arguments would be possible in our case as well but we stick to the now standard formulation of ( Liero et al. , 2018 ; Chizat et al. , 2018 ) for simplicity . In the next section , we will first leverage a classical family of algorithms in inverse problems and NMF , namely MM , to obtain new algorithms for KL and ` 2-penalized UOT ( possibly with entropic regularization in the first case ) . Second , we will leverage results about non-negative Lasso to design an efficient algorithm to compute the regularization path of ` 2-penalized UOT .

Summary: Authors in this paper flattened the matrix-form unbalanced OT and reformulated it as a linear regression regularized by either the KL divergence or $\ell_2$ between the transport plan and two marginals. Then the problem converts to an MM problem and the authors derived the optimization steps for solving them, especially the one with $\ell_2$ in which case the authors proposed an algorithm to regress the optimal combinations of the transport plan and the regularization strength. Numerical results show that the new formulation leads to faster computation.