Encoded Prior Sliced Wasserstein AutoEncoder for learning latent manifold representations

1 INTRODUCTION . Generative models have the potential to capture rich representations of data and use them to generate realistic outputs . In particular , Variational AutoEncoders ( VAEs ) ( Kingma & Welling , 2014 ) can capture important properties of high-dimensional data in their latent embeddings , and sample from a prior distribution to generate realistic images . Whille VAEs have been very successful in a variety of tasks , the use of a simplistic standard normal prior is known to cause problems such as under-fitting and over-regularization , and fails to use the network ’ s entire modeling capacity ( Burda et al. , 2016 ) . Gaussian or Gaussian mixture model ( GMM ) priors are also limited in their ability to represent geometric and topological properties of the underlying data manifold . High-dimensional data can typically be modeled as lying on or near an embedded low-dimensional , nonlinear manifold ( Fefferman et al. , 2016 ) . Learning improved latent representations of this nonlinear manifold is an important problem , for which a more flexible prior may be desirable . Conventional variational inference uses Kullback-Leibler ( KL ) divergence as a measure of distance between the posterior and the prior , restricting the prior distribution to cases that have tractable approximations of the KL divergence . Many works such as Guo et al . ( 2020 ) ; Tomczak & Welling ( 2018 ) ; Rezende & Mohamed ( 2015 ) etc . have investigated the use of more complicated priors ( notably GMMs ) which lead to improved latent representation and generation compared to a single Gaussion prior . Alternate approaches such as adversarial training learn arbitrary priors by using a discriminator network to compute a divergence ( Wang et al. , 2020 ) , however they are reported to be harder to train and are computationally expensive . In this work , we introduce the Encoded Prior Sliced Wasserstein AutoEncoder ( EPSWAE ) , which consists of a conventional autoencoder architecture and an additional prior-encoder network that learns an unconstrained prior distribution that encodes the geometry and topology of any data manifold . We use a type of Sliced Wasserstein ( SW ) distance ( Bonnotte , 2013 ; Bonneel et al. , 2015 ) , a concept from optimal transport theory that is a simple and convenient alternative to the KL divergence for any sampleable distributions . A Sliced Wasserstein AutoEncoder ( SWAE ) that regularizes an autoencoder using SW distance was proposed in Kolouri et al . ( 2018a ) . Several works improve the SW distance through additional optimizations ( Deshpande et al. , 2019 ; Chen et al. , 2020b ; Deshpande et al. , 2018 ) , and show improved generation , however involve additional training and use a fixed ( usually Gaussian ) prior . Kolouri et al . ( 2019 ) presents a comparison between max-SW distance , polynomial generalized SW distances , and their combinations . In contrast , we use a simple and efficient nonlinear shearing which requires no additional optimization . Additionally , we introduce a structural consistency term that encourages the latent space to be isometric to the feature space , which is typically measured at the output of the convolutional layers of the data encoder . Variants of this penalty have previously been used to encourage isometry between the latent space and data space ( Yu et al. , 2013 ; Benaim & Wolf , 2017 ; Sainburg et al. , 2018 ) . The structural consistency term further encourages the prior to match the encoded data manifold by preserving feature-isometry , which in turn is expected to assist with encoding the geometry of the data manifold , thus leading to improved latent representations . A key contribution of our work is the graph-based geodesic-interpolation algorithm . Conventionally , VAEs use Euclidean interpolation between two points in latent space . However , since manifolds typically have curvature , this is an unintuitive distance metric that can lead to unrealistic intermediate points . Our goal is to learn a true representation of the underlying data manifold , hence it is natural to interpolate along the manifold geodesics in latent space . Several works such as Shao et al . ( 2018 ) ; Miolane & Holmes ( 2020b ) endow the latent space with a Riemannian geometry and measure corresponding distances , however these are difficult and involve explicitly solving expensive ordinary differential equations . In this work , we introduce ‘ network-geodesics ’ , a graph-based method for interpolating along a manifold in latent space , that maximizes sample density along paths while minimizing total energy . This involves first generating a distance graph between samples from the prior . Then this network is non-uniformly thresholded such that the set of allowable paths from a given sample traverse high density regions through short hops . Lastly , we use a shortest path algorithm like Dijkstra ’ s algorithm ( Dijkstra , 1959 ) to identify the lowest ‘ energy ’ path between two samples through the allowable paths . Since the prior is trained to learn the data manifold , the resulting network-geodesic curves give a notion of distance on the manifold and can be used to generate realistic interpolation points with relatively few prior samples . The novel contributions of this work are : • We introduce a novel architecture , EPSWAE , that consists of a prior-encoder network that is efficiently trained ( without expensive adversarial methods ) to generate a prior that encodes the geometric and topological properties of data . • We introduce a novel graph-based method for interpolating along network-geodesics in latent space through maximizing sample density while minimizing total energy . We show that it generates natural interpolations through realistic images . • Improvements to the latent space representation are obtained by using a structural consistency term in the loss that encourages isometry between feature space and latent space and by using a simple and efficient nonlinear variant of the SW distance . 2 BACKGROUND AND RELATED WORK . Several works have attempted to increase the complexity of the prior in order to obtain better latent representations . Most data can mathematically be thought of as living on a high dimensional manifold . In an image dataset , for instance , if images in high dimensional pixel space are effectively parametrized using a small number of continuous variables , they will lie on or near a low dimensional manifold ( Lu et al. , 1998 ) . Many works such as Weinberger & Saul ( 2006 ) ; Rahimi et al . ( 2005 ) have investigated image and video manifolds . Encoding and exploiting the topological and geometric properties of such data is a question of increasing interest . Some representative works are Dilokthanakul et al . ( 2016 ) which shows improved unsupervised clustering of latent space by using GMM priors , Takahashi et al . ( 2019 ) which uses the density ratio trick to calculate the KL divergence using implicit priors , Rezende & Mohamed ( 2015 ) , which maps a Gaussian prior by a series of explicit transformations , Guo et al . ( 2020 ) which learns a GMM prior using an approximate ELBO , Goyal et al . ( 2017 ) which uses a hierarchical Bayesian framework , and VampPrior Tomczak & Welling ( 2018 ) which learns a two-layer hierarchical GMM from aggregated posteriors . Yin & Zhou ( 2018 ) ; Molchanov et al . ( 2019 ) ; Liutkus et al . ( 2019 ) expand the variational family to incorporate implicit variational distributions while providing exact theoretical support . The KL divergence is tractable only for Gaussian distributions , however the Wasserstein distance of 1D projections ( also called Sliced Wasserstein ( SW ) distance ) has a closed-form for any arbitrary distribution ( Kolouri et al. , 2018b ) . Wasserstein distances ( Villani , 2003 ) have several of the same properties as KL divergence , but often lead to better sampling ( Gulrajani et al. , 2017 ; Tolstikhin et al. , 2018 ) , and have been used in several machine learning applications ( e.g . Arjovsky et al . ( 2017 ) ; Tolstikhin et al . ( 2018 ) ; Kolouri et al . ( 2018b ) ) . A Sliced Wasserstein AutoEncoder ( SWAE ) that regularizes an autoencoder using SW distance was proposed in Kolouri et al . ( 2018a ) . There exist several variants of SW distance that have been successful at improving generative ability - for example , Deshpande et al . ( 2018 ) ; Chen et al . ( 2020b ) train discriminator-like networks to perform nonlinear transformations ( instead of purely linear projections ) whereas Deshpande et al . ( 2019 ) introduces max-SW distance , which looks for the best linear projection , generalized in Nguyen et al . ( 2020 ) to finding optimal distributions of projections . Wu et al . ( 2019 ) uses a different optimization approach to computing the SW distance based on the Kantorovich dual formulation . All of the above use some form of additional training . In contrast , we use a simple and efficient nonlinear shearing which requires no additional optimization and is sufficient for improving latent representation , however , our method can easily incorporate other SW distances if desired . Alternatively , many works such as Wang et al . ( 2020 ) ; Makhzani et al . ( 2015 ) ; Arjovsky et al . ( 2017 ) ; Sainburg et al . ( 2018 ) replace the KL divergence with an adversarial loss , however adversarial methods tend to be significantly more expensive and difficult to optimize . In higher dimensions , using a discriminator network as in adversarial approaches is a natural way of implicitly computing an equivalent of the Wasserstein-1 distance ( Arjovsky et al. , 2017 ; Tolstikhin et al. , 2018 ) . However , using the SW distance is much simpler and more efficient . Adversarial training for interpolation in Sainburg et al . ( 2018 ) uses a structural consistency term that encourages relative distances in data space to be preserved in latent space . Similar distance-preserving terms have also been used successfully in Yu et al . ( 2013 ) ; Benaim & Wolf ( 2017 ) , however Euclidean distances in data space can be a poor measure of the natural geometry of the data . In our work , we preserve relative distances in latent space and feature space to improve latent encoding , where features are extracted at the output of the last convolutional layer in the data encoder . Several manifold learning techniques such as Dollár et al . ( 2007 ) ; Weinberger & Saul ( 2006 ) compute embeddings of high-dimensional data but are less general and lack generative or interpolation abilities . The hyperspherical VAE Davidson et al . ( 2018 ) outperforms the standard VAE for data residing on a hyperspherical manifold . Miolane & Holmes ( 2020a ) formulates a Riemannian VAE , however computing its ELBO is challenging . Along similar lines , Arvanitidis et al . ( 2017 ) shows that under certain conditions , a Riemannian metric is naturally induced in the latent space , and uses it to compute geodesics . In Chen et al . ( 2020a ) , a flat manifold is approximated by penalizing curvature in latent space , and geodesics are defined through Euclidean distances on the flat manifold . There exists limited work on integrating graphical structures with generative models ( Kipf & Welling , 2016 ) . Hadjeres et al . ( 2017 ) studies monophonic music using a latent space geodesic regularization that allows interpolations in the latent space to be more meaningful , giving a notion of geodesic distance . Several approaches such as Tenenbaum et al . ( 2000 ) ; Bernstein et al . ( 2000 ) ; Klein & Zachmann ( 2004 ) ; Mémoli & Sapiro ( 2005 ) ; Luo & Hu ( 2020 ) approximate geodesics on point clouds by , for instance , building a nearest-neighbor network on the manifold and applying a shortest path algorithm , however , these often require many samples in order generate reasonable approximations and aren ’ t robust to noise ( Sober et al. , 2020 ) . Inspired by this , we introduce an energy-based network algorithm to identify network-geodesics in latent space with relatively few , noisy samples .

Summary: This paper addresses the issues of representation learning with VAEs and propose EPSWAE as a solution. EPSWAE applies a prior encoder to construct an implicit prior, which is more flexible. Moreover, the authors apply the sliced Wasserstein distance for the matching between the posterior and the prior, enhance the conventional SWD with non-linear transformations and make the latent space similar to the feature space with a structural consistency loss. This paper also proposes a graph-based algorithm for minimizing the pathwise energy to achieve the manifold walking to improve the interpolation in the latent space.