ENCODED PRIOR SLICED WASSERSTEIN AUTOEN- CODER FOR LEARNING LATENT MANIFOLD REPRESEN- TATIONS
Year: 2020

Abstract:
While variational autoencoders have been successful in a variety of tasks, the use of conventional Gaussian or Gaussian mixture priors are limited in their ability to encode underlying structure of data in the latent representation. In this work, we introduce an Encoded Prior Sliced Wasserstein AutoEncoder (EPSWAE) wherein an additional prior-encoder network facilitates learns an embedding of the data manifold which preserves topological and geometric properties of the data, thus improving the structure of latent space. The autoencoder and prior-encoder networks are iteratively trained using the Sliced Wasserstein (SW) distance, which efficiently measures the distance between two arbitrary sampleable distributions without being constrained to a specific form as in the KL divergence, and without requiring expensive adversarial training. To improve the representation, we use (1) a structural consistency term in the loss that encourages isometry between feature space and latent space and (2) a nonlinear variant of the SW distance which averages over random nonlinear shearing. The effectiveness of the learned manifold encoding is best explored by traversing the latent space through interpolations along geodesics which generate samples that lie on the manifold and hence are advantageous compared to standard Euclidean interpolation. To this end, we introduce a graph-based algorithm for interpolating along network-geodesics in latent space by maximizing the density of samples along the path while minimizing total energy. We use the 3D-spiral data to show that the prior does indeed encode the geometry underlying the data and to demonstrate the advantages of the network-algorithm for interpolation. Additionally, we apply our framework to MNIST, and CelebA datasets, and show that outlier generations, latent representations, and geodesic interpolations are comparable to the state of the art.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND AND RELATED WORK
text: 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.

heading: 3 EPSWAE

heading: 3.1 NONLINEAR SLICED WASSERSTEIN DISTANCE
text: Wasserstein distances are a natural metric for measuring distances between probability distributions, however they xare difficult to compute in dimensions two and higher. The Sliced Wasserstein distance dSW averages over the Wasserstein distance of 1D projections (see Appendix A for derivation). Several works (e.g. Kolouri et al. (2019); Deshpande et al. (2019)) discuss why a conventional SW distance may be sub-optimal as a large number of linear projections may be required to distinguish two distributions. In this work we use a Nonlinear Sliced Wasserstein distance (dNSW ), an averaging procedure over (random) nonlinear 1D projections, between two distributions µ and ν defined as:
dNSW (µ, ν) = EdSW (N ζ,γ∗ µ,N ζ,γ∗ ν) ≈ 1
L L∑ `=1 dSW (N ζ`,γ`∗ µ,N ζ`,γ`∗ ν) (1)
where L is the number of nonlinear transformations, and ζ, γ are chosen to be normal random variables with mean 0, and variance matching that of µ, and N is defined below. We define the “push-forward” (also known as “random variable transform”) as follows: given any function f : X → Z and probability measure µ, we define f∗µ(A) = µ(f−1(A)) for all A ⊂ Z measurable. Most importantly, to generate samples of z ∼ f∗µ one simply generates x ∼ µ and defines z = f(x).
The nonlinear transformations are related to a special case of the generalized SW distance (Kolouri et al., 2019) (with the difference that we choose distribution-dependent shear frequencies that are most likely to produce deformations which highlight differences in the measures, thus breaking the homogeneity condition H2 in Kolouri et al. (2019)), and are given by
N ζ`,γ`(x) = x+ ζ` sin(γ` · x). (2) In all our experiments, we use L = 5 nonlinear transformations with 50 linear 1D projections each. Several works such as Deshpande et al. (2018; 2019); Wu et al. (2019); Nguyen et al. (2020); Chen et al. (2020b) improve the SW distance through some form of optimization or training an additional discriminator-like network, and show improved generative results (with Deshpande et al. (2018) being the most similar to our method). We did not implement these other variants, however they could easily be used in conjunction with our method. Our choice of nonlinearity is motivated by certain tail-behavior considerations (such as boundedness and non-saturation of nonlinearity), along with computational efficiency, and for our goals, a simple nonlinearity was sufficient. See section 4.2 and Appendix E for computational time and loss comparisons with some other nonlinearities discussed in Kolouri et al. (2019). For further discussion on NSW distance and derivation see Appendix A.

heading: 3.2 ALGORITHM DETAILS
text: In the Encoded Prior Sliced Wasserstein AutoEncoder (EPSWAE) (see Fig. 1 for schematic), let’s define the data encoder as ΨE , the decoder ΨD, and the prior-encoder as ΨPE , each with parameters φE , φD, φPE respectively. Input samples x(j) ∼ PX , where PX is the probability distribution of the input data, are passed through ΨE to generate posterior samples z(j) ∼ (ΨE)∗PX by setting z(j) = ΨE(x
(j)). Similarly, prior-encoder input samples ξ(j) ∼ µ, where µ is the probability distribution of the input to the prior-encoder (chosen to be a mixture of Gaussians) from a distribution, are passed through ΨPE to generate prior samples y(j) ∼ (ΨPE)∗µ by setting y(j) = ΨPE(ξ(j)). The priorencoder network and the autoencoder (data encoder and decoder) network are trained iteratively in a two step process:
1. Given a minibatch, parameters of the autoencoder (φE , φD) are trained for k1 steps while parameters of the prior encoder (φPE) are fixed. The loss function for the autoencoder consists of the reconstruction error, the NSW distance, and a Feature Structural Consistency (FSC) term LFSC (given in Eqn. 4):
LAE = αEx∼PXLrec (x, x̄) + βdNSW ((ΨE)∗PX , (ΨPE)∗µ) + κLFSC , (3)
where x̄ = ΨD(ΨE(x)) is the autoencoder output. Note that dNSW is efficiently computed just from samples of the distributions. In this measure theoretic notation, (ΨE)∗PX is the posterior distribution (denoted as qφE (z|x) in Bayesian literature), and (ΨPE)∗µ is the prior distribution (denoted as pφPE (z) in Bayesian literature). Similar to β-VAE ((Higgins et al., 2016)), the hyperparameters α, β, κ are tuneable; in our experiments they are fixed, but it could be advantageous to have them vary over training epochs. LFSC encourages relative distances in feature space to be preserved in latent space, where features are extracted at the output of the last convolutional layer in the encoder (or at input data if no convolutional layers are present). LFSC (adapted from Sainburg et al. (2018)) for two point clouds F = [f1, . . . fN ] in feature space and the corresponding points Z = [z1, . . . zN ] in latent space in a minibatch of size N is given by:
LFSC = 1
N2 N∑ `,j=1
( log ( 1 +
‖fj − f`‖2 1 N2 ∑ m,n ‖fm − fn‖2
) − log ( 1 +
‖zj − z`‖2 1 N2 ∑ m,n ‖zm − zn‖2
))2 .
(4) 2. In the second step of minimization, the parameters of the prior-encoder (ΨPE) are trained
for k2 steps while parameters of the autoencoder (φE , φD) are fixed. The loss function for the prior-encoder consists of the NSW distance between the prior and posterior:
LPE(x) = dNSW ((ΨE)∗PX , (ΨPE)∗µ). (5)
The pseudocode is outlined in Appendix B. Additional architecture and training details for the experiments in this paper are in Appendix C.

heading: 3.3 INTERPOLATION AND APPROXIMATE GEODESICS
text: EPSWAE attempts to learn a representation of the embedded manifold geometry that the data lies along. The question arises: how does one make use of this representation? A natural way of interpolating data which lies on a manifold is through geodesics or paths on the manifold which traverse dense regions. Here we present a network algorithm for efficiently approximating networkgeodesics with relatively few, noisy samples from the prior, by encouraging connections through high density regions of latent space while minimizing total energy.
1. Gather samples of the posterior (ΨE)∗PX , i.e., ΨE(x) with x is a minibatch of data. Additional prior samples can be used to supplement this if desired;
2. For each sample j, compute the average Euclidean distance cj of the k-nearest neighbors; 3. Generate a thresholded network: Let samples be represented as nodes. For a given thresh-
old value t, sample j is connected to sample i with edge weight dh(i, j) if d(i, j) < t · cj . Here h is the energy parameter in the edge weight (chosen to be 1 or 2 in experiments, with h = 2 encouraging shorter hops). Sample-specific thresholding, i.e., thresholding dependent on cj , increases the probability of a ‘central’ node having a connection. Thus, this encourages (1) paths through high density regions, i.e., high-degree nodes, and (2) traversal of latent space through short hops as a consequence of localization of paths, i.e., paths existing only between nearby points due to thresholding.
4. Continue increasing t until the graph is connected, i.e., there exists a path, direct or indirect, from every sample to every other sample. Then, use Dijkstra’s algorithm (Dijkstra, 1959) to identify network-geodesics with least total energy through allowable paths on the thresholded network.

heading: 4 EXPERIMENTS
text: We run experiments using EPSWAE on three datasets. First, we use a 3D spiral data, where the latent space can be visualized easily, to demonstrate that the learned prior captures the geometry of the 3D spiral. We also present here the advantages of the network-geodesic interpolation over linear interpolation. Then we present interpolation and generation results on the MNIST (LeCun et al., 2010) and CelebA (Liu et al., 2018) datasets.
4.1 ARCHITECTURE
In all of the experiments, the prior-encoder ΨPE consists of three fully connected hidden layers with ReLU activation. A sampleable distribution with dimension larger than the dimension of latent space is input to the prior-encoder which yields the prior. In principle, samples from any distribution can be the input µ to the prior-encoder. In the artificial data set we use a normal distribution, whereas for MNIST and CelebA we use a mixture of 10 Gaussians with random i.i.d means (from Gaussian with σ = 2). The schematic of the architecture is presented in Fig. 1(a). Note that the specifics of the layers shown in this schematic are for the image datasets MNIST and CelebA, the 3D Spiral data does not have convolutional layers. Details of the architecture, training, and hyperparameters for each dataset are given in Appendix C. The optimizer Adam (Kingma & Ba, 2014) with a learning rate of 0.001 was used for learning in both networks: the prior-encoder and the autoencoder.

heading: 4.2 LEARNING A LATENT MANIFOLD
text: We consider a 3D spiral randomly embedded to 40D space with noise as follows: (a) The formula for the spiral is given by (x(t), y(t), z(t)) = (t cos 7tπ/4, t sin 7tπ/4, 2t), (b) a random 40 × 3 matrix is generated with i.i.d normal entries to map the spiral into R40, and (c) Gaussian noise with standard deviation 0.1 is added. The input to the prior-encoder is a Gaussian in R40, and the latent space is R3. For these tests, the data encoder, prior-encoder, and decoder all consist of three Fully Connected (FC) layers with 40 nodes each and ReLU activations.
As seen in Figure 2, after training, the posterior matches the spiral almost exactly and the prior learns the geometry of the spiral. A Nonlinear Sliced Wasserstein (NSW) distance and structural consistency term (where latent space is isometric with data space since convolutional layers aren’t used) are used in the loss in Eqns. 3 and 5. A comparison highlighting the effects of these terms is shown in Fig. 5 in the Appendix D. We observe that the use of the NSW distance (as opposed to linear SW distance) improves the accuracy of the learned prior.
Manifold interpolation (in red in Fig. 2) uses the network-algorithm outlined in section 3.3. The interpolation is seen to have the desired form on the manifold, i.e., it approximates geodesics on the manifold. The larger the number of samples of the prior, the smoother the corresponding interpolation. In contrast, linear interpolation (brown line) between two points on the spiral does not capture the geometry at all, and goes through regions that are largely empty and untrained. In higher dimensional datasets, this may result in unrealistic interpolation. In order to traverse the manifold, both the prior-encoder (that allows the prior to capture the geometry of the spiral), and the networkinterpolation algorithm are equally essential. Comparisons with baselines SWAE (Kolouri et al., 2018a) and VAE (Kingma & Welling, 2014) are shown in Appendix Fig. 6, and EPSWAE is seen to significantly outperform them in learning an improved latent representation.
Computational cost The choice of nonlinearity was, in part, motivated by computational efficiency and simplicity, i.e., not requiring additional optimizations every evaluation or the training of a discriminator-like network to select optimal transformations. Our sinusoidal shear nonlinearity is closely related to a special case of the generalized SW distance defined in Kolouri et al. (2019) (see discussion in section 3.1 for differences ). Here, we present a comparison of computational time with two polynomial-type nonlinearities discussed in Kolouri et al. (2019) - cubic and quintic. Computations done on a laptop with Intel Core i7-10710 CPU, 16GB RAM, and computed over 1000 runs. In practice, the losses were indistinguishable with choice of nonlinearity (see Appendix E), whereas the sine shear is slightly less expensive (see Table 1). However, the complexity of cubic and quintic nonlinearities are exponential in dimension and hence are impractical for larger data.

heading: 4.3 GENERATION AND OUTLIERS IN LATENT SPACE
text: In order to demonstrate one of the advantages of using a prior encoder in improving latent structure, we consider generation with an increased probability of sampling from outliers, i.e., by increasing the standard deviation of the distribution that feeds into the prior encoder. Studying generation from outliers may seem nonstandard, however, it is informative in the study of how data is encoded in the latent space. In Fig. 3 compare our results with the baseline SWAE (Kolouri et al., 2018a) in order to assess the effect of the prior encoder in improving latent structure. It is worth noting that while there exist several recent works that build upon SWAE for better generation, they do not attempt to solve the same problem as ours (learning latent manifolds). As seen in Fig. 3, the behavior of SWAE and EPSWAE at outliers is very different. For large σ, SWAE tends to generate unrealistic faces with distorted colorations, whereas EPSWAE is more likely to generate realistic faces (albiet with increased mode collapse at large σ). We see that EPSWAE encodes coherent information in a large region of latent space as a consequence of nonlinear transformations through the prior encoder. Since we are not competing with state of the art in generation here, we use downsized images and a small, unsophisticated model compared to state of the art methods.
Both SWAE and EPSWAE use the same data encoder and decoder, and are independently optimized: the encoder is (after downsizing CelebA images to 64x64) Conv (3,16,3) → BatchNorm → ReLu → MaxPool (2,2) → Conv (16,32,3) → BatchNorm → ReLu → MaxPool (2,2) → Conv (32,64,3) → BatchNorm → ReLu followed by two FC layers of 512 and 256 nodes respectively with a leaky ReLu nonlinearity and the decoder is the reverse (replace Conv layers with Conv-transpose and MaxPool with Upsample). The prior-encoder consists of three FC layers from 186D → 128D latent space. See Appendix C for more details and hyperparameter values.
EPSWAE
SWAE
σ = 1 σ = 2 σ = 3 σ = 4
σ = 1 σ = 2 σ = 3 σ = 4
We also present a comparison with SWAE on CelebA generation from samples around the mean in Appendix G, and find that EPSWAE sees a marginal improvement in the quality of faces. Generation on MNIST using EPSWAE (and comparisons with SWAE (Kolouri et al., 2018a) and VAE (Kingma & Welling, 2014)) are shown in Appendix F. MNIST results are obtained with 5D latent space. EPSWAE generation is more natural, whereas both other baselines generate some bloated and unidentifiable numbers.
For completeness, in Table 4.3 we present Frechet Inception Distance (FID) (Salimans et al., 2016) scores with 10000 images with σ = 1 (code from Seitzer (2020)). It is important to note that one can have good generation, despite learning a poor latent representation (and vice-versa), so scores like FID may not be the best way to evaluate the latent structure. However, comparison between SWAE and EPSWAE-baseline(bl) indicate that the prior-encoder significantly improves generation.
Computational cost : The computational cost of the different terms in the EPSWAE loss on the CelebA dataset is as follows: (1) FSC computation: 0.00066s ± 0.0001s, (2) linear SW distance computation: 0.01561s±0.0005s (3) sine-shear NSW computation: 0.0150s±0.0007s (with same number of total projections). All computations for CelebA were done on a workstation with one NVIDIA Tesla P100 GPU, two Intel Xeon E5-26660v4 CPUs, and 64Gb RAM.

heading: 4.4 INTERPOLATION
text: In this paper, we introduce a graph-based algorithm that interpolates over network-geodesics in curved latent manifolds, thus capturing properties of the geometry. Here we show that we obtain smooth interpolation using the network-geodesics algorithm, in contrast with linear interpolations that often tend to generate unrealistic intermediate images.
The start and end points correspond to real images (posterior samples). Figure 4(a) presents a comparison of linear interpolation (top) vs interpolation through network-geodesics (bottom) described
in section 3.3 on CelebA. One can see that linear interpolations often go through unrealistic images (this corresponds to regions in latent space where training is limited), whereas geodesic interpolations go through more realistic and less blurry faces. This also serves as evidence, that the latent space for CelebA contains some natural structure that can be exploited.
Figure 4 (b) shows examples of interpolations on MNIST and CelebA datasets. MNIST interpolations are smooth and intuitive, for instance, in the bottom left MNIST panel in Fig. 4(b), the top part of a ‘7’ first changes to a ‘9’ naturally, followed by transformation of the ‘9’ to a ‘4’. Interpolation along network-geodesics ensures that reconstructions of intermediate samples are realistic. Interpolations on CelebA are smooth and pass through intermediate images that could arguably pass for celebrities these days. The state of the art interpolations along network-geodesics on the manifold indicate that the learned prior does indeed encode the data manifold. For all interpolations shown in Fig. 4, we use energy parameter h = 2. Comparisons between interpolations corresponding to h = 1 and h = 2 are presented in Appendix 10. In experiments, energy parameter is not found to have a significant impact on the quality of interpolations. Additionally, comparisons of linear interpolation between equivalent networks of EPSWAE and SWAE are presented in Appendix I.

heading: 5 CONCLUSION
text: We introduce the Encoded Prior Sliced Wasserstein AutoEncoder (EPSWAE) that learns improved latent representations through training an encoded prior to approximate an embedding of the data manifold that preserves geometric and topological properties. The learning of an arbitrary shaped prior is facilitated by the use of the Sliced Wasserstein distance, which can be efficiently computed from samples only. We use a nonlinear variant of SW distance to capture differences between two distributions more efficiently, and employ a feature structural consistency term to improve the latent space representation. Finally, we introduce an energy-based algorithm to identify networkgeodesics in latent space that maximize path density while minimizing total energy. We demonstrate EPSWAE’s ability to learn the geometry and topology of a 3D spiral from a noisy 40D embedding. We also show that our model embeds information in large regions of latent space, leading to better outlier generation. Lastly, we show that our geodesic interpolation results on MNIST and CelebA are efficient and comparable to state of the art techniques. Our code is publicly available.

heading: A BACKGROUND ON THE SLICED WASSERSTEIN DISTANCE
text: Wasserstein distances provide a natural metric for measuring distances between probability distributions based on optimal transport, i.e., the cost of deforming one probability distribution into another. For a measurable cost function c(x, y) for x, y ∈ Rd, and probability distributions µ and ν on Rd, we define the p-Wasserstein distance between the distributions as
dW,p(µ, ν) = inf Γ∈Π(µ,ν) (ˆ Rd×Rd cp(x, y)dΓ(x, y) )1/p , (6)
where Π is the set of all joint distributions with marginals µ and ν (see Villani (2003)).
This Wasserstein distance between probability distributions is extremely difficult and computationally intensive in dimensions two and higher, i.e., d ≥ 2. However, in dimension one, there exists a closed-form solution (Villani, 2003). A simple algorithm for computing the 1D Wasserstein distance is given as follows (see for example (Kolouri et al., 2018a)): (a) Generate N one dimensional i.i.d samples xj ∼ µ, yj ∼ ν; (b) sort each list X = [x1, . . . , xN ], Y = [y1, . . . , yN ] into ascending order denoted by X̃, Ỹ respectively; (c) define the approximation to the 1D p-Wasserstein distance
dW,p(µ, ν) ≈  1 N N∑ j=1 cp(x̃j , ỹj) 1/p , (7) where x̃j ∈ X̃, ỹj ∈ Ỹ . The Sliced Wasserstein (SW) distance defines a metric on probability measures (Bonnotte, 2013) which provides an alternative to Eqn. 6 by exploiting the computational feasibility of the 1D Wasserstein distance in Eqn. 7. It involves averaging over one-dimensional orthogonal projections πθx := (θ · x)θ as follows:
dSW (µ, ν) = ( Sd−1 dpW,p(π θ ∗µ, π θ ∗ν)dS(θ) )1/p , (8)
The SW distance has seen a variety of implementations (Bonneel et al., 2015; Kolouri et al., 2018b; Lee et al., 2019).
The integral in Eqn. 8 can be easily approximated by sampling M random one-dimensional vectors θk uniformly on Sd−1 and computing
dSW (µ, ν) ≈
( 1
M M∑ k=1 dpW,p(π θk ∗ µ, π θk ∗ ν)
)1/p . (9)
Several works (for example Nguyen et al. (2020); Deshpande et al. (2019)) have discussed reasons why this often does not result in the best computational method. Linear projections may be suboptimal for extracting information about the differences between µ and ν, since a large number of linear projections may be required to get an accurate approximation for dSW . Several works have suggested possible methods for improving the effectiveness of the SW distance (Chen et al., 2020b; Kolouri et al., 2019; Nguyen et al., 2020; Deshpande et al., 2019). In contrast, we use a Nonlinear Sliced Wasserstein (NSW) distance, an averaging procedure over (random) nonlinear transformations. For our goals, the choice of nonlinearity was motivated by the following considerations: (1) a bounded non-linearity would be beneficial since unbounded non-linearities (such as cubic polynomials) have a pronounced deformation on the tails of a measure and may excessively weight outliers. (2) A sigmoid is another potential candidate, but it saturates at high values and we want the non-linearity to be similarly effective everywhere. (3) The use of polynomials has been discussed in Kolouri et al. (2019), however, a full set of all higher order polynomials have exponential complexity and may be prohibitively expensive in high latent dimensions. We compared our method to cubic and quintic polynomials for the simple 3D spiral and found that the choice of the nonlinearity did not have a significant effect on performance, with our choice being the fastest (see section 4.2 and Appendix E for computational time and loss comparisons). In principle, other ensembles of nonlinear transformations could be used.

heading: B PSEUDOCODE
text: Algorithm 1 Training EPSWAE 1: while not converged do 2: Update the autoencoder ΨE ,ΨD: 3: for k1 substeps do 4: Sample minibatch from data {x(1), ...,x(N)} with x(j) ∼ PX 5: Compute feature extractor samples f (j) = ΨFE(x(j)) 6: Compute posterior samples z(j) = ΨE(x(j)) 7: Compute decoded output x(j)recon = ΨD((z) j )
8: Generate samples {ξ(1), ...ξ(J)} with ξ(j) ∼ µ 9: Compute prior samples y(j) = ΨPE(ξ(j))
10: Compute reconstruction error: Lrecon = 1J ∑J j=1 d(x (j),x (j) recon)
11: Compute NSW distance: LNSWdistance = dNSW ( 1N ∑N j=1 δz(j) , 1 N ∑N j=1 δy(j)) 12: Compute FSC loss: LFSC = dFSC({f (1), ..., f (J)}, {z(1), ..., z(J)}) 13: Compute autoencoder loss LAE = αLrecon + βLNSWdistance + κLFSC 14: Compute gradients of LAE wrt to φE , φD 15: Update φE , φD 16: end for 17: Update the prior-encoder ΨPE : 18: for k2 substeps do 19: Sample minibatch from data {x(1), ...,x(J)} with x(j) ∼ PX 20: Compute posterior samples z(j) = ΨE(x(j)) 21: Generate samples {ξ(1), ...ξ(J)} with ξ(j) ∼ µ 22: Compute prior samples y(j) = ΨPE(ξ(j)) 23: Compute prior-encoder loss: LPE = dNSW ( 1N ∑N j=1 δz(j) , 1 N ∑N j=1 δy(j)) 24: Compute gradients of LPE wrt to φPE 25: Update φPE 26: end for 27: end while

heading: C ARCHITECTURE AND TRAINING DETAILS
text: For all datasets, the prior-encoder ΨPE consists of three fully connected hidden layers and ReLU activations. For all datasets, the autoencoder and prior-encoder losses (given in Eqns 3 and 5) are trained iteratively using the optimizer Adam (Kingma & Welling, 2014) with a learning rate of 0.001. We experimented with both p = 1 and p = 2 (corresponding to p-Wasserstein) in the SW distance and did not find any significant differences; all results in this paper use p = 2. For each calculation of the NSW distance, L = 5 random nonlinear transformations were taken followed by M = 50 one dimensional projections per transformation. Data-specific model and parameter details are given below.
3D Spiral dataset: The input to the prior-encoder is a 40D Gaussian, and the latent space is 3D. The input to the autoencoder is a 40D embedding of a 3D spiral manifold with 10% noise. The dataset consists of 10000 samples, and a batch size of 100 was used. The prior-encoder, the data encoder, and the decoder consist of three Fully Connected (FC) layers with 40 nodes each and ReLU activations. The reconstruction loss is given by the Mean Square Error and α = 1, β = 0.1, κ = 0.01 in Eqn 3. In the absence of convolutional layers, the FSC term encourages the pairwise distances of the minibatch in latent space to be similar to the pairwise distances of the minibatch in the data space. The prior-encoder is trained k1 = 2 times for each training of the autoencoder k1 = 1. Power of distance h = 2 is used to compute edge weights for computing network-geodesics.
MNIST dataset: The input to the prior-encoder is a 40 dimensional mixture of 10 Gaussians, and the latent space is 5 dimensional. The data encoder takes MNIST images as input using a batch size of 100, and consists of the following layers Conv (1,10,3) → BatchNorm → ReLu → MaxPool
(2,2) → Conv (10,16,3) → BatchNorm → ReLu followed by two FC layers of 512 and 256 nodes respectively with a leaky ReLu nonlinearity. The encoder outputs a 5 dimensional latent representation. The decoder consists of the reverse, i.e., three fully connected layers of size 256, 512, and 1936 nodes respectively. This is followed by ConvTranspose (16,10,3) → LeakyReLu → Upsample(2,2) → ConvTranspose (10,1,3) → Sigmoid. The decoder output has size 28 × 28. The prior-encoder is trained k1 = 1 times for each training of the autoencoder k2 = 1. The reconstruction loss is given by the Binary Cross Entropy and α = 1, β = 0.1, κ = 0.001 in Eqn. 3. The FSC loss encourages the pairwise distances of the minibatch in latent space to be similar to the pairwise distances of the minibatch in feature space computed at the output of the last convolutional layer in the data encoder. Energy parameter h = 2 is used to compute edge weights for computing network-geodesics.
CelebA dataset: The input to the prior-encoder is a 186 dimensional mixture of 10 Gaussians, and the latent space is 128 dimensional. We ran experiments with different latent dimensions and found that the generation ability didn’t vary significantly as a function of latent dimension. We trade-off image quality of the input images for computational speed by downsizing the input and employ a fairly simple network compared to state of the art computer vision architectures that use CelebA. The data encoder takes CelebA images (of size 218× 178× 3) and downsizes them to size 64 × 64 × 3. The data encoder consists of the following layers Conv (3,16,3) → BatchNorm → ReLu → MaxPool (2,2) → Conv (16,32,3) → BatchNorm → ReLu → MaxPool (2,2) → Conv (32,64,3) → BatchNorm → ReLu followed by two FC layers of 512 and 256 nodes respectively with a leaky ReLu nonlinearity. The encoder outputs a 128 dimensional latent representation. As in the case of MNIST, the decoder consists of the reverse, with convolutions replaced by Convolution Transpose, and MaxPool replaced by Upsample. The output of the decoder is passed through a sigmoid nonlinearity and is of size 64 × 64 × 3. The prior-encoder was trained k1 = 1 times for each training of the autoencoder. The reconstruction loss is given by the Binary Cross Entropy and α = 500, β = 50, κ = 0.05 in Eqn. 3. The FSC loss encourages the pairwise distances of the minibatch in latent space to be similar to the pairwise distances of the minibatch in feature space, i.e., computed at the output of the last convolutional layer in the encoder. Energy parameter h = 2 is used to compute edge weights for computing network-geodesics.

heading: D SPIRAL BASELINE COMPARISONS
text: Figure 5 shows the effects of the different loss terms in EPSWAE on the geometry of the learned prior and posterior. The position and orientation of the spiral in the 3D plots are random and the views
in the image are hand-chosen to be equivalent. (a) shows latent space and interpolations (red) using EPSWAE with a linear Sliced Wasserstein distance in the loss, and no structural consistency term. (b) shows that adding a structural consistency term doesn’t remarkably improve the quality of the manifold learned, however, consistent with other experiments, it seems to improves interpolation slightly. Note here that since we don’t use convolutional layers for the 3D spiral, the stuctural consistency term preserves distance in latent space corresponding to distances in data space. (c) shows that employing the NSW distance term significantly improves the learned structure in latent space. The improvements resulting from incorporation of the NSW and structural consistency terms as seen in these visualizations of the 3D spiral lead us to use both loss terms (as in (d)) on all results in the main paper.
Figure 6 compares the EPSWAE model with baselines VAE and SWAE. The VAE uses KL divergence in the loss, which constrains the prior to be sampled from a Gaussian distribution. While the vanilla SWAE uses the SW distance, the prior remains a Gaussian, leading to an unnatural embedding of the data manifold. In contrast, EPSWAE results in a significantly better learned prior as a consequence of the nonlinear SW distance and the prior-encoder network trained explicitly to improve the latent representation.

heading: E COMPARISON WITH OTHER SLICED WASSERSTEIN NONLINEARITIES
text: Here, we compare our choice of nonlinearity in the SW loss (sinusoidal shear) with some other common nonlinearities. While there exist several methods to improve the SW distance (e.g. maxSW (Deshpande et al., 2019) which has seen several subsequent variations), these involve additional training, and hence aren’t considered (although many such versions of SW distance could be used in conjunction with out method). Instead we consider the polynomial generalized Radon transforms described in (Kolouri et al., 2019). The computational cost of a single polynomial generalized Radon transform is O(dk) for dimension d and kth order of polynomial, and hence we limit our investigation to cubic and quintic polynomials. We show results on the artificial 3D Spiral dataset trained upto 100 epochs with α = 1, β = 0.1, κ = 0.01. We see in Fig. 7 that the choice of nonlinearity does not have a significant effect on the loss, however, as seen in Table4.2, the sineshear has slightly lower computational cost.

heading: F MNIST GENERATION RESULTS
text: (a) EPSWAE (b) SWAE (c) VAE
Figures 8 (a,b,c) shows generation on the MNIST dataset after 100 epochs on EPSWAE, baseline SWAE(Kolouri et al., 2018a), and baseline VAE (Kingma & Welling, 2014) respectively. All networks EPSWAE and baselines SWAE and VAE use equivalent architectures (outlined in section of Appendix C), and hyperparameters are optimized individually. As seen in the figure, we observe that EPSWAE generated samples were consistently found to have a lower fraction of ‘false’ digits, i.e., digits that are unrealistic.

heading: G CELEBA GENERATION RESULTS
text: CelebA raw images are downsized to 64 × 64 pixels, and a simple architecture is employed (see Appendix C). This serves as a proof of principle for a better latent representation and interpolation, without a high computational cost. Naturally, using a more sophisticated network (for instance ResNet, VGG etc.) without downsizing the data would yield higher quality images at a computational cost.
Figure 9 shows generation on the CelebA dataset after 100 epochs with (a) EPSWAE and (b) baseline SWAE (right) respectively. Both employ equivalent architectures and take downsized images as input. As seen in the figure, EPSWAE generated images are more realistic. Note that this comparison is for equivalent training epochs and architectural details, and does not make claims about the quality of SWAE images with longer training.

heading: H EFFECT OF ENERGY PARAMETER ON CELEBA INTERPOLATIONS
text: Additional CelebA interpolations are shown in Fig. 10. The length of interpolations is automatically selected by the network algorithm. In some cases, a linearly interpolated point is added between every two samples on the network-geodesic to make smoother transitions; this would be unnecessary with enough samples of the prior, however it could be impractical in high latent dimension to generated sufficiently many prior samples. Note that even if only one intermediate point is selected by the network geodesic, the generated interpolation could be significantly different from a purely linear interpolation (an illustration is provided in Figure 4.1 - the direct linear interpolation from A to B is significantly different from the interpolation through the point C). Fig. 10 presents a comparison between energy parameter h = 1 and h = 2. The energy parameter determines the power of the distance metric used to compute the edge weight. One can think of the higher value (h = 2) as corresponding to stronger connections between points, and encouraging shorter hops. In practice, as seen from the figure, there is no clear advantage to choosing a specific value of energy parameter h.

heading: I LINEAR INTERPOLATION COMPARISON EPSWAE AND SWAE
text: We compare here linear interpolations on the CelebA dataset for MNIST and CelebA. The autoencoder networks used are identical and outlined in section C. In contrast, EPSWAE uses the additional prior encoder. Both models are independently optimized, and the interpolations are linear. As seen in Fig. 11, interpolations in EPSWAE are more realistically identifiable as ’faces’ than SWAE, which mixes up the features and generates blurry intermediate images. This is suggestive that the prior-encoder may play a role in improving latent representation.

Review summaries:
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. 

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. 

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.