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: The paper introduces an additional prior-encoder network to autoencoders to learn an unconstrained prior. The autoencoder and prior-encoder networks are iteratively trained with the sliced Wasserstein distance (SWD). To strengthen SWD, this paper further applies nonlinear transformations with a structural consistency term for better match between two distributions. For better interpolation on the latent space, it also introduces a graph-based algorithm.