Unbalanced Optimal Transport through Non-negative Penalized Linear Regression
Year: 2022

Abstract:
This paper addresses the problem of Unbalanced Optimal Transport (UOT) in which the marginal conditions are relaxed (using weighted penalties in lieu of equality) and no additional regularization is enforced on the OT plan. In this context, we show that the corresponding optimization problem can be reformulated as a non-negative penalized linear regression problem. This reformulation allows us to propose novel algorithms inspired from inverse problems and nonnegative matrix factorization. In particular, we consider majorization-minimization which leads in our setting to efficient multiplicative updates for a variety of penalties. Furthermore, we derive for the first time an efficient algorithm to compute the regularization path of UOT with quadratic penalties. The proposed algorithm provides a continuity of piece-wise linear OT plans converging to the solution of balanced OT (corresponding to infinite penalty weights). We perform several numerical experiments on simulated and real data illustrating the new algorithms, and provide a detailed discussion about more sophisticated optimization tools that can further be used to solve OT problems thanks to our reformulation.

heading: 1 Introduction
text: 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)].

heading: 2 Reformulation of UOT as non-negative penalized linear regression

heading: 2.1 Background on Optimal Transport
text: 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.

heading: 2.2 Reformulation of UOT
text: 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.

heading: 3 Novel numerical solvers for UOT

heading: 3.1 Majorization-Minimization (MM) for UOT
text: General MM framework. MM algorithms have been around a long time in inverse problems and NMF to solve problems of form (3). Classical algorithms for NMF such as (Lee and Seung, 2001) have built on seminal MM algorithms for inverse problems such as (Richardson, 1972; De Pierro, 1993). Subsequent works in NMF such as (Dhillon and Sra, 2005; Févotte and Idier, 2011; Yang and Oja, 2011) have further contributed novel MM algorithms for larger classes of problems, including larger families of divergences. In a nutshell, MM consists in iteratively building and minimizing an upper bound of the objective function which is tight at the current parameter estimate (and referred to as auxiliary function), see Hunter and Lange (2004); Sun et al. (2017) for tutorials. In NMF, a common approach consists of alternating the updates of the dictionary H and of the embeddings. In our case,H is fixed and we may use the results of (Dhillon and Sra, 2005) to build an auxiliary function for term Dϕ(Ht,y), to which we may simply add the linear term c>t/λ to obtain a valid auxiliary function for Fλ(t). Let t̃ denote the current estimate of t, Z̃i,j =
Hi,j t̃j∑ lHi,l t̃l and
Gλ(t, t̃) = ∑ i,j Z̃i,jϕ
( Hi,jtj
Z̃i,j ) + ∑ j [ cj λ − ∑ i Hi,jϕ ′(yi) ] tj + cst, (4)
where cst = ∑ i[ϕ ′(yi)yi − ϕ(yi]. Then, Gλ(t, t̃) is an auxiliary function for Fλ(t), i.e., ∀t, Gλ(t, t̃) ≥ Fλ(t) and Gλ(t̃, t̃) = Fλ(t̃). Let t(k+1) = argmint≥0Gλ(t, t(k)), then Fλ(t(k)) = Gλ(t
(k), t(k)) ≥ Gλ(t(k+1), t(k)) ≥ Fλ(t(k+1)), producing a descent algorithm over F . The trick to obtain G is to apply Jensen inequality to ϕ( ∑ j Hi,jtj) = ϕ( ∑ j Z̃i,j Hi,j Z̃i,j tj) ≤ ∑ j Z̃i,jϕ( Hi,j Z̃i,j
tj), thanks to the convexity of ϕ, see details in (Dhillon and Sra, 2005). We provide below the resulting algorithms for the KL and `2 penalizations, with detailed computations available in Section ?? of the supplementary.
MM for KL-penalized UOT. The KL divergence is obtained with ϕ(y) = y log y−y. Minimizing Gλ(t, t (k)) in that case leads to following multiplicative update:
t (k+1) j = t (k) j exp
( [H> log(y)−H> log ( Ht(k) ) ]j − 1λcj
[H>1]j
) . (5)
Owing to the structure of t andH , the update can be re-written in the following matrix form:
T (k+1) = diag ( a
T (k)1m
) 1 2 ( T (k) exp ( −C
2λ
)) diag ( b
T (k)>1n
) 1 2
, (6)
where is entrywise multiplication and divisions are taken entrywise as well. The multiplicative update (6) is remarkably similar to the well-known Sinkhorn-Knopp algorithm that has been used in
numerous OT problems involving KL regularization. But instead of two separate steps for the left and right scaling, Eq. (6) applies these scalings simultaneously in a unique update using the diagonal matrices (and a form of geometrical average). Also note how the scaling factor exp ( − C2λ ) penalizes along iterations the coefficients of the transport plan with large costs.
MM for `2-penalized UOT. The quadratic loss is obtained with ϕ(y) = y 2
2 . In that case, minimizing Gλ(t, t(k)) s.t. non-negativity leads to following multiplicative update:
T (k+1) = T (k) max
( 0,a1>m + 1nb > − 1λC )
T (k)Om + OnT (k)
with O` = 1`1>` . (7)
Interestingly enough, update (7) prunes any coefficient Ti,j in T such that ai + bj − 1λCi,j < 0 from the very first iteration, providing a useful certificate on the support of the solution.
3.2 Regularization path for `2-penalized UOT
Let us focus on the case where Dϕ is a quadratic divergence. As mentioned in Section 2.2, Eq. (3) is then a positive weighted Lasso problem, allowing us to derive the first regularization path algorithm for computing the whole set of solutions for a varying λ from 0 to +∞. Note that the path’s extreme point recovers the balanced OT solution. We show that the path is piecewise linear in 1/λ between changes in the active set A = supp(tλ), where tλ = vec(T λ) and T λ is the OT plan for given hyperparameter λ. The main steps of the algorithm are roughly as follows: given a current solution (λk,T
λk) and a current active set Ak, we look for the next value λk+1 > λk such that the active set changes (i.e., Ak+1 6= Ak), either because one component enters or leaves the active set. We describe our algorithm below.
KKT conditions of the `2-penalized UOT problem. The Lagrangian for problem (3) writes:
Lλ(t,γ) = 1
λ c>t+
1 2 (Ht− y)>(Ht− y)− γ>t (8)
where γ represents the Lagrange parameters. We denote m = H>y = vec(a1>m + 1nb >). KKT optimality conditions state that i)∇tLλ(t, λ) = 1λc+H >Ht−m−γ = 0 (stationarity condition), ii) γ t = 0 (complementary condition) and iii) γ ≥ 0 (feasibility condition).
Piecewise linearity of the path. Assume that, at iteration k, we know the current active setA = Ak and we look for tλA (the other values of tA being 0). LetHA,mA and cA denote the corresponding sub-matrix and vectors (see Appendix ?? for rigorous definitions). Because of the complementary condition, we have γA = 0. Using λ = λk + , with > 0 small enough to ensure that the active set remains the same, the stationarity condition writes:
H>AHAt λ A = mA −
1 λ cA ⇒ tλA = m̃A − 1 λ c̃A (9)
with m̃A = (H>AHA) −1mA and c̃A = (H>AHA) −1cA . Eq. (9) shows that the optimal tλA (and hence tλ) can be solved for any λ ∈ [λk, λk+1], i.e., when the active set A remains the same, by solving a linear problem. It also reveals the piecewise linearity in λ−1 of the path when A is fixed. As expected, balanced OT is recovered when λ→∞.
Finding (λk+1,Ak+1) given (λk,Ak). Given a current solution (λk, tλk) and λ = λk + , we increase the until we reach a change in the set of active components. This happens whenever the first of the following two situations occurs.
•One component inA becomes inactive. In that case, we remove the index i ∈ A with the smallest λr > λk that violates the constraint. In such case, [m̃A]i = [c̃A]i/λ and we may write
λr = min >λk ( c̃A m̃A ) (10)
where min>λk indicates the minimum value in the vector greater than λk and the division is entrywise.
Algorithm 1 Regularization path of `2-penalized UOT Require: a, b, C, λ0 = 0, t0 = 0, A = A0 = ∅, k = 1 λ1 = min
cĀ mĀ , A = A1 = arg min cĀ mĀ ,H>AHA = 2
tλ1A1 = mA 2 − 1λ1 cA 2 while (Htλk − y)>(Htλk − y) 6= 0 do λr, λa← Compute as in Eq. (10) and Eq. (11) λk+1 ← min(λr, λa) t λk+1 A ← (H>AHA)−1mA − 1λk+1 (H > AHA)
−1cA A = Ak+1 ← Update active set for next iteration. (H>AHA)
−1 ← Update from (H>AkHAk)−1 with Schur complement (see supplementary ??) k ← k + 1
end while return (λk, tλk)k
• One component in Ā becomes active. This occurs when the KKT positivity constraint γĀ ≥ 0 becomes violated. Assume this happens at index i ∈ Ā for the smallest value λa > λk of λ. In such case, the stationarity condition outside the active set can be rewritten:
[ 1
λ cĀ +
[ H>H ( m̃+ 1 λ c̃ )] Ā −mĀ ] i = [γĀ]i ⇒ λa = min >λk ( cĀ − [ H>Hc̃ ] Ā mĀ − [ H>Hm̃ ] Ā ) , (11)
where m̃ (resp. c̃) equals m̃A (resp. c̃A) on A and zero on Ā. In practice, at each step of the path, we compute both λr and λa and set λk+1 = min{λr, λa}. The active set Ak+1 is obtained by either removing the index i ∈ A correponding of the arg min of eq. (10) (case λk+1 = λr) or by adding the index i ∈ Ā corresponding to the arg min of eq. (11) (case λk+1 = λa).
Numerical computation of the entire path. Eq. (9) involves the computation of the matrix (H>AHA)
−1, which is of size |A| × |A|. As only one index leaves or enters the active set at each iteration, we can use the Schur complement of the matrix to compute its value from (H>AkHAk)
−1, alleviating the computational burden of the algorithm as it only involves matrix-vector computations (see Section ?? of supplementary). Algorithm 1 sums up the different steps of the full path computation. At each iteration, we compute λa, λr, update the inverse matrix (H>AkHAk)
−1 and estimate the solution tλk+1 with a complexity of O(nm).
Regularization path of the semi-relaxed `2-penalized UOT. As a side result, let us consider the semi-relaxed OT problem SROTλ(a, b) = minT≥0,T>1n =b 〈C,T 〉 + λ‖T1m − a‖2. The main difference with UOT is that the equality constraint T>1n = b (equivalent toHct = b) must always be met. This leads to the following Lagrangian:
Lλ(t,γ,u) = 1
λ c>t+
1 2 (Hrt− a)>(Hrt− a) + (Hct− b)>u− γ>t, (12)
where u ∈ Rm contains the Lagrange parameters associated to the m equality constraints. The KKT optimality conditions now dictate that i)∇tLλ(t,γ,u) = 1λc+H > r Hrt−H>r a+H>c u−γ = 0, ii) γ t = 0, iii) γ ≥ 0 and Hct − b = 0. We can use the same reasoning than previously to compute the entire path. Details are provided in Section ?? of the supplementary. The main difference lies in solving, at each iteration, a linear system of size (m + |A|) to comply with the marginal equality constraint. The path is initialized as follows: the jth column of T 0 for λ0 = 0 is set to the weighted canonical vector bi?ei? , where i? = argmin{Ci,j}i.

heading: 4 Numerical experiments
text: In this section, we first show the solutions obtained with our solvers2 on simple and interpretable examples. We then evaluate the computational complexity of the different algorithms and finally we show an application where the regularization path can be used on a domain adaptation problem.
Illustration of the algorithms. We first illustrate the regularization path for `2-penalized UOT on a simple example between two distributions containing 3 points each, with different masses and a cost matrix C given in Fig. 1 (left). We can see on Fig. 1 (right) that, starting from λ0 = 0 and T = 0, we successively add or remove components in the active set A when increasing the λ values. When λ =∞, we recover the balanced OT solution. Recall that the path is linear in 1/λ (and not λ). We then illustrate the path for both `2-penalized UOT and semi-relaxed UOT on two 2D distributions with n = m = 100 samples. We can see in Fig. 2 the difference between the two regularization paths for specific values of λ. UOT starts with an empty plan for λ = 0 and then activates samples from both source and target from the closest to the farthest ones until convergence to the balanced OT plan. Semi-relaxed UOT starts with all target samples active due to marginal constraints and progressively activates the source samples.
Comparison of the performances of the algorithms. We now provide an empirical evaluation of the running times of the proposed algorithms, using 2 sets of 10-dimensional points with n = m and drawn according to IID Gaussian distributions. The cost matrix C is computed using a squared `2
2Our implementation of the regularization path has been contributed to POT Flamary et al. (2021) and the MM algorithms are provided in the repository https://github.com/lchapel/ UOT-though-penalized-linear-regression.
norm. We first study the running times of the regularization path algorithm, for n = m ranging from 100 to 1000, averaging the results over 5 runs, see Fig. 3 (left). We empirically observe that log-log plot is near-linear, with an empirical complexity O(n3.27) in this example.
Using n = m = 500, we compare the running times of the current state-of-the-art BFGS algorithm (Blondel et al., 2018)3 using SciPy (Virtanen et al., 2020) and those of our algorithms: the `2- penalized UOT formulated as a Lasso problem (with both the Celer algorithm (Massias et al., 2018) and the coordinate descent solvers from Scikit-learn (Pedregosa et al., 2011)), the multiplicative algorithm for both the `2 and the KL penalties and the regularization path algorithm (see Section ?? of the supplemental material for more details about the solvers and their parameters, together with a comparison of the results of the MM algorithms computed on both CPU and GPU). Figure 3 (middle and right) shows the average running time for all algorithms. For `2-penalized UOT, we observe that, for large λ values, the Lasso solvers are the fastest and that, whatever the value λ, BFGS is the slowest. We also notice that, for large λ, the running times for computing the path remain constant: when the last active set is found, computing the OT plan only involves a weighted sum. As for KL-penalized UOT, the BFGS algorithm is more efficient when large values of λ are considered. One can also notice that, similarly to Sinkhorn which is fast for large regularization values, the multiplicative algorithms for both penalties are also fast for high 1/λ values.
Regularization path for unbalanced domain adaptation. We demonstrate the interest of having the entire regularization path in a classification context where some of the data collection may be polluted by outliers. We consider a setup similar to Mukherjee et al. (2020). Let the sourceX be a set of 400 MNIST digits sampled from the digits 0, 1, 2, 3 (100 points per class) and let the target Y be a set of digits 0, 1 of MNIST (LeCun et al., 2010) and of digits 8, 9 from Fashion MNIST (Xiao et al., 2017). Our setting is simple classification: we classify a sample of the target dataset by propagating
the label of the source sample it is the most transported to, provided that the transported mass of the target point is greater than 0.25bj . Note that similarly to Mukherjee et al. (2020) a validation set can be used here to select the best λ. Figure 4 shows the overall accuracy, defined as the number of samples that are correctly classified divided by the total number of points, and the current accuracy, which is the proportion of well-classified points among the points that are classified, i.e., that are receiving mass. One can notice that, as the number of classified points increases (with λ), the overall accuracy increases as more and more points are well classified while the current accuracy remains stable until outliers are included in the labeled set. This suggests that UOT can be used not only for classification but also as an automated outlier detection method.
3Note that we cannot compare our result with solvers for entropic formulation of UOT with λreg → 0 (that should provide a sparse transport plan) as in that case, their algorithm becomes a fixed point method that cannot converge to a solution of the problem.

heading: 5 Discussion and perspectives
text: We showed that UOT can be recast as a non-negative penalized linear regression problem, encouraging us to dig into this well-established field of research in order to adapt existing algorithmic solutions to the structure of the UOT problem. In this section, we discuss the relation between the proposed algorithms and classical solvers used in OT, and also investigate some research directions that can widen the scope of proposed methods.
Multiplicative algorithms for UOT. As discussed in Section 3.1, the multiplicative updates for the KL divergence obtained from MM resemble the Sinkhorn algorithm from Chizat et al. (2018), except for the joint scaling and the weighting matrix exp(−C/2). Interestingly, this scaling matrix also appears in the Inexact Proximal Point OT (IPOT) algorithm of Xie et al. (2020) to solve balanced OT. As a matter of fact, we show in Section ?? of the supplementary that IPOT is a MM algorithm. The idea is to re-write the OT objective as [〈C,T 〉 + λDϕ(T ,ab>)] − λDϕ(T ,ab>) and upper bound the concave term by its tangent. This further supports the interest of MM for OT and UOT, and highlight an important feature of one of our contributions: designing the first Sinkhorn-like multiplicative algorithm for UOT that can be applied when the OT plan is not entropy-regularized.
More efficient solvers. Despite the positive experimental results of Section 4, multiplicative and regularization path algorithms can be slow, especially for large values of λ. Various accelerations can be envisaged. Regarding path algorithms, the approach of Mairal and Yu (2012) can compute a regularization path with precision in o(1/ ) iterations. This would lead in our setting to a full complexity of O(mn/ ) that is even interesting to approximate balanced OT. Another way to speed up computations is to use screening. In sparse regression, this consists of eliminating during optimization components that will not belong to the support of the solutions thanks to safe screening tests. Methods such as (El Ghaoui et al., 2012; Wang et al., 2015; Dantas et al., 2021) can readily be adapted to our `2 or KL-penalized UOT algorithms. Finally, an other line of improvement is to consider stochastic optimization methods such as (Defazio et al., 2014). Given the particular structure of H , the complexity of stochastic updates shall be small and can lead to very efficient implementations (Nesterov, 2014).
General case and entropy-regularized UOT. Following (Frogner et al., 2015; Chizat et al., 2018; Séjourné et al., 2019), general regularized UOT can be expressed as:
RUOTλ(a, b) = min T≥0 〈C,T 〉+ λ1Dϕ(T1m,a) + λ2Dϕ(T>1n, b) + λregDϕ(T ,ab>). (13)
As it turns out, this general problem involving different regularization weights (λ1, λ2, λreg) can easily be addressed in our framework as well using two simple tricks. The first one consists of absorbing the regularization weights into the divergences. Indeed, many divergences are homogeneous, i.e., satisfy a relation of the form λDϕ(x|y) = Dϕ(λαx|λαy) where α is divergence-specific. This holds in particular for the KL divergence (α = 1) and the squared `2 norm (α = 1/2). The second one consists of complementingH and y with suitable terms to account for the regularization term. In the end, we may re-write Eq. (13) into Eq. (3) with λ = 1, H = [λα1H > r , λ α 2H > c , λ α regI]
> and y> = [λα1a >, λα2 b >, λαregvec(ab
>)>]. In particular, we obtain the following multiplicative update in the case of entropy-regularized KL-penalized UOT:
T (k+1) = diag ( a
T (k)1m
) λ1 λall (( T (k) )λ1+λ2 λall K ) diag ( b
T (k)>1n
) λ2 λall
(14)
where K = ( ab> )λreg λall exp ( − 1λallC ) and λall = λ1 + λ2 + λreg. This multiplicative update is slightly more complex than the Sinkhorn algorithms of Frogner et al. (2015); Chizat et al. (2018) and as such, it might have limited practical interest but is conceptually interesting and novel. Note that balanced UOT as of Eq. (2) is simply obtained with λreg = 0.
Non-linear UOT. Finally, we discuss how our proposed reformulation of UOT can accommodate non-linear variants in which the linear term 〈C,T 〉 is replaced by a sparsity/robustness-promoting
term, leading to problems of the form
NLUOTλ(a, b) = min T≥0 ∑ i,j g(Ci,jTi,j) + λ1Dϕ(T1m,a) + λ2Dϕ(T >1n, b) (15)
where g(·) is a usually concave function, see, e.g., (Candes et al., 2008; Gasso et al., 2009). Our MM setting can readily accommodate such a formulation by majorizing the concave terms by their tangent. The non-linearity may improve robustness w.r.t outliers and better model realistic OT problems. For instance, in real life, the costs of transporting some goods between two places can be nonlinear due to economies of scale.
Broad and potential negative societal impact. The contributions in this paper are methodological and focus on a reformulation of a fundamental OT problem and adapting existing algorithms to solve it. In this sense, we bring more efficient solvers that run on GPU but this computational advantage can be counterbalanced by the possibility that it brings to be applied on larger datasets. The application of OT in domain adaptation has shown that it can be used to infer labels on samples/individuals when no labels are available, suggesting a capacity for violating user privacy. A potential application of UOT is the case where two datasets of users acquired by different methods contain some shared users. UOT can be used here to find correspondences between the users in the two datasets and also identify unique users in each dataset (those that do not receive mass).

heading: 6 Conclusion
text: In this paper, we reformulate the UOT problem as a non-negative penalized linear regression, allowing us to propose two new classes of algorithms. We first derive multiplicative algorithms for both KL and `2-penalized UOT, providing numerical solutions that are fast and easy to implement. For the specific case of `2-penalized UOT, we provide the first regularization path algorithm that computes the whole set of solutions for all the regularization parameter values. We finally build on the extensive literature in inverse problem and NMF to draw some fruitful perspectives on even more efficient algorithmic solutions or the definition of new OT problems.
Acknowledgments and Disclosure of Funding
The authors want to thank Hicham Janati for interesting discussions and providing us with the experiments of convergence for the MM algorithm in the supplemental. This work is partially funded by the French National Research Agency (ANR; grants OATMIL ANR-17-CE23-0012, RAIMO ANR-20-CHIA-0021-01, MULTISCALE ANR-18-CE23-0022-01, E4C ANR-18-EUR0006-02, 3IA Côte d’Azur ANR-19-P3IA-0002, 3IA ANITI ANR-19-PI3A-0004) and the European Research Council (ERC; grant FACTORY-CoG-6681839). Furthermore, this research was produced within the framework of Energy4Climate Interdisciplinary Center (E4C) of IP Paris and Ecole des Ponts ParisTech. This action benefited from the support of the Chair “Challenging Technology for Responsible Energy” led by l’X - Ecole Polytechnique and the Fondation de l’Ecole Polytechnique, sponsored by TOTAL.

Review summaries:
This paper proposes a new method for Unbalanced Optimal Transport (UOT) without a regularization term. The solution consists of two steps; (i) formulate the UOT as a non-negative penalized linear regression problem (ii) solve it using the Majorization-Minimization (MM) algorithm. The authors derive specific algorithms for several penalty functions. Furthermore, a regularization path tracking method based on the relationship between the proposed formulation and lasso is proposed. The authors validate the effectiveness of the proposed method by numerical experiments using synthetic datasets and the MNIST dataset.

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.

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.