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 extends the variational autoencoder framework with a richer prior distribution to model more complex correlations in the latent variable distribution. They start with a Gaussian mixture distribution as the prior for the latent variables, and add an encoder network to allow richer correlation structure in the latent variables.  Training the prior distribution requires an optimization between the prior distribution and the latent encoded distribution of the training data set. The paper starts with an existing method of optimizing the prior by computing an approximation of the Wasserstein distance between prior and encoded training distribution that uses an average over slices through the prior and encoded training distribution. The paper replaces linear projections used in prior work with a non-linear projection. The paper also employs a structural consistency term which has been used in prior work, however, the paper employs this term differently than prior work by applying it between encoder features and latent variables rather than inputs and latent variables. Since the latent variable space is now a complex and possibly a nonconvex submanifold, points in the latent space R^D  lying between points corresponding to training data points may not actually fall in the training distribution. The paper therefore proposes a method of interpolating between points in the manifold by constructing a graph between points sampled from the manifold and then choosing points lying along lines in the graph. The paper tests the method on three datasets, a synthetic 40 dimensional spiral dataset, the venerable MNIST dataset and a scaled down CELEB A dataset. Plots of the latent space trained on the spiral dataset shows that the latent space can in fact have complex internal structure.