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: In this work, the authors propose to phrase the Unbalanced Optimal Transport problem under general Bregman divergence penalty as a non-negative penalized linear regression problem. They exploit an existing algorithm in Non-negative Matrix Factorization problem (NMF), namely Majorization-Minimization, in order to solve the obtained problem. In the case of KL-penalized UOT, the obtained algorithm is a variation of the Sinkhorn algorithm. However, this algorithm solves a different algorithm than Sinkhorn, namely KL-penalized UOT as opposed to OT with KL penalized transport plan.  For the l2-penalized UOT, the authors design an algorithm, based on LARS, that solves the problem for all values of lambda, from 0 to infinity.  The running time of different variants of their proposed algorithm are compared on a toy problem versus the current state-of-the-art method that uses BFGS. Overall, for l2-penalized UOT, the proposed method is significantly faster than BFGS. For KL-penalized UOT, it depends on the value of lambda.