Generative Modeling with Optimal Transport Maps

1 INTRODUCTION . Since the discovery of Generative Adversarial Networks ( GANs , Goodfellow et al . ( 2014 ) ) , there has been a surge in generative modeling ( Radford et al. , 2016 ; Arjovsky et al. , 2017 ; Brock et al. , 2019 ; Karras et al. , 2019 ) . In the past few years , Optimal Transport ( OT , Villani ( 2008 ) ) theory has been pivotal in addressing important issues of generative models . In particular , the usage of Wasserstein distance has improved diversity ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ) , convergence ( Sanjabi et al. , 2018 ) , and stability ( Miyato et al. , 2018 ; Kim et al. , 2021 ) of GANs . Generative models based on OT can be split into two classes depending on what OT is used for . First , the optimal transport cost serves as the loss for generative models , see Figure 1a . This is the most prevalent class of methods which includes WGAN ( Arjovsky et al. , 2017 ) and its modifications : WGAN-GP ( Gulrajani et al. , 2017 ) , WGAN-LP ( Petzka et al. , 2018 ) , and WGAN-QC ( Liu et al. , 2019 ) . Second , the optimal transport map is used as a generative model itself , see Figure 1b . Such approaches include LSOT ( Seguy et al. , 2018 ) , AE-OT ( An et al. , 2020a ) , ICNN-OT ( Makkuva et al. , 2020 ) , W2GN ( Korotin et al. , 2021a ) . Models of the first class have been wellstudied , but limited attention has been paid to the second class . Existing approaches of the second class primarily consider OT maps in latent spaces of pre-trained autoencoders ( AE ) , see Figure 3 . The performance of such generative models depends on the underlying AEs , in which decoding transformations are often not accurate ; as a result this deficiency limits practical applications in high-dimensional ambient spaces . For this reason , using OT in the latent space does not necessarily guarantee superior performance in generative modeling . The focus of our paper is the second class of OT-based models using OT map as the generative map . Finding an optimal mapping is motivated by its ability to preserve specific attributes of the input samples , a desired property in unpaired learning . For example , in unpaired image-to-image translation , the learner has to fit a map between two data distributions which preserves the image content . CycleGAN-based models ( Zhu et al. , 2017 ) are widely used for this purpose . However , they typically have complex optimization objectives consisting of several losses ( Amodio & Krishnaswamy , 2019 ; Lu et al. , 2019 ) in order to make the fitted map preserve the required attributes . The main contributions of this paper are as follows : 1 . We propose an end-to-end algorithm ( M4.3 ) to fit OT maps for the quadratic cost ( Wasserstein-2 distance ) between distributions located on the spaces of equal dimensions ( M4.1 ) and extend the method to unequal dimensions as well ( M4.2 ) . We prove error bounds for the method ( M4.4 ) . 2 . We demonstrate large-scale applications of OT maps in popular computer vision tasks . We consider image generation ( M5.1 ) and unpaired image restoration ( M5.2 ) tasks . Our strict OT-based framework allows the theoretical analysis of the recovered transport map . The OT map obtained by our method can be directly used in large-scale computer vision problems which is in high contrast to previous related methods relying on autoencoders and OT maps in the latent space . Importantly , the performance and computational complexity of our method is comparable to OT-based generative models using OT cost as the loss . Notations . In what follows , X and Y are two complete metric spaces , µ ( x ) and ν ( y ) are probability distributions on X and Y , respectively . For a measurable map T : X → Y , T # µ denotes the pushforward distribution of µ , i.e. , the distribution for which any measurable set E ⊂ Y satisfies T # µ ( E ) = µ ( T −1 ( E ) ) . For a vector x , ‖x‖ denotes its Euclidean norm . We use 〈x , y〉 to denote the inner product of vectors x and y . We use Π ( µ , ν ) to denote the set of joint probability distributions on X ×Y whose marginals are µ and ν , respectively ( couplings ) . For a function f : RD → R∪ { ±∞ } its Legendre–Fenchel transform ( the convex conjugate ) is f ( y ) = supx∈RD { 〈x , y〉 − f ( x ) } . It is convex , even if f is not . 2 BACKGROUND ON OPTIMAL TRANSPORT . Consider a cost of transportation , c : X × Y → R defined over the product space of X and Y . Monge ’ s Formulation . The optimal transport cost between µ and ν for ground cost c ( · , · ) is Cost ( µ , ν ) def= inf T # µ=ν ∫ X c ( x , T ( x ) ) dµ ( x ) , ( 1 ) where the infimum is taken over all measurable maps T : X → Y pushing µ to ν , see Figure 2 . The map T ∗ on which the infimum in ( 1 ) is attained is called the optimal transport map . Monge ’ s formulation does not allow splitting . For example , when µ is a Dirac distribution and ν is a non-Dirac distribution , the feasible set of equation ( 1 ) is empty . Kantorovich ’ s Relaxation . Instead of asking to which particular point y ∈ Y should all the probability mass of x be moved , Kantorovich ( 1948 ) asks how the mass of x should be distributed among all y ∈ Y . Formally , a transport coupling replaces a transport map ; the OT cost is given by : Cost ( µ , ν ) def= inf π∈Π ( µ , ν ) ∫ X×Y c ( x , y ) dπ ( x , y ) , ( 2 ) where the infimum is taken over all couplings π ∈ Π ( µ , ν ) of µ and ν . The coupling π∗ attaining the infimum of ( 2 ) is called the optimal transport plan . Unlike the formulation of ( 1 ) , the formulation of ( 2 ) is well-posed , and with mild assumptions on spaces X , Y and ground cost c ( · , · ) , the minimizer π∗ of ( 2 ) always exists ( Villani , 2008 , Theorem 4.1 ) . In particular , if π∗ is deterministic , i.e. , π∗ = [ idX , T ∗ ] # µ for some T ∗ : X → Y , then T ∗ minimizes ( 1 ) . Duality . The dual form of ( 2 ) is given by ( Kantorovich , 1948 ) : Cost ( µ , ν ) = sup ( u , v ) { ∫ X u ( x ) dµ ( x ) + ∫ Y v ( y ) dν ( y ) : u ( x ) + v ( y ) ≤ c ( x , y ) } , ( 3 ) with u ∈ L1 ( µ ) , v ∈ L1 ( ν ) called Kantorovich potentials . For u : X → R and v : Y → R define their c-transforms by uc ( y ) = infx∈X { c ( x , y ) − u ( x ) } and vc ( x ) = infy∈Y { c ( x , y ) − v ( y ) } respectively . Using c-transform , ( 3 ) is reformulated as ( Villani , 2008 , M5 ) Cost ( µ , ν ) =sup v { ∫ X vc ( x ) dµ ( x ) + ∫ Y v ( y ) dν ( y ) } =sup u { ∫ X u ( x ) dµ ( x ) + ∫ Y uc ( y ) dν ( y ) } . ( 4 ) Primal-dual relationship . For certain ground costs c ( · , · ) , the primal solution T ∗ of ( 1 ) can be recovered from the dual solution u∗ of ( 3 ) . For example , if X = Y = RD , c ( x , y ) = h ( x− y ) with strictly convex h : RD → R and µ is absolutely continuous supported on the compact set , then T ∗ ( x ) = x− ( ∇h ) −1 ( u∗ ( x ) ) , ( 5 ) see ( Santambrogio , 2015 , Theorem 1.17 ) . For general costs , see ( Villani , 2008 , Theorem 10.28 ) . 3 OPTIMAL TRANSPORT IN GENERATIVE MODELS . OPTIMAL TRANSPORT COST AS THE LOSS ( Figure 1a ) . Starting with the works of Arjovsky & Bottou ( 2017 ) ; Arjovsky et al . ( 2017 ) , the usage of OT cost as the loss has become a major way to apply OT for generative modeling . In this setting , given data distribution ν and fake distribution µθ , the goal is to minimize Cost ( µθ , ν ) w.r.t . the parameters θ . Typically , µθ is a pushforward distribution of some given distribution , e.g. , N ( 0 , I ) , via generator network Gθ . The Wasserstein-1 distance ( W1 ) , i.e. , the transport cost for ground cost c ( x , y ) = ‖x − y‖ , is the most practically prevalent example of such a loss . Models based on this loss are known as Wasserstein GANs ( WGANs ) . They estimateW1 ( µθ , ν ) based on the dual form as given by ( 4 ) . For W1 , the optimal potentials u∗ , v∗ of ( 4 ) satisfy u∗ = −v∗ where u∗ is a 1-Lipschitz function ( Villani , 2008 , Case 5.16 ) . As a result , to computeW1 , one needs to optimize the following simplified form : W1 ( µθ , ν ) = sup ‖u‖L≤1 { ∫ X u ( x ) dµθ ( x ) − ∫ Y u ( y ) dν ( y ) } . ( 6 ) In WGANs , the potential u is called the discriminator . Optimization of ( 6 ) reduces constrained optimization of ( 4 ) with two potentials u , v to optimization of only one discriminator u . In practice , enforcing the Lipschitz constraint on u is challenging . Most methods to do this are regularizationbased , e.g. , they use gradient penalty ( Gulrajani et al. , 2017 , WGAN-GP ) and Lipschitz penalty ( Petzka et al. , 2018 , WGAN-LP ) . Other methods enforce Lipschitz property via incorporating certain hard restrictions on the discriminator ’ s architecture ( Anil et al. , 2019 ; Tanielian & Biau , 2021 ) . General transport costs ( other than W1 ) can also be used as the loss for generative models . They are less popular since they do not have a dual form reducing to a single potential function similar to ( 6 ) for W1 . Consequently , the challenging estimation of the c-transform uc is needed . To avoid this , Sanjabi et al . ( 2018 ) consider the dual form of ( 3 ) with two potentials u , v instead form ( 4 ) with one u and softly enforce the condition u ( x ) + v ( y ) ≤ c ( x , y ) via entropy or quadratic regularization . Nhan Dam et al . ( 2019 ) use the dual form of ( 4 ) and amortized optimization to compute uc via an additional neural network . Both methods work for general c ( · , · ) , though the authors test them for c ( x , y ) = ‖x − y‖ only , i.e. , W1 distance . Mallasto et al . ( 2019 ) propose a fast way to approximate the c-transform and test the approach ( WGAN- ( q , p ) ) with several costs , in particular , the Wasserstein-2 distance ( W2 ) , i.e. , the transport cost for the quadratic ground cost c ( x , y ) = 12‖x−y‖ 2 . Specifically forW2 , Liu et al . ( 2019 ) approximate the c-transform via a linear program ( WGAN-QC ) . A fruitful branch of OT-based losses for generative models comes from modified versions of OT cost , such as Sinkhorn ( Genevay et al. , 2018 ) , sliced ( Deshpande et al. , 2018 ) and minibatch ( Fatras et al. , 2019 ) OT distances . They typically have lower sample complexity than usual OT and can be accurately estimated from random mini-batches without using dual forms such as ( 3 ) . In practice , these approaches usually learn the ground OT cost c ( · , · ) . The aforementioned methods use OT cost in the ambient space to train GANs . There also exist approaches using OT cost in the latent space . For example , Tolstikhin et al . ( 2017 ) ; Patrini et al . ( 2020 ) use OT cost between encoded data and a given distribution as an additional term to reconstruction loss for training an AE . As the result , AE ’ s latent distribution becomes close to the given one . OPTIMAL TRANSPORT MAP AS THE GENERATIVE MAP ( Figure 1b ) . Methods to compute the OT map ( plan ) are less common in comparison to those computing the cost . Recovering the map from the primal form ( 1 ) or ( 2 ) usually yields complex optimization objectives containing several adversarial terms ( Xie et al. , 2019 ; Liu et al. , 2021 ; Lu et al. , 2020 ) . Such procedures require careful hyperparameter choice . This needs to be addressed before using these methods in practice . Primal-dual relationship ( M2 ) makes it possible to recover the OT map via solving the dual form ( 3 ) . Dual-form based methods primarily considerW2 cost due to its nice theoretical properties and relation to convex functions ( Brenier , 1991 ) . In the semi-discrete case ( µ is continuous , ν is discrete ) , An et al . ( 2020a ) and Lei et al . ( 2019 ) compute the dual potential and the OT map by using the Alexandrov theory and convex geometry . For the continuous case , Seguy et al . ( 2018 ) use the entropy ( quadratic ) regularization to recover the dual potentials and extract OT map from them via the barycenteric projection . Taghvaei & Jalali ( 2019 ) , Makkuva et al . ( 2020 ) , Korotin et al . ( 2021a ) employ input-convex neural networks ( ICNNs , see Amos et al . ( 2017 ) ) to parametrize potentials in the dual problem and recover OT maps by using their gradients . The aforementioned dual form methods compute OT maps in LATENT SPACES for problems such as domain adaptation and latent space mass transport , see Figure 3 . OT maps in high-dimensional ambient spaces , e.g. , natural images , are usually not considered . Recent evaluation of continuous OT methods forW2 ( Korotin et al. , 2021b ) reveals their crucial limitations , which negatively affect their scalability , such as poor expressiveness of ICNN architectures or bias due to regularization .

Summary: The paper proposes to learn the optimal transport map from the latent distribution to the data distribution directly by optimizing the W2 distance. To find the OT map, the authors replace $y$ with $T(x)$ in the Kantorovich dual problem. Experiments show that the proposed method works well and learns the image distribution successfully.