Optimal Transport Maps For Distribution Preserving Operations on Latent Spaces of Generative Models

1 INTRODUCTION & RELATED WORK . Generative models such as Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) and Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) have emerged as popular techniques for unsupervised learning of intractable distributions . In the framework of Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , the generative model is obtained by jointly training a generator G and a discriminator D in an adversarial manner . The discriminator is trained to classify synthetic samples from real ones , whereas the generator is trained to map samples drawn from a fixed prior distribution to synthetic examples which fool the discriminator . Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) are also trained for a fixed prior distribution , but this is done through the loss of an Autoencoder that minimizes the variational lower bound of the data likelihood . For both VAEs and GANs , using some data X we end up with a trained generator G , that is supposed to map latent samples z from the fixed prior distribution to output samples G ( z ) which ( hopefully ) have the same distribution as the data . In order to understand and visualize the learned model G ( z ) , it is a common practice in the literature of generative models to explore how the output G ( z ) behaves under various arithmetic operations on the latent samples z . However , the operations typically used so far , such as linear interpolation ( Goodfellow et al. , 2014 ) , spherical interpolation ( White , 2016 ) , vicinity sampling and vector arithmetic ( Radford et al. , 2015 ) , cause a distribution mismatch between the latent prior distribution and the results of the operations . This is problematic , since the generator G was trained on a fixed prior and expects to see inputs with statistics consistent with that distribution . To address this , we propose to use distribution matching transport maps , to obtain analogous latent space operations ( e.g . interpolation , vicinity sampling ) which preserve the prior distribution of the latent space , while minimally changing the original operation . In Figure 1 we showcase how our proposed technique gives an interpolation operator which avoids distribution mismatch when interpolating between samples of a uniform distribution . The points of the ( red ) matched trajectories are obtained as minimal deviations ( in expectation of l1 distance ) from the the points of the ( blue ) linear trajectory . 1.1 DISTRIBUTION MISMATCH AND RELATED APPROACHES . In the literature there are dozens of papers that use sample operations to explore the learned models ( Bengio et al . ( 2013 ) ; Goodfellow et al . ( 2014 ) ; Dosovitskiy et al . ( 2015 ) ; Reed et al . ( 2016 ) ; Brock et al . ( 2016 ) ; Reed et al . ( 2016 ) to name a few ) , but most of them have ignored the problem of distribution mismatch . Kingma & Welling ( 2013 ) and Makhzani et al . ( 2015 ) sidestep the problem when visualizing their models , by not performing operations on latent samples , but instead restrict the latent space to 2-d and uniformly sample the percentiles of the distribution on a 2-d grid . This way , the samples have statistics that are consistent with the prior distribution . However , this approach does not scale up to higher dimensions - whereas the latent spaces used in the literature can have hundreds of dimensions . White ( 2016 ) experimentally observe that there is a distribution mismatch between the norm for points drawn from uniform or Gaussian distribution and points obtained with linear interpolation ( SLERP ) , and ( heuristically ) propose to use a so-called spherical linear interpolation to reduce the mismatch , obtaining higher quality interpolated samples . While SLERP has been subjectively observed to produce better looking samples than linear interpolation and is now commonly , its heuristic nature has limited it from fully replacing the linear interpolation . Furthermore , while perhaps possible it is not obvious how to generalize it to other operations , such as vicinity sampling , n-point interpolation and random walk . In Section 2 we show that for interpolation , in high dimensions SLERP tends to approximately perform distribution matching the approach taken by our framework which can explain why it works well in practice . Kilcher et al . ( 2018 ) further analyze the ( norm ) distribution mismatch observed by White ( 2016 ) ( in terms of KL-Divergence ) for the special case of Gaussian priors , and propose an alternative prior distribution with dependent components which produces less ( but still nonzero ) distribution mismatch for linear interpolation , at the cost of needing to re-train and re-tune the generative models . In contrast , we propose a framework which allows one to adapt generic operations , such that they fully preserve the original prior distribution while being faithful to the original operation . Thus the KL-Divergence between the prior and the distribution of the results from our operations is zero . The approach works as follows : we are given a ‘ desired ’ operation , such as linear interpolation y = tz1 + ( 1− t ) z2 , t ∈ [ 0 , 1 ] . Since the distribution of y does not match the prior distribution of z , we search for a warping f : Rd → Rd , such that ỹ = f ( y ) has the same distribution as z . In order to have the modification ỹ as faithful as possible to the original operation y , we use optimal transform Published as a conference paper at ICLR 2019 Operation Expression 2-point interpolation y = tz1 + ( 1− t ) z2 , t ∈ [ 0 , 1 ] n-point interpolation y = ∑n i=1 tizi with ∑ i ti = 1 maps ( Santambrogio , 2015 ; Villani , 2003 ; 2008 ) to find a minimal modification of y which recovers the prior distribution z . This is illustrated in Figure 1a , where each point ỹ of the matched curve is obtained by warping a corresponding point y of the linear trajectory , while not deviating too far from the line . 2 FROM DISTRIBUTION MISMATCH TO OPTIMAL TRANSPORT . With implicit models such as GANs ( Goodfellow et al. , 2014 ) and VAEs ( Kingma & Welling , 2013 ) , we use the data X , drawn from an unknown random variable x , to learn a generator G : Rd 7→ Rd′ with respect to a fixed prior distribution pz , such that G ( z ) approximates x . Once the model is trained , we can sample from it by feeding latent samples z through G. We now bring our attention to operations on latent samples z1 , · · · , zk from pz , i.e . mappings κ : Rd × · · · × Rd → Rd . ( 1 ) We give a few examples of such operations in Table 1 . Since the inputs to the operations are random variables , their output y = κ ( z1 , · · · , zk ) is also a random variable ( commonly referred to as a statistic ) . While we typically perform these operations on realized ( i.e . observed ) samples , our analysis is done through the underlying random variable y . The same treatment is typically used to analyze other statistics over random variables , such as the sample mean , sample variance and test statistics . In Table 1 we show example operations which have been commonly used in the literature . As discussed in the Introduction , such operations can provide valuable insight into how the trained generator G changes as one creates related samples y from some source samples . The most common such operation is the linear interpolation , which we can view as an operation yt = tz1 + ( 1− t ) z2 , ( 2 ) where z1 , z2 are latent samples from the prior pz and yt is parameterized by t ∈ [ 0 , 1 ] . Now , assume z1 and z2 are i.i.d , and let Z1 , Z2 be their ( scalar ) first components with distribution pZ . Then the first component of yt is Yt = tZ1 + ( 1− t ) Z2 , and we can compute : Var [ Yt ] = Var [ tZ1 + ( 1− t ) Z2 ] = t2Var [ Z1 ] + ( 1− t ) 2Var [ Z2 ] = ( 1 + 2t ( t− 1 ) ) Var [ Z ] . ( 3 ) Since ( 1+ 2t ( t− 1 ) ) 6= 1 for all t ∈ [ 0 , 1 ] \ { 0 , 1 } , it is in general impossible for yt to have the same distribution as z , which means that distribution mismatch is inevitable when using linear interpolation . A similar analysis reveals the same for all of the operations in Table 1 . This leaves us with a dilemma : we have various intuitive operations ( see Table 1 ) which we would want to be able to perform on samples , but their resulting distribution pyt is inconsistent with the distribution pz we trained G for . Due to the curse of dimensionality , as empirically observed by White ( 2016 ) , this mismatch can be significant in high dimensions . We illustrate this in Figure 2 , where we plot the distribution of the squared norm ‖yt‖2 for the midpoint t = 1/2 of linear interpolation , compared to the prior distribution ‖z‖2 . With d = 100 ( a typical dimensionality for the latent space ) , the distributions are dramatically different , having almost no common support . Kilcher et al . ( 2018 ) quantify this mismatch for Gaussian priors in terms of KL-Divergence , and show that it grows linearly with the dimension d. In Appendix A ( see Supplement ) we expand this analysis and show that this happens for all prior distributions with i.i.d . entries ( i.e . not only Gaussian ) , both in terms of geometry and KL-Divergence . 2.1 DISTRIBUTION MATCHING WITH OPTIMAL TRANSPORT . In order to address the distribution mismatch , we propose a simple and intuitive framework for constructing distribution preserving operators , via optimal transport : Strategy 1 ( Optimal Transport Matched Operations ) . 1 . We construct an ’ intuitive ’ operator y = κ ( z1 , · · · , zk ) . 2 . We analytically ( or numerically ) compute the resulting ( mismatched ) distribution py 3 . We search for a minimal modification ỹ = f ( y ) ( in the sense that Ey [ c ( ỹ , y ) ] is minimal with respect to a cost c ) , such that distribution is brought back to the prior , i.e . pỹ = pz . The cost function in step 3 could e.g . be the euclidean distance c ( x , y ) = ‖x − y‖ , and is used to measure how faithful the modified operator , ỹ = f ( κ ( z1 , · · · , zk ) ) is to the original operator k. Finding the map f which gives a minimal modification can be challenging , but fortunately it is a well studied problem from optimal transport theory . We refer to the modified operation ỹ as the matched version of y , with respect to the cost c and prior distribution pz . For completeness , we introduce the key concepts of optimal transport theory in a simplified setting , i.e . assuming probability distributions are in euclidean space and skipping measure theoretical formalism . We refer to Villani ( 2003 ; 2008 ) and Santambrogio ( 2015 ) for a thorough and formal treatment of optimal transport . The problem of step ( 3 ) above was first posed by Monge ( 1781 ) and can more formally be stated as : Problem 1 ( Santambrogio ( 2015 ) Problem 1.1 ) . Given probability distributions px , py , with domains X , Y respectively , and a cost function c : X × Y → R+ , we want to minimize inf { Ex∼px [ c ( x , f ( x ) ) ] ∣∣∣f : X → Y , f ( x ) ∼ py } ( MP ) We refer to the minimizer f∗X → Y of ( MP ) ( if it exists ) , as the optimal transport map from px to py with respect to the cost c. However , the problem remained unsolved until a relaxed problem was studied by Kantorovich ( 1942 ) : Problem 2 ( Santambrogio ( 2015 ) Problem 1.2 ) . Given probability distributions px , py , with domains X , Y respectively , and a cost function c : X × Y → R+ , we want to minimize inf { E ( x , y ) ∼px , y [ c ( x , y ) ] ∣∣∣ ( x , y ) ∼ px , y , x ∼ px , y ∼ py } , ( KP ) where ( x , y ) ∼ px , y , x ∼ px , y ∼ py denotes that ( x , y ) have a joint distribution px , y which has ( previously specified ) marginals px and py . We refer to the joint px , y which minimizes ( KP ) as the optimal transport plan from px to py with respect to the cost c. The key difference is to relax the deterministic relationship between x and f ( x ) to a joint probability distribution px , y with marginals px and py for x and y . In the case of Problem 1 , the minimization might be over the empty set since it is not guaranteed that there exists a mapping f such that f ( x ) ∼ y . In contrast , for Problem 2 , one can always construct a joint density px , y with px and py as marginals , such as the trivial construction where x and y are independent , i.e . px , y ( x , y ) : = px ( x ) py ( y ) . Note that given a joint density px , y ( x , y ) over X × Y , we can view y conditioned on x = x for a fixed x as a stochastic function f ( x ) from X to Y , since given a fixed x do not get a specific function value f ( x ) but instead a random variable f ( x ) that depends on x , with f ( x ) ∼ y|x = x with density py ( y|x = x ) : = px , y ( x , y ) px ( x ) . In this case we have ( x , f ( x ) ) ∼ px , y , so we can view the Problem KP as a relaxation of Problem MP where f is allowed to be a stochastic mapping . While the relaxed problem of Kantorovich ( KP ) is much more studied in the optimal transport literature , for our purposes of constructing operators it is desirable for the mapping f to be deterministic as in ( MP ) ( see Appendix C for a more detailed discussion on deterministic vs stochastic operations ) . To this end , we will choose the cost function c such that the two problems coincide and where we can find an analytical solution f or at least an efficient numerical solution . In particular , we note that the operators in Table 1 are all pointwise , such that if the points zi have i.i.d . components , then the result y will also have i.i.d . components . If we combine this with the constraint for the cost c to be additive over the components of x , y , we obtain the following simplification : Theorem 1 . Suppose px and py have i.i.d components and c over X × Y = Rd × Rd decomposes as c ( x , y ) = d∑ i=1 C ( x ( i ) , y ( i ) ) . ( 4 ) Consequently , the minimization problems ( MP ) and ( KP ) turn into d identical scalar problems for the distributions pX and pY of the components of x and y : inf { EX∼pX [ C ( X , T ( X ) ) ] ∣∣∣T : R→ R , T ( X ) ∼ pY } ( MP-1-D ) inf { E ( X , Y ) ∼pX , Y [ C ( X , Y ) ] ∣∣∣ ( X , Y ) ∼ pX , Y , X ∼ pX , Y ∼ pY } , ( KP-1-D ) such that an optimal transport map T for ( MP-1-D ) gives an optimal transport map f for ( MP ) by pointwise application of T , i.e . f ( x ) ( i ) : = T ( x ( i ) ) , and an optimal transport plan pX , Y for ( KP-1-D ) gives an optimal transport plan px , y ( x , y ) : = ∏d i=1 pX , Y ( x ( i ) , y ( i ) ) for ( KP ) . Proof . See Appendix . Fortunately , under some mild constraints , the scalar problems have a known solution : Theorem 2 ( Theorem 2.9 in Santambrogio ( 2015 ) ) . Let h : R → R+ be convex and suppose the cost C takes the form C ( x , y ) = h ( x− y ) . Given an continuous source distribution pX and a target distribution pY on R having a finite optimal transport cost in ( KP-1-D ) , then TmonX→Y ( x ) : = F [ −1 ] Y ( FX ( x ) ) , ( 5 ) defines an optimal transport map from pX to pY for ( MP-1-D ) , where FX ( x ) : = ∫ x −∞ pX ( x ′ ) dx′ is the Cumulative Distribution Function ( CDF ) of X and F [ −1 ] Y ( y ) : = inf { t ∈ R|FY ( t ) ≥ y } is the pseudo-inverse of FY . Furthermore , the joint distribution of ( X , TmonX→Y ( X ) ) defines an optimal transport plan for ( KP-1-D ) . The mapping TmonX→Y ( x ) in Theorem 2 is non-decreasing and is known as the monotone transport map from X to Y . It is easy to verify that TmonX→Y ( X ) has the distribution of Y , in particular FX ( X ) ∼ Uniform ( 0 , 1 ) and if U ∼ Uniform ( 0 , 1 ) then F [ −1 ] Y ( U ) ∼ Y . Now , combining Theorems 1 and 2 , we obtain a concrete realization of the Strategy 1 outlined above . We choose the cost c such that it admits to Theorem 1 , such as c ( x , y ) : = ‖x − y‖1 , and use an operation that is pointwise , so we just need to compute the monotone transport map in ( 5 ) . That is , if z has i.i.d components with distribution pZ , we just need to compute the component distribution pY of the result y of the operation , the CDFs FZ , FY and obtain TmonY→Z ( y ) : = F [ −1 ] Z ( FY ( y ) ) ( 6 ) as the component-wise modification of y , i.e . ỹ ( i ) : = TmonY→Z ( y ( i ) ) . In Figure 3 we show the monotone transport map for the linear interpolation y = tz1 + ( 1− t ) z2 for various values of t. The detailed calculations and examples for various operations are given in Appendix B , for both Uniform and Gaussian priors .

Summary: This paper considers the issue of distribution mismatch between the input data used for training generative models and the new data for new instance generation. Given a sample operation, the authors propose to use the so-called optimal transport to map the distribution of the new data to that of the input data that were used training. The optimal transport is essentially a monotonic transformation as the composite of the inverse of the target distribution and the source distribution.