Symmetric Wasserstein Autoencoders

1 INTRODUCTION . Deep generative models have emerged as powerful frameworks for modelling complex data . Widely used families of such models include Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , Variational Autoencoders ( VAEs ) ( Rezende et al. , 2014 ; Kingma & Welling , 2014 ) , and autoregressive models ( Uria et al. , 2013 ; Van Oord et al. , 2016 ) . The VAE-based framework has been popular as it yields a bidirectional mapping , i.e. , it consists of both an inference model ( from data to latent space ) and a generative model ( from latent to data space ) . With an inference mechanism VAEs can provide a useful latent representation that captures salient information about the observed data . Such latent representation can in turn benefit downstream tasks such as clustering , classification , and data generation . In particular , the VAE-based approaches have achieved impressive performance results on challenging real-world applications , including image synthesizing ( Razavi et al. , 2019 ) , natural text generation ( Hu et al. , 2017 ) , and neural machine translation ( Sutskever et al. , 2014 ) . VAEs aim to maximize a tractable variational lower bound on the log-likelihood of the observed data , commonly called the ELBO . Since VAEs focus on modelling the marginal likelihood of the data instead of the joint likelihood of the data and the latent representation , the quality of the latent is not well assessed ( Alemi et al. , 2017 ; Zhao et al. , 2019 ) , which is undesirable for learning useful representation . Besides the perspective of maximum-likelihood learning of the data , the objective of VAEs is equivalent to minimizing the KL divergence between the encoding and the decoding distributions , with the former modelling the joint distribution of the observed data and the latent representation induced by the encoder and the latter modelling the corresponding joint distribution induced by the decoder . Such connection has been revealed in several recent works ( Livne et al. , 2019 ; Esmaeili et al. , 2019 ; Pu et al. , 2017b ; Chen et al. , 2018 ) . Due to the asymmetry of the KL divergence , it is highly likely that the generated samples are of a low probability in the data distribution , which often leads to unrealistic generated samples ( Li et al. , 2017b ; Alemi et al. , 2017 ) . A lot of work has proposed to improve VAEs from different perspectives . For example , to enhance the latent expressive power VampPrior ( Tomczak & Welling , 2018 ) , normalizing flow ( Rezende & Mohamed , 2015 ) , and Stein VAEs ( Pu et al. , 2017a ) replace the Gaussian distribution imposed on the latent variables with a more sophisticated and flexible distribution . However , these methods are all based on the objective of VAEs , which therefore are unable to alleviate the limitation of VAEs induced by the objective . To improve the latent representation ( Zhao et al. , 2019 ) explicitly includes the mutual information between the data and the latent into the objective . Moreover , to address the asymmetry of the KL divergence in VAEs ( Livne et al. , 2019 ; Chen et al. , 2018 ; Pu et al. , 2017b ) leverage a symmetric divergence measure between the encoding and the decoding distributions . Nevertheless , these methods typically involve a sophisticated objective function that either depends on unstable adversarial training or challenging approximation of the mutual information . In this paper , we leverage Optimal Transport ( OT ) ( Villani , 2008 ; Peyré et al. , 2019 ) to symmetrically match the encoding and the decoding distributions . The OT optimization is generally challenging particularly in high dimension , and we address this difficulty by transforming the OT cost into a simpler form amenable to efficient numerical implementation . Owing to the symmetric treatment of the observed data and the latent representation , the local structure of the data can be implicitly preserved in the latent space . However , we found that with the symmetric treatment only the performance of the generative model may not be satisfactory . To improve the generative model we additionally include a reconstruction loss into the objective , which is shown to significantly benefit the quality of the generation and reconstruction . Our contributions can be summarized as follows . Firstly , we propose a new family of generative autoencoders , called Symmetric Wasserstein Autoencoders ( SWAEs ) . Secondly , we adopt a learnable latent prior , parameterized as a mixture of the conditional priors given the learnable pseudo-inputs , which prevents SWAEs from over-regularizing the latent variables . Thirdly , we empirically perform an ablation study of SWAEs in terms of the KNN classification , denoising , reconstruction , and sample generation . Finally , we empirically verify , using benchmark tasks , the superior performance of SWAEs over several state-of-the-art generative autoencoders . 2 SYMMETRIC WASSERSTEIN AUTOENCODERS . In this section we introduce a new family of generative autoencoders , called Symmetric Wasserstein Autoencoders ( SWAEs ) . 2.1 OT FORMULATION . Denote the random vector at the encoder as e , ( xe , ze ) ∈ X×Z , which contains both the observed data xe ∈ X and the latent representation ze ∈ Z . We call the distribution p ( e ) ∼ p ( xe ) p ( ze|xe ) the encoding distribution , where p ( xe ) represents the data distribution and p ( ze|xe ) characterizes an inference model . Similarly , denote the random vector at the decoder as d , ( xd , zd ) ∈ X × Z , which consists of both the latent prior zd ∈ Z and the generated data xd ∈ X . We call the distribution p ( d ) ∼ p ( zd ) p ( xd|zd ) the decoding distribution , where p ( zd ) represents the prior distribution and p ( xd|zd ) characterizes a generative model . The objective of VAEs is equivalent to minimizing the ( asymmetric ) KL divergence between the encoding distribution p ( e ) and the decoding distribution p ( d ) ( see Appendix A.1 ) . To address the limitation in VAEs , first we propose to treat the data and the latent representation symmetrically instead of asymmetrically by minimizing the pth Wasserstein distance between p ( e ) and p ( d ) leveraging Optimal Transport ( OT ) ( Villani , 2008 ; Peyré et al. , 2019 ) . OT provides a framework for comparing two distributions in a Lagrangian framework , which seeks the minimum cost for transporting one distribution to another . We focus on the primal problem of OT , and Kantorovich ’ s formulation ( Peyré et al. , 2019 ) is given by : Wc ( p ( e ) , p ( d ) ) , inf Γ∈P ( e∼p ( e ) , d∼p ( d ) ) E ( e , d ) ∼Γ c ( e , d ) , ( 1 ) where P ( e ∼ p ( e ) , d ∼ p ( d ) ) , called the coupling between e and d , denotes the set of the joint distributions of e and d with the marginals p ( e ) and p ( d ) , respectively , and c ( e , d ) : ( X , Z ) × ( X , Z ) → [ 0 , +∞ ] denotes the cost function . When ( ( X , Z ) × ( X , Z ) , d ) is a metric space and the cost function c ( e , d ) = dp ( e , d ) for p ≥ 1 , Wp , the p-th root of Wc is defined as the p-th Wasserstein distance . In particular , it can be proved that the p-th Wasserstein distance is a metric hence symmetric , and metrizes the weak convergence ( see , e.g. , ( Santambrogio , 2015 ) ) . Optimization of equation 1 is computationally prohibitive especially in high dimension ( Peyré et al. , 2019 ) . To provide an efficient solution , we restrict to the deterministic encoder and decoder . Specifically , at the encoder we have the latent representation ze = E ( xe ) with the function E : X → Z , and at the decoder we have the generated data xd = D ( zd ) with the function D : Z → X . It turns out that with the deterministic condition instead of searching for an optimal coupling in high dimension , we can find a proper conditional distribution p ( zd|xe ) with the marginal p ( zd ) . Theorem 1 Given the deterministic encoder E : X → Z and the deterministic decoder D : Z → X , the OT problem in equation 1 can be transformed to the following : Wc ( p ( e ) , p ( d ) ) = inf p ( zd|xe ) Ep ( xe ) Ep ( zd|xe ) c ( e , d ) , ( 2 ) where the observed data follows the distribution p ( xe ) and the prior follows the distribution p ( zd ) . The proof of Theorem 1 extends that of Theorem 1 in ( Tolstikhin et al. , 2018 ) , and is provided in Appendix A.2 . If X × Z is the Euclidean space endowed with the Lp norm , then the expression in equation 2 equals the following : Wc ( p ( e ) , p ( d ) ) = inf p ( zd|xe ) Ep ( xe ) Ep ( zd|xe ) ‖xe −D ( zd ) ‖ p p + ‖E ( xe ) − zd‖pp , ( 3 ) where in the objective we call the first term the x-loss and the second term the z-loss . With the above transformation , we decompose the loss in the joint space into the losses in both the data and the latent spaces . Such decomposition is crucial and allows us to treat the data and the latent representation symmetrically . The x-loss , i.e. , ‖xe−D ( zd ) ‖pp , represents the discrepancy in the data space , and can be interpreted from two different perspectives . Firstly , since D ( zd ) represents the generated data , the x-loss essentially minimizes the dissimilarity between the observed data and the generated data . Secondly , the x-loss is closely related to the objective of Denoising Autoencoders ( DAs ) ( Vincent et al. , 2008 ; 2010 ) . In particular , DAs aim to minimize the discrepancy between the observed data and a partially destroyed version of the observed data . The corrupted data can be obtained by means of a stochastic mapping from the original data ( e.g. , via adding noises ) . By contrast , the x-loss can be explained in the same way with the generated data being interpreted as the corrupted data . This is because the prior zd in D ( zd ) is sampled from the conditional distribution p ( zd|xe ) , which depends on the observed data xe . Consequently , the generated data D ( zd ) , obtained by feeding zd to the decoder , is stochastically related to the observed data xe . With this insight , the same as the objective of DAs , the x-loss can lead to the denoising effect . The z-loss , i.e. , ‖E ( xe ) − zd‖pp , represents the discrepancy in the latent space . The whole objective in equation 3 hence simultaneously minimizes the discrepancy in the data and the latent spaces . Observe that in equation 3 E ( xe ) is the latent representation of xe at the encoder , while zd can be thought of as the latent representation of D ( zd ) at the decoder . With such connection , the optimization of equation 3 can preserve the local data structure in the latent space . More specifically , since xe and D ( zd ) are stochastically dependent , roughly speaking , if two data samples are close to each other in the data space , their corresponding latent representations are also expected to be close . This is due to the symmetric treatment of the data and the latent representation . In Figure 1 we illustrate this effect and compare SWAE with VAE . Comparison with WAEs ( Tolstikhin et al. , 2018 ) The objective in equation 3 minimizes the OT cost between the joint distributions of the data and the latent , i.e. , Wc ( p ( e ) , p ( d ) ) , while the objective of WAEs ( Tolstikhin et al. , 2018 ) minimizes the OT cost between the marginal distributions of the data , i.e. , Wc ( p ( xe ) , p ( xd ) ) , where p ( xd ) is the marginal data distribution induced by the decoding distribution p ( d ) . The problem of WAEs is first formulated as an optimization with the constraint p ( ze ) = p ( zd ) , where p ( ze ) is the marginal distribution induced by the encoding distribution p ( e ) , and then relaxed by adding a regularizer . With the deterministic decoder , the final optimization problem of WAEs is as follows : inf p ( ze|xe ) Ep ( xe ) Ep ( ze|xe ) c ( xe , D ( ze ) ) + λD ( p ( ze ) , p ( zd ) ) , ( 4 ) where D ( , ) denotes some divergence measure . Comparing equation 4 to equation 3 , we can see that both methods decompose the loss into the losses in the data and the latent spaces . Differently , in equation 4 the first term reflects the reconstruction loss in the data space and the second term represents the distribution-based dissimilarity in the latent space ; while in equation 3 the x-loss is closely related to the denoising and the generation quality and the z-loss measures the sample-based dissimilarity . Moreover , equation 4 is optimized over the posterior p ( ze|xe ) with a fixed prior p ( zd ) , while equation 3 is optimized over the conditional prior p ( zd|xe ) with a potentially learnable prior .

Summary: This works proposes an new auto-encoder variant based on an Optimal Transport (OT) penalty.  While there are many such previous works of OT and auto-encoders, this work proposes a joint OT penalty on data and latent space. As the scalability of computing OT penalties in high dimensions is a concern, the authors address this by restricting to deterministic encoders and decoders in Theorem 1, an extension to joint distributions of Theorem 1 of Tolstikhin 2018. The resulting algorithm amounts to a loss involving L2 penalties for (1) the reconstruction loss (2) decoded latents (conditional on "pseudo-inputs") and real samples (3) encoded samples and the conditional latents. Next experimental results are shown on small-scale datasets (MNIST, Fashion-MNIST, Coil20, subest of CIFAR-10) and compared against the VAE, WAE-{GAN,MMD}, VampPrior, and MIM.