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: Noticing that widely used latent code interpolations for exploring the generative capabilities of VAEs and GANs have distribution mismatch problems, this paper proposes to utilize monotone transport map to exactly eliminate the distribution mismatch between modified interpolated codes and a prior distribution, assuming I.I.D. code components and a L1 code distance. More precisely, a transformation of the latent space operation is learnt with the objective that the distribution of the transformed variable match the prior distribution used in training the generative models. Optimal transport is used as a measure to minimize the two distributions. By restricting the class of cost functions used in the optimal transport formulation, the solution to the optimal transport problem (and hence the transformation function) has been shown to take a simple form (closed form in cases where cdf has a analytical form). Experiments on CIFAR-10, LLD-icon, LSUN, CelebA datasets show that, the minimally modified interpolated codes for several different interpolations produce samples with higher Inception Scores and better visual effects under an improved Wasserstein GAN than the original interpolated codes.