Generative adversarial networks (GANs) are an expressive class of neural generative models with tremendous success in modeling high-dimensional continuous measures. In this paper, we present a scalable method for unbalanced optimal transport (OT) based on the generative-adversarial framework. We formulate unbalanced OT as a problem of simultaneously learning a transport map and a scaling factor that push a source measure to a target measure in a cost-optimal manner. We provide theoretical justification for this formulation, showing that it is closely related to an existing static formulation by Liero et al. (2018). We then propose an algorithm for solving this problem based on stochastic alternating gradient updates, similar in practice to GANs, and perform numerical experiments demonstrating how this methodology can be applied to population modeling. 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 BID34 .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 BID24 , which seeks the optimal probabilistic coupling between measures and can be solved using linear programming methods for discrete measures. Recently, BID13 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 BID19 . Optimal transport has been applied to many areas, such as computer graphics BID17 BID36 and domain adaptation (Courty et al., 2014; .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) BID20 , which push a source distribution to a target distribution by training against an adversary. Numerous transport problems in image translation BID30 BID38 BID37 , natural language translation BID22 , domain adaptation BID5 and biological data integration (Amodio & (a) (b)Figure 1: (a) Illustration of the problem of modeling the transformation of a source population µ to a target population ν. In this example, one sub-population is growing more rapidly than the others. (b) Schematic of Monge-like formulations of unbalanced optimal transport. The objective is to learn a transport map T (for transporting mass) and scaling factor ξ (for mass variation) to push the source µ to the target ν, using a deterministic transport map (top) BID7 or a stochastic transport map (bottom).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 cannot handle mass variation.Several formulations have been proposed for extending the theory of OT to the setting where the measures can have unbalanced masses BID7 BID26 BID27 BID18 . In terms of numerical methods, a class of scaling algorithms BID8 ) 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 BID27 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 BID8 , tumor growth modeling BID6 and computational biology BID34 . However, while optimal entropy-transport allows mass variation , it cannot 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 BID7 .• By relaxing this problem, we obtain an alternative form of the optimal entropy-transport problem by BID27 , 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 BID20 ).• 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 BID27 in the continuous setting. The algorithm extends the work of to unbalanced OT and is a scalable alternative to the algorithm of BID8 for very large or continuous datasets.
We propose new methodology for unbalanced optimal transport using generative adversarial networks. The authors consider the unbalanced optimal transport problem between two measures with different total mass using a stochastic min-max algorithm and local scaling The authors propose an approach to estimate unbalanced optimal transport between sampled measures that scales well in the dimension and in the number of samples. The paper introduces a static formulation for unbalanced optimal transport by learning simultaneously a transport map T and scaling factor xi.