Scalable Unbalanced Optimal Transport using Generative Adversarial Networks

1 INTRODUCTION . We consider the problem of unbalanced optimal transport : given two measures , find a cost-optimal way to transform one measure to the other using a combination of mass variation and transport . Such problems arise , for example , when modeling the transformation of a source population into a target population ( Figure 1a ) . In this setting , one needs to model mass transport to account for the features that are evolving , as well as local mass variations to allow sub-populations to become more or less prominent in the target population ( Schiebinger et al. , 2017 ) . Classical optimal transport ( OT ) considers the problem of pushing a source to a target distribution in a way that is optimal with respect to some transport cost without allowing for mass variations . Modern approaches are based on the Kantorovich formulation ( Kantorovich , 1942 ) , which seeks the optimal probabilistic coupling between measures and can be solved using linear programming methods for discrete measures . Recently , Cuturi ( 2013 ) showed that regularizing the objective using an entropy term allows the dual problem to be solved more efficiently using the Sinkhorn algorithm . Stochastic methods based on the dual objective have been proposed for the continuous setting ( Genevay et al. , 2016 ; Seguy et al. , 2017 ; Arjovsky et al. , 2017 ) . Optimal transport has been applied to many areas , such as computer graphics ( Ferradans et al. , 2014 ; Solomon et al. , 2015 ) and domain adaptation ( Courty et al. , 2014 ; 2017 ) . In many applications where a transport cost is not available , transport maps can also be learned using generative models such as generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) , which push a source distribution to a target distribution by training against an adversary . Numerous transport problems in image translation ( Mirza & Osindero , 2014 ; Zhu et al. , 2017 ; Yi et al. , 2017 ) , natural language translation ( He et al. , 2016 ) , domain adaptation ( Bousmalis et al. , 2017 ) and biological data integration ( Amodio & Krishnaswamy , 2018 ) have been tackled using variants of GANs , with strategies such as conditioning or cycle-consistency employed to enforce correspondence between original and transported samples . However , all these methods conserve mass between the source and target and therefore can not handle mass variation . Several formulations have been proposed for extending the theory of OT to the setting where the measures can have unbalanced masses ( Chizat et al. , 2015 ; 2018 ; Kondratyev et al. , 2016 ; Liero et al. , 2018 ; Frogner et al. , 2015 ) . In terms of numerical methods , a class of scaling algorithms ( Chizat et al. , 2016 ) that generalize the Sinkhorn algorithm for balanced OT have been developed for approximating the solution to optimal entropy-transport problems ; this formulation of unbalanced OT by Liero et al . ( 2018 ) corresponds to the Kantorovich OT problem in which the hard marginal constraints are relaxed using divergences to allow for mass variation . In practice , these algorithms have been used to approximate unbalanced transport plans between discrete measures for applications such as computer graphics ( Chizat et al. , 2016 ) , tumor growth modeling ( Chizat & Di Marino , 2017 ) and computational biology ( Schiebinger et al. , 2017 ) . However , while optimal entropy-transport allows mass variation , it can not explicitly model it , and there are currently no methods that can perform unbalanced OT between continuous measures . Contributions . Inspired by the recent successes of GANs for high-dimensional transport problems , we present a novel framework for unbalanced optimal transport that directly models mass variation in addition to transport . Concretely , our contributions are the following : • We propose to solve a Monge-like formulation of unbalanced OT , in which the goal is to learn a stochastic transport map and scaling factor to push a source to a target measure in a cost-optimal manner . This generalizes the unbalanced Monge OT problem by Chizat et al . ( 2015 ) . • By relaxing this problem , we obtain an alternative form of the optimal entropy-transport problem by Liero et al . ( 2018 ) , which confers desirable theoretical properties . • We develop scalable methodology for solving the relaxed problem . Our derivation uses a convex conjugate representation of divergences , resulting in an alternating gradient descent method similar to GANs ( Goodfellow et al. , 2014 ) . • We demonstrate in practice how our methodology can be applied towards population modeling using the MNIST and USPS handwritten digits datasets , the CelebA dataset , and a recent single-cell RNA-seq dataset from zebrafish embrogenesis . In addition to these main contributions , for completeness we also propose a new scalable method ( Algorithm 2 ) in the Appendix for solving the optimal-entropy transport problem by Liero et al . ( 2018 ) in the continuous setting . The algorithm extends the work of Seguy et al . ( 2017 ) to unbalanced OT and is a scalable alternative to the algorithm of Chizat et al . ( 2016 ) for very large or continuous datasets . 2 PRELIMINARIES . Notation . Let X , Y ⊆ Rn be topological spaces and let B denote the Borel σ-algebra . Let M 1+ ( X ) , M+ ( X ) denote respectively the space of probability measures and finite non-negative measures over X . For a measurable function T , let T # denote its pushforward operator : if µ is a measure , then T # µ is the pushforward measure of µ under T . Finally , let πX , πY be functions that project onto X and Y ; for a joint measure γ ∈M+ ( X × Y ) , πX # γ and πY # γ are its marginals with respect to X and Y respectively . Optimal transport ( OT ) addresses the problem of transporting between measures in a cost-optimal manner . Monge ( 1781 ) formulated this problem as a search over deterministic transport maps . Specifically , given µ ∈M 1+ ( X ) , ν ∈M 1+ ( Y ) and a cost function c : X × Y → R+ , Monge OT seeks a measurable function T : X → Y minimizing inf T ∫ X c ( x , T ( x ) ) dµ ( x ) ( 1 ) subject to the constraint T # µ = ν . While the optimal T has an intuitive interpretation as an optimal transport map , the Monge problem is non-convex and not always feasible depending on the choices of µ and ν . The Kantorovich OT problem is a convex relaxation of the Monge problem that formulates OT as a search over probabilistic transport plans . Given µ ∈ M 1+ ( X ) , ν ∈ M 1+ ( Y ) and a cost function c : X × Y → R+ , Kantorovich OT seeks a joint measure γ ∈M 1+ ( X × Y ) subject to πX # γ = µ and πY # γ = ν minimizing W ( µ , ν ) : = inf γ ∫ X×Y c ( x , y ) dγ ( x , y ) . ( 2 ) Note that the conditional probability distributions γy|x specify stochastic maps from X to Y and can be considered a “ one-to-many ” version of the deterministic map from the Monge problem . In terms of numerical methods , the relaxed problem is a linear program that is always feasible and can be solved in O ( n3 ) time for discrete µ , ν. Cuturi ( 2013 ) recently showed that introducing entropic regularization results in a simpler dual optimization problem that can be solved efficiently using the Sinkhorn algorithm . Based on the entropyregularized dual problem , Genevay et al . ( 2016 ) and Seguy et al . ( 2017 ) proposed stochastic algorithms for computing transport plans that can handle continuous measures . Unbalanced OT . Several formulations that extend classical OT to handle mass variation have been proposed ( Chizat et al. , 2015 ; 2018 ; Kondratyev et al. , 2016 ) . Existing numerical methods are based on a Kantorovichlike formulation known as optimal-entropy transport ( Liero et al. , 2018 ) . This formulation is obtained by relaxing the marginal constraints of ( 2 ) using divergences as follows : given two positive measures µ ∈ M+ ( X ) and ν ∈ M+ ( Y ) and a cost function c : X × Y → R+ , optimal entropy-transport finds a measure γ ∈M+ ( X × Y ) that minimizes Wub ( µ , ν ) : = inf γ ∫ X×Y c ( x , y ) dγ ( x , y ) +Dψ1 ( π X # γ|µ ) +Dψ2 ( πY # γ|ν ) , ( 3 ) whereDψ1 , Dψ2 areψ-divergences induced byψ1 , ψ2 . Theψ-divergence between non-negative finite measures P , Q over T ⊆ Rd induced by a lower semi-continuous , convex entropy function ψ : R→ R ∪ { ∞ } is Dψ ( P |Q ) : = ψ′∞P⊥ ( T ) + ∫ T ψ ( dP dQ ) dQ , ( 4 ) where ψ′∞ : = lims→∞ ψ ( s ) s and dP dQQ+ P ⊥ is the Lebesgue decomposition of P with respect to Q . Note that mass variation is allowed since the marginals of γ are not constrained to be µ and ν . In terms of numerical methods , the state-of-the-art in the discrete setting is a class of iterative scaling algorithms ( Chizat et al. , 2016 ) that generalize the Sinkhorn algorithm for computing regularized OT plans ( Cuturi , 2013 ) . There are no practical algorithms for unbalanced OT between continuous measures , especially in high-dimensional spaces . 3 SCALABLE UNBALANCED OT USING GANS . In this section , we propose the first algorithm for unbalanced OT that directly models mass variation and can be applied towards transport between high-dimensional continuous measures . The starting point of our development is the following Monge-like formulation of unbalanced OT , in which the goal is to learn a stochastic transport map and scaling factor to push a source to a target measure in a cost-optimal manner . Unbalanced Monge OT . Let c1 : X × Y → R+ be the cost of transport and c2 : R+ → R+ the cost of mass variation . Let the probability space ( Z , B ( Z ) , λ ) be the source of randomness in the transport map T . Given two positive measures µ ∈M+ ( X ) and ν ∈M+ ( Y ) , we seek a transport map T : X ×Z → Y and a scaling factor ξ : X → R+ minimizing L ( µ , ν ) : = inf T , ξ ∫ X ∫ Z c1 ( x , T ( x , z ) ) dλ ( z ) ξ ( x ) dµ ( x ) + ∫ X c2 ( ξ ( x ) ) dµ ( x ) , ( 5 ) subject to the constraint T # ( ξµ × λ ) = ν . Concretely , the first and second terms of ( 5 ) penalize the cost of mass transport and variation respectively , and the equality constraint ensures that ( T , ξ ) pushes µ to ν exactly . A special case of ( 5 ) is the unbalanced Monge OT problem by Chizat et al . ( 2015 ) , which employs a deterministic transport map ( Figure 1b ) . We consider the more general case of stochastic ( i.e . one-to-many ) maps because it is a more suitable model for many practical problems . For example , in cell biology , it is natural to think of one cell in a source population as potentially giving rise to multiple cells in a target population . In practice , one can take Z = Rn and λ to be the standard Gaussian measure if a stochastic map is desired ; otherwise λ can be set to a deterministic distribution . The following are examples of problems that can be modeled using unbalanced Monge OT . Example 3.1 ( Figure 2a ) . Suppose the objective is to model the transformation from a source measure ( Column 1 ) to the target measure ( Column 2 ) , which represent a population of interest at two distinct time points . The transport map T models the transport/movement of points from the source to the target , while the scaling factor ξ models the growth ( replication ) or shrinkage ( death ) of these points . Different models of transformation are optimal depending on the relative costs of mass transport and variation ( Columns 3-6 ) . Example 3.2 ( Figure 2b ) . Suppose the objective is to transport points from a source measure ( 1st panel , color ) to a target measure ( 1st panel , grey ) in the presence of class imbalances . A pure transport map would muddle together points from different classes , while an unbalanced transport map with a scaling factor is able to ameliorate the class imbalance ( 2nd panel ) . In this case , the scaling factor tells us explicitly how to downweigh or upweigh samples in the source distribution to balance the classes with the target distribution ( 3rd panel ) . Relaxation . From an optimization standpoint , it is challenging to satisfy the constraint T # ( ξµ× λ ) = ν . We hence consider the following relaxation of ( 5 ) using a divergence penalty in place of the equality constraint : Lψ ( µ , ν ) : = inf T , ξ ∫ X ∫ Z c1 ( x , T ( x , z ) ) dλ ( z ) ξ ( x ) dµ ( x ) + ∫ X c2 ( ξ ( x ) ) dµ ( x ) +Dψ ( T # ( ξµ× λ ) |ν ) , ( 6 ) using an appropriate choice of ψ that satisfies the requirements of Lemma C.2 in the Appendix1.This relaxation is the Monge-like version of the optimal-entropy transport problem ( 3 ) by Liero et al . ( 2018 ) . Specifically , ( T , ξ ) specifies a joint measure γ ∈M+ ( X × Y ) given by γ ( C ) : = ∫ X ∫ Z 1C ( x , T ( x , z ) ) dλ ( z ) ξ ( x ) dµ ( x ) , ∀C ∈ B ( X ) × B ( Y ) , and by reformulating ( 6 ) in terms of γ instead of ( T , ξ ) , one obtains the objective function for optimal-entropy transport . The main difference between the formulations is their search space , since not all joint measures γ ∈M+ ( X × Y ) can be specified by some choice of ( T , ξ ) . For example , if T is a deterministic transport map , then γ is necessarily restricted to the set of deterministic couplings . Even if T is sufficiently random , it is generally not possible to specify all joint measures γ ∈M+ ( X × Y ) : in the asymmetric Monge formulation ( 6 ) , all the mass transported to Y must come from somewhere within the support of µ , since the scaling factor ξ allows mass to grow but not to materialize outside of its original support . Therefore equivalence can be established in general only when restricting the support of γ to supp ( µ ) × Y as described in the following lemma , whose proof is given in the Appendix . Lemma 3.3 . Let G be the set of joint measures supported on supp ( µ ) × Y , and define W̃c , ψ1 , ψ2 ( µ , ν ) : = inf γ∈G ∫ c dγ +Dψ1 ( π X # γ|µ ) +Dψ2 ( πY # γ|ν ) . If ( Z , B ( Z ) , λ ) is atomless and c2 is an entropy function , then Lψ ( µ , ν ) = W̃c1 , c2 , ψ ( µ , ν ) . Based on the relation between ( 3 ) and ( 6 ) , several theoretical results for ( 6 ) follow from the analysis of optimal entropy-transport by Liero et al . ( 2018 ) . Importantly , one can show the following theorem , namely that for an appropriate and sufficiently large choice of divergence penalty , solutions of the relaxed problem ( 6 ) converge to solutions of the original problem ( 5 ) . The proof is given in the Appendix . Theorem 3.4 . Suppose c1 , c2 , ψ satisfy the existence assumptions of Proposition B.1 in the Appendix , and let ( Z , B ( Z ) , λ ) be an atomless probability space . Furthermore , let ψ be uniquely minimized at ψ ( 1 ) = 0 . Then for a sequence 0 < ζ1 < · · · < ζk < · · · diverging to∞ indexed by k , limk→∞ Lζkψ ( µ , ν ) = L ( µ , ν ) . Additionally , let γk be the joint measure specified by a minimizer of Lζkψ ( µ , ν ) . If L ( µ , ν ) < ∞ , then up to extraction of a subsequence , γk converges weakly to γ , the joint measure specified by a minimizer of L ( µ , ν ) . Algorithm . Using the relaxation of unbalanced Monge OT in ( 6 ) , we now show that the transport map and scaling factor can be learned by stochastic gradient methods . While the divergence term can not easily be minimized using the definition in ( 4 ) , we can write it as a penalty witnessed by an adversary function f : Y → ( −∞ , ψ′∞ ] using the convex conjugate representation ( see Lemma B.2 ) : Dψ ( T # ( ξµ× λ ) |ν ) = sup f ∫ X ∫ Z f ( T ( x , z ) ) dλ ( z ) ξ ( x ) dµ ( x ) − ∫ Y ψ∗ ( f ( y ) ) dν ( y ) , ( 7 ) where ψ∗ is the convex conjugate of ψ . The objective in ( 6 ) can now be optimized using alternating stochastic gradient updates after parameterizing T , ξ , and f with neural networks ; see Algorithm 1 2 . The optimization 1Note that Dψ ( T # ( ξµ× λ ) |ν ) = 0 ; T # ( ξµ× λ ) = ν in general since the total mass of the transported measure is not constrained ; for the relaxation to be valid , ψ ( s ) should attain a unique minimum at s = 1 ( see Lemma C.2 ) . 2We assume that one has access to samples from µ , ν , and in the setting where µ , ν are not normalized , then samples to the normalized measures µ̃ , ν̃ as well as the normalization constants cµ , cν . These are reasonable assumptions for practical applications : for example , in a biological assay , one might collect cµ cells from time point 1 and cν cells from time point 2 . In this case , the samples are the measurements of each cell and the normalization constants are cµ , cν . procedure is similar to GAN training and can be interpreted as an adversarial game between ( T , ξ ) and f : • T takes a point x ∼ µ and transports it from X to Y by generating T ( x , z ) where z ∼ λ . • ξ determines the importance weight of each transported point . • Their shared objective is to minimize the divergence between transported samples and real samples from ν that is measured by the adversary f . • Additionally , cost functions c1 and c2 encourage T , ξ to find the most cost-efficient strategy . Algorithm 1 Generative-Adversarial Framework for Unbalanced Monge OT Input : Initial parameters θ , φ , ω ; step size η ; normalized measures µ̃ , ν̃ , constants cµ , cν . Output : Updated parameters θ , φ , ω. while ( θ , φ , ω ) not converged do Sample x1 , · · · , xn from µ̃ , y1 , · · · , yn from ν̃ , z1 , · · · , zn from λ ; ` ( θ , φ , ω ) : = 1 n n∑ i=1 [ cµc1 ( xi , Tθ ( xi , zi ) ) ξφ ( xi ) + cµc2 ( ξφ ( xi ) ) ( 8 ) + cµξφ ( xi ) fω ( Tθ ( xi , zi ) ) − cνψ∗ ( fω ( yi ) ) . ] Update ω by gradient descent on − ` ( θ , φ , ω ) . Update θ , φ by gradient descent ` ( θ , φ , ω ) . end while Table 1 in the Appendix provides some examples of divergences with corresponding entropy functions and convex conjugates that can be plugged into ( 7 ) . Further practical considerations for implementation and training are discussed in Appendix C. Relation to other approaches . The probabilistic Monge-like formulation ( 6 ) is similar to the Kantorovichlike entropy-transport problem ( 3 ) in theory , but they result in quite different numerical methods in practice . Algorithm 1 solves the non-convex formulation ( 6 ) and learns a transport map T and scaling factor ξ parameterized by neural networks , enabling scalable optimization using stochastic gradient descent . The networks are immediately useful for many practical applications ; for instance , it only requires a single forward pass to compute the transport and scaling of a point from the source domain to the target . Furthermore , the neural architectures of T , ξ imbue their function classes with a particular structure , and when chosen appropriately , enable effective learning of these functions in high-dimensional settings . Due to the non-convexity of the optimization problem , however , Algorithm 1 is not guaranteed to find the global optimum . In contrast , the scaling algorithm of Chizat et al . ( 2016 ) based on ( 3 ) solves a convex optimization problem and is proven to converge , but is currently only practical for discrete problems and has limited scalability . For completeness , in Section A of the Appendix , we propose a new stochastic method based on the same dual objective as Chizat et al . ( 2016 ) that can handle transport between continuous measures ( Algorithm 2 in the Appendix ) . This method generalizes the approach of Seguy et al . ( 2017 ) for handling transport between continuous measures and overcomes the scalability limitations of Chizat et al . ( 2016 ) . However , the output is in the form of the dual solution , which is less interpretable for practical applications compared to the output of Algorithm 1 . In particular , while one can compute a deterministic transport map known as a barycentric projection from the dual solution , it is unclear how best to obtain a scaling factor or a stochastic transport map that can generate samples outside of the target dataset . In the numerical experiments of Section 4 , we show the advantage of directly learning a transport map and scaling factor using Algorithm 1 . The problem of learning a scaling factor ( or weighting factor ) that “ balances ” measures µ and ν also arises in causal inference . Generally , µ is the distribution of covariates from a control population and ν is the distribution from a treated population . The goal is to scale the importance of different members from the control population based on how likely they are to be present in the treated population , in order to eliminate selection biases in the inference of treatment effects . Kallus ( 2018 ) proposed a generative-adversarial method for learning the scaling factor , but they do not consider transport .

Summary: The authors propose a novel approach to estimate unbalanced optimal transport between sampled measures that scales well in the dimension and in the number of samples. This formulation is based on a formulation of the entropy-transport problems of Liero et al. where the transport map, growth maps and Lagrangian multipliers are parameterized by neural networks. The effectiveness of the approach is shown on some tasks.