Encoded Prior Sliced Wasserstein AutoEncoder for learning latent manifold representations
Authors: Sanjukta Krishnagopal
Jacob Bedrossian
Venue:  (2021)
Keywords: VAE
sliced Wasserstein distance
latent representation
interpolation
manifold embedding
geodesics
network algorithm

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 \textit{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 \textit{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.

Under review as a conference paper at ICLR 2021
ENCODED PRIOR SLICED WASSERSTEIN AUTOEN- CODER FOR LEARNING LATENT TATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
While variational autoencoders have been successful in variety of tasks, the use conventional Gaussian or Gaussian mixture priors are limited 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 additional prior-encoder network facilitates learns embedding of the data manifold which preserves topological and geometrie properties of the data, thus improving the structure of latent space. The autoencoder and prior-encoder net- works are iteratively trained using the Sliced Wasserstein (SW) distance, which efficiently measures the distance between two arbitrary sampleable distributions without being constrained to specific form as in the KL divergence, and without requiring expensive adversarial training. To improve the representation, we use (1) structural consistency term in the loss that encourages isometry between fea- ture space and latent space and (2) a nonlinear variant of the SW distance which averages over random nonlinear shearing. The effectiveness of the learned man- ifold encoding is best explored by traversing the latent space through interpola- tions along geodesics which generate samples that lic 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-geodeste latent space by maximizing the density of samples along the path while minimiz- ing total energy. We use the 3D-spiral data show that the prior does indeed encode the geometry underlying the data and to demonstrate the advantages of the algooott for interpolation. Additionally, we apply our framework MNIST, and CelebA datasets, and show that outlier generations, latent represen- tations, and geodesic interpolations are comparable to the state of the art.
in
INTRODUCTION
Generative models have the potential to capture rich representations of data and use them to gen- erate realistic outputs. In particular, Variational AutoEncoders (VAEs) (Kingma & Welling, 2014) can capture important properties of bigh-dimensional data in their latent embeddings, and sample from prior distribution to generate realistic images. Whille VAEs have been very successful in variety of tasks, the use of a simplistic standard normal prior is known to cause problems such as under-fitting and over-regularrizaioo. and fails 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-dimensions data can typically be modeled as lying or near an embedded low-dimensional. nonlinear manifold (Fefferman et al., 2016). Learning improved latent representations of this nonlinear manifold 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 cases that have tractable approximations of the KL divergence. Many works such as Guo et al. (200);; Tomezak & 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. Altemate approaches such adversarial training learn arbitrary priors by using

---

Under review as a conference paper at ICLR 2021
discriminator network to compute 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 conventional autoencoder architecture and an additional prior-encoder network that learns unconstrained prior distribution that encodes the geometry and topology of any data man- ifold. We use type of Sliced Wasserstein (SW) distance (Bonnotte) 2013 Bonneel et al. (013). a concept from optimal transport theory that is a simple and convenient alternative the KL diver- gence 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. Chen et al.] [2020b; Desh- pande 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 dis- tance, 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 iso- metric 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 pre- serving 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 work is the graph-based geodesic-interpolatii algorithm. Conventionally, VAEs use Euclidean interpolation between two points in latent space. However, since manifolds typ- ically 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 geodesies in latent space. Several works such as (Shao et al. (2018); Miolane & Holmes (2020b) endow the latent space with a Riemannian geometry and measure cor- responding distances, however these are difficult and involve explicitly solving expensive ordinary differential equations. In this work, we introduce 'network-geodesies", 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-ume non-uniformly thresholded such that the set of allowable paths from a given sample traverse high density regions through short hops. Lastly, gorithm like Dijkstra's 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 result- ing network-geodest curves give notion of distance on the manifold and can be used to generate realistic interpolation points with relatively few prior samples.
shortest path al-
The novel contributions of this work are:
We introduce novel architecture, EPSWAE, that consists of 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-geodesies 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 consis- tency 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.
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 man- ifold. In 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 dimen- sional manifold (Lu et al. 1998). Many works such as Weinberger & Saul 2006); Rahimi et al.
2

---

Under review as a conference paper at ICLR 2021
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 KI. divergence using implicit priors, Rezende & Mohamed (205)), which maps a Gaussian prior by a series of explicit transformations, [uoo et al. (2020) which learns a GMM prior using an approximate ELBO, Goyal et al. (2017) which uses hierarchical Bayesian framework, and VampPrior Tomczak & Welling 2018)) which learns a two-layer hierarchical GMM from aggregated posteriors. Yin & Zhou| (2018); Molchanov 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 closed-fome for any arbitrary distribution (Kolouri et al.] (2018b). Wasserstein distances (Villani, (003) have several of the same properties as KL divergence, but often lead to better sampling (Gulrajani et al., 2017; TTostikhiin ct al. 2018), and have been used in several machine learning applications (e.g. Arjovsky et al.
(2017) Tolstikhin et al. (2018); Kolouri et al. (0188)).
Sliced Wasserstein AutoEncoder (SWAE) that regularizes an autoencoder using SW distance was proposed in Kolouri al. (2018a). There exist several variants of SW distance that have been suc- cessful at improving generative ability for example, Deshpande et al. (2018); [Chen et al. (220bb) train discriminator-like networks perform nonlinear transformations (instead of purely linear pro- jections) whereas Deshpande et al. (2019) introduces max-SW distance, which looks for the best linear projection, generalized in Nnuuee et al. (2020) 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 optiiization 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); Makhhzaní et al. (2015); Arjovsky Arovsyyy al. (2017) 7);
Sainburg et al. (2018) replace the KL divergence with an adversarial loss, however adversarial meth- ods tend to significantly more expensive and difficult to optimize. In higher dimensions, using discriminator network as in adversarial approaches is a natural way of implicitly computing an cquivalent of the Wasserstein-I 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. (0018) uses a structural consistency term that encourages relative distances in data space to be preserved in latent space. Similar distance preserrin terms have also been used suc- cessfully in Yu et al. (2013); Benaim & Wolf (201).. 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 Dollár et al. (2077); Weinberger & Saul (2006) com- pute embeddings of -ighniimeniional 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 hyperspherical manifold. Miolane & Holmes (2020a) formulates Riemannian VAE, however computing its ELBO is challenging. Along similar lines, Arvanitidis et al. (2017) shows that under certain conditions, a Riemannian metric naturally induced in the latent space, and uses it to compute geodesics. In Chen et al. (2020a), a flat manifold is approximated by penalizing cur- vature in latent space, and geodesics are defined through Euclidean distances on the flat manifold.
There exists limited work 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 be more meaningful, giving notion of geodesic distance. Several approaches such as Tenenbaum et al. (2000): [Bernstein et al. (2000);
Klein & Zachmann (2004): Mé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 't robust to noise (Sober et al. 2020). Inspired by this, we introduce an
3

---

Under review as a conference paper at ICLR 2021
energy-based network algorithm to identify network-geodesien in latent space with relatively few, noisy samples.
EPSWAE
3.1 NONLINEAR SLICED WASSERSTEIN DISTANCE
Wasserstein distances are a natural metric for measuring distances between probability distributions, however they xare difficult compute in dimensions two and higher. The Sliced Wasserstein dis- tance dsw averages over the Wasserstein distance of ID projections (see Appendix(Affor derivation). derrivitton Several works (e.g. Kolouri et al. (2019): Deshpande et al. (2099) discuss why a conventional SW distance may be sub-optimal as large number of linear projections may be required to distinguish two distributions. In this work we use a Nonlinear Sliced Wasserstein distance (dysw), an averag- ing procedure over (random) nonlinear ID projections, between two distributions μ and y defined as:
dNsw = WEdww (NS ~ I
dsw (, NC
(1)
8=1
where L is the number of nonlinear transformations, and S, 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 J4, we define (.,u(A) = (()=)(A) for all A C Z measurable. Most importantly, to generate samples of fap one simply generates x μ and defines fz))
The nonlinear transformations are related to a special case of the generalized SW distance (Kolouri et al. 2019) (with the difference that we choose ietributiin-depennde most likely to produce deformations which highlight differences the measures, thus breaking the homogeneity condition H2 in Kolouri et al. (2019)), and are given by
shear frequencies that are
A (x) + Çe sin(nv w.).
In all our experiments, we use Several works such as Deshpande et al.) (2018 2019); Wu et al. (2019); Nguyen al. (2020); Chen let 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) be- ing the most similar to our method). We did not implement these other variants, however they could easily be used in conjunction with lail-behavior considerations (such as boundedness and non-saturation of nonlinearity), along with computational efficiency, and for our goals, 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
nonlinear transformations with 50 linear ID projections each.
method. Our choice of nonlinearity is motivated by certain
3.2 ALGORITHM DETAILS
In the Encoded Prior Sliced Wasserstein AutoEncoder (EPSWAE) (see Fig. I for schematic), let's define the data encoder as V E. the decoder V W. and the prior-encoder as prg. each with parameters
Φg, @DD.PE &rg respectively. Input samples x Px. where Px is the probability distribution of the input data, are passed through W generate posterior samples z(d)
V. by setting 20) V (x()))) Similarly, prior-encoder input samples g(i) μl, where μ is the probability distribution of the input to the prior-encoder (chosen to be a mixture of Gaussians) from distribution, are passed through V pg to generate prior samples y( ~ (I PE)+ by setting y(i) = P (g()). The prior- encoder network and the autoencoder (data encoder and decoder) network are trained iteratively in a two step process:
1. Given minibatch, parameters of the autoencoder (dhE, (DD) are trained for kn steps while parameters of the prior encoder (#Pg) 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 ofEx-Py Px free (x, x) + Bdd ((Vr). Px, (VrE).H) + &Lrsc,
4

---

3. Generate thresholded network: Let samples be represented nodes. For a given thresh- old value i, sample j is connected to sample with edge weight (i, if d(i..) < t L.Cg. Cj- Under review as a conference paper at ICLR 2021
where - ( is the autoencoder output. Note that dy sw is efficiently computed just from samples of the distributions.
In this measure theoretic notation, (V/ (Wg), Px is the posterior distribution (denoted as 40 (2|2r in Bayesian literature), and (g) is the prior distribution (denoted as Pore (z) in Bayesian literature). Similar to 3-VAE ((iggiin et al. (016)), the hyperparam- eters a, B, k. are tuneable; our experiments they are fixed, but it could be advantageous to have them vary over training epochs.
CFsC 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 (2018) for two point clouds F = =ff [fi.... fN] in feature space and the corresponding points Z 21.... 2N in latent space in minibatch of size N is given by:
present). Lesc (adapted from (Sainburg et al.
Lrsc
1
Σ 5-=1
log (1
1+
IIf, fell?
log (
1
12 N2 m,
zel12
2
N2
N2 Em,n 1l - fnll?
2m
2n
112
2. In the second step of minimization, the parameters of the prior-encoder (g) are trained for li2 steps while parameters of the autoencoder (øg, ΦD) are fixed. The loss function for the prior-encoder consists of the NSW distance between the prior and posterior:
CPg x dxsw Px (B))))).
(5)
The pseudocode is outlined in Appendix periments in this paper are in Appendix
Additional architecture and training details for the ex-
3.3 INTERPOLATION AND APPROXIMATE GEODESICS
EPSWAE attempts to learn representation of the embedded manifold geometry that the data lies along. The question arises: how does one make use of this representation? natural way of interpolating data which lies manifold is through geodesics or paths on the manifold which traverse dense regions. Here we present a network algorithm for efficiently approximating network- geodesics 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 F.) Px, i.e., (E(() with is a minibatch of data. Ad- ditional prior samples can be used to supplement this if desired;
2. For each sample j. compute the average Euclidean distance ey of the k-nearest neighbors:
Here h is the energy parameter in the edge weight (chosen be I or in experiments, with h = encouraging shorter hops). Sample-specific thresholding, i.e., thresholding depen- dent on Cj. increases the probability of 'central' node having a connection. Thus, this encourages (1) paths through high density regions, i.c., high-degree nodes, and (2) traver- sal of latent space through short hops as consequence of localization of paths, i.e., paths existing only between nearby points due to thresholding.
4. Continue increasing until the graph is connected, i.e., there exists a path, direct or indi- rect, from every sample every other sample. Then, use Dijkstra's algorithm (Dijkstra 1959) to identify network-geodesicc with least total energy through allowable paths on the thresholded network.
4 EXPERIMENTS
We run experiments using EPSWAE on three datasets. First, we use 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-geode interpolation. Then we present interpolation and generation results on the MNIST (LeCun et al, 2010) and CelebA (Liu et al. 2018) datasets.
interpolation over linear
5

---

Under review as a conference paper at ICLR 2021
4.1
ARCHITECTURE
GMM F
(a)
Conv Batch Norm Maxpool
(b)
Proo-Ennodder
Conv Transpose Upsample
FC layers
Y
Data Encoder x
plyl
Decader v,
Input x
Reconstruction
piz xx)
pG
Latent
flo
Space
FSC
NSW distance
Reconstruction Loss
Figure 1: (a) is schematic of the EPSWAE architecture. The red arrows indicate calculation of the loss terms. The prior-encoder generates a prior in latent space. (b) is simple example of shape where interpolating between points A and B 'through the manifold' (red dashed line through point C) desired, since linear interpolation (black line) leads out the concave hull.
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 y to the prior-encoder. In the artificial data set we use normal distribution, whereas for MNIST and CelebA we use mixture of 10 Gaussians with random i.i.d means (from Gaussian with a 2). The schematic of the architecture is presented Fig. Ia). Note that the specifics of for the image datasets MNIST and CclebA, the 3D Spiral data does not have convolutional layers. Details of the architecture, training, and hyperparameters for each dataset are given in Appendix [ The optimizer Adam (Kingma & Ba 2014) with learning rate 0.001 was used for learning in both networks: the prior-encoder and the autoencoder.
the layers shown in this schematic
4.2 LEARNING LATENT MANIFOLD
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), ((t)),=()) (t cos 7t /4, sin 7t /4, 2f). (b) random 40 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- space is R For these tests, the data encoder, prior-encoder, and decoder all consist of three Fully Connected (FC) layers with nodes each and ReLU activations.
is Gaussian in R40 and the latent
geodesic interpolation
(a)
linear interpolation
(b)
Posterior Prior
Figure 2: (a) side view and (b) top view of samples of the prior (green) and posterior (blue) in 3D latent space generated by EPSWAE after 100 epochs on the high dimensional input. The red line shows the interpolation along network-geodesse between two prior samples. The brown line shows trajectory of linear interpolation between two points on the spiral.

---

Under review as a conference paper at ICLR 2021
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 '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 [ 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-allorithm 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 in- terpolation. 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 network- interpolation algorithm are equally essential. Comparisons with baselines SWAE (Kolouri et al.) 2018a) and VAE (Kingma & Welling 2014) are shown in Appendix Fig. 0 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 effi- ciency and simplicity, i.e., not requiring additional optimizations every evaluation or the training of a discriminator-lik network to select optimal transformations. Our sinusoidal shear nonlinearity is closely related to special case of the generalized SW distance defined in Kolouri al. (2019) (see discussion in section [. for differences ). Here, we present a comparison of computational time with two polynomial-type nonlinearities discussed in Kolouri al.] al.] (2019)- cubic and quintic. Com- putations done on laptop with Intel Core i7-10710 10 CPU, 16GB RAM, and computed over 1000 runs. In practice, the losses were indistinguishable with choice of nonlinearity (see Appendix B.
whereas the sine shear is slightly less expensive (see Table DT. However, the complexity of cubic and quintic nonlinearities are exponential in dimension and hence are impractical for larger data.
Table 1: Comparison of computational time for nonlinearities in SW distance. sine-shear NSW 0.0050s + 0.0006s 0.0054s + 0.0007s 0.0060s + 0.0008s
cubic NSW
quintic NSW
Computational Cost
4.3 GENERATION AND OUTLIERS IN LATENT SPACE
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. (2188) 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 a, SWAE tends to generate unrealistic faces with distorted colorations, whereas EPSWAE is more likely 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 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 opti- mized: the encoder is (after downsizing CelebA images to 64r64) Conv (3.16.3) - BatchNorm BatchNorm + ReLu - MaxPool (2,2) + Conv (32,64,3) + BatchNorm + ReLu followed by two FC layers of 512 and 256 nodes re- spectively with leaky ReLu nonlincarity 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.
+ ReLu + MaxPool (2,2) + Conv (16.32.3)
7

---

Under review as a conference paper at ICLR 2021
EPSWAE
SWAE
Figure 3:
Images generated after 100 epochs from prior samples in (a) EPSWAE (b) Baseline SWAE at increasing standard deviations o, i.e., progressively more 'outlying'.
We also present a comparison with SWAE on CelebA generation from samples around the mean in Appendix G and find that EPSWAE sees marginal improvement in the quality of faces. Gen- cration on MNIST using EPSWAE (and comparisons with SWAE (Kolouri et al.) (018a) and VAE (Kingma & Welling 2014)) shown in Appendix F MNIST results obtained with 5/ latent space. EPS WAE generation is more natural, whereas both other baselines generate some bloated and unidentifiable numbers.
Table 2: FID score comparison generated with 10000 samples (lower is better)
SWAE 178.04
EPSWAE (bl) EPSWAE (FSC only) EPSWAE (NSW only) ESPWAE (FSC NSW) 161.55
162.14
161.06
157.86
we present Frechet Inception Distance (FID) (Sllimans et al. 2016) scores with 10000 images with a = (code from Seitzer (2020)). is important to note that one 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
For completeness, in Table 4.3
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.
4.4 INTERPOLATION
In this paper, we introduce graph-based algorithm that interpolates over network-gcodesics in curved latent manifolds, thus capturing properties of the geometry. Here we show that obtain smooth interpolation using the network- -eeoeesseeoessscs 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 4a) presents a com- parison of linear interpolation (top) interpolation through network-geodesies (bottom) described
8

---

Under review as a conference paper at ICLR 2021
(a) Linear (L) vs Network-geodesic (NG) interpolation
NG
L
NG
NG
[b) Geodesic interpolation
7 7 9
9 9
4
4
4
4
4
4
2
2
2
2
8
8
6
6
6
S
S
5
Figure 4: (a) Comparison of linear vs geodesic interpolation using EPSWAE. The top row of each row-pair is linear, and the bottom is geodesie-inlerpolation. (b) Randomly chosen instances for geodesic interpolations for MNIST and CelebA. The first and last images are reconstructions of real data, and the interpolations traverse through samples of the prior using the network-geodesi algorithm. A total of 400 samples are used all cases.
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 interpola- tions 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 g (b) shows examples of interpolations on MNIST and CelebA datasets. MNIST interpola- tions are smooth and intuitive, for instance, in the bottom left MNIST panel in Fig. (bb) the top part of '7' first changes to a '9' naturally, followed by transformation of the '9" to '4' Interpolation along network-geodest
ensures that reconstructions of intermediate samples are realistic. Inter- polations CelebA are smooth and pass through intermediate images that could arguably pass for celebritics these days. The state the art interpolations along stwork-geodessccs on the manifold indicate that the learned prior does indeed encode the data manifold. For all interpolations shown in Fig. we use energy parameter h = 2. Comparisons between interpolations corresponding to = and h 2 are presented in Appendix I0] In experiments, energy parameter 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 !
CONCLUSION
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- algorithm to identify network- geodesics in latent space that maximize path density while minimizing total energy. We demonstrate EPSWAE's ability to learn the geometry and topology 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 efficient and comparable to state of the art techniques. Our code is publicly available.
3D spiral from a noisy 40D embedding.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. arXiv preprint SEI::1701.00875. 2017.
Georgios Arvanitidis. Lars Kai Hansen, and Søren Hauberg. Latent space oddity: on the curvature of deep generative models. arXiv preprint 17X1v:1710.11379, 2017.
Sagie Benaim and Lior Wolf. One-sided unsupervised domain mapping. In Advances in neural information processing systems, pp. 752-762, 2017.
Mira Bernstein, Vin De Silva, John Langford, and Joshua B Tenenbaum. Graph approximations to geodesics on embedded manifolds. Technical report, Citeseer, 2000.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51(1):22 2015.
Nicolas Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD thesis, Paris 11, 2013.
Yuri Burda, Roger B. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. CoRR, abs/1509.005 .,.000599,,0 2016.
Nutan Chen, Alexej Klushyn, Francesco Ferroni, Justin Bayer, and Patrick van der Smagt. Learning flat latent manifolds with vaes. arXiv preprint arXiv:2002.04888/ 2020a.
Xiongjie Chen, Yongxin Yang, and Yunpeng Li. Augmented sliced wasserstein distances. arXiv preprint arXiv:2006.0883 08812, 2020b.
Tim R. Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M. Tomczak. Hyperspheri- variational auto-encoders. 34th Conference on Uncertainty in Artificial Intelligence (UAI-18), 2018.
Ishan Deshpande, Ziyu Zhang, and Alexander G Schwing. Generative modeling using the sliced wasserstein distance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3483-3491, 2018
Ishan Deshpande, Yuan- Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 10648-10656, 2019.
Edsger W Dijkstra. A note on two problems in connexion with graphs. Numerische mathematik, (1):269-271, 1959.
Nat Dilokthanakul, Pedro AM Mediano, Marta Garnlo,, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture varia- tional autoencoders. arXiv preprint drXiv.1611.02648. 2016.
Piotr Dollár, Vincent Rabaud, and Serge J Belongie.
Learning to traverse image manifolds. In Advances in neural information processing systems, pp. 361-368, 2007.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 29(4):983-104) 2016.
Prasoon Goyal, Zhiting Hu, Xiaodan Liang, Chenyu Wang, and Eric Xing. Nonparametric vari- ational auto-encoders for hierarchical representation learning. In Proceedings of the IEEE Inter- national Conference on Computer Vision, pp. 5094-5102, 2017.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. In Advances in neural information processing systems, pp- 5767-5777, 2017.
Chunsheng Guo, Jialuo Zhou, Huahua Chen, Na Ying, Jianwu Zhang, and Di Zhou. Variational autoencoder with optimizing gaussian mixture model priors. IEFE Access, 8:43992-4400 2020.
10

---

Under review as a conference paper at ICLR 2021
Gaëtan Hadjeres, Frank Nielsen, and François Pachet. Glsr-vae: Geodesic latent space regularization for variational autoencoder architectures. In 2017 IEEE Symposium Series on Computational Intelligence (SSCI), pp. 1-7. IEEE, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerehner. beta-VAE. Learning basic visual concepts with a constrained variational framework. 2016.
D. P. Kingma and M. Welling. Auto-encoding variational bayes. In International Conference Learning Representations. 2014.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Conference Learning Representations, 12 2014.
International
Thomas N Kipf and Max Welling. Variational graph auto-encoders. NIPS Workshop on Bayesian
Deep Learning, 2016.
Jan Klein and Gabriel Zachmann. Point cloud surfaces using geometric proximity graphs. Comput- & Graphics, 28(6):839 8500 2004.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced wasserstein auto- encoders. In International Conference on Learning Representations, 2018a.
Soheil Kolouri, Gustavo K Rohde, and Heiko Hoffmann. Sliced wasserstein distance for learning gaussian mixture models. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 3427-3436 2018b.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced wasserstein distances. In Advances in Neural Information Processing Systems, pp. 261- 272, 2019.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http:/yyam.lecu..commeeeee
2010.
Chen- Yu Lee, Tanmay Batra, Mohammad Haris Baig, and Daniel Ulbricht. Sliced wasserstein discrepancy for unsupervised domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 10285-10295, 2019.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba)
dataset. Retrieved August, 15:2018, 2018.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Rober Stöter. wassertww flows: Nonparametric generative modeling via optimal transport and diffu- sions. In International Conference on Machine Learning, pp. 4104 4113. PMLR, 2019.
Haw-Minn Lu, Yeshaiahu Fainman, and Robert Hecht-Niclsen. Image manifolds. In Applications of Artificial Neural Networks in Image Processing III, volume 3307, pp. 52-63. International Society for Optics and Photonics, 1998.
Shitong Luo and Wei Hu. Diffcrentiable manifold reconstruction for point cloud denoising. arXiv preprint arXiv:2007. 2020.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, lan Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprint arXiv: 1tf:::551..0000 1.05644, 2015.
Facundo Mémoli and Guillermo Sapiro. Distance functions and geodesics on submanifolds of \rdd and point clouds. SIAM Journal Applied Mathematics, 65(4):1227-1260) 2005.
Nina Miolane and Susan Holmes. Learning weighted submanifolds with variational autoencoders and riemannian variational autoencoders. Proceedings of the IEEE/CVF Conference on Com- puter Vision and Pattern Recognition, pp. 14503-14511, 2020a.
Nina Miolane and Susan Holmes. Learning weighted submanifolds with variational autoencoders and riemannian variational autoencoders. Proceedings of the IEEE/CVF Conference on Com- pufer Vision and Pattern Recognition, pp. 14503-14511, 2020b.

---

Under review as a conference paper at ICLR 2021
Dmitry Molchanov, Valery Kharitonov, Artem Sobolev, and Dmitry Vetrov. Doubly semi-implict variational inference. In The 22nd International Conference on Artificial Intelligence and Statis- fics, pp.
2019.
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional slised-wasserstein and applica- tions to generative modeling. arXiv preprint arXiv:2002.07367 2020.
Ali Rahimi, T Darrell, and B Recht. Learning appearance manifolds from video. In 2005 IEEE Com- puter Society Conference on Computer Vision and Pattern Recognition (CVPR'0S) 00), volume pp- 868-875. IEEE, 2005.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional Conference on Machine Learning, pp. 1530-1538, 2015.
Tim Sainburg, Marvin Thielk, Brad Theilman, Benjamin Migliori, and Timothy Gentner. Gener- ative adversarial interpolative autoencoding: adversarial training on latent space interpolations encourage convex latent distributions. arXiv preprint arXiv:1807.06650 2018.
Tim Salimans. Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Advances in neural information processing systems, pp. 2234-2242, 2016.
Maximilian Seitzer. pytorch- fid: FID Score for PyTorch. http:s://gittub..co/mmmess pytorch-fid August 2020. Version 0.1.1.
Hang Shao, Abhishek Kumar, and P Thomas Fletcher. The riemannian geometry of deep generative models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 315-323, 2018.
Barak Sober, Ingrid Daubechies, and Robert Ravier. Approximating the riemannian metric from point clouds via manifold moving least squares. arXiv preprint arXiv:2007.0988. 09885, 2020.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka, Masanori Yamada, and Satoshi Yagi. Vari-
ational autoencoder with implicit optimal priors. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 5066-5073, 2019.
Joshua Tenenbaum, Vin De Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, 290(5500):2319-2322 2000.
I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schölkopf. Wasserstein auto-encoders. ArXiv, abs/1711.01 2018.
Jakub Tomczak and Max Welling, Vae with a vampprior. In International Conference Artificial Intelligence and Statistics, pp. 1214-1223, 2018.
Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
Hui-Po Wang, Wen-Hsiao Peng, and Wei-Jan Ko. Learning priors for adversarial autoencoders.
APSIPA Transactions on Signal and Information Processing, 9, 2020.
Kilian Weinberger and Lawrence K Saul. Unsupervised learning of image manifolds by semidef-
inite programming. International journal of computer vision, 70(1):77-90, 2006.
Jiqing Wu, Zhiwu Huang, Dinesh Acharya, Wen Li, Janine Thoma, Danda Pani Paudel, and Luc Van Gool. Sliced wasserstein generative models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3713-3722. 2019.
Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. volume 80 of Proceedings of Machine Learning Research, pp. 5660-5669, Stockholmsmssann Stockholm Sweden, 10-15 Jul 2018. PMLR. BRLLitttp://rroceddnng wtt.://wwwwee///////nnnl.. 3tmll
Wenehao Yu, Guangxiang Zeng. Ping Luo, Fuzhen Zhuang, Qing He, and Zhongzhi Shi. Embed- ding with autoencoder regularization. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 208-223. Springer, 2013.
12

---

Under review as a conference paper at ICLR 2021
A BACKGROUND ON THE SLICED WASSERSTEIN DISTANCE
Wasserstein distances provide a natural metric for measuring distances between probability distribu- tions based on optimal transport, i.e., the cost of deforming one probability distribution into another. For measurable cost function c(æ, y) for x, y E IR', and probability distributions and v on IR we define the p- Wasserstein distance between the distributions as
Verrze (e vdd rell(u.)
1/P
dw dw,)(p.v)=
inf
(6)
where II is the set of all joint distributions with marginals 1 and (see Villani (003)).
This Wasserstein distance between probability distributions is extremely difficult and computation ally intensive in dimensions two and higher, i.c., d 2 2. However, in dimension one, there exists closed-form solution (Villani 2003). simple algorithm for computing the ID Wasserstein distance is given follows (see for example (Kolouri et al. 2018a)): (a) Generate N one dimensional i.i.d samples 2j ~ μ., ~ v; (b) sort each list X #1.... Ey Y = I,, UUv] into ascending order denoted by X, Y respectively; (c) define the approximation to the ID p- Wasserstein distance
1/2
N ( 055
dw v) 2
N
where ē; X Ý. The Sliced Wasserstein (SW) distance defines metric on probability
measures (Bonnotte) 2013) which provides an alternative to exploiting the computational
feasibility of the TD Wasserstein distance in Eqn. 7 It involves averaging over one-dimensional orthogonal projections x" := (0 x)0 as follows:
dsw (M, = (for ( 1/P
(8)
The SW distance has seen a variety of implementations (Bonneel et al. 2015 Kolouri et al.
2018 1b
Lee et al. 2019).
The integral in Eqn.
g can be easily approximated by sampling M random one-dimensional vectors
0g uniformly Sd and computing
1/P
dsw (1, v)
Bee (m pl p,, mb(v)
M
Several works (for example Nguyen al. (2020): Deshpande et al. (2199) have discussed reasons why this often does not result in the best computational method. Linear projections may be sub- optimal for extracting information about the differences between 1a and v, 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. 2020bb 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 transforma- tions. 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 poly- nomials) have a pronounced deformation on the tails of a measure and may excessively weight outliers. (2) A sigmoid is another potential candidate, but saturates at high values and we want the non-liny 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 com- plexity 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 significant effect on performance, with choice being the fastest (see section 4.2 and Appendix for computational time and loss comparisons). In principle, other ensembles of nonlinear transformations could be used.
13

---

Under review as a conference paper at ICLR 2021
B PSEUDOCODE
Algorithm Training EPSWAE
1: while not converged do
2: Update the autoencoder V We D:
3: for k1 substeps do
4: Sample minibatch from data (x(1, ,.., ,.., x with x66 ~ Px
5:
Compute feature extractor samples p() = (x)(x(1)
6:
Compute posterior samples z(ا) (x((())
7:
Compute decoded output Xrecon (n()))
8:
Generate samples (g(1), ...,(())) with g(t) 1
9:
Compute prior samples y()) = V PEE(E()
Compute reconstruction error: Lrecon
Σラ=1 d(x(1), x
Compute NSW distance: LNSW distance dw ( Σ 82(5), 23=18
12:
Compute FSC loss: LESC desc((f(),
f()) } fz( ((.,),,,,'(,)
Compute autoencoder loss LAE = alrreoo + B NL SWisttance + K&Lsc
14:
Compute gradients of LAE wrt to $E ØD
15:
Update ΦD.6D end for
16:
17:
Update the prior-encoder Ppg
18:
for l2 substeps do
19:
Sample minibatch from data (x(1) x())} with x(5) ~ Px
20:
Compute posterior samples z zi6 V (x(())
21:
Generate samples ((()),..((')) with g(i) 1
22:
Compute prior samples y()) = V (F((()
23:
Compute prior-encoder loss: LPF dnsw ( -ÿ=1 O262, $ Eÿ=1
24:
Compute gradients of LPÆ ΦPE
25:
Update ΦrE
end for
27:
end while
C ARCHITECTURE AND TRAINING DETAILS
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 B and 5 are trained iteratively using the optimizer Adam (Kingma & Welling 2014) with a learning rate of 0.001. We experimented with both p and = (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 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 40D Gaussian, and the latent space 3D. The input to the autoencoder is 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 given by the Mean Square Error and a 1, = 0.1, K 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 i = 2 times for each training of the autoencoder k 1. Power of distance 2 is used to compute edge weights for computing network -ewwork-eeddsscss....
MNIST dataset: The input to the prior-encoder is 40 dimensional mixture of 10 Gaussians, and the latent space is 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
14

---

Under review as a conference paper at ICLR 2021
(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 Cony' Transpose (16,10,3) -> LeakyRel.u Upsample(2.2) - ConvTrannposse (10,1,3) - Sigmoid. The decoder output has size 28 x 28. The prior-encoder is trained l 1 times for each training of the autoencoder kn 1. The reconstruction loss is given by the Binary Cross Entropy and a = 3 0.1, K 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.
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 trade-off image quality of the input images for computational speed by downsizing the input and employ fairly simple network compared to state of the art computer vision architectures that use CelebA. The data encoder takes CelebA images (of size 218 x 178 x 3) and downsizes them to size 64 x 64 x 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
CelebA dataset:
a function of latent dimension. We
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 x 64 3. The prior-encoder was trained k 1 times for each training of the autoencoder. The reconstruction loss is given by the Binary Cross Entropy and 500, 3 = 50, x = 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, ie.. computed at the output of the last convolutional layer in the encoder. Energy parameter = used to compute edge weights for computing network-geodess
D SPIRAL BASELINE COMPARISONS
Linear SW distance No structure consistency loss
Linear SW distance structure consistency loss
(c) Nonlinear SW distance No structure consistency loss
(d) Nonlinear SW distance structure consistency loss
Fastrior
Paderis
Figure 5: Comparisons of the EPSWAE model with different loss terms. Top panels present top views, and bottom panels present the corresponding side views. The red curves show interpola- tion between two randomly selected samples using the network-geodesie algorithm. All figures are generated after 100 epochs with a 0r0001,, and batch size =100. k 1, k2 2.
Figure 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
15

---

Under review as a conference paper at ICLR 2021
(a) EPSWAE
(b) SWAE
(c) VAE
Interpaation
Posteria Price
Figure 6: Comparison of EPSWAE with baselines SWAE (Kolouri et al. 2018a) and VAE (Kingma & Welling 2014). All figures are generated after 100 epochs with Ir-0.01, and batch size =100.
For EPSWAE, k = 1, k2 2, a = 1, = 0.1, K = 0.001.
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 structural consistency term doesn't remarkably improve the quality of the manifold learned, however, consistent with other experiments, 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 seen in these visualizations of the 3D spiral lead us both loss terms (as in (d)) all results in the main paper.
Figure 6 compares the EPSWAE model with baselines VAE and SWAE The VAE uses KI. diver- gence 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 Gaussian, leading to an unnatural em- bedding of the data manifold. In contrast, EPSWAE results in significantly better learned prior as a consequence of the nonlinear SW distance and the prior-encoder network trained explicitly to improve the latent representation.
E COMPARISON WITH OTHER SLICED WASSERSTEIN NONLINEARITIE,
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. max- SW (Deshpande et al.] 2019) which has seen several subsequent variations), these involve additional training. and hence aren't 't considered (although many such versions of SW distance could be used in conjunction with out method), Instead consider the polynomial generalized Radon transforms described in (Kolouri et al. 2019). The computational cost of a single polynomial generalized Radon transform is ( for dimension and t order of polynomial, and hence we limit investigation to cubic and quintic polynomials. We show results on the artificial 3D Spiral dataset 3 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 sine- shear has slightly lower computational cost.
trained upto 100 epochs with
16

---

Under review as a conference paper at ICLR 2021
Cubic poly
10a
Quintic poly
Sine shear
100
Total Loss
Linear
10-1
10-2
10
0
20
40
60
80
100
Epochs
Figure 7: Loss as function of training epochs for linear SW, cubic nonlinearity, quintic nonlinearity and our (sine-shear) NSW distance. A full set of cubic and quintic functions terms are considered for computing the generalized Radon transform in the SW distance. nonlinear transforms with M = 50 linear 1D projections each taken for all cases.
MNIST GENERATION RESULTS
(a) EPSWAE
(b) SWAE
(c) VAE
5
3
2
9
5
2
q
6
7
4
5 7
7
2
3
6
4
6
O
5
3
8
9
2
8
7
0
3
6
4
9
7
9
9
6
4
S
5 1
2
9
1
9
1
5
3
8
3
7
6
5
7
9
9
9
S
7
O
3
2
0
4
7
6
(
2
0 6 /
6
/
3
0
3
9
4
7
/
0
0
6
6
3
9
2
2
a
5
3
7
3
3
4
1
6
/
3
4
2
0
6
3
64
S
/
3
5
6
9
9
1
0
0
9
2
3
4
6
5
6
0
2
4
2
6
7
4
6
5
2
9
0
8
5
S
3
5 7
$
9
5
3
2 1 3
4
8
0
3
3
0
6
0
5
5
6
Figure 8: Comparison between (a) EPSWAE (b) SWAE and (c) VAE for generation MNIST. for
all three, batch size 200, e0ochs=100, 0.001, for EPSWAE k = 2, k = 1, a = 1,3
0.1, K = 0.001.
Figures 8 (a.b.c shows generation on the MNIST dataset after 100 epochs on EPSWAE, baseline SWAE(Kolours et al. 2018a), and baseline VAE (Kingma & Welling 2014) respectively. All net- works EPSWAE and baselines SWAE and VAE use equivalent architectures (outlined section of Appendix C) and hyperparameters that EPSWAE generated samples were consistently found to have a lower fraction of 'false' digits, ie.. digits that unrealistie.
optimized individually. As in the figure, observe
17

---

Under review as a conference paper at ICLR 2021
G CELEBA GENERATION RESULTS
(a) EPSWAE
(b) SWAE
Figure 9: Images generated from prior samples in (a) EPSWAE (b) Baseline SWAE. Batch size 200, epochs=100, Ir 0.001. 500, B 50, K 0.05.
CelebA raw images downsized to 64 x 64 pixels, and simple architecture is employed (see Appendix C This serves as a proof of principle for a better latent representation and interpola- sophisticated network (for instance ResNet, VGG etc.) without downsizing the data would yield higher quality images at computa- tional cost.
tion, without high computational cost. Naturally, using
FigureDshows 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 for equivalent training epochs and architectural details, and does not make claims about the quality of SWAE images with longer training.
18

---

Under review as a conference paper at ICLR 2021
H EFFECT OF ENERGY PARAMETER ON CELEBA INTERPOLATIONS
Additional CelebA interpolations are shown in Fig. 10 The length of interpolations is automatically selected by the network algorithm. In some cases, linearly interpolated point is added between every two samples on the network-geodesic to make smoother transitions; this would be unneces- sary with enough samples of the prior, however it could be impractical 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 purely linear interpolation (an illustration is provided in Figure - the direct linear interpolation from A to B significantly different from the interpolation through the point C). Fig. T presents a compar- ison between energy parameter h 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.
(a) h=1
(b) h=2
Figure 10: Additional CelebA interpolations. The first and last images are reconstructions of real data, and the interpolations traverse through samples of the prior using the network- algo- rithm. (a) Top three panels show interpolations for energy parameter h 1, and (b) bottom three panels show interpolations for energy parameter h = 2. Latent space is 128 dimensional. Hyper- parameters are the same as those used in the paper. A total of 400 samples in latent space are used.
19

---

Under review as a conference paper at ICLR 2021
I
LINEAR INTERPOLATION COMPARISON EPSWAE AND SWAE
(a)
EPSWAE
SWAE
EPSWAE
SWAE
(b)
EPSWAE
SWAE
EPSWAE
SWAE
Comparisons of linear interpolation between EPSWAE and SWAE. The first and last images are reconstructions of randomly picked real data (a) shows long interpolation (over two lines), (b.c) show shorter interpolations. Hyperparameters are the same those used the paper. Latent space is 128 dimensional. A total of 400 samples in latent space are used.
Figure 11:
We compare here linear interpolations on the CelebA dataset for MNIST and CelebA. The autoen- coder networks used are identical and outlined in section C In contrast, EPSWAE uses the addi- tional prior encoder. Both models are independently optimized, and the interpolations are linear. As seen in Fig. IT 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.
20


Reviews:
[{"review_id": "bIENVj0QHo", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nWhile the paper cites some works that also aims at addressing the drawback of SWD, it still misses some important related works like [a, b, c]. The paper is highly expected to make more discussion and suitable empirical comparison with them.\n\n[a] Deshpande et al., Generative modeling using the sliced wasserstein distance, CVPR 2018.\n\n[b] Wu et al., Sliced wasserstein generative models, CVPR 2019.\n\n[c] Liutkus et al., Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions, ICML 2019.\n\nThe motivation of using nonlinear transformation sounds not convincing. It is indeed known that traditional SW approximation generally requires a large amount of linear transformations. The bottleneck has been overcome by some existing works (Chen et al., 2020b; Kolouri et al., 2019; Nguyen et al., 2020; Deshpande et al., 2019, 2018). Why can the suggested nonlinear transformation avoid suffering from this issue? It is also not clear why to choose Eq.(6) for the nonlinear transformation. Are there any more excellent properties of the suggested nonlinear transformation compared to the existing methods? It is also necessary to make more discussions on the application of other nonlinear transformations.\n\nThe motivation of using the additional prior-encoder is not clear to me? The introduction states that it learns an unconstrained prior distribution that matches any data manifold topology. Unfortunately, I cannot find any clear explanation about this in the proposed method part. \n\nThe evaluation is highly insufficient. The paper merely compares  the proposed method with SWAE which was published in 2018. More recent methods like [Deshpande et al. 2019, Wu et al. 2019, Liutkus et al. 2019] should be compared for a more complete study. Moreover, new generative modeling methods should be evaluated quantitatively using popular metrics like inception score, FID etc. Unfortunately, this paper does study this at all. In addition, the visual results of the proposed method seems not comparable with the state of the art on the CelebA dataset.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "POuXoPrMjAK", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\n-------------------\n\nThe paper is well written. The pipeline is clear and easy to understand. The representation learning with VAEs is a widely studied topic. Using a flexible implicit prior will boost the learning of the latent codes. The usage of the encoder (or generator in the adversarial cases) in the latent space is widely discussed in previous works, such as vampprior, semi-implicit VI, doubly semi-implicit VI, etc. Thus the contribution in this part is limited.  \n\nThe sliced Wasserstein distance is an efficient approximation of the Wasserstein distance for the distribution matching. To avoid the projections that contain useless information, as cited in this paper, a lot of papers generalize the typical random linear transform to non-linear transformation (Chen et al., 2020b; Kolouri et al., 2019; Nguyen et al., 2020; Deshpande et al., 2019). Here the authors propose the non-linear sliced Wasserstein (NSW) distance in Equation 5 with a transformation shown in Equation 6, which appears to be a special case that satisfies the four conditions discussed in (Kolouri et al., 2019). Is there any difference between NSW distance and the generalized SWD in Kolouri et al., 2019? \n\nThe graph-based method is an interesting way for the manifold walking and is much better than conventional ways such as linear interpolation. I think this part should be discussed more in the paper. \n\n-------------------\n\nSome detailed questions about the technique\n\n* The authors claim the usage of FSC encourages pairwise distances of the latent code to be similar to the pairwise distances of the data features. I am curious if FSC is necessary for manifold learning. In figure 5, it does not show much difference with/without the FSC loss. \n\n* As claimed in the paper, the adversarial methods in latent space are expensive, while in the experiment part, there is no computation (such as time per update step) comparison with the adversarial methods. The usage of FSC also needs to compute the pairwise distance. Is that expensive too?\n\n* In the abstract and the pseudocode shown in the appendix, the prior encoder is trained with sliced Wasserstein loss. Why we choose SWD for the training instead of NSW? \n\n* It is also confusing that in the method part, the prior encoder is trained with NSW loss, which is not consistent as claimed in the other parts of the paper.\n\n* In the experiments, the authors claim the generations are better, while the quantitative results, such as fid, are not provided. \n\n* For the spiral toy dataset: views in Fig. 5 seem not consistent and are hard to compare. How about we compare in 2d cases?\n\nIn conclusion, the paper has the merits, and these investigations may be helpful for this problem, but it is not enough and need to dig out more to be an ICLR publication.  \n\n---\n\nUpdate after the discussions\n\nI appreciate the efforts that the authors make in their responses, some of which address my concerns and improve the quality of the paper.\n\nI have raised my rating. However, taking into account all information during the discussion phase, I stick to my original review that this paper still needs to explore more to be a mature publication. For example, if the main contribution comes from the prior encoder, as I said the contribution is limited since the usage of the encoder (or generator in the adversarial cases) in the latent space is widely discussed in previous works, such as vampprior, semi-implicit VI, doubly semi-implicit VI, etc. This also seems to make the contribution of the sliced Wasserstein part incremental. Plus, as this paper has several components, their relations need to be discussed in a more clear way. Thus, more ablation studies are needed to help this paper to present its insight in a much more clear way.\n\nThanks again for the efforts that the authors make and I hope my reviews could help them to polish this paper to be a nice publication. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Iinm5DjqIzK", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Paper Summary\n\nThe 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. \n\n# Pros and Cons\n\nImproving the ability of generative models to capture high-dimensional empirical distributions accurately is a key problem in machine learning and central to the representation learning theme of the ICLR community. \n\nThe paper clearly states contributions up front, namely: using an encoder to generate richer priors for the latent variable distribution, non-linear projections for sliced Wasserstein approximation, and a graph based interpolation method.  They also alter the structural consistency term so that it applies to more abstract features instead of inputs. \n\nThe paper does a thorough and clear job of covering prior work and technical background such as the Wasserstein metric, perhaps even excessively so.  \n\nThe particular choice of a sinusoidal non-linear projection, a key contribution according to the paper, is not motivated in the text. On first glance, the sinusoidal term seems like an odd choice for a non-linearity. After looking at the results which include a test on spirals, it is clearer why this might have be chosen, but is a sinusoidal term likely to be helpful for non periodic data? It might be possible to shed some light on this by investigating whether the coefficients in the non-linear term, zeta and gamma, are significantly different from zero after training on MNIST or the CELEB A data set used in the paper.  \n\nFigure 3 comparing EPSWAE and SWAE doesn’t clearly illustrate the benefits of the EP component. In the print version, both grids of images seem blurry and prone to oddities such as overly large and dark eyes. Even when one blows up the image to 3X size using the digital version the advantage does not jump out. While I recognize that evaluating generative models is hard, the observations do not clearly support the author's hypothesis that the EP component provides an advantage. Maybe they could evaluate Freschet Inception Distance over the whole data set? Use blind human reviewers to choose between EPSWAE and SWAE based on realism and report preference scores? Interestingly, there is one duplicate image in the EPSWAE grid: row 1, col 5 and row 5, col 8 look identical. Seems odd to get identical images: is this because of sampling from a discrete graph structure? The highly similar but not identical images row 1, col 1 and row 1, col 8 are more what I would expect. Maybe the advantage could be made clearer by helping us focus on relevant features. For instance, it might be that EPSWAE is a little bit less likely to generate large black eyes? I can't tell from this small sample, but a grid that focuses on this might make the point. Hair versus background also seems to be a challenge. SWAE image row 3, col 4 seems to have two hair regions for the same face, but, EPSWAE images such as row 4, col 7 also look like they have two distinct hair regions. There are a couple of SWAE images that seem particularly ill formed: row 5, col 5  and maybe row 5, col 8, and row 7, col 4. It might be worth focusing on a few of the worst examples from both EPSWAE and SWAE to show differences in tails rather than the mean? I wonder if you could do a leave-one-out kind of analysis where you check the probability of held out training data points under the prior for EPSWAE vs. SWAE to assess if the prior is capturing the empirical distribution better? You would have to invert the prior network with gradient descent to do this ... but it might work. \n\nIt might make sense to compress some of the tutorial material up front to make room for more results demonstrating the efficacy of the model. The MNIST results, fig 7 in appendix D, shows some advantage for EPSWAE: For instance, figure 7b and 7c both have more bloated numbers … especially 8, 3 and 0.  Also figure 7b has one degenerate number in position 1,2 … possibly a 2?  \n\nThe paper argues that the structured latent space improves generative power. Is there any evidence of structure in the latent space after training on CELEB A? If we plot 2 or 3D projections, or plot projections of structure preservering factorizations such as PCA, do we see structure in the encoding training data points or are they distributed independently? One might also try independence tests between variables in the encoded latent space to see if pairs of variables are being encoded with correlations. This could be compared easily between EPSWAE and SWAE. \n\nFigure 4 in section 5.4 on interpolation shows smooth interpolation between two points A and B in latent space presumably drawn from training data. The interpolations are pretty smooth. Nice! However, the authors claim is that the graph embedding gives better interpolations than linear interpolation between points in the latent space. To make this point, we would also need to see interpolations between points using linear interpolations. The plots on spiral might make this point, but it is not clear.  \n\nThe ability of VAE’s to disentangle the dimensions of an empirical distribution into indpendent latent variables is sometimes seen as a feature, not a bug. For instance, if the training data truly lie along a spiral, isn’t this really a 1D latent space and not a 3D space? While I can see the appeal of improving generated distribution realism, some discussion by the paper on the merits of improving encoder and decoder vs. complexifying the latent space would help to motivate this approach. \n\nIt isn't clear to me that the structural consistency term is a good idea in general. Ideally we want the latent space to capture something fundamental about the underlying structure of the data and not features of the input. Moving the structural consistency from input to more abstract features addresses this concern somewhat, but aren't the latent values themselves the ultimate goal? \n\n# Recommendation\n\n I recommend a rejection of the paper. The hypotheses (that richer priors and geodesic interpolation generate better images on realistic images) are not clearly supported by the experimental results provided.\n\n# Questions\n\n For the spiral training, what was the dimensionality of the latent space? Was it 3D?   \n\nFigure 4 caption contains statement \"through an intermediate sample corresponding to the midpoint in latent space\". Are you actually literally using the midpoint? I thought graph embedding was being used to avoid using midpoints? Maybe the sentence is just ambiguous? \n\nHow many samples are used in the Wasserstein approximation? How were the coefficients in the multi-term loss function defined (alpha, beta and kappa)? Oh - I see these are in the appendix... Seems like appendices are becoming pretty integral to papers these days ... Probably worth including these for the main results in section 5 for the two results presented. \n\nStep 3 in the graph embedding didn't make sense to me after reading it a couple of times. It wasn't quite clear how this sample specific weighting works. It would probably be worth expanding this a bit at the expense of background material, as it is one of the contributions of the paper. \n\n# Other Feedback\n\nPage 2 \"Adversarial methods are harder to train\" ... also adversarial methods are implicit distributions -- you can sample from them but you cannot easily calculate the likelihood of an image under the adversarial model (although you can use gradient descent to try to find the latent parameters). This makes things like outlier detection difficult. \n\nPage 7, Section 5.2, last sentence refers to Fig 6, but I think this should be Appendix D, Fig 7? \n\nIf you flipped the name around from EPSWAE to SWEP-AE, you would have a much more memorable acronym for people to take away from your paper/talk/poster although I recognize this doesn’t have the same “build” on previous work dynamic.  ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All three referees have provided detailed comments, both before and after the author response period. While the authors have carefully revised the paper and provided detailed responses, leading to clearly improved clarity and quality, there remain clear concerns on novelty (at least not sufficiently supported with ablation study) and experiments (neither strong enough nor sufficient to support the main hypotheses). The authors are encouraged to further improve their paper for a future submission.
Decision: Reject