Title: SLICED-WASSERSTEIN AUTO-ENCODERS

Abstract: In this paper we use the geometric properties of the optimal transport (OT) problem and the Wasserstein distances to define a prior distribution for the latent space of an auto-encoder. We introduce Sliced-Wasserstein Auto-Encoders (SWAE), that enable one to shape the distribution of the latent space into any samplable probability distribution without the need for training an adversarial network or having a likelihood function specified. In short, we regularize the auto-encoder loss with the sliced-Wasserstein distance between the distribution of the encoded training samples and a samplable prior distribution. We show that the proposed formulation has an efficient numerical solution that provides similar capabilities to Wasserstein Auto-Encoders (WAE) and Variational Auto-Encoders (VAE), while benefiting from an embarrassingly simple implementation. We provide extensive error analysis for our algorithm, and show its merits on three benchmark datasets.
Scalable generative models that capture the rich and often nonlinear distribution of high-dimensional data, (i.e., image, video, and audio), play a central role in various applications of machine learning, including transfer learning Isola et al.

:

Learning such generative models boils down to minimizing a dissimilarity measure between the data distribution and the output distribution of the generative model. To this end, and following the work of Arjovsky et al. (2017) and Bousquet et al. (2017) , we approach the problem of generative modeling from the optimal transport point of view. The optimal transport problem Villani (2008) ; Kolouri et al. (2017) provides a way to measure the distances between probability distributions by transporting (i.e., morphing) one distribution into another. Moreover, and as opposed to the common information theoretic dissimilarity measures (e.g., f -divergences), the p-Wasserstein dissimilarity measures that arise from the optimal transport problem: 1) are true distances, and 2) metrize a weak convergence of probability measures (at least on compact spaces). Wasserstein distances have recently attracted a lot of interest in the learning community Frogner et al. (2015) ; Gulrajani et al. (2017) ; Bousquet et al. (2017) ; Arjovsky et al. (2017) ; Kolouri et al. (2017) due to their exquisite geometric characteristics Santambrogio (2015) . See the supplementary material for an intuitive example showing the benefit of the Wasserstein distance over commonly used f -divergences.

In this paper, we introduce a new type of auto-encoders for generative modeling (Algorithm 1), which we call Sliced-Wasserstein auto-encoders (SWAE), that minimize the sliced-Wasserstein distance between the distribution of the encoded samples and a samplable prior distribution. Our work is most closely related to the recent work by Bousquet et al. (2017) and more specifically the follow-up work by Tolstikhin et al. (2017) . However, our approach avoids the need to perform adversarial training in the encoding space and is not restricted to closed-form distributions, while still benefiting from a Wasserstein-like distance measure in the latent space. Calculating the Wasserstein distance can be computationally expensive, but our approach permits a simple numerical solution to the problem. Finally, we note that there has been several concurrent papers, including the work by Deshpande et al. (2018) and Şimşekli et al. (2018) , that also looked into the application of sliced-Wasserstein distance in generative modeling. Regardless of the concurrent nature of these papers, our work remains novel and is distinguished from these methods. Deshpande et al. (2018) use the sliced-Wasserstein distance to match the distributions of high-dimensional reconstructed images, which require large number of slices, O(10 4 ), while in our method and due to the distribution matching in the latent space we only need O(10) slices. We also note that Deshpande et al. (2018) proposed to learn discriminative slices to mitigate the need for a very large number of random projections that is in essence similar to the adversarial training used in GANs, which contradicts with our goal of not using adversarial training. Şimşekli et al. (2018) , on the other hand, take an interesting but different approach of parameter-free generative modeling via sliced-Wasserstein flows.

1. NOTATION AND PRELIMINARIES:

Let X denote the compact domain of a manifold in Euclidean space and let x n ∈ X denote an individual input data point. Furthermore, let ρ X be a Borel probability measure defined on X. We define the probability density function p X (x) for input data x to be:

dρ X (x) = p X (x)dx

Let φ : X → Z denote a deterministic parametric mapping from the input space to a latent space Z (e.g., a neural network encoder). To obtain the density of the push forward of ρ X with respect to φ, i.e., ρ Z = φ * (ρ X ), we use Random Variable Transformation (RVT) Gillespie ( 1983)). In short, the probability density function of the encoded samples z can be expressed in terms of φ and p X by:

EQUATION

where δ denotes the Dirac distribution function. Similar to variational Auto-Encoders (VAEs) Kingma & Welling ( 2013) and the Wasserstein Auto-Encoders (WAE) Tolstikhin et al. (2017) , our main objective is to encode the input data points x ∈ X into latent codes z ∈ Z such that: 1) x can be recovered/approximated from z, and 2) the probability density function of the encoded samples, p Z , follows a prior distribution q Z . Let ψ : Z → X be the decoder that maps the latent codes back to the original space such that

EQUATION

where y denotes the decoded samples. It is straightforward to see that when ψ = φ -1 (i.e. ψ(φ(•)) = id(•)), the distribution of the decoder p Y and the input distribution p X are identical. Hence, in its most general form, the objective of such auto-encoders simplifies to learning φ and ψ, so that they minimize a dissimilarity measure between p Y and p X , and between p Z and q Z . In what follows, we briefly review the existing dissimilarity measures for these distributions.

1.1 MINIMIZING DISSIMILARITY BETWEEN p X AND p Y We first emphasize that the VAE often assumes stochastic encoders and decoders Kingma & Welling (2013), while we consider the case of only deterministic mappings. Although, we note that, similar to WAE, SWAE can also be formulated with stochastic encoders. Different measures have been used previously to compute the dissimilarity between p X and p Y . Most notably, Nowozin et al. (2016) showed that for the general family of f -divergences, D f (p X , p Y ), (including the KL-divergence, Jensen-Shannon, etc.), using the Fenchel conjugate of the convex function f and minimizing D f (p X , p Y ) leads to a min-max problem that is equivalent to the adversarial training widely used in the generative modeling literature Goodfellow et al. (2014) ; Makhzani et al. (2015) ; Mescheder et al. (2017) .

Others have utilized the rich mathematical foundation of the OT problem and Wasserstein distances Arjovsky et al. (2017) ; Gulrajani et al. (2017) ; Bousquet et al. (2017) ; Tolstikhin et al. (2017) to define a distance between p X and p Y . In Wasserstein-GAN, Arjovsky et al. (2017) utilized the Kantorovich-Rubinstein duality for the 1-Wasserstein distance, W 1 (p X , p Y ), and reformulated the problem as a min-max optimization that is solved through an adversarial training scheme.

Inspired by the work of Bousquet et al. (2017) and Tolstikhin et al. (2017) , it can be shown that (see supplementary material for a proof):

W c (p X , p Y ) ≤ W ‡ c (p X , p Y ) := E p X (c(x, ψ(φ(x)))) (3) = X c(x, ψ(φ(x)))p X (x)dx,

Furthermore, the r.h.s. of equation 3 supports a simple implementation where for i.i.d samples of the input distribution, {x n } N n=1 , the upper bound can be approximated as:

EQUATION

The r.h.s of equation 3 and equation 4 take advantage of the existence of pairs x n and y n = ψ(φ(x n )), which make f (•) = ψ(φ(•)) a transport map between p X and p Y (but not necessarily the optimal transport map). In this paper, we minimize W ‡ c (p X , p Y ) following equation 4 to minimize the discrepancy between p X and p Y . Next, we focus on the discrepancy measures between p Z and q Z . 1.2 MINIMIZING DISSIMILARITY BETWEEN p Z AND q Z If q Z is a known distribution with an explicit formulation (e.g. Normal distribution) the most straightforward approach for measuring the (dis)similarity between p Z and q Z is the log-likelihood of z = φ(x) with respect to q Z , formally:

sup φ X p X (x)log(q Z (φ(x)))dx

(5) maximizing the log-likelihood is equivalent to minimizing the KL-divergence between p Z and q Z , D KL (p Z , q Z ) (see supplementary material for more details and derivation of Equation equation 5). This approach has two major limitations: 1) The KL-Divergence and in general f -divergences do not provide meaningful dissimilarity measures for distributions supported on non-overlapping lowdimensional manifolds Arjovsky et al. (2017) ; Kolouri et al. (2018) (see supplementary material), which is common in hidden layers of neural networks, and therefore they do not provide informative gradients for training φ, and 2) we are limited to distributions q Z that have known explicit formulations, which is restrictive as it eliminates the ability to use the much broader class of samplable distributions.

Various alternatives exist in the literature to address the above-mentioned limitations. These methods often sample Z = {z j } N j=1 from q Z and Z = {z n = φ(x n )} N n=1 from p X and measure the discrepancy between these sets (i.e. point clouds). Note that there are no one-to-one correspondences between zj s and z n s. In their influential WAE paper, Tolstikhin et al. (2017) proposed two different approaches for measuring the discrepancy between Z and Z, namely the GAN-based and the maximum mean discrepancy (MMD)-based approaches. The GAN-based approach proposed in Tolstikhin et al. (2017) defines a discriminator network, D Z (p Z , q Z ), to classify zj s and z n s as coming from 'true' and 'fake' distributions correspondingly, and proposes a min-max adversarial optimization for learning φ and D Z . The MMD-based approach, utilizes a positive-definite reproducing kernel k : Z × Z → R to measure the discrepancy between Z and Z. The choice of the kernel and its parameterization, however, remain a data-dependent design parameter.

An interesting alternative approach is to use the Wasserstein distance between p Z and q Z . Following the work of Arjovsky et al. (2017) , this can be accomplished utilizing the Kantorovich-Rubinstein duality and through introducing a min-max problem, which leads to yet another adversarial training scheme similar to the GAN-based method in Tolstikhin et al. (2017) . Note that, since elements of Z and Z are not paired, an approach similar to equation 4 could not be used to minimize the discrepancy. In this paper, we propose to use the sliced-Wasserstein metric, Rabin & Peyré (2011) ; Rabin et al. (2011) ; Bonneel et al. (2015) ; Kolouri et al. (2016b) ; Carriere et al. (2017) ; Kolouri et al. (2018) , to measure the discrepancy between p Z and q Z . We show that using the sliced-Wasserstein distance ameliorates the need for training an adversary network or choosing a data-dependent kernel (as in WAE-MMD), and provides an efficient, stable, and simple numerical implementation.

Before explaining our proposed approach, it is worthwhile to point out the major difference between learning auto-encoders as generative models and GANs. In GANs, one needs to minimize a distance between {ψ(z j )|z j ∼ q Z } M j=1 and {x n } M n=1 , which are high-dimensional point clouds for which there are no correspondences between ψ(z j )s and x n s. For the auto-encoders, on the other hand, there exists correspondences between the high-dimensional point clouds {x n } M n=1 and {y n = ψ(φ(x n ))} M n=1 , and the problem simplifies to matching the lower-dimensional point clouds {φ(x n )} M n=1 and {z j ∼ q Z } M j=1 . In other words, the encoder performs a nonlinear dimensionality reduction, that enables us to solve a simpler problem compared to GANs. Next we introduce the details of our approach.

2. PROPOSED METHOD:

In what follows we first provide a brief review of the necessary equations to understand the Wasserstein and sliced-Wasserstein distances and then present our Sliced Wasserstein auto-encoder (SWAE).

2.1. WASSERSTEIN DISTANCES:

The Wasserstein distance between probability measures ρ X and ρ Y , with corresponding densities dρ X = p X (x)dx and dρ Y = p Y (y)dy is defined as:

EQUATION

where Γ(ρ X , ρ Y ) is the set of all transportation plans (i.e. joint measures) with marginal densities p X and p Y , and c : X × Y → R + is the transportation cost. equation 6 is known as the Kantorovich formulation of the optimal mass transportation problem, which seeks the optimal transportation plan between p X and p Y . If there exist diffeomorphic mappings, f : X → Y (i.e. transport maps) such that y = f (x) and consequently,

EQUATION

where det(D•) is the determinant of the Jacobian, then the Wasserstein distance could be defined based on the Monge formulation of the problem (see Villani (2008) and Kolouri et al. (2017) ) as:

EQUATION

where M P is the set of all diffeomorphisms that satisfy equation 7. As can be seen from equation 6and equation 8, obtaining the Wasserstein distance requires solving an optimization problem. We note that various efficient optimization techniques have been proposed in the past (e.g. Cuturi (2013) ; Solomon et al. (2015) ; Oberman & Ruan ( 2015)) to solve this optimization. For one-dimensional probability densities, p X and p Y , however, the Wasserstein distance has a closed-form solution. Let P X and P Y be the cumulative distributions of one-dimensional probability distributions p X and p Y , correspondingly. The Wassertein distance can then be calculated as below (see Kolouri et al. (2017) for more details):

EQUATION

This closed-form solution motivates the definition of sliced-Wasserstein distances.

2.2 SLICED-WASSERSTEIN DISTANCES Sliced-Wasserstein distance has similar qualitative properties to the Wasserstein distance, but it is much easier to compute. The sliced-Wasserstein distance was used in Rabin & Peyré (2011) ; Rabin et al. (2011) to calculate barycenter of distributions and point clouds. Bonneel et al. (2015) provided a nice theoretical overview of barycenteric calculations using the sliced-Wasserstein distance. Kolouri et al. (2016b) used it to define positive definite kernels for distributions and Carriere et al. (2017) to define a kernel for persistence diagrams. Sliced-Wasserstein was recently used for learning Gaussian mixture models in Kolouri et al. (2018) , and it was also used as a measure of goodness of fit for GANs in Karras et al. (2017) .

The main idea behind the sliced-Wasserstein distance is to slice (i.e., project) higher-dimensional probability densities into sets of one-dimensional marginal distributions and compare these marginal distributions via the Wasserstein distance. The slicing/projection process is related to the field of Integral Geometry and specifically the Radon transform (see Helgason (2011) ). The relevant result to our discussion is that a d-dimensional probability density p X can be uniquely represented as the set of its one-dimensional marginal distributions following the Radon transform and the Fourier slice theorem Helgason (2011) . These one dimensional marginal distributions of p X are defined as:

Rp X (t; θ) = X p X (x)δ(t -θ • x)dx, ∀θ ∈ S d-1 , ∀t ∈ R (10)

where S d-1 is the d-dimensional unit sphere. Note that for any fixed θ ∈ S d-1 , Rp X (•; θ) is a one-dimensional slice of distribution p X . In other words, Rp X (•; θ) is a marginal distribution of p X that is obtained from integrating p X over the hyperplane orthogonal to θ.Utilizing these marginal distributions in equation 10, the sliced Wasserstein distance could be defined as:

SW c (p X , p Y ) = S d-1 W c (Rp X (•; θ), Rp Y (•; θ))dθ (11)

Given that Rp X (•; θ) and Rp Y (•; θ) are one-dimensional, the Wasserstein distance in the integrand has a closed-form solution (see equation 9). Moreover, it can be shown that SW c is a true metric (Bonnotte (2013) and Kolouri et al. (2016a) ), and it induces the same topology as W c , at least on compact sets Santambrogio (2015) . A natural transportation cost that has extensively studied in the past is the 2 2 , c(x, y) = x -y 2 2 , for which there are theoretical guarantees on existence and uniqueness of transportation plans and maps (see Santambrogio (2015) and Villani (2008) ). When c(x, y) = x -y p p for p ≥ 2, the following upper bound hold for the SW distance:

EQUATION

where, Bonnotte (2013) proves this inequality. In our paper, we are interested in p = 2, for which α p,d = 1 d , and we have:

α d,p = 1 d S d-1 θ p p dθ ≤ 1. Chapter 5 in

EQUATION

In the Numerical Implementation Section, we provide a numerical experiment to compare W 2 and SW 2 , that confirms the above equation.

2.3. SLICED-WASSERSTEIN AUTO-ENCODER (SWAE):

Our proposed formulation for the SWAE is as follows:

argmin φ,ψ W ‡ c (p X , p Y ) + λSW c (p Z , q Z ) (14)

where φ is the encoder, ψ is the decoder, p X is the data distribution, p Y is the data distribution after encoding and decoding ( equation 2), p Z is the distribution of the encoded data ( equation 1), q Z is a predefined samplable distribution, and λ indicates the relative importance of the loss functions. To further clarify why we use the sliced-Wasserstein distance to measure the difference between p Z and q Z , we reiterate that due to the lack of correspondences between zi s and z j s, one cannot minimize the upper-bound in equation 4, and calculation of the Wasserstein distance requires an additional optimization step to obtain the optimal coupling between p Z and q Z . To avoid this additional optimization, while maintaining the favorable characteristics of the Wasserstein distance, we use the sliced-Wasserstein distance to measure the discrepancy between p Z and q Z .

3. NUMERICAL IMPLEMENTATION:

We now describe the numerical details of our approach.

3.1. NUMERICAL IMPLEMENTATION OF THE WASSERSTEIN DISTANCE IN 1D:

The Wasserstein distance between two one-dimensional probability densities p X and p Y is obtained from equation 9. The integral in equation 9 can be numerically estimated using the midpoint Riemann sum,

1 M M m=1 a m , where a m = c(P -1 X (τ m ), P -1 Y (τ m ))

and τ m = 2m-1 2M (see Fig. 1 ). In scenarios where only samples from the distributions are available, x m ∼ p X and y m ∼ p Y , the empirical densities can be estimated as

p X ≈ p X,M = 1 M M m=1 δ xm and p Y ≈ p Y,M = 1 M M m=1 δ ym ,

where δ xm is the Dirac delta function centered at x m . Therefore the corresponding empirical distribution function of p X is P X (t) ≈ P X,M (t) = 1 M M m=1 u(t-x m ) where u(.) is the step function (P Y,M (t) is defined similarly). From Glivenko-Cantelli Theorem we have that sup t |P X,M (t) -P X (t)| a.s.

--→ 0, where the convergence behavior is achieved via Dvoretzky-Kiefer-Wolfowitz inequality bound: P rob(sup t |P X,M (t) -P X (t)| > ) ≤ 2 exp (-2M 2 ). Calculating the Wasserstein distance with the empirical distribution function is computationally attractive. Sorting x m s in an ascending order, such that x i[m] ≤ x i[m+1] and where i[m] is the index of the sorted x m s, it is straightforward to see that P -1 X,M (τ m ) = x i[m] (see Fig. 1 for a visualization). The Wasserstein distance can be approximated by first sorting x m s and y m s and then calculating: The problem of calculating the Wasserstein distance between samples from one-dimensional densities simplifies to solving two sorting problems (solved in O(M )/O(M log(M )) best/worst case).

EQUATION

We need to address one final question here. How well does equation 15 approximate the Wasserstein distance, W c (p X , p Y )? We first note that the rates of convergence of empirical distributions, for the p-Wasserstein metric (i.e., c(x, y) = |x -y| p ) of order p ≥ 1, have been extensively studied in the mathematics and statistics communities (see for instance Bobkov & Ledoux (2014) and Dedecker et al. (2015) ). A detailed description of these rates is, however, beyond the scope of this paper, especially since these rates are dependent on the choice of p. In short, for p = 1 it can be shown that E(W1(pX,M , p X ) ≤ C √ M where C is an absolute constant. Similar results are achieved for E(Wp(pX,M , p X )) and (E(W p p (p X,M , p X )))

1 p , although under more strict assumptions on p X (i.e., slightly stronger assumptions than having a finite second moment). Using the triangle inequality together with the convergence rates of empirical distributions with respect to the p-Wasserstein distance, see Bobkov & Ledoux (2014) , for W 1 (p X,M , p X ) (or more generally W p (p X,M , p X )) we can show that (see supplementary material):

EQUATION

for some absolute constant, C. We reiterate that similar bounds could be found for W p although with slightly more strict assumptions on p X and p Y .

3.2. SLICING EMPIRICAL DISTRIBUTIONS:

In scenarios where only samples from the d-dimensional distribution, p X , are available, x m ∼ p X , the empirical density can be estimated as p X,M = 1 M M m=1 δ xm . Following equation 10 it is straightforward to show that the marginal densities (i.e. slices) are obtained from:

Rp X (t, θ) ≈ Rp X,M (t, θ) = 1 M M m=1 δ(t -x m • θ), ∀θ ∈ S d-1 , and ∀t ∈ R (17)

see the supplementary material for a proof. The Dvoretzky-Kiefer-Wolfowitz upper bound holds for Rp X (t, θ) and Rp X,M (t, θ).

3.3. MINIMIZING SLICED-WASSERSTEIN VIA RANDOM SLICING:

Minimizing the sliced-Wasserstein distance (i.e., as in the second term of 14) requires an integration over the unit sphere in R d , i.e., S d-1 . In practice, this integration is approximated by using a simple Monte Carlo scheme that draws uniform samples from S d-1 and replaces the integral with a finite-sample average, Moreover, the global minimum for SW c (p Z , q Z ) is also a global minimum for each

SW c (p Z , q Z ) ≈ 1 |Θ| θ l ∈Θ W c (Rp Z (•; θ l ), Rq Z (•; θ l ))

W c (Rp Z (•; θ l ), Rq Z (•; θ l )). Note that SW c (p Z , q Z ) = E S (d-1) (W c (Rp Z (•; θ), Rq Z (•; θ))).

A fine sampling of S d-1 , however, is required for a good approximation of SW c (p Z , q Z ). Intuitively, if p Z and q Z are similar, then their projections with respect to any finite subset of S d-1 would also be similar. This leads to a stochastic gradient descent scheme where in addition to the random sampling of the input data, we also random sample the projection angles from S d-1 .

A natural question arises on the effect of the number of random slices, L = |Θ|, on the approximation of the SW distance. Here, we devised a simple experiment that demonstrates the effect of L on approximating the SW distance. We generated two random multi-variate Gaussian distributions in a d-dimensional space, where d ∈ {2 n } 10 n=1 , to serve as p X = N (µ X , Σ X ) and p X = N (µ Y , Σ Y ). The Wasserstein distance for the two Gaussian distributions has a closed form solution,

EQUATION

2 ), which served as the ground-truth distance between the distributions. We then measured the SW distance between M = 1000 samples generated from the two Gaussian distributions using L ∈ {1, 10, 50, 100, 500, 1000} random slices. We repeated the experiment for each L and d, a thousand times and report the means and standard deviations in Figure 2 . Following equation 13 we scaled the SW distance by √ d. Moreover we found out empirically that 1.22

√ d E(SW2(pX,M , p Y,M )) ≈ W 2 (p X , p Y ).

It can be seen from Figure 2 that the expected value of the scaled SW -distance closely follows the true Wasserstein distance. A more interesting observation is that the variance of estimation increases for higher dimensions d and decreases as the number of random projections, L, increases. Hence, calculating the SW distance in the image space, as in Deshpande et al. (2018) , requires a very large number of projections L to get a less variant approximation of the distance.

3.4. PUTTING IT ALL TOGETHER:

To optimize the proposed SWAE objective function in equation 14 we use a stochastic gradient descent scheme as described here. In each iteration, let {x m ∼ p X } M m=1 and {z m ∼ q Z } M m=1 be i.i.d random samples from the input data and the predefined distribution, q Z , correspondingly. Let {θ l } L l=1 be randomly sampled from a uniform distribution on S d-1 . Then using the numerical approximations described in this section, the loss function in equation 14 can be rewritten as:

EQUATION

)

where i[m] and j[m] are the indices of sorted θ l •z m s and θ l •φ(x m ) with respect to m, correspondingly. The steps of our proposed method are presented in Algorithm 1. It is worth pointing out that sorting is by itself an optimization problem (which can be solved very efficiently), and therefore the sorting followed by the gradient descent update on φ and ψ is in essence a min-max problem, which is being solved in an alternating fashion. Finally, we point out that each iteration of SWAE costs O(LM log(M )) operations. 

S K-1 Sort θ l • zM such that θ l • zi[m] ≤ θ l • zi[m+1] Sort θ l • φ(xm) such that θ l • φ(x j[m] ) ≤ θ l • φ(x j[m+1]

) Update φ and ψ by descending:

M m=1 c(xm, ψ(φ(xm))) + λ L l=1 M m=1 c(θ l • zi[m] , θ l • φ(x j[m]

)) end while

4. EXPERIMENTS:

In our experiments we used three image datasets, namely the MNIST dataset by LeCun (1998), the CelebFaces Attributes Dataset (CelebA) by Liu et al. (2015) , and the LSUN Bedroom Dataset by Yu et al. (2015) . For the MNIST dataset we used a simple auto-encoder with mirrored classic deep convolutional neural networks with 2D average poolings, leaky rectified linear units (Leaky-ReLu) as the activation functions, and upsampling layers in the decoder. For the CelebA and LSUN datasets we used the DCGAN Radford et al. (2015) architecture similar to Tolstikhin et al. (2017) .

To test the capability of our proposed algorithm in shaping the latent space of the encoder, we started with the MNIST dataset and trained SWAE to encode this dataset to a two-dimensional latent space (for the sake of visualization) while enforcing a match between p X and p Y and p Z and q Z . We chose four different samplable distributions as shown in Figure 3 . It can bee seen that SWAE can successfully embed the dataset into the latent space while enforcing p Z to closely follow q Z . In addition, we sample the two-dimensional latent spaces on a 25 × 25 grid in [-1, 1] 2 and decode these points to visualize their corresponding images in the digit/image space.

To get a sense of the convergence behavior of SWAE, and similar to the work of Karras et al. (2017) , we calculate the Sliced Wasserstein distance between p Z and q Z as well as p X and p Y at each batch iteration where we used p-LDA Wang et al. (2011) to calculate projections (See supplementary material). We compared the convergence behavior of SWAE with the closest related work, WAE Tolstikhin et al. (2017) (specifically WAE-GAN) where an adversarial training is used to match p Z to q Z , while the loss function for p X and p Y remains exactly the same between the two methods. We repeated the experiments 100 times and report the summary of results in Figure 4 . We mention that the exact same models and optimizers were used for both methods in this experiment. An interesting observation, here is that while WAE-GAN provides good or even slightly better generated random samples for MNIST (lower sliced-Wasserstein distance between p X and p Y ), it fails to provide a good match between p Z and q Z for the choice of the prior distribution reported in Figure 4 . This phenomenon seems to be related to the mode-collapse problem of GANs, where the adversary fails to sense that the distribution is not fully covered. Finally, in our experiments we did not notice a significant difference between the computational time for SWAE and WAE-GAN. For the MNIST experiment and on a single NVIDIA Tesla P 100 GPU, each batch iteration (batchsize=500) of WAE-GAN took 0.2571 ± 0.0435(sec) while SWAE (with L = 50 projections) took 0.2437 ± 0.0391(sec). Table 1 : Quantitative comparison of the SWAE and WAE-GAN using the sliced-Wasserstein distance with discriminant slices in the latent space, SW (p Z , q Z ), and the output space, SW (p X , p Y ). The distribution in the 64-dimensional latent space, q Z , was set to Normal. We also report the negative log-likelihood of {z i = φ(x i )} with repect to q Z for 1000 testing samples for both datasets. We did not use Nowizin's trick for the GAN models.

Model FID -CelebA FID -LSUN Bedroom SWAE 79 ± 6 225 ± 7 WAE-GAN 53 ± 2 232 ± 2 WAE-MMD(IMQ) 55 ± 1 226 ± 2 WAE-MMD(RBF) 363 ± 17 378 ± 12

True Data 2 3

Table 2 : FID score statistics (N = 5) at final iteration of training. Lower is better. Scores were computed with 10 4 random samples from the testing set against an equivalent amount of generated samples.

The CelebA face and the LSUN bedroom datasets contain higher degrees of variations compared to the MNIST dataset and therefore a two-dimensional latent-space does not suffice to capture the variations in these datasets (See supplementary material for more details on the dimensionality of the latent space). We used a K = 64 dimensional latent spaces for both the CelebA and the LSUN Bedroom datasets, and also used a larger auto-encoder (i.e., DCGAN, following the work of Tolstikhin et al. (2017) ). For these datasets SWAE was trained with q Z being the Normal distribution to enable the calculation of the negative log likelihood (NLL). Table 1 shows the comparison between SWAE and WAE for these two datasets. We note that all experimental parameters were kept the same to enable an apples to apples comparison. Finally, Figure 5 demonstrates the interpolation between two sample points in the latent space, i.e. ψ(tφ(I 0 ) + (1 -t)φ(I 1 )) for t ∈ [0, 1], for all three datasets.

5. CONCLUSIONS:

We introduced Sliced Wasserstein auto-encoders (SWAE), which enable one to shape the distribution of the encoded samples to any samplable distribution without the need for adversarial training or having a likelihood function specified. In addition, we provided a simple and efficient numerical scheme for this problem, which only relies on few inner products and sorting operations in each SGD iteration. We further demonstrated the capability of our method on three image datasets, namely the MNIST, the CelebA face, and the LSUN Bedroom datasets, and showed competitive performance, in the sense of matching distributions p Z and q Z , to the techniques that rely on additional adversarial trainings. Finally, we envision SWAE could be effectively used in transfer learning and domain adaptation algorithms where q Z comes from a source domain and the task is to encode the target domain p X in a latent space such that the distribution follows the distribution of the target domain. Figure 4 : Sample convergence behavior for our method compared to the WAE-GAN, where q Z is set to a ring distribution (Figure 3 , top left). The columns represent batch iterations (batchsize= 500).

The top half of the table shows results of ψ(z) for z ∼ q Z , and the bottom half shows z ∼ q Z and φ(x) for x ∼ p X . It can be seen that the adversarial loss in the latent space does not provide a full coverage of the distribution, which is a similar problem to the well-known 'mode collapse' problem in the GANs. It can be seen that SWAE provides a superior match between p Z and q Z while it does not require adversarial training.

Figure 6 : These plots show W 1 (p, q τ ) and JS(p, q τ ) where p is a uniform distribution around zero and q τ (x) = p(x -τ ). It is clear that JS divergence does not provide a usable gradient when distributions are supported on non-overlapping domains.

SUPPLEMENTARY MATERIAL COMPARISON OF DIFFERENT DISTANCES:

Following the example by Arjovsky et al. (2017) and later Kolouri et al. (2018) here we show a simple example comparing the Jensen-Shannon divergence with the Wasserstein distance. First note that the Jensen-Shannon divergence is defined as,

JS(p, q) = KL(p, p + q 2 ) + KL(q, p + q 2 )

where KL(p, q) = X p(x)log( p(x) q(x) )dx is the Kullback-Leibler divergence. Now consider the following densities, p(x) be a uniform distribution around zero and let q τ (x) = p(x -τ ) be a shifted version of the p. Figure 6 show W 1 (p, q τ ) and JS(p, q τ ) as a function of τ . As can be seen the JS divergence fails to provide a useful gradient when the distributions are supported on non-overlapping domains.

LOG-LIKELIHOOD:

To maximize (minimize) the similarity (dissimilarity) between p Z and q Z , we can write :

argmax φ Z p Z (z)log(q Z (z))dz = Z X p X (x)δ(z -φ(x))log(q Z (z))dxdz = X p X (x)log(q Z (φ(x)))dx

where we replaced p Z with equation 1. Furthermore, it is straightforward to show:

argmax φ Z p Z (z)log(q Z (z))dz = argmax φ Z p Z (z)log( q Z (z) p Z (z) )dz = argmin φ D KL (p Z , q Z ) PROOF OF EQUATION 3

The Wasserstein distance between the two probability measures ρ X and ρ Y with respective densities p X and p Y , can be measured via the Kantorovich formulation of the optimal mass transport problem:

EQUATION

y)} is the set of all transportation plans (i.e., couplings or joint distributions) over p X and p Y . Now, note that the two step process of encoding p X into the latent space Z and decoding it to p Y , provides a unique decomposition of γ as γ 0 (x, y) = δ(y -ψ(φ(x)))p X (x) ∈ Γ. Therefore we can write:

W c (p X , p Y ) = inf γ∈Γ X Y c(x, y)γ(x, y)dxdy ≤ W ‡ c (p X , p Y ) := X Y c(x, y)γ 0 (x, y)dxdy = X c(x, ψ(φ(x)))p X (x)dx

which proves equation 3. Finally, taking the infimum of the two sides of the inequality, with respect to φ and ψ, we have:

inf ψ,φ W c (p X , p Y ) = inf ψ,φ inf γ∈Γ ψ,φ X Y c(x, y)γ(x, y)dxdy ≤ inf ψ,φ W ‡ c (p X , p Y ) = inf ψ,φ X c(x, ψ(φ(x)))p X (x)dx where Γ ψ,φ := {γ| Y γ(x, y)dy = p X (x), X γ(x, y)dx = X p X (x)δ(y -ψ(φ(x)))dx}. Fig- ure 7 demonstrates a simple scenario were the Wasserstein distance, W c (p X , p Y ), is zero how- ever, W ‡ c (p X , p Y ) is non-zero. Finally, we note that ψ(φ(•)) = id(•) is a global optima for both W c (p X , p Y ) and W ‡ c (p X , p Y ).

Following equation 10 a distribution can be sliced via: 

Rp X (t, θ) = X p X (x)δ(t -θ • x)dx

Rp X (t, θ) = 1 M M m=1 X δ(x -x m )δ(t -θ • x)dx = 1 M M m=1 δ(t -θ • x m )

DIMENSIONALITY OF THE LATENT SPACE:

Figure 9 demonstrates the outputs of trained SWAEs with K = 2 and K = 128 for sample input images. The input images were resized to 64 × 64 and then fed to our auto-encoder structure. This effect can also be seen for the MNIST dataset as shown in Figure 10 . When the dimensionality of the latent-space (i.e. information bottleneck) is too low the latent space will not contain enough information to reconstruct crisp images. Increasing the dimensionality of the latent space leads to crisper images.

CALCULATING THE SLICED WASSERSTEIN DISTANCE AS A MEASURE OF GOODNESS OF FIT:

In this paper we also used the sliced Wasserstein distance as a measure of goodness of fit (for convergence analysis). To provide a fair comparison between different methods, we avoided random projections for this comparison. Instead, we calculated a discriminant subspace to separate ψ(z) from ψ(φ(x)) for z ∼ q Z and x ∼ p X , and set the projection parameters θs to the calculated discriminant components. This will lead to only slices that contain discriminant information. We point out that the linear discriminant analysis (LDA) is not a good choice for this task as it only leads to one discriminant component (because we only have two classes). We used the penalized linear discriminant analysis (p-LDA) that utilizes a combination of LDA and PCA. In short, p-LDA solves the following objective function: for some absolute constant C, where the last line comes from the empirical convergence bounds of distributions with respect to the Wasserstein distance, see Bobkov & Ledoux (2014) .

Figure 11 : Sample convergence behavior for our method compared to the WAE- GAN Tolstikhin et al. (2017) , where q

Z is set to the ring distribution. The columns represent batch iterations (batchsize= 500).

The top half of the table shows results of ψ(z)

for z ∼ q Z , and the bottom half shows z ∼ q Z and φ(x) for x ∼ p X . It can be seen that SWAE provides a superior match between φ(x) (i.e. p Z ) and q Z while being less computationally expensive and not requiring adversarial training.

Figure 12 : Sample convergence behavior for our method compared to the WAE- GAN Tolstikhin et al. (2017) , where q

Z is set to the uniform polar distribution. The columns represent batch iterations (batchsize= 500).

The top half of the table shows results of ψ(z)

for z ∼ q Z , and the bottom half shows z ∼ q Z and φ(x) for

x ∼ p X .It can be seen that SWAE provides a superior match between φ(x) (i.e. p Z ) and q Z while being less computationally expensive and not requiring adversarial training.



Figures and Tables:

Figure 1: The Wasserstein distance for one-dimensional probability distributions p X and p Y (top left) is calculated based on equation 9. For a numerical implementation, the integral in equation 9 is substituted with 1 M

Figure 2: SW approximations (scaled by 1.22 √ d) of the Wdistance in different dimensions, d ∈ {2 n } 10 n=1 , and different number of random slices, L.Such Monte Carlo estimation was used inRabin & Peyré (2011), and later used inBonneel et al. (2015);Kolouri et al. (2018);Şimşekli et al. (2018);Deshpande et al. (2018).Moreover, the global minimum for SW c (p Z , q Z ) is also a global minimum for eachW c (Rp Z (•; θ l ), Rq Z (•; θ l )). Note that SW c (p Z , q Z ) = E S (d-1) (W c (Rp Z (•; θ), Rq Z (•; θ))).A fine sampling of S d-1 , however, is required for a good approximation of SW c (p Z , q Z ). Intuitively, if p Z and q Z are similar, then their projections with respect to any finite subset of S d-1 would also be similar. This leads to a stochastic gradient descent scheme where in addition to the random sampling of the input data, we also random sample the projection angles from S d-1 .

Figure 5: Interpolation in the latent space, ψ(tφ(I 0 ) + (1 -t)φ(I 1 )) for t ∈ [0, 1].

Figure 3: The results of SWAE on the MNIST dataset with a two-dimensional embedding space for four different distributions as ,q Z , namely the ring distribution (top left), the uniform distribution (bottom left), the uniform polar distribution (top right), and a custom polar distribution (bottom right). Note that the far right visualization demonstrates the decoding of a 25 × 25 grid in [-1, 1] 2 .

Figure 7: The optimal coupling (i.e., transport plan) between p X and p Y could be equal or different from γ(x, y) = δ(y -ψ(φ(x)))p X (x). This leads to the scenario on the right where W c (p X , p Y ) = 0 but W ‡ c (p X , p Y ) > 0.

Figure 8 visualizes two sample slices for an example distribution p X . Here we calculate a Radon slice of the empirical distribution p X (x) = 1 M

Figure 8: Visualization of the slicing process defined in equation 10

Figure 10: Interpolation results for on the MNIST dataset with various dimensions of the latent space. The parameter t ∈ [0, 1] indicates the interpolation parameter.

Algorithm 1 Sliced-Wasserstein Auto-Encoder (SWAE) Regularization coefficient λ, and number of random projections, L. Initialize the parameters of the encoder, φ, and decoder, ψ while φ and ψ have not converged do Sample {x1, ..., xM } from training set (i.e. pX ) Sample {z1, ..., zM } from qZ Sample {θ1, ..., θL} from

Table content: <table/>

Peer review (iclr2019):
 This paper proposes training generative models with Wasserstein auto-encoders. It uses the sliced-Wasserstein distance to measure the dissimilarity between p_z and q_z. Strengths: 1. This paper is easy to read. 2. Concepts are introduced clearly. My major comments are the following: 1. The innovation is a bit on the incremental level, especially given the results from WAE (Tolstikhin, ICLR18). The training objective is the same as Eq(4) in the WAE paper. The only difference is that the dissimilarity measure between p_z and q_z used in this paper is the sliced- Wasserstein distance, while WAE used GAN/MMD-based penalties. The advantage of using sliced-Wasserstein distance is not clear to me either. 2. The empirical results are fairly weak. The authors may consider reporting the sample qualities (e.g. FID) for all the methods. 3. The results of WAE-MMD are not reported.