AE-OT: A NEW GENERATIVE MODEL BASED ON EXTENDED SEMI-DISCRETE OPTIMAL TRANSPORT

1 INTRODUCTION . Generative adversarial networks ( GANs ) ( Goodfellow et al . ( 2014 ) ) and variational autoencoders ( VAEs ) ( Kingma & Welling ( 2013 ) ) emerge as the dominant approaches for unconditional image generation . When trained on appropriate datasets , they are able to produce realistic and visual appealing samples . GAN methods train an unconditional generator that regresses real images from random noise and a discriminator that measures the difference between generated samples and real images . Despite GANs ’ advantages , they have critical drawbacks . 1 ) Training of GANs are tricky and sensitive to hyperparameters . 2 ) GANs suffer from mode collapse , in which the generator only learns to generate few modes of data distribution while missing others , although samples from the missing modes occur throughout the training data ( see e.g . Goodfellow ( 2016 ) ) . While for the VAEs , the encoder is used to map the data distribution to a Gaussian latent distribution , which is then mapped back to the data distribution by the decoder . While standard VAEs tend to capture all modes , they often generate ambiguous images on multimodal real data distributions . We propose that these phenomena relates deeply with the singularities of distribution transport maps . Manifold Distribution Hypothesis In deep learning , the manifold distribution hypothesis is well accepted , which assumes the distribution of a specific class of natural data is concentrated on a low dimensional manifold embedded in the high dimensional data space Tenenbaum et al . ( 2000 ) . Therefore , GANs and VAEs implicitly aim to accomplish two major tasks : 1 ) manifold embedding : to find the encoding/decoding maps between the data manifold embedded in the image space and the latent space ; 2 ) probability distribution transport : to transport a given white noise distribution to the data distribution , either in the latent or in the image space . Distribution Transformation The generator of GAN model and the decoder of VAE model are trained to compute a transport map that transforms a known continuous distribution ( e.g . Gaussian white noise ) to the real data distribution . Namely , the transport map pushes forward the given noise distribution to a generated distribution to approximate the real data distribution , the similarity between the two distributions determines the generalization ability of the generator Ben-David et al . ( 2010 ) . Discontinuity and Mode Collapse/Mixture It is a common practice among GAN/VAE models that the generators/decoders are expressed by deep neural networks , which can only represent continuous mappings . Unfortunately , as pointed out by works Nagarajan & Kolter ( 2017 ) ; Khayatkhoei et al . ( 2018 ) ; Xiao et al . ( 2018 ) , the transport maps may be discontinuous when there are multiple modes in the data distribution . This intrinsic conflict can cause mode collapse or mode mixture . The later means that the generated samples are mixtures of multiple modes and look spurious or ambiguous . 1 Furthermore , even when the real data distribution has a single mode , ambiguous data ( e.g . a human face image with mismatched eye colors ) can still present . This can be explained by Brenier ’ s polar factorization theorem Brenier ( 1991b ; 1987 ; 1991a ) and Figalli ’ s regularity theorem Figalli ( 2010 ) ; Chen & Figalli ( 2017 ) ( Thm . 5 in Appendix B ) , which asserts that if the support of the target distribution is not convex , then there will be singularity sets on the support of the source distribution , such that the transport map is discontinuous on these sets . This shows the intrinsic difficulties of conventional GANs/VAEs can not be eliminated , as shown in Fig . 1 . Conquering Mode Collapse/Mixture However , according to Brenier ( 1987 ; 1991a ) theorem , the optimal transport map can be represented as the gradient map of the Brenier potential . At the regular points , the Brenier potential is differentiable , its gradient map ( the transport map ) is continuous ; at the singularities , the Brenier potential is continuous but not differentiable , and its gradient map is discontinuous . Conventional GANs and VAEs model the gradient map directly and encounter the trouble of discontinuity . In contrast , we propose to model the globally continuous Brenier potential to avoid mode collapse/mixture . More specifically , our proposed AE-OT model separates the manifold embedding step and the probability distribution transformation step , the former is carried out by an autoencoder ( AE ) , the latter is accomplished by a convex optimization framework ( OT ) . The OT step computes the Brenier potential explicitly and is able to locate the singularity set ( the discontinuous points of the gradient map ) based on Figalli ’ s theory . Our experimental results demonstrate that the proposed method can not only cover all of the modes , but also avoid generating spurious samples ( mode mixture ) . Contributions ( i ) From theoretical aspect , this work gives a thorough explanation of mode collapse and mode mixture by the regularity theory of optimal transportation developed by Figalli ( 2018 1For example , a generator generates obscure digits mixing 0 and 8 but neither 0 nor 8 on the MNIST dataset . Fields medalist ) and the reasons why standard GANs/VAEs can not solve this problem perfectly . ( ii ) From practical aspect , this work proposes a novel model called AE-OT , which first encodes the data manifold into the latent space , then compute the Brenier potential to represent the optimal transportation map in the latent space . The Figalli ’ s singularity set can be located efficiently and avoided when generating new samples . In this way , our model eliminates mode collapse and mode mixture successfully . ( iii ) The algorithm for finding the Brenier potential and the optimal transportation map can be accelerated with GPU based convex optimization algorithm . The method converges to the unique global optimum with bounded error estimate . ( iv ) Our experiment results demonstrate the efficiency and efficacy of the proposed method . 2 RELATED WORK . Optimal Transport Optimal transport ( OT ) has been successfully applied in the areas like economy , optics and machine learning , as surveyed in Solomon ( 2018 ) and Peyré & Cuturi ( 2018 ) . In Gu et al . ( 2016 ) , the intrinsic connection between Brenier theory in OT and Alexandroff theory in convex geometry was established , and applied for deep learning in Lei et al . ( 2017 ) by an convex optimization . Figalli and the collaborators Figalli ( 2010 ) ; Chen & Figalli ( 2017 ) proposed that when the support of the data distribution is non-convex , the transport map will be discontinuous . Generative models Generative model is one of the main tasks in the machine learning field . One of the most used image generation methods is Variational Autoencoders ( VAEs ) , where the decoders approximate real data distributions from a Gaussian distribution in a variational approach ( e.g Kingma & Welling ( 2013 ) and Rezende et al . ( 2014 ) ) . Later , Adversarial Autoencoders ( AAEs ) Makhzani et al . ( 2015 ) and Wasserstein Autoencoders ( WAEs ) Tolstikhin et al . ( 2018 ) are proposed following the similar scheme . Although VAEs are relatively simple to train , images they generate look blurry . Generative Adversarial Networks ( GANs ) Goodfellow et al . ( 2014 ) proposed by Goodfellow et.al can produce better quality images . While being a powerful tool in generating realistically looking samples , GANs suffer from the mode collapse problem . To solve this problem , a huge number of methods , including changing the loss function ( e.g . Wasserstein GAN Arjovsky et al . ( 2017 ) ) , regularizing the discriminators to be Lipschitz ( clipping Arjovsky et al . ( 2017 ) , gradient regularization Gulrajani et al . ( 2017 ) , Mescheder et al . ( 2018 ) or spectral normalization Miyato et al . ( 2018 ) ) , were proposed . Besides , various non-adversarial methods has also been proposed recently . GLO Bojanowski et al . ( 2017 ) adopts the “ encoder-less autoencoder '' method to generate new images with a non-adversarial loss function . IMLE Li & Malik ( 2018 ) proposed an ICP related generative model training approach . Later GLANN Hoshen & Malik ( 2019 ) combines the advantages of GLO and IMLE , an embedding from image space to latent space was first found using GLO and then a transformation between an arbitrary distribution and latent code was computed using IMLE . Mitigating Mode Collapsing Recently , Nagarajan & Kolter ( 2017 ) ; Khayatkhoei et al . ( 2018 ) ; Xiao et al . ( 2018 ) also realize the training difficulties of GANs come from the approximation of discontinuous functions with continuous DNNs . By the gradient-based regularization , GDGAN Nagarajan & Kolter ( 2017 ) do relieve the mode collapse phenomenon of GANs , but mode mixture still exists . Khayatkhoei et al . ( 2018 ) proposes to use multiple GANs to overcome the mode collapse . Xiao et al . ( 2018 ) proposed to embed the images into a latent space according to Bourgain ’ s theorem , and train the generator by sampling a Gaussian mixture distribution in the latent space instead of a unimodal Gaussian . The recently introduced normalized diversification by Liu et al . ( 2018 ) can also help overcome mode collapse successfully . However , all of them can not solve the mode mixture well . All these works Nagarajan & Kolter ( 2017 ) ; Khayatkhoei et al . ( 2018 ) ; Xiao et al . ( 2018 ) explain that if the target data distribution has multiple modes , the transport map is discontinuous , but DNNs can only represent continuous mappings , the intrinsic conflict causes mode collapse . 3 COMPUTATIONAL ALGORITHMS . Overview of AE-OT Model Our AE-OT model is summarized in Fig . 2 , it has two major components : i ) ( AE ) An autoencoder is trained to encode ( fθ ) the data manifold from the image space X to the latent space Z , and map the data distribution to the latent code distribution ; then the decoder gξ decodes the latent code back to the data manifold . ii ) ( OT ) This module computes the optimal transportation map T from the noise distribution to the latent code distribution . First , the Brenier potential is found by a convex optimization process according to Gu et al . ( 2016 ) , whose gradient is the semi-discrete optimal transport map , where the target is the discrete set of latent codes of training samples ; then the transport map is piece-wise linearly extended to a global continuous map T̃ , where the image domain becomes a simplicial complex obtained by triangulating the above latent codes . Finally , the singularity set in the source domain is located and avoided when generating new samples . As a result , given a random noise x , we can get the generated image by gξ ◦ T̃ ( x ) . Semi-Discrete OT Map Suppose the source measure µ ( Gaussian or uniform distribution ) is absolutely continuous defined on a convex domain Ω ⊂ Rd , the target domain is a discrete set , Y = { y1 , y2 , · · · , yn } , yi ∈ Rd , the target measure is a Dirac measure , ν = ∑n i=1 νiδ ( y − yi ) , i = 1 , 2 , . . . , n , with the equal total mass as the source measure , µ ( Ω ) = ∑n i=1 νi . Under a semi- discrete transport map T : Ω → Y , a cell decomposition is induced Ω = ⋃n i=1Wi , such that every x in each cell Wi is mapped to the target yi , T : x ∈ Wi 7→ yi . The map T is measure preserving , denoted as T # µ = ν , if the µ-volume of each cell Wi equals to the ν-measure of the image T ( Wi ) = yi , µ ( Wi ) = νi . The cost function is given by c : Ω × Y → R , where c ( x , y ) represent the cost for transporting a unit mass from x to y . The total cost of T is given by ∫ Ω c ( x , T ( x ) ) dµ ( x ) = ∑n i=1 ∫ Wi c ( x , yi ) dµ ( x ) . Semi-discrete optimal transport map is the measure-preserving map that minimizes the total cost , T ∗ : = arg minT # µ=ν ∫ Ω c ( x , T ( x ) ) dµ ( x ) . When the cost function is the L2 distance c ( x , y ) = 1/2‖x − y‖2 , Brenier ’ s theorem claims that the semi-discrete OT map is given by the gradient map of a piece-wise ( PL ) convex function , the so-called Brenier potential uh : Ω→ R , uh ( x ) : = maxni=1 { πh , i ( x ) } , where πh , i ( x ) = 〈x , yi〉+ hi is the hyperplane corresponding to yi ∈ Y . As shown in Fig . 3 ( a ) , the projection of the graph of uh decomposes Ω into cells Wi ( h ) , each cell Wi ( h ) is the projection of the supporting plane πh , i ( x ) . The height vector h is the unique optimizer of the following convex energy under the condition that∑ i hi = 0 , E ( h ) = ∫ h 0 n∑ i=1 wi ( η ) dηi − n∑ i=1 hiνi , ( 1 ) where wi ( η ) is the µ-volume ofWi ( η ) . The convex energyE ( h ) can be optimized simply by gradient descend method with∇E ( h ) = ( wi ( h ) − νi ) T . The key is to compute the µ-volume wi ( h ) of each cell Wi ( h ) , which can be estimated using conventional Monte Carlo method . We draw N random samples from µ distribution , { xj } ∼i.i.d . µ , ∀j ∈ J , the estimated µ-volume of each cell is ŵi ( h ) = # { j ∈ J | xj ∈ Wi ( h ) } /N . Given xj , we can find Wi in which xj ∈ Wi by i = arg maxi { 〈xj , yi〉 + hi } , i = 1 , 2 , . . . , n. When N is large enough , ŵi ( h ) converges to wi ( h ) . Then the gradient of the energy is approximated as ∇E ≈ ( ŵi ( h ) − νi ) T . Once the gradient is estimated , we can use Adam algorithm Kingma & Ba ( 2015 ) to minimize the energy . Sampling of x is independent of each other and finding the cell that x is located only involves matrix multiplication and sorting . Hence the Monte Carlo method has a natural parallel computation implementation on GPUs . The approximation error is proportional to the inverse of the square root the amount of Monte Carlo samples . Asymptotically the number of Monte Carlo samples increases exponentially with respect to the dimension d ( see e.g . Weed & Bach ( 2017 ) ) . This brings huge computational burdens . To find a good balance between precision and speed , we adaptively adjust the number of random samples . In practice , we apply the following strategy : if the energy E ( h ) ceases decreasing for a number of consecutive steps , we double the amount of Monte Carlo samples . The algorithmic details of semi-discrete OT map are summarized in Alg . 1 . Algorithm 1 Semi-discrete OT Map 1 : Input : Latent codes Y = { yi } i∈I , empirical latent code dis- tribution ν = 1|I| ∑ i∈I δyi , number of Monte Carlo samples N , positive integer s. 2 : Output : Optimal transport map T ( · ) . 3 : Initialize h = ( h1 , h2 , . . . , h|I| ) ← ( 0 , 0 , . . . , 0 ) . 4 : repeat 5 : GenerateN uniformly distributed samples { xj } Nj=1 . 6 : Calculate∇h = ( ŵi ( h ) − νi ) T . 7 : ∇h = ∇h−mean ( ∇h ) . 8 : Update h by Adam algorithm with β1 = 0.9 , β2 = 0.5 . 9 : ifE ( h ) has not decreased for s steps then 10 : N ← N × 2 . 11 : end if 12 : until Converge 13 : OT map T ( · ) ← ∇ ( maxi〈· , yi〉+ hi ) . Algorithm 2 Generate latent code 1 : Input : Optimal transport map T ( · ) , number of samples to gen- erate n , angle threshold θ̂ . 2 : Output : Generated latent code P . 3 : Compute ĉi by Monte Carlo method . 4 : repeat 5 : Sample x ∼ µ , Find the smallest d + 1 vertex around x as { d ( x , ĉi0 ) , d ( x , ĉi1 ) , . . . , d ( x , ĉid ) } . 6 : Compute dihedral angles θik between πi0 and πik . 7 : Select θik with θik ≤ θ̂ , result in îk = 0 , 1 , . . . , d1 . 8 : if ∀k , θik > θ̂ then Abandon x 9 : else Generate latent code T̃ ( x ) = ∑d1 k=0 λkT ( ĉîk ) with λk = d −1 ( x , ĉîk ) / ∑d1 j=0 d −1 ( x , ĉîj ) . 10 : end if 11 : until Generate n new latent code Piece-wise Linear Extension The semi-discrete OT map ∇uh : Ω → Y maps all x ∈ Ω to the latent codes of training samples { yi } ’ s and won ’ t generate new samples . Therefore , we extend the semi-discrete OT map T = ∇uh to a piecewise linear ( PL ) mapping T̃ as follows . The projection of uh in the source domain induces a cell decomposition of Ω , of which each cell is of µ-volume νi and is mapped to the corresponding yi . By representing the cells by their µ-mass centers as ci : = ∫ Wi ( h ) xdµ ( x ) , we can get the point-wise map t : ci 7→ yi . The Poincaré of the cell decomposition induces a triangulation of the centers C = { ci } : if Wi ∩Wj 6= ∅ , then ci is connected with cj to form an edge [ ci , cj ] . Similarly , if Wi0 ∩Wi1 · · · ∩Wik 6= ∅ , then there is a k-dimensional simplex [ ci0 , ci1 , . . . , cik ] . All these simplices form a triangulation of C ( a simplicial complex ) , denoted as T ( C ) ( the green triangles in the left of Fig . 3 ( b ) ) . We can triangulate Y in the same way to obtain the triangulation T ( Y ) ( the green triangles in the right of Fig . 3 ( b ) ) . Once a random sample x is drawn from the distribution µ , we can find the simplex σ in T ( C ) containing x . Assume the simplex σ has d+ 1 vertices { ci0 , ci1 , . . . , cid } , the bary-centric coordinates of x in σ is defined as x = ∑d k=0 λkcik , and ∑d k=0 λk = 1 with all λk non-negative . Then the generated latent code of x under this piece-wise linear map is given by T̃ ( x ) = ∑d k=0 λkyik ( In Fig . 3 ( b ) , the green dot x1 is mapped to be T̃ ( x1 ) ) . Because all of the yis are used to construct the simplicial complex T ( Y ) in the support of the target distribution , we can guarantee that no mode is lost . In practice , the µ-mass center ci is approximated by the mean value of all the Monte-Carlo samples inside Wi ( h ) , ĉi = ∑ xj∈Wi xj/ # { xj ∈ Wi } , where xj ∼ µ . The connectivity information T ( C ) is too complicated to construct and to store in high dimensional space , thus T ( C ) is not explicitly built . Instead , we find the simplex σ ∈ T ( C ) containing x as follows : given a random point x ∈ Ω , evaluate and sort its Euclidean distances to the centers d ( x , ĉi ) , i = 1 , 2 , . . . , n in the ascending order . Suppose the first d+ 1 items are { d ( x , ĉi0 ) , d ( x , ĉi1 ) , . . . , d ( x , ĉid ) } , then σ is formed by { ĉik } . The bary-centric coordinates λ̂ik are estimated as λ̂ik = d −1 ( x , ĉik ) / ∑d k=0 d −1 ( x , ĉik ) . However , this may generate some spurious samples . To overcome it , we need further to detect the singular set . Singular Set Detection According to Figalli ’ s theory Figalli ( 2010 ) ; Chen & Figalli ( 2017 ) , if there are multiple modes or the support of the target distribution is concave , there will be singular sets Σ ⊂ Ω , where the Brenier potential is continuous but not differentiable , making its gradient map , i.e . the transport map , discontinuous . As shown in Fig . 3 ( a ) , the source distribution is uniformly defined on Ω , and the target empirical distribution has two modes . There is one ridge ( the red line ) on the Brenier potential uh , whose projection is the singular set Σ ( the red line in Ω ) . Ω \ Σ consists of two connected components , each of them is mapped onto a single mode . Σ consists of codimension 1 facets of cells . If Wi ( h ) ∩ Wj ( h ) ⊂ Σ , then the dihedral angle between two supporting planes πh , i and πh , j of uh is prominently large . Therefore , on the graph of Brenier potential , we pick the pairs of facets whose dihedral angles are larger than a given threshold , the projection of their intersection gives a co-dimension 1 cell in the singular set Σ . During the generation process , if a random sample x is around Σ , it will be mapped by T̃ to the gaps among the modes . When generating new latent codes , we just abandon such samples . And this help our method prevent the mode mixture phenomenon . Given the extended OT map T̃ ( x ) , some of the polyhedrons transverse the singular set ( the red lines of Fig .3 ( b ) ) , which means that different vertices of the polyhedron belongs to different mode . If the sample x falls into such a polyhedron ( the dotted red triangle ) , we just abandon it ( as shown in Fig . 3 ( b ) , the red dot x2 is just abandoned ) . Specifically , given x , we can detect if it belongs to the singular set by checking the angles θik between πi0 and πik , k = 1 , 2 , . . . , d as θik = 〈yi0 , yik〉/‖yi0‖ · ‖yik‖ . If all of the angles θik is larger than a threshold θ̂ , we say x belongs to the singular set and just abandon it . Or we just select a subset { πik } with θik ≤ θ̂ , denoted as { πîk , k = 0 , 1 , . . . , d1 } . Then we can compute λk = d−1 ( x , ĉîk ) / ∑d1 j=0 d −1 ( x , ĉîj ) and T̃ ( x ) = ∑d1 k=0 λkT ( ĉîk ) . Intuitively , T̃ ( · ) smooths the discrete function T ( · ) in regions where latent codes are dense and keep the discontinuity of T ( · ) where latent codes are very sparse . In this way we avoid generating spurious latent code and thus improve the generation quality . The algorithm to generate new code is shown in Alg . 2 and the effect of threshold filtering is further investigated in Appendix 4.1 .

Summary: General Comments:  The generator in Generative Adversarial Networks (GANS) computes an optimal transportation from the noise distribution to the data distribution.  However, such maps are in general discontinuous.  Since deep neural networks can only represent continuous maps, this brings two problems: mode collapse and mode mixture. This paper approaches both problems using Figalli's regularity theory. They separate the manifold embedding (here an autoencoder maps input data to a latent space) from the optimal transportation (this map is found by convex optimization). Composing these two steps yields the proposed method. Their method basically avoids representing discontinuous maps by the generator. Empirically, the proposed method performs similar or better than state-of-the-art.