Symmetric Wasserstein Autoencoders
Authors: Sun Sun
Hongyu Guo
Venue:  (2021)
Keywords: generative models
variational autoencoders

Abstract: Leveraging the framework of Optimal Transport, we introduce a new family of generative autoencoders with a learnable prior, called Symmetric Wasserstein Autoencoders (SWAEs). We propose to symmetrically match the joint distributions of the observed data and the latent representation induced by the encoder and the decoder. The resultant algorithm jointly optimizes the modelling losses in both the data and the latent spaces with  the loss in the data space leading to the denoising effect. With the symmetric treatment of the data and the latent representation, the algorithm implicitly preserves the local structure of the data in the latent space. To further improve the latent representation, we incorporate a reconstruction loss into the objective, which significantly benefits both the generation and reconstruction. We empirically show the superior performance of SWAEs over several state-of-the-art generative autoencoders in terms of classification, reconstruction, and generation.

Under review as a conference paper at ICLR 2021
SYMMETRIC WASSERSTEIN AUTOENCODERS
Anonymous authors
Paper under double-blind review
ABSTRACT
Leveraging the framework of Optimal Transport, we introduce new family of generative autoencoders with a learnable prior, called Symmetric Wasserstein Au- toencoders (SWAEs). We propose to symmetrically match the joint distributions of the observed data and the latent representation induced by the encoder and the decoder. The resulting algorithm jointly optimizes the modelling losses both the data and the latent spaces with the loss in the data space leading to the denoising effect. With the symmetric treatment of the data and the latent representation, the algorithm implicitly preserves the local structure of the data in the latent space. To further improve the latent representation, we incorporate a reconstruction loss into the objective. which significantly benefits both the generation and reconstruction. We empirically show the superior performance of SWAEs over the state-of-the-e generative autoencoders in terms of classification, reconstruction, and generation.
INTRODUCTION
Deep generative models have emerged as powerful frameworks for modelling complex data. Widely used families of such models include Generative Adversarial Networks (GANs) (Goodfellww al,
2014). Variational Autoencoders (VAEs) (Rezende et al. 2014 Kingma & Welling (014). and autoregressiive models (Uria et al. 2013] Van Oord al. 2016). The VAE-based framework has been popular as it yields a bidirectional mapping, i.e., it consists of both an inference model (from data to latent space) and generative model (from latent data space). With an inference mechanism VAEs can provide a useful latent representation that captures salient information about the observed data. Such latent representation can in turn benefit downstream tasks such as clustering, classification, and data generation. In particular, the VAF-based approaches have achieved impressive performance results on challenging real-worke applications, including image synthesizing (Razavi et al 2019). natural text generation (Hu et al.) 2017). and neural machine translation (Sutskever et al. 2014).
VAES aim to maximize a tractable variational lower bound on the log-likelihood of the observed data, commonly called the ELBO. Since VAEs focus modelling the marginal likelihood of the data instead of the joint likelihood of the data and the latent representation, the quality of the latent is not well assessed (Almi et al. 2017 Zhao et al.) 2019). which is undesirable for learning useful representation. Besides the perspective of maximum-likese likelihood learning of the data, the objective of VAEs is equivalent to minimizing the KI. divergence between the encoding and the decoding distributions, with the former modelling the joint distribution of the observed data and the latent representation induced by the encoder and the latter modelling the corresponding joint distribution induced by the decoder. Such connection has been revealed in several recent works (Livne et al.) 2019 Esmacili al. 2019; Pu al. 20176; Chen al. 2018). Due the asymmetry of the KL divergence, it is highly likely that the generated samples are of low probability in the data distribution, which often leads to unrealistic generated samples (Li ei al. 20176 Alemi et al., 2017).
A lot of work has proposed to improve VAEs from different perspectives. For example, to enhance the latent expressive power VampPrior (Tomezak Welling 2018), normalizing flow (Rezende & Mohamed 2015), and Stein VAEs (Pu et 20177)))) replace the Gaussian distribution imposed on the latent variables with a more sophisticated and flexible distribution. However, these methods are all based the objective of VABs, which therefore are unable to alleviate the limitation of VAEs induced by the objective. To improve the latent representation (Zhao et al. 2019) explicitly includes the mutual information between the data and the latent into the objective. Moreover, to address the asymmetry of the KL divergence in VAEs (Livne et al.) 2019] Chen et al., 2018 Pu et al.) (017b) leverage symmetric divergence measure between the encoding and the decoding distributions.

---

Under review as a conference paper at ICLR 2021
Nevertheless, these methods typically involve sophisticated objective function that either depends on unstable adversarial training or challenging approximation of the mutual information.
In this paper, leverage Optimal Transport (OT) (Villani 2008 Peyré al. 2019) to symmetri- cally match the encoding and the decoding distributions. The OT optimization is generally challeng- ing particularly in high dimension, and we address this difficulty by transforming the OT cost into a simpler form amenable to efficient numerical implementation. Owing to the symmetric treatment of the observed data and the latent representation, the local structure of the data can be implicitly preserved in the latent space. However, we found that with the symmetrie treatment only the per- formance of the generative model may not be satisfactory. To improve the generative model we additionally include a reconstruction loss into the objective, which is shown to significantly benefit the quality of the generation and reconstruction.
Our contributions can be summarized as follows. Firstly, we propose a new family of generative au- toencoders, called Symmetric Wasserstein Autoencoders (SWAEs). Secondly, we adopt a learnable latent prior, parameterized as mixture of the conditional priors given the learnable pseudo-inputs, which prevents SWAES from over-regularizing the latent variables. Thirdly, we empirically perform an ablation study of SWAEs in terms of the KNN classification, denoising, reconstruction, and sam- ple generation. Finally, we empirically verify, using benchmark tasks, the superior performance of SWAEs over several state-of-the afrt generative autoencoders.
2
SYMMETRIC WASSERSTEIN AUTOENCODERS
In this section we introduce new family of generative autoencoders, called Symmetric Wasserstein Autoencoders (SWAEs).
2.1 OT FORMULATION
Denote the random vector the encoder as _ (Xe, 2.)) X'x Z, which contains both the observed data X. E X' and the latent representation z, E Z. We call the distribution p(e) ~ p(xe) )p(z, (xe) the encoding distribution, where p(x.) represents the data distribution and p(ze xx)) characterizes an inference model. Similarly, denote the random vector the decoder as d a (x,, E X x Z, which consists of both the latent prior E Z and the generated data x, E X. We call the distribu- tion p(d) ~ p(z.) )p( xa)za) the decoding distribution, where p(za) represents the prior distribution and p(xa Za) characterizes generative model. The objective of VAEs is equivalent to minimizing the (asymmetric) KL divergence between the encoding distribution p(e) and the decoding distri- bution p(d) (see Appendix I. To address the limitation in VAEs, first we propose treat the data and the latent representation symmetrically instead of asymmetrically by minimizing the p- th Wasserstein distance between p(c) and p(d) leveraging Optimal Transport (OT) (Villani 2008 Peyré et al. 2019).
OT provides a framework for comparing two distributions in a Lagrangian framework, which seeks the minimum cost for transporting one distribution to another. We focus on the primal problem of OT, and Kantorovich's 's formulation (Peyré et al. 2019) is given by:
inf (P(e))) dep
B(a.d)-r c(e,d).
(1)
where P(e p (e), d ~ p(d)), called the coupling between e and d, denotes the set of the joint
distributions of and d with the marginals p(e) and p(d), respectively, and c(c, d) (X' Z) x (X + [0, +oo denotes the cost function. When ((Y x (X' 2). is metric space and the cost function c(e, d) (e, d) for 2 1, W,. the p-th root of W. is defined as the p-th Wasserstein distance. In particular, can be proved that the p-th Wasserstein distance metric hence symmetric, and metrizes the weak convergence (see, e.g., (Santambrogio 2015)).
Optimization of equation iss computationally prohibitive especially in high dimension (Peyré et al.) 2019). To provide an efficient solution, we restrict to the deterministic encoder and decoder. Specif- ically, at the encoder we have the latent representation Z, = E(x.) with the function E E:-' + Z, and at the decoder we have the generated data xa D(za) with the function D 2 + X. It turns out that with the deterministic condition instead of searching for an optimal coupling in high dimension, can find proper conditional distribution p(za| (x.) with the marginal y p(z.).
2

---

Under review as a conference paper at ICLR 2021
Data index
Data Index
0.0
1
60
-100
-1
-20 -2-
-20
(a) SWAE
(b) VAE
Figure I: Latent representations of 100 GMM samples (mode 5 and dimension 10) with dim-z = 2. The indexes of these latent representations are sorted based on the distance to target sample in the data space, i.e., Index 0 is associated with the target sample and Index 100 is associated with the furthest sample the target in the data space. With SWAE (left) data samples that are close in the data space are also close in the latent space, while VAE (right) cannot preserve such correspondence.
Theorem Given the deterministic encoder :
> Z and the deterministic decoder D : 2 X, the Of problem in equation can be transformed the following:
Welpls)) = inf
ce.d))
p(aa(x.)
where the observed data follows the distribution p(xc) and the prior follows the distribution p(z.)
The proof of Theorem extends that of Theorem I in (Tolstikhin et al., 2018), and provided in Appendix A.2 If X' x 2 is the Euclidean space endowed with the Lp norm, then the expression in equation equals the following:
p((p(c).p(d)) = inf p(aa|xx.)
||xe - D(za)lB E(xx)) - zalle,
(3)
where in the objective we call the first term the x x-loss and the second term the z-loss. With the above transformation, we decompose the loss in the joint space into the losses both the data and the latent spaces. Such decomposition is crucial and allows us to treat the data and the latent representation symmetrically.
The x-loss, i.e., ||x - D(za) Ipp. represents the discrepancy in the data space, and can be interpreted from two different perspectives. Firstly, since D(za) represents the generated data, the x-loss es- sentially minimizes the dissimilarity between the observed data and the generated data. Secondly, the x-loss is closely related to the objective of Denoising Autoencoders (DAs) |Vincent al.) 2008 2010). In particular, DAs aim to minimize the discrepancy between the observed data and a partially destroyed version of the observed data. The corrupted data can be obtained by means of stochastic mapping from the original data (e.g., via adding noises). By contrast, the x-loss can be explained in the same way with the generated data being interpreted the corrupted data. This is because the prior Zu in D(za) is sampled from the conditional distribution p(za(xe), which depends on the observed data xe- Consequently, the generated data D(za), obtained by feeding zg the decoder, is stochastically related to the observed data x.. With this insight, the same as the objective of DAs, the x-loss can lead to the denoising effect.
The zlloss. i.e., [xxx) - 2e represents the discrepancy in the latent space. The whole objective in equation 3 hence simultaneously minimizes the discrepancy in the data and the latent spaces. Observe that in equation 3 Exx.) is the latent representation of x, at the encoder, while z, can be thought of as the latent representation of D(za) at the decoder. With such connection, the optimiza- tion of equation 3 can preserve the local data structure in the latent space. More specifically, since x. and D(z.) are stochastically dependent, roughly speaking, if two data samples are close to each other in the data space, their corresponding latent representations also expected to be close. This due to the symmetric treatment of the data and the latent representation. In Figure I we illustrate this effect and compare SWAE with VAE.
3

---

Under review as a conference paper at ICLR 2021
Reconstructed:
Input:
Xe
Latent:
E(x))
(
Encoder
Decoder
Prior: Zd
Generated: D(za)
Find the closest pseudo-input
Conditional prior p,(*u(n)
Sampling
Sampling
Figure 2: Network architecture of SWAEs. To generate new data latent samples are first drawn from the marginal prior p(za) based on the conditional priors p(z ug). and then fed the decoder.
Comparison with WAEs (Tosstikhin et al. 2018) The objective in equation minimizes the OT cost between the joint distributions the data and the latent, i.e., (p(p)e), p(d)). while the objective of WAEs (Tolstikhin et al. 2018) minimizes the OT cost between the marginal distributions of the data, i.e., (x(eexe), ((xx)) where p(xa) is the marginal data distribution induced by the decoding distribution p(dl (d). The problem of WAEs is first formulated as an optimization with the constraint p(z.) - p(za), where p(z.) is the marginal distribution induced by the encoding distribution p(e), and then relaxed by adding regularizer. With the deterministic decoder, the final optimization problem WAEs as follows:
inf Bp(x.) Pp... c(Xe, D(ze)) XD(p(za), p(a, (x,)
divergence measure. Comparing equation that both methods decompose the loss into the losses in the data and the latent spaces. Differently, in equation 4 the first term reflects the reconstruction loss in the data space and the second term represents the distribution-basee dissimilarity in the latent space; while in equation 3 the x-loss is closely related to the denoising and the generation quality and the z-loss measures the sample-based dissimilarity. Moreover, equation is optimized over the posterior ((2)(xe) with a fixed prior p( (Zz), while equation is optimized over the conditional prior p(za(x,) with potentially learnable prior.
where D(,) denotes
equation 3 can
2.2 IMPROVEMENT OF LATENT REPRESENTATION
The objective in equation 3 only seeks to match the encoding and the decoding distributions. Besides the encoder and the decoder structures, there is no explicit constraint on the correlation between the data and the latent representation within each joint distribution. Lacking of such constraint typically results in a low quality of reconstruction (Dumoulin et al.) 2017 Li et al.) (017a). Therefore, we incorporate a reconstruction- based loss into the objective associated with a controllable coefficient. Additionally, since the dimension of the latent space is usually much smaller than that of the data space, we associate weighting parameter to balance these two types of losses. Overall, the objective function can be represented as follows:
inf Ep(x..) Ex(es p(Ra(x.)
B||x DE + + (1- D(2,) + E().) - zalip (5)
where ||xe D (Zee)) 1 denotes the reconstruction loss, and 3(0 < 3 < 1) and a(a > 0) are the weighting parameters. The weighting parameter controls the trade-off between the x-loss and the reconstruction loss, and a smaller value of / generally leads to better reconstruction. To achieve better trade-off between the generation and reconstruction needs to carefully chosen. We will perform an ablation study of SWAEs and show the importance of including the reconstruction loss into the objective for the generative model in Section
2.3 ALGORITHM
Similar to many VAE-based generative models, assume that the encoder, the decoder, and the conditional prior are parameterized by deep neural networks. Unlike the canonical VAEs, where the prior distribution is simple and given in advance, the proposed method adopts learnable prior. The
4

---

Under review as a conference paper at ICLR 2021
benefits of a learnable prior, e.g., avoiding over-regularization and hence improving the quality of the latent representation, have been revealed in several recent works (Hoffman & Johnson 2016 Tom- czak & Welling 2018 Atanov et al. 2019 Klushyn et al. 2019). Obviously, the conditional pror is related to the marginal prior via 1xx. P x.) )= = p(za). This indicates a way to design the prior as a mixture of the conditional distributions, i.e., p" pz(z) N > p(z. Xe.n), where Xe.1,'' Xe,N are the training samples. To avoid over-fitting, similar to (Tomezak & Welling 2018), we replace the training samples with learnable pseudo-inputs and parameterize the prior distribution p(za) as py = t E P-y (Zu))u))), where 7 denotes the parameters of the conditional prior network, u; E X' are the learnable pseudo-inputs, and K is the number of the pseudo-inputs We empha- size that the conditional prior p(za(xe) (or approximated p(zalu:)) is used to obtain the marginal prior p(za); while the posterior p(z,]x.) is used for inference. In experiment, we parameterize the conditional prior as a Gaussian distribution.
We call the proposed generative model Symmetric Wasserstein Autoencoders (SWAFs) as we treat the observed data and the latent representation symmetrically. We summarize the training algorithm in Algorithm D and show the network architecture in Figure 2 function c(,) as the squared L2 norm.
As example, we define the cost
Algorithm 1: Symmetric Wasserstein Autoencoders (SWABs)
Require: The number of the pseudo-inputs K. The weighting parameters / and a. Initialize the parameters Φ, 0, and y of the encoder network, the decoder network, and the conditional prior network, respectively.
while (Ø, 0, 7: {f not converged do
1. Sample {xe,1.-.. .x.N) from the training dataset.
2. Find the closest pseudo-input (1(*) of each training sample from the set {ui.... ,
3. Sample Z,,, from the conditional prior P., (za) (n))) for n ,,.,,,NN
Update (Ø) 0, {u&}) by descending the cost function
N >> 8l1xe. (Z )1 + (1 B)|lx..n (((E(x))))) + all al)E(xenn) Zd.
Since we use the greudo-inputs instead of the training samples in the conditional prior, given each training sample we need to find the closest pseudo-input in Step 2. To measure the similarity, we can use, e.g., the T.2 norm or the cosine similarity. Since the dimension of the latent space is usually much smaller than that of the data space, to reduce the searching time we can alternatively perform Step 2 in the latent space sentation of E(xe.n) from the set { {E(u), ..... E(uk)} so as to obtain the corresponding closest innut. From the experiment found that such approximation results in little performance degradation, and we attribute it to the preservation of the local structure as explained before.
approximation. Specifically, we can find the closest latent repre-
3 EXPERIMENTAL RESULTS
In this section, we compare the performance of the proposed SWAE with several contemporary generative autoencoders, namely VAE (Kingma & Welling 2014). WAE-GAN (Tolstikhin et al. 2018), WAE-MMID (Tolstikhin et al.} 2018), VampPrior (Toeeza & Welling 2018). and MIM (Livne al. 2019), using four benchmark datasets: MNIST, MNSSS, Coil20, and CIFAR10 with a subset of classes (denoted as CIFARIO-sub).
3.1 EXPERIMENTAL SETUP
The design of neural network architectures is orthogonal to that of the algorithm objective, and can greatly affect the algorithm performance (Vahdat & Kautz, (020). Since MIM has the same network architecture as that of VampPrior, for fair comparison we also build SWAE as well as VAE based on the VampPrior network architecture. In particular, VampPrior adopts the hierarchical latent structure with the convolutional layers (i.e., convHVAE (L = 2)). where the gating mechanism is non-lincar,.y. The building block of the network structure of VAE and
utilized an element-wing SWAE is the same as that of VampPior except that the latent structure is non- non-hirarrchcal. Different from SWAE, the prior of VampPrior and MIM is designed as mixture of the posteriors (instead
5

---

Under review as a conference paper at ICLR 2021
Table Classification accuracy of 5-NN (averaged over 5 trials). The standard deviation is generally less than 0.01 and is omitted in the table.
Dataset
dim-z
SWAE
SWAE
SWAE
VAF
WAF-GAN
WAE-MMD
VampPrior
MIM
(B
0.5)
(3 0)
0.96
0.97
0.97
0.96
0.87
0.97
0.97
0.97
MNIST
40
0.97
0.97
0.97
0.80
0.68
0.96
0.93
0.97
80
0.97
0.97
0.97
0.60
0.90
0.94
0.86
0.96
8
0.82
0.81
0.83
0.80
0.71
0.80
0.81
0.82
Fashion-
40
0.84
0.83
0.84
0.54
0.62
0.82
0.81
0.82
MNIST
80
0.84
0.83
0.83
0.37
0.54
0.76
0.74
0.81
8
0.95
0.97
0.97
0.95
0.98
0.78
0.91
0.89
Coil20
40
0.97
0.98
0.97
0.90
0.99
0.99
0.96
0.96
80
0.97
0.98
0.98
0.97
0.98
0.98
0.96
0.98
80
0.69
0.67
0.65
0.67
0.61
0.68
0.68
0.66
CIFARIO-
256
0.70
0,66
0.61
0.62
0.61
0.68
0.65
0.65
sub
512
0.70
0.66
0.62
0.55
0.60
0.68
0.64
0.66
of a mixture of the conditional priors as in SWAE) conditioned the learnable pseudo-inputs. The pseudo-inputs in SWAE, VampPrior, and MIM are initialized with the training samples. For VampPrior and MIM, the number of the pseudo-inputs K is carefully chosen via the validation set. Unlike these two algorithms, for SWAE we found that increasing K improves the algorithm performance. The setup of K for SWAE. VampPrior, and MIM on all datasets Appendix A.3 For SWAF, we set the weighting parameter a to all cases; Step 2 we use the 1.2 norm the similarity measure in the data space. WAE-GAN and WAE-MMD based models, where the divergence measure in the latent space is based on GAN and the maximum mean discrepancy (MMD), respectively. The network structure of WAB-GAN and WAE-MMD is the same as that used in (Tolstikhin et al., 2018). The prior of VAE, WAE-GAN, and WAE-MMD is set as isotropic Gaussian. A detailed description of the datasets, the applied network architectures, and the training parameters can be found in Appendix A.3
be found in
the WAE-
3.2 LATENT REPRESENTATION
The latent representation is expected to capture salient features of the observed data and be useful for the downstream applications. The considered datasets are all associated with the labels. In the experiment we use the latent representation for the K-Nearest Neighbor (KNN) classification and compare the classification accuracy of 5-NN in Table I where dim-z z denotes the dimension of the latent space. The results of 3-NN and 10-NN are similar to those of 5-NN and thus omitted. We found that the classification results of all algorithms on CIFAR10 10 are unsatisfactory based on the current networks (accuracy was around 0.3 0.4; this may due the limited expressive power of the shallow network architectures used), so instead we create a subset of CIFARIO (CIFARIO-sub) which contains classes: bird, cat, and ship.
Since the prior of VAE, WAE-GAN, and WAE-MMD is an isotropic Gaussian, setting dim-z greater than the intrinsic dimensionality of the observed data would force p(z.) be in manifold in the latent space (Tolstikhin et al. 2018). This makes it impossible to match the marginal p(z.) with the prior p(za) and thus leads to unsatisfactory latent representation. Such concern can be verified particularly on Fashion-MNIST where the classification accuracy of VAE and WAE-GAN drops dramatically when dim-z is increased. For SWAE, we consider two cases: / (i.e., without the reconstruction loss) and 0.5. The classification accuracy of SWAE (8 1) is comparable to SWAE (B = 0.5) and is generally superior for different values of dim-z to the benchmarks.
To further show the structure of the latent representation, we project the latent representation to 2D using t-SNE (Maaten & Hinton, 2008) as the visualization tool. As an example, we show the projection of the latent representation structure of the observed data in the latent space and lead to tight clusters, which is consistent to our expectation explained in Section 2.1
MNIST in Figure We can that SWAEs keep the local

---

Under review as a conference paper at ICLR 2021
(a) SWAE (3 1)
(b) SWAE (3 0.5)
(c) VAE
(d) WAE-GAN
(e) WAE-MMD
(1) VampPrior
MIM
Figure 3: Projection of the latent representation to 2D via t-SNE MNIST. dim-z 80 for all methods.
Table 2: Fréchet Inception Distance (FID) on generated images (smaller is better).
Dataset
dim-z
SWAE
SWAE
SWAE
SWAR
VAE WAE-GAN WAE-MMD VampPrior
MIM
(5
(B
0.5)
(3
(5r)
MNIST
50
40
24
21
24
17
24
74
Fashion-MNIST
65
57
48
47
60
41
100
51
83
Coil20
80
97
89
102
89
278
278
320
97
113
CIFARIO-sub
512
105
44
183
44
242
114
341
68
59
3.3 GENERATION AND RECONSTRUCTION
To generate data, latent samples first drawn from the marginal prior distribution p (Zza) based on the conditional priors p(zz ug), and then fed to the decoder. We put the generated images of all methods in Appendix A.4 and show the Fréchet Inception Distance (FID) (Heusel et al., 2017), which is commonly used for evaluating the quality of generated images, in Table 2 For SWAEs, we observe that the reconstruction loss term is crucial for improving the generation quality as SWAE (3 = 1) generally cannot lead to the lowest FID. On MNIST and Fashion-MNIST, the FID of the best SWAE (indicated as 3*) is slightly higher than that of WAE-GAN, but lower than all the other benchmarks. The visual difference between SWAE (5*) and WAE-GAN on MNIST and MesshonnMNISTT is however negligible. In Section 2.T we compare the formulation of SWAES (3 = 1) with WAF. In particular, the objective of WAE includes iistribution-basee dissimilarity the latent space while the z-loss SWAEs measures the sample-based dissimilarity. On Coil20 and CIFARI0-sule SWAE (8*) achieves the lowest FID and generates new images that are visually much better than those generated by the benchmarks.
In Table we compare the reconstruction loss, defined as ||x, ((2.) the four datasets. As expected, icreasing the value of dim-z can reduce the reconstruction loss but the reduction becomes marginal when dim-z is large enough. Additionally, since smaller value of 3 leads more em- phasis on the reconstruction based loss the quality of reconstruction is generally better. We observe that SWAE (3 0.5) results in the lowest reconstruction loss in all cases. The reconstructed images of all methods are provided in Appendix AA.4 for reference. Without including the reconstruction loss into the objective, the reconstruction quality of SWAE (3 1) can be unsatisfactory (e.g., on CIFAR 10-sub).
3.4 DENOISING EFFECT WITH SWAE (8 = 1)
As discussed in Section 2.1 the x-loss has close relationship to the objective of Denoising Autoen- coders (DAs). After training, we feed the noisy images, which are obtained by adding the Gaussian random samples with mean zero and standard deviation 0.3 the clean test samples, to the encoder. In Figure 4 as an example, we show the reconstructed images on Pashion-MNIST Since the recon- struction loss is highly related to the dimension of the latent space, for fair comparison set dim-z

---

Under review as a conference paper at ICLR 2021
Table 3: Reconstruction loss (averaged over 5 trials).
Datasct
dim-z
SWAE
SWAE
VAE
WAE-GAN
WAE-MMD
VampPrior
MIM
(3 = 1)
(3 = 0.5)
30.11 = 0.14
23.20 = 0.04
24.34 = 0.07
26.86
0.37
24.76 0.31
24.05
24.04
MNIST
40
26.29 = 0.13
6,93 # 0.05
16.06 18.40
0.15
13.78
17.92
0.33
0.09
10.78 0.11
0.05
17.29 0.04
0.20
74.74.0000
71.03
72.36
18.17
7..30
0.02
7.034
Fashion-
40
73.39 0.08
57.90 0.25
69.85
0.04
74.84 = 0.23
75.86 0.41
68.67
MNIST
80
73.35 + 0.08
44.30 * 0.71
69.90 0.08
70.74 + 1.16
71.28 + 3.80
68.54
0.13
7.07 + 0.64
5.69 * 0.51
8.14 0.36
21.20 15.30
8,17
1.02
+
Coil20
5.52 0.40
4.27 0.66
5.67
5.82
8.07
6.31 0.62
0.77
3,56
4.33 0.40
5.62 1.26
5.83
5.67 0.33
CIFARIO-sub
512
50.8243.78
6,5020.08
9.41=0.27
13.3711622
13,09E192
12,06 #0.91
10.37
(a) Noisy real images
SWAE C 1)
SWAE (3 0.5)
VAE
(c) WAE-GAN
ff) WAE-MMD
g
VampPrior
(h) MIM
Figure 4: Denoising effect: reconstructed images Fashion-MNIST. dim-z 80 for all methods.
to S0 for all methods. We observe that only SWAE (3 = 1) can recover clean images. This observa- tion confirms the denoising effect induced by the x-loss, and thus the resultant latent representation is robust to partial destruction of the observed data.
RELATED WORK
The objective of VAEs uses the asymmetric KI. divergence between the encoding and the decoding distributions (see Appendix [.I]. To improve VAEs (Livne et al. 2019 Chen et al., 2018 Pu et al.
2017b) propose symmetric divergence measures instead of the asymmetric KL divergence in VAE- based generative models. For example, MIM (Livne et al. 2019) adopts the Jensen-Shannon (JS) divergence between the encoding and the decoding distributions together with a regularizer maxi- mizing the mutual information between the data and the latent representation. Due to the difficulty of estimating the mutual information and the unavailability of the data distribution, an upper bound of the desired loss is proposed. AS-VAE (Pu et al.) (017b) and the following work (Chen et al.) 2018) propose a symmetric form of the KL divergence optimized with adversarial training. These methods typically involve a difficult objective either depending on (unstable) adversarial training or containing the mutual information that requires further approximation. In contrast, the proposed SWAEs yield a simple expression of objective and do not involve adversarial training.
Compared to VAES, GANs lack an efficient inference model thus are incapable of providing the corresponding latent representation given the observed data. To bridge the gap between VAEs and GANs, recent works attempt to integrate an inference mechanism into GANs by symmetrically treat- ing the observed data and the latent representation, i.e., the discriminator is trained to discriminate the joint samples in both the data and the latent spaces. In particular, the JS divergence between the encoding and the decoding distributions is deployed in ALI (Dumoulin et al.) 2017) and BiGANs (Donahue al. 2017). To address the non-identifiability issue in ALI (e.g., unfaithful reconstruc- tion), later ALICE (Li et al. 2017a) proposes to regularize ALI using conditional entropy.
Generative modelling is closely related minimizing dissimilarity measure between two distribu- tions. As opposed to many other commonly adopted dissimilarity measures, e.g., the and the KL divergences, the Wasserstein distances that arise from the OT problem provide weaker distance
8

---

Under review as a conference paper at ICLR 2021
between probability distributions (see (Santambrogio 2015 Peyré al. 2019 Kolouri et al. 2017 for more background on OT). This is crucial as in many applications the observed data are essen- tially supported on a low dimensional manifold. In such cases, common dissimilarity measures may fail to provide a useful gradient for training. Consequently, the Wasserstein distances have received
surge of attention for learning generative models (Arjovsky et al. [017 Balaji et al. 2019: Sanjabi et al. 2018 Kolouri et al., 2019 Patrini et al. 2019 Tolstikhin et al., 20188 Deshpande et al. 2019 Nguyen et al. 2220). Particularly, the VAE-based models (Tolstikhin
al., 2018 Kolouri al,
2019 Patrini et al. 2019) are all based on minimizing the OT cost of the marginal distributions in the data space with the difference of how to measure the divergence in the latent space: (Tolstikhin) al. 2018) proposes the GAN-based and the MMD-based divergences, (Kolouri et al. 2019) adopts Siced-Wasserstein distance, and (Patrini al. 2019) exploits the Sinkhorn divergence. Unlike these works, our proposed SWAFS directly minimize the OT cost of the joint distributions of the observed data and the latent representation with the inclusion of reconstruction loss for further improving the generative model.
CONCLUSION AND FUTURE WORK
We contributed novel family of generative autoencoders, termed Symmetric Wasserstein Autoen- coders (SWAEs) under the framework of OT, We proposed to symmetrically match the encoding and the decoding distributions with the inclusion of a reconstruction loss for further improving the generative model. We conducted empirical studies on benchmark tasks to confirm the superior per- formance of SWAES over state-of-the-ie generative autoencoders.
We believe that symmetrically aligning the encoding and the decoding distributions with proper regularizer crucial improving the performance of generative models. To further enhance the performance of SWAEs, it is worthwhile to exploit other methods for the prior design, e.g., the flow- based approaches (Rezende & Mohamed 2015 Dinh et al., 2014: 2016), and other forms of the reconstruction loss, e.g., the cross entropy.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Alexander A Alemi, Ben Poole, lan Fischer, Joshua V Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken ELBO. arXiv preprint 1711 1.00664, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICMI., 2017.
Andrei Atanov, Arsenii Ashukha, Kirill Struminsky, Dmitry Vetrov, and Max Welling. The deep weight prior. ICLR, 2019.
Yogesh Balaji, Hamed Hassani, Rama Chellappa, and Soheil Feizi. Entropic GANs meet VAEs: A statistical approach compute sample likelihoods in GANs. In JCML, 2019.
Liqun Chen, Shuyang Dai, Yunchen Pu, Erjin Zhou, Chunyuan Li, Qinliang Su, Changyou Chen, and Lawrence Carin. Symmetric variational autoencoder and connections to adversarial learning. In AISTATS, 2018.
Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced wasserstein distance and its use for GANs. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 10648-10656, 2019.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516, 2014.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. arXiv preprint a605v 08803, 2016.
Jeff Donahue, Philipp Krähenbühl, and Trevor Darrell. Adversarial feature learning. In ICLR, 2017.
Vincent Dumoulin, Ishmael Belghazi, Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Ar- jovsky, and Aaron Courville. Adversarially learned inference. In ICLR, 2017.
Babak Esmaeili, Hao Wu, Sarthak Jain, Alican Bozkurt, Narayanaswamy Siddharth, Brooks Paige, Dana H Brooks, Jennifer Dy, and Jan- Billeilee van de Meent. Structured disentangled representa- tions. In AISTATS, 2019.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. NIPS, 2014.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to local Nash equilibrium. In Advances in neural information processing systems, 2017.
Matthew D Hoffman and Matthew Johnson. ELBO surgery: yet another way carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS, 2016.
Zhiting Hu, Zichao Yang, Xiaodan Liang, Ruslan Salakhutdinov, and Eric P. Xing. Toward con- trolled generation of text. In ICML, 2017.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014.
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning hierarchical priors in VAES. In NewrlPS, 2019.
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K. Rohde. Optimal mass transport: Signal processing and machine-learning applications. IEEE signal processing magazine, 34(4):43-59 2017.
Soheil Kolouri, Phillip E Pope, Charles Martin, and Gustavo K Rohde. Sliced-Wasserstein Wasserstein auto- encoders. In ICLR, 2019.
10

---

Under review as a conference paper at ICLR 2021
Chunyuan Li, Hao Liu, Changyou Chen, Yuchen Pu, Liqun Chen, Ricardo Henao, and Lawrence Carin. Alice: Towards understanding adversarial learning for joint distribution matching. In NIPS, 2017a.
Jerry Li, Aleksander Madry, John Peebles, and Ludwig Schmidt. Towards understanding the dy- namics of generative adversarial networks. In ICLR, 2017b.
Micha Livne, Kevin Swersky, and David Fleet. MIM: Mutual information machine. arXiv preprint 1:X.V:1910.03115 2019.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research, 9Novv) 2008.
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced-wasserstein and applica- tions to generative modeling. arXiv preprint arXiv:2002.07367, 2020.
Giorgio Patrini, Rianne van den Berg. Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. UAI, 2019.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 11(5-6):355-077 2019.
Yuchen Pu, Zhe Gan, Ricardo Henao, Chunyuan Li, Shaobo Han, and Lawrence Carin. VAE learn- ing via Stein variational gradient descent. NIPS, 2017a
Yuchen Pu, Weiyao Wang, Ricardo Henao, Liqun Chen, Zhe Gan, Chunyuan Li, and Lawrence Carin. Adversarial symmetric variational autoencoder. In NIPS, 2017b.
Razavi, Aäron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ-VAB-2. In NeurIPS, 2019.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In JCML, 2014.
Maziar Sanjabi. Jimmy Ba, Meisam Razaviyayn, and Jason D Lee. On the convergence and robust- ness of training GANs with regularized optimal transport. In NeurlPS, 2018.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Ilya Sutskever, Oriol Vinyals, and Quoe V. Le. Sequence to sequence learning with neural networks. In NIPS, 2014.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. In ICLR, 2018.
Jakub Tomczak and Max Welling. VAE with VampPrior. In AISTATS, 2018.
Benigno Uria, lain Murray, and Hugo Larochelle. RNADE: The real-valued neural autoregressive density-estimator. In NIPS, 2013.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. arXiv preprint arXiv:2007.03880 2020.
Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In ICML, 2016.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoime Manzagol. Extracting and composing robust features with denoising autoencoders. In ICMI, 2008
11

---

Under review as a conference paper at ICLR 2021
Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in a deep network with local denoising criterion. Journal of machine learning research, LI(Dec):3371-340) 2010.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Info VAE: Balancing learning and inference in variational autoencoders. In AAAI, 2019.
A
APPENDIX
A.1
OBJECTIVE OF VAES
The objective of VAES is to maximize a tractable variational lower bound on the data log-likelihood, called the Evidence Lower Bound (ELBO):
Bx(x.) (Pp(.. p( - DEL. (p(x.. (xx.(
can be also shown that the objective of VAEs is equivalent to minimizing the KL divergence (or maximizing the negative KL divergence) between the encoding and the decoding distributions
(Livne et al. 2019 Esmaeili et al., 2019 Pu et al. 2017b Chen et al. 2018):
p(xa. Zd) px (lg p(x))
(p 2e) [[(xx 22)) Exp(.. a.)
log
plzel x
The right hand side of equation 7 only different from equation i terms of constant, which is the entropy of the observed data.
A.2 PROOF OF THEOREM
The proof extends that of Theorem 1 in (Tolstikhin et al. 2018). In particular, Tolstikhin al. 2018) aims to minimize the OT cost of the marginal distributions x(x)) and p(xa). and the proof there is based the joint probability of three random variables: the observed data, the generated data, and the latent representation. In contrast, we propose to minimize the OT cost the joint distributions of the observed data and the latent representation induced by the encoder and the decoder. As a result our proof is based on the joint distribution of four random variables (xe, Ze, Xd. za) E X x Z x X Z. We assume that the joint distribution pl X Ze Xd, Za) satisfies the following three conditions:
e (Xe, be) ~ ((ze)))(e))).....
2. d
(xx, za) ~ (xzx)p(xx zz); and
3. X, IL x. Zd (conditional independence).
The first two conditions specify the encoder and the decoder respectively, and the last condition indicates that given the latent prior the generated data and the observed data are independent.
Denote the of the above joint distributions as P(x,, Ze, Xx, Za). Obviously, we have P(xx,, Ze, Xd. za) C Peee
(e), ~ p(d)) due to the third condition. If the decoder is deter- ministic, p(xalza) is a Dirac distribution thus P(xe; Ze; Xa, Za) P(e ~ p (e), d p(d)). With this result, we can rewrite the objective of the underlying OT problem as follows:
Wc(p(e).p(d))
inf DEP(x. inf rEP(x. inf p(z,(xe), p(zalx..2.)
Eledyrr c(c,
E Exx, c(e,
(8)
E
c(edd)
inf Ppp(..) E (xpiaax))) c(e,d), p(a.(x))
(10)
where in equation
8 P(x Ze, Za) denotes the set of the joint distributions of (x,, Ze, za) induced
by P(x,, Ze. Xa, z and it holds due to the deterministic decoder, and equation 10 holds due to the deterministic encoder.
12

---

Under review as a conference paper at ICLR 2021
A.3 DATASETS AND NETWORK ARCHITECTURES
In this section, we briefly describe the datasets, the network architectures, and the hyperparameters that are used in our training algorithm.
MNIST: The dataset includes 70, 000 binarized images of numeric digits from 0 to 9, each the size 28 x 28. There are 7, 000 images per class. The training set contains 50, 000 images, the validation set contains 10, 000 images for choosing the best model based on the loss function, and the test set contains 10, 000 images.
Fashion-MNISS MNIST:
The dataset includes 70, 000 binarized images of fashion products in 10 classes. This dataset has the same image size and the split of the training, validation, and test sets as MNIST.
Coil20: The dataset includes gray-scale images of 20 objects, each image of the size 32 32. The training set contains 1040 images, the validation set contains 200 images for choosing the best model based on the loss function, and the test contains 200 images.
CIFARIO-sub: The CIFAR-10 dataset consists of 60, 000 32 32 colour images in 10 classes with 6, 000 images per class. There are 40, 000 training, 10, 000 validation, and 10, 000 test images. We randomly select three classes to form the CIFARIO-sub dataset, namely bird, car, and ship.
Celeba: The Celeba dataset is resized to 64 x 64 resolution. The training set contains 162, 770 images, the validation set contains 19, 867 images, and the contains 19, 962 images.
Network architecture of SWAE: The building block of the network structure of SWAE is based on VampPrior, called GatedConv2d. GatedConv2d contains two convolutional layers with the gating mechanism utilized as an clement-wise non-linearity. The parameters in the function GatedConv2d( represent the number of the input channels, the number of the output channels, kernel size, stride, and padding, respectively. The conditional prior network outputs the mean and the log- of Gaussian distribution, based which the latent prior is sampled.
The structure of the encoder network: GatedConv2d(1332.77. (ateeoo22) 6,,,,,,,,4, GatedCComn2ad(6446.3....111. followed by one fully-connected layer with no activation function.
The structure of the conditional prior network: The layers of GatedConv2d are the same as those in the encoder network, which are followed by two fully-connected layers. One produces the mean, and the other produces the log-variance with the activation function Hardtanh.
The structure of the decoder network: Two fully-connected layers with the 1... followed by convolutional layer
gating mechanism, followed by GatedConv2d(644.44 (.A.J...))CaaeeeCoo2226664. with the activation function Sigmoid.
The algorithm is trained by Adam with the learning rate 0.001, 31, = 0.9, and 83 0.999.
Setup of the number of the pseudo-inputs K: As suggested in Tomezak Welling 2018 Livne et al. 2019) we set the value of K in VampPrior and MIM on MNIST and Fashion-M to 500. We found K = 500 is also suitable for VampPrior and MIM on Coil20, CIFAR10-sub, and Celeba. Unlike VampPrior and MIM, for SWAE we found that increasing K improves the performance and we set K to 4000 on MNIST, Moshion-MNIST" CIFAR 10-sub, and Celeba. Coil20 is relatively small dataset and set K to 500 for SWAE, VampPrior, and MIM.
A.4 MORE EXPERIMENTAL RESULTS
In this section, we show more experimental results based on the comparison with the benchmarks.
13

---

Under review as a conference paper at ICLR 2021
(a) SWAE (3 = 1)
(b) SWAE (3 0.5)
(c) VAE
(d) WAE-GAN
(C) WAE-MMD
VampPrior
MIM
Figure 5: Latent representation MNIST; dim-z 2 for all methods. With more compressed latent representation, the classification accuracy generally decreases except VAE (5-NN accuracy with dim-z = 2: SWAE(3 1)(0.55), SWAE($ = 0.5)(..86) VAE(0.84), WANO-.,,,.8I., WAE- MMD(0.81), VampPrior(00.2)) and MIM(0.87)). Also, the quality of reconstruction and generation decreases when dim-z 2.
/
4
3
3
9
3
5
9
5
7 7
9
7
3
5
/
7
7
9
5
L
9
2
8
5
/
2
0
2
4
9
2
7
7
4
9
5
2
0
3
4
O
7
3
4
6
6
7
0
3
3
0
6
0
9
4
4
2
0
4
2
8
2 8
4
0
4
9
O
2
8
4
5
2
6
2
3
4
0
2
5
4
6
9
4
7
8
5
Y
6
3
4
8
(a) SWAE (3"
0.02)
(b) SWAE (8
1)
(c)
SWAE
(3
0)
7
3
3
3
8
C
5
9
7
2
7
5
8
6
3
7
2
0
9
8
4
6
6
2
2
4
5
8
6
0
3
7
?
7
5
5
7
9
J
/
7
2
6
7
4
1
4
7
8
5
6
3
9
0
1
5
4
G
8
0
6
2
6
9
3
2
6
2
2
0
5
0
1
J
1
5
3
2
5
2
7
2
7
2
0
3
2
2
2
(d)
VAE
(e)
WAE-GAN
(I)
WAE-MMD
/
6
3
9
6
2
/ 0
7
0
3
0
2
0
/
6
5
2
3
9
1
5
7
/
5
O
7
2
9
2
/
4
7
9
6
9
9
2
2 7
7
4
2
6
5
0
2
3
3
9
0
4
9
0
3 7
9
5
2
4
6
0
0
7 2
VampPrior
(h) MIM
Figure 6: Generated new samples MNIST. dim-z 8 for all methods.
14

---

Under review as a conference paper at ICLR 2021
SWAE (3" 0.05)
(b) SWAE (3
(c) SWAE (B
(d) VAE
WAE-GAN
(I) WAE-MMD
(g)
VampPrior
(h) MIM
Figure 7: Generated new samples on Fashion-NNIST. dim-z = 8 for all methods.
(a) SWAE (8" 0.5)
(b) SWAE (3
(c) SWAE (B 0)
(d) VAE
ee) WAE-GAN
(I) MMD
VampPrior
(h) MIM
Figure 8: Generated new samples on Coil20. dim-z = 80 for all methods.
15

---

Under review as a conference paper at ICLR 2021
(a) SWAE (8" 0.5)
(b) SWAE (3 = 1)
(c) SWAE (B 0)
(d) VAB
(c) WAE-GAN
f WAE-MMD
VampPrior
(h) MIM
Figure 9: Generated new samples on CIFARIO-sub. dim-z = 512 for all methods.
(a) SWAE (6* 0.5)
(b) SWAE (3 1)
(c) SWAE (8 0)
(d) VAE
(e) WAE-GAN
( WAE-MMD
VampPrior
(h) MIM
Figure 10: Generated new samples Celeba. dim-z 512 for all methods. VampPrior has the best generation quality visually. SWAE (8* = 0.5) is better than SWAE (8 1) and SWAE (3 0).
16

---

Under review as a conference paper at ICLR 2021
7
2
O
4 5
9 3
S
2
2
4
G
3 4
0
O
G
S
S
4
O
O
R
O
7
6
4
O
7
4
O
O
4
O
7
O
O
4
3
3
2
L
2
I
3
3
2
4
(a) Real images
(b) SWAE (8
1)
(c) SWAE (3 = 0.5)
(d) SWAE (8
7
7
7
5
O
7
5
4
O
9
7
O
7
D
7
7
4
2
2
I
3
7
4
4
(c) VAE
(f
WAE-GAN
WAE-MMD
(h) Vamp Prior
2
O
/
S
0
O
9
0
9
7 4
3
7
6
6
5
O
O
3
7
5
ii) MIM
Figure 11: Reconstructed images MNIST. dim-z = so for all methods. As expected, for SWAEs
a smaller 3 leads to a higher quality of reconstruction.
Real images
SWAE (
SWAE C 0.5)
(d) SWAE (8
(c) VAE
(I) WAE-GAN
WAE-MMD
(h) VampPrior
ii) MIM
Figure 12: Reconstructed images on Fashion-MNIST. dim-z = 80 for all methods.
17

---

Under review as a conference paper at ICLR 2021
(a) Real images
(b) SWAE (8
(c) SWAE (8 = 0.5)
(d) SWAE (B
(c) VAB
(f) WAE-GAN
WMD
(h) VampPrior
MIM
Figure 13: Reconstructed images Coil20. dim-z 80 for all methods. For SWAEs, the difference of the reconstruction error for different values of is insignificant, and the reconstructed images look visually the same.
Real images
(b) SWAE (B 1)
SWAE ( 0.5)
(d) SWAE (8
(c) VAE
(i) WAE-GAN
WAE-MMD
(h) VampPrior
MIM
Figure 14: Reconstructed images CIFARIO-sub. dim-z 512 for all methods. Excluding the reconstruction loss in the objective, the reconstruction of SWAE (,3 = is blurry.
18

---

Under review as a conference paper at ICLR 2021
(a) Real images
(b) SWAE (B 1)
(c) SWAE (8 0.5)
(d) SWAE (B
AE
BES
DO
&
O
ES
OR
(c) VAE
WAE-GAN
WAE-MMD
(h) VampPrior
ABAR
O
ii MIM
Figure 15: Reconstructed images on Celeba. dim-z 512 for all methods. The average reconstruc-
tion error over three seeds is follows: SWAE (B 1) 301.11 #4.85, SWAE (8 0.5): 37.78 +6.71, SWAE (B = O): 30.14 #0.18, VAE: 38.64 6445.39, WAE-GAN: 23.63+5.79, WAE-MMD: 28.64 VampPrior: 37.96+0.71, and MIM: 36.40+2.84. As expected, for SWAE smaller value 3 leads to a lower reconstruction loss.
7
0
9
9
0
6
9
0
5
9
7
3
7
4
5
O
7
4
O
3
2
7
2
4
(a) Noisy real images
(b) SWAE (B
(c) SWAE (8 0.5)
(d) VAE
0
4
Y
C
2
4
9 7 &
O
9
D
6
9
0
l
O
4
D
9
6
L
O
4
O
2
3
7
2.
Y
4
4
(e) WAE-GAN
(I WAE WAE-MMD
VampPrior
(h) MIM
Figure 16: Denoising effect: reconstructed images MNIST. dim-z 80 for all methods. SWAE
(6 = 1), WAE-GAN, and WAE-MMD can recover clean images. However, for WAE-GAN and WAE-MMD, we can still see some noisy dots around the digits.
19

---

Under review as a conference paper at ICLR 2021
Noisy real images
(b) SWAE (8
SWAE C8 0.5)
(d) VAE
(c) WAE-GAN
(f) MMD
VampPrior
(h) MIM
Figure 17: Denoising effect: reconstructed images on Coil20. dim-z so for all methods. Except WAE-GAN and WAE-MMD, the other methods can produce clear images.
20


Reviews:
[{"review_id": "S3LWKJ_t6XN", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 proposed symmetric Wasserstein autoencoders (SWAE), which is a new type of generative autoencoders within the framework of optimal transport (OT). This work is based on Wasserstein autoencoders (WAE), but leverage a symmetric distance between the encoding distribution and the decoding distribution to better preserve the local structure of the data in the latent space. A reconstruction loss based regularization is suggested for better performance. The paper is well written and easy to follow. My only concern is that the improvement over WAE seems to be marginal. Below are some minor issues that need to be addressed. \n\n1. Please provide the appendix.\n\n2. The data sets used in the experiment are kind of too simple. Can the author provide results for full CIFAR10? \n\n3. The FID of WAE-GAN and WAE-MMD in Table 2 does not look right, can the author provide implementation info? Did you use your own code?\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "QQyE01dU6n5", "reviewer": "AnonReviewer3", "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": "##########################################################################\n\nSummary:\n \nThis works proposes an new auto-encoder variant based on an Optimal Transport (OT) penalty.  While there are many such previous works of OT and auto-encoders, this work proposes a joint OT penalty on data and latent space. As the scalability of computing OT penalties in high dimensions is a concern, the authors address this by restricting to deterministic encoders and decoders in Theorem 1, an extension to joint distributions of Theorem 1 of Tolstikhin 2018. The resulting algorithm amounts to a loss involving L2 penalties for (1) the reconstruction loss (2) decoded latents (conditional on \"pseudo-inputs\") and real samples (3) encoded samples and the conditional latents. Next experimental results are shown on small-scale datasets (MNIST, Fashion-MNIST, Coil20, subest of CIFAR-10) and compared against the VAE, WAE-{GAN,MMD}, VampPrior, and MIM.\n\n##########################################################################\n\nReasons for score: \n\nOverall I vote for rejection. Given the numerous prior works (properly cited in the paper), the novelty of the proposed loss appears limited to me. Theorem 1 also appears to be limited novelty over the Theorem 1 of Tolstikhin 2018. Additionally, the empirical evaluation of the method is only limited to small-scale datasets. \n \n\n##########################################################################\n\nPros: \n* Easy to read\n* Reasonable selection of prior models to compare against\n\n##########################################################################\n\nCons: \n* Only small-scale datasets are evaluated. Thus the empirical advantage is unclear\n* Theorem 1 appears to be of limited novelty over Theorem 1 in Tolstikin 2018\n\n\n##########################################################################\n\nQuestions during rebuttal period: \n \nWhy is your theorem novel over Theorem 1 in Tolstikin 2018?\n\nWhat is the performance of your method on more complex datasets such as CelebA and LSUN Bedroom?\n\n#########################################################################\n\nPOST-REBUTTAL RESPONSE:\n\nThanks for clarification on Theorem 1 of this paper, i.e. that the novelty is in interpretation. I agree that the \"denoising\" between observed and generated data is an interesting idea.\n\nI read the author's additional experiments on CelebA. In Figure 10, VampPrior is still qualitatively superior to the author's best result $SWAE(\\beta^*=0.5)$. I have some skepticism over the reported results for WAE-{GAN/MMD}, which are much worse than the results in the original paper (Tolstikhin 2018). The authors appear to have used different encoder/decoder architectures, which complicates the comparison. Is WAE's decreased performance due to choice of architecture or algorithm? FID scores on CelebA would be also helpful.\n\nAll told, I raise my score, but still harbor some doubts over the empirical advantage of this work.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "N6HBl1ivLwF", "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": "In the paper, the authors propose a new family of generative auto-encoders, named Symmetric Wasserstein Autoencoders (SWAEs), based on replacing the KL divergence between the encoding and decoding distributions on the traditional VAE framework into the Wasserstein metric between these distributions. Then, they also carried out experiments to demonstrate the favorable performance of their auto-encoder over previous base-line autoencoders.\n\nI think the SWAEs is quite interesting but lacks novelty. Here are my comments with the paper:\n\n(1) The result of Theorem 1 is under the assumption that both the encoder and decoder are deterministic, which is quite restrictive. Can the authors provide some understandings when either one of them is random? The result of Theorem 1 is also quite similar to the main result in the Wasserstein autoencoder work. I think the authors should at least cite this point properly.\n\n(2) Simply replacing KL divergence between the encoding and decoding distributions by Wasserstein metric to have symmetric properties sounds not novel. The KL divergence has a good interpretation in terms of ELBO. What can we interpret the Wasserstein metric between the encoding and decoding distributions? Theorem 1 does not seem convincing to me.\n\n(3) In the experiments, the choice of $\\beta = 1/2$ seems to yield best results, i.e., we should balance the reconstruction loss and the discrepancy in the data space in the objective function (4). Can the authors provide some intuition behind that?\n\n(4) The related work with generative modeling based on optimal transport lacks several relevant recent works; see for examples [1], which achieves SOTA among sliced-based Wasserstein distances in deep generative models, or [2], which proposes max-sliced Wasserstein distance for deep generative models.\n\nReferences:\n\n[1] K. Nguyen, N. Ho, T. Pham, H. Bui. Distributional sliced-Wasserstein and applications to deep generative modeling. Arxiv preprint Arxiv: 2002.07367, 2020.\n\n[2] I. Deshpande, Y. Hu, R. Sun, A. Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, D. Forsyth, A. Schwing. Max-sliced Wasserstein distance and its use for GANs. CVPR, 2019.", "extra_scores": {}, "extra_text": {}}, {"review_id": "C9gaiZC5oj", "reviewer": "AnonReviewer2", "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 proposes to treat the encoding and the decoding pairs symmetrically as a solution to OT problems. SWAE minimizes $p(x_d, z_d)$ and $p(x_e, z_e)$ in a jointly manner and shows better latent representation learning and generation. Moreover, the symmetric treatment for encoding and decoding shows an advantage in data denoising. \n\n\n----------\n\nIt is interesting to minimize the distance between $p(x_d, z_d)$ and $p(x_e, z_e)$ with the OT formulation.  The usage of the deterministic encoder and decoder could solve the problem of $W(p_{x_e}, p_{e_d})$ minimization, while it is difficult for the latent code since the latent codes from the prior and the posterior are not paired. Here it is good to see the authors apply aggregated posterior methods to solve this. The usage of the closest pseudo-inputs in constructing the pair of $z_e$ and $z_d$ could be a good approximation for $W(p_{z_e}, p_{z_d})$.\n\nMy questions mainly lie in the experiment part:\n\nThe authors report SWAE with three configurations, where $\\beta=1$, $\\beta=0.5$ and $\\beta=0$ respectively. When $\\beta \\rightarrow 1$, SWAE is only minimizing the $d( p(x_d, z_d), p(x_e, z_e) ) $; when $\\beta \\rightarrow 0$, SWAE behaves like an autoencoder with regularization in latent space.  The classification results and the reconstruction results of the case $\\beta = 0$  are missing. I am curious how these two cases perform, are they perform like VAE and Vampprior VAE? \n\nThe authors visualize the latent space with t-sne where the dimensionality of $z$ is 80 before the dimension reduction. As far as I know, when the dimension of $z$ is large enough, the t-sne results of VAEs and WAEs are similar to figure 3a, 3b and 3g. The results shown in c,d,e and f seem degenerated.  Especially 3c, it does not look like what we have seen in previous papers. For example, figure 2 in Makhzani et al, 2015 (https://arxiv.org/pdf/1511.05644.pdf) shows that the latent representations produced by VAE have several modes.  I doubt this could be the problem of dimension reduction. It will be more convincing if the author could show the visualization of these models where $dim(z) = 2$ and without any dimension reduction. \n\nFor the reconstruction results (not the denoising reconstruction) shown in Figure 9 and 10, the images seem to be dynamically binarized. However, in the description of the experiment details, this part is not mentioned. Considering SWAE is not applying a Bernoulli decoder, the dynamic binarization is not necessary. Moreover, the reconstruction from SWAE ($\\beta=0.5$) and SWAE ($\\beta=0$) seem to be binarized, which is not reasonable.\n\nWhen $\\beta=1$, a good baseline for SWAE will be ALI (Dumoulin et al., 2017), since ALI is minimizing $JS(p(x_d, z_d)|| p(x_e, z_e))$. I am curious how SWAE ($\\beta=1$) outperforms/underperforms compared to ALI and some corresponding analysis from the authors will be appreciated.\n\nIn the random generation results, only SWAE ($\\beta^\\star$) is shown. The generation results by SWAE ($\\beta=1$) and SWAE ($\\beta=0$) are missing while the quantitative results are shown in table 2.\n\n\n---------------\n\nIn general, the idea of symmetrically treating encoding and decoding to solve with the Wasserstein distance is interesting and is worthy of study. However, some details in the experiments remain unclear; some experiments results and the corresponding analysis are missing. The authors might need to dig out more to support their method. \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors proposed a new variant of the Wasserstein autoencoder (WAE), which matches the joint distribution of data and the latent codes induced by the encoder and the joint distribution induced by the decoder in the framework of optimal transport. Because of matching the distributions that are not considered by existing autoencoders like WAEs or VAEs, I agree with the authors that the proposed method is novel to some degrees. 

However, the experimental part does not support the superiority of the method well. For example, some reviewers (including me) think the results of the baselines shown in Figures 8-10 are underestimated. According to my personal experience, the WAE should perform much better on CelebA than that shown in Figure 10. The experiments in Figures 11-17 provide more reasonable results, but the advantage of the proposed method is not convincing.

Here is my suggestion:
1) Because the proposed method can achieve flexible prior, besides randomly generating data, the authors can consider adding some experiments on conditional generation, i.e., generating data from a single modality of the learned prior. I believe the proposed method will be more convincing if it can show some advantages in the conditional generation task.
2) The runtime comparison for the method and the baselines in the training phase should be discussed. 
3) The short name "SWAE" is in conflict with an earlier work "Sliced Wasserstein Autoencoder", which is also called "SWAE".
Decision: Reject