Year: 2019

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. (2017); Murez et al. (2018), super-resolution Ledig et al. (2016); Kolouri & Rohde (2015), image inpainting and completion Yeh et al. (2017), and image retrieval Creswell & Bharath (2016), among many others. The recent parametric generative models, including Generative Adversarial Networks (GANs) Goodfellow et al. (2014); Radford et al. (2015); Arjovsky et al. (2017); Berthelot et al. (2017) and Variational auto-encoders (VAE) Kingma & Welling (2013); Mescheder et al. (2017); Bousquet et al. (2017) enable an unsupervised and end-to-end modeling of the high-dimensional distribution of the training data. 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

text: 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. (2017); Murez et al. (2018), super-resolution Ledig et al. (2016); Kolouri & Rohde (2015), image inpainting and completion Yeh et al. (2017), and image retrieval Creswell & Bharath (2016), among many others. The recent parametric generative models, including Generative Adversarial Networks (GANs) Goodfellow et al. (2014); Radford et al. (2015); Arjovsky et al. (2017); Berthelot et al. (2017) and Variational auto-encoders (VAE) Kingma & Welling (2013); Mescheder et al. (2017); Bousquet et al. (2017) enable an unsupervised and end-to-end modeling of the high-dimensional distribution of the training data.
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 Şimş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(104), 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. Şimşekli et al. (2018), on the other hand, take an interesting but different approach of parameter-free generative modeling via sliced-Wasserstein flows.

heading: 1 NOTATION AND PRELIMINARIES
text: Let X denote the compact domain of a manifold in Euclidean space and let xn ∈ X denote an individual input data point. Furthermore, let ρX be a Borel probability measure defined on X . We define the probability density function pX(x) for input data x to be:
dρX(x) = pX(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 pX by:
pZ(z) = ∫ X pX(x)δ(z − φ(x))dx, (1)
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, pZ , follows a prior distribution qZ . Let ψ : Z → X be the decoder that maps the latent codes back to the original space such that
pY (y) = ∫ X pX(x)δ(y − ψ(φ(x)))dx, (2)
where y denotes the decoded samples. It is straightforward to see that when ψ = φ−1 (i.e. ψ(φ(·)) = id(·)), the distribution of the decoder pY and the input distribution pX 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 pY and pX , and between pZ and qZ . In what follows, we briefly review the existing dissimilarity measures for these distributions.
1.1 MINIMIZING DISSIMILARITY BETWEEN pX AND pY 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 pX and pY . Most notably, Nowozin et al. (2016) showed that for the general family of f -divergences, Df (pX , pY ), (including the KL-divergence, JensenShannon, etc.), using the Fenchel conjugate of the convex function f and minimizing Df (pX , pY ) 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 pX and pY . In Wasserstein-GAN, Arjovsky et al. (2017) utilized the Kantorovich-Rubinstein duality for the 1-Wasserstein distance, W1(pX , pY ), 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):
Wc(pX , pY ) ≤W ‡c (pX , pY ) := EpX (c(x, ψ(φ(x)))) (3)
= ∫ X c(x, ψ(φ(x)))pX(x)dx,
Furthermore, the r.h.s. of equation 3 supports a simple implementation where for i.i.d samples of the input distribution, {xn}Nn=1, the upper bound can be approximated as:
W ‡c (pX , pY ) ≈ 1
N N∑ n=1 c(xn, ψ(φ(xn))) (4)
The r.h.s of equation 3 and equation 4 take advantage of the existence of pairs xn and yn = ψ(φ(xn)), which make f(·) = ψ(φ(·)) a transport map between pX and pY (but not necessarily the optimal transport map). In this paper, we minimize W ‡c (pX , pY ) following equation 4 to minimize the discrepancy between pX and pY . Next, we focus on the discrepancy measures between pZ and qZ .
1.2 MINIMIZING DISSIMILARITY BETWEEN pZ AND qZ If qZ is a known distribution with an explicit formulation (e.g. Normal distribution) the most straightforward approach for measuring the (dis)similarity between pZ and qZ is the log-likelihood of z = φ(x) with respect to qZ , formally:
supφ ∫ X pX(x)log(qZ(φ(x)))dx (5)
maximizing the log-likelihood is equivalent to minimizing the KL-divergence between pZ and qZ , DKL(pZ , qZ) (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 qZ 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}Nj=1 from qZ and Z = {zn = φ(xn)}Nn=1 from pX and measure the discrepancy between these sets (i.e. point clouds). Note that there are no one-to-one correspondences between z̃js and zns. 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,DZ(pZ , qZ), to classify z̃js and zns as coming from ‘true’ and ‘fake’ distributions correspondingly, and proposes a min-max adversarial optimization for learning φ and DZ . 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 pZ and qZ . 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 pZ and qZ . 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 ∼ qZ}Mj=1 and {xn}Mn=1, which are high-dimensional point clouds for which there are no correspondences between ψ(z̃j)s and xns. For the auto-encoders, on the other hand, there exists correspondences between the high-dimensional point clouds {xn}Mn=1 and {yn = ψ(φ(xn))}Mn=1, and the problem simplifies to matching the lower-dimensional point clouds {φ(xn)}Mn=1 and {z̃j ∼ qZ}Mj=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.

heading: 2 PROPOSED METHOD
text: 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).

heading: 2.1 WASSERSTEIN DISTANCES
text: The Wasserstein distance between probability measures ρX and ρY , with corresponding densities dρX = pX(x)dx and dρY = pY (y)dy is defined as:
Wc(pX , pY ) = infγ∈Γ(ρX ,ρY ) ∫ X×Y c(x, y)dγ(x, y) (6)
where Γ(ρX , ρY ) is the set of all transportation plans (i.e. joint measures) with marginal densities pX and pY , 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 pX and pY . If there exist diffeomorphic mappings, f : X → Y (i.e. transport maps) such that y = f(x) and consequently,
pY (y) = ∫ X pX(x)δ(y − f(x))dx When f is−−−−−−−−−→ a diffeomorphism pY (y) = det(Df −1(y))pX(f −1(y)) (7)
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:
Wc(pX , pY ) = minf∈MP ∫ X c(x, f(x))dρX(x) (8)
where MP is the set of all diffeomorphisms that satisfy equation 7. As can be seen from equation 6 and 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, pX and pY , however, the Wasserstein distance has a closed-form solution. Let PX and PY be the cumulative distributions of one-dimensional probability distributions pX and pY , correspondingly. The Wassertein distance can then be calculated as below (see Kolouri et al. (2017) for more details):
Wc(pX , pY ) = ∫ 1 0 c(P−1X (τ), P −1 Y (τ))dτ, (9)
This closed-form solution motivates the definition of sliced-Wasserstein distances.

heading: 2.2 SLICED-WASSERSTEIN DISTANCES
text: 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 pX 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 pX are defined as:
RpX(t; θ) = ∫ X pX(x)δ(t− θ · x)dx, ∀θ ∈ Sd−1, ∀t ∈ R (10)
where Sd−1 is the d-dimensional unit sphere. Note that for any fixed θ ∈ Sd−1, RpX(·; θ) is a one-dimensional slice of distribution pX . In other words,RpX(·; θ) is a marginal distribution of pX
that is obtained from integrating pX over the hyperplane orthogonal to θ.Utilizing these marginal distributions in equation 10, the sliced Wasserstein distance could be defined as:
SWc(pX , pY ) = ∫ Sd−1 Wc(RpX(·; θ),RpY (·; θ))dθ (11)
Given thatRpX(·; θ) andRpY (·; θ) are one-dimensional, the Wasserstein distance in the integrand has a closed-form solution (see equation 9). Moreover, it can be shown that SWc is a true metric (Bonnotte (2013) and Kolouri et al. (2016a)), and it induces the same topology as Wc, at least on compact sets Santambrogio (2015). A natural transportation cost that has extensively studied in the past is the `22, c(x, y) = ‖x− y‖22, 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‖pp for p ≥ 2, the following upper bound hold for the SW distance:
SW pp (pX , pY ) ≤ αd,pW pp (pX , pY ) (12) where, αd,p = 1d ∫ Sd−1 ‖θ‖ p pdθ ≤ 1. Chapter 5 in Bonnotte (2013) proves this inequality. In our paper, we are interested in p = 2, for which αp,d = 1d , and we have:
SW2(pX , pY ) ≤ 1√ d W2(pX , pY ) (13)
In the Numerical Implementation Section, we provide a numerical experiment to compare W2 and SW2, that confirms the above equation.

heading: 2.3 SLICED-WASSERSTEIN AUTO-ENCODER (SWAE)
text: Our proposed formulation for the SWAE is as follows:
argminφ,ψW ‡ c (pX , pY ) + λSWc(pZ , qZ) (14)
where φ is the encoder, ψ is the decoder, pX is the data distribution, pY is the data distribution after encoding and decoding ( equation 2), pZ is the distribution of the encoded data ( equation 1), qZ 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 pZ and qZ , we reiterate that due to the lack of correspondences between z̃is and zjs, 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 pZ and qZ . 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 pZ and qZ .

heading: 3 NUMERICAL IMPLEMENTATION
text: We now describe the numerical details of our approach.

heading: 3.1 NUMERICAL IMPLEMENTATION OF THE WASSERSTEIN DISTANCE IN 1D
text: The Wasserstein distance between two one-dimensional probability densities pX and pY is obtained from equation 9. The integral in equation 9 can be numerically estimated using the midpoint Riemann sum, 1M ∑M m=1 am, where am = 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, xm ∼ pX and ym ∼ pY , the empirical densities can be estimated as pX ≈ pX,M = 1M ∑M m=1 δxm and pY ≈ pY,M = 1 M ∑M m=1 δym , where δxm is the Dirac delta function centered at xm. Therefore the corresponding empirical distribution function of pX is PX(t) ≈ PX,M (t) = 1M ∑M m=1 u(t−xm) where u(.) is the step function (PY,M (t) is defined similarly). From Glivenko-Cantelli Theorem we have that supt |PX,M (t)−PX(t)| a.s.−−→ 0, where the convergence behavior is achieved via Dvoretzky–Kiefer–Wolfowitz inequality bound: Prob(supt |PX,M (t)− PX(t)| > ) ≤ 2 exp (−2M 2). Calculating the Wasserstein distance with the empirical distribution function is computationally attractive. Sorting xms in an ascending order, such that xi[m] ≤ xi[m+1] and where i[m] is the index of the sorted xms, it is straightforward to see that P−1X,M (τm) = xi[m] (see Fig. 1 for a visualization). The Wasserstein distance can be approximated by first sorting xms and yms and then calculating:
Wc(pX , pY ) ≈Wc(pX,M , pY,M ) = 1
M M∑ m=1 c(xi[m], yj[m]) (15)
The problem of calculating the Wasserstein distance between samples from one-dimensional densities simplifies to solving two sorting problems (solved in O(M)/O(Mlog(M)) best/worst case). We need to address one final question here. How well does equation 15 approximate the Wasserstein distance, Wc(pX , pY )? 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 , pX) ≤ C√M where C is an absolute constant. Similar results are achieved for E(Wp(pX,M , pX)) and (E(W pp (pX,M , pX))) 1 p , although under more strict assumptions on pX (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 W1(pX,M , pX) (or more generally Wp(pX,M , pX)) we can show that (see supplementary material):
E(W1(pX , pY )−W1(pX,M , pY,M )) ≤ C√ M
(16)
for some absolute constant, C. We reiterate that similar bounds could be found for Wp although with slightly more strict assumptions on pX and pY .

heading: 3.2 SLICING EMPIRICAL DISTRIBUTIONS
text: In scenarios where only samples from the d-dimensional distribution, pX , are available, xm ∼ pX , the empirical density can be estimated as pX,M = 1M ∑M m=1 δxm . Following equation 10 it is straightforward to show that the marginal densities (i.e. slices) are obtained from:
RpX(t, θ) ≈ RpX,M (t, θ) = 1
M M∑ m=1 δ(t− xm · θ), ∀θ ∈ Sd−1, and ∀t ∈ R (17)
see the supplementary material for a proof. The Dvoretzky–Kiefer–Wolfowitz upper bound holds for RpX(t, θ) andRpX,M (t, θ).

heading: 3.3 MINIMIZING SLICED-WASSERSTEIN VIA RANDOM SLICING
text: Minimizing the sliced-Wasserstein distance (i.e., as in the second term of 14) requires an integration over the unit sphere in Rd, i.e., Sd−1. In practice, this integration is approximated by using a simple Monte Carlo scheme that draws uniform samples from Sd−1 and replaces the integral with a
finite-sample average,
SWc(pZ , qZ) ≈ 1 |Θ| ∑ θl∈Θ Wc(RpZ(·; θl),RqZ(·; θl))
√ d) of the W-
Such Monte Carlo estimation was used in Rabin & Peyré (2011), and later used in Bonneel et al. (2015); Kolouri et al. (2018); Şimşekli et al. (2018); Deshpande et al. (2018). Moreover, the global minimum for SWc(pZ , qZ) is also a global minimum for each Wc(RpZ(·; θl),RqZ(·; θl)). Note that SWc(pZ , qZ) = ES(d−1)(Wc(RpZ(·; θ),RqZ(·; θ))). A fine sampling of Sd−1, however, is required for a good approximation of SWc(pZ , qZ). Intuitively, if pZ and qZ are similar, then their projections with respect to any finite subset of Sd−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 Sd−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 ∈ {2n}10n=1, to serve as pX = N (µX ,ΣX) and pX = N (µY ,ΣY ). The Wasserstein distance for the two Gaussian distributions has a closed form solution,
W 22 (pX , pY ) = ‖µX − µY ‖22 + trace(ΣX + ΣY − 2(Σ 1 2 XΣY Σ 1 2 X) 1 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 √ dE(SW2(pX,M , pY,M )) ≈ W2(pX , pY ). 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.

heading: 3.4 PUTTING IT ALL TOGETHER
text: To optimize the proposed SWAE objective function in equation 14 we use a stochastic gradient descent scheme as described here. In each iteration, let {xm ∼ pX}Mm=1 and {z̃m ∼ qZ}Mm=1 be i.i.d random samples from the input data and the predefined distribution, qZ , correspondingly. Let {θl}Ll=1 be randomly sampled from a uniform distribution on Sd−1. Then using the numerical approximations described in this section, the loss function in equation 14 can be rewritten as:
L(φ, ψ) = 1 M M∑ m=1 c(xm, ψ(φ(xm))) + λ LM L∑ l=1 M∑ m=1 c(θl · z̃i[m], θl · φ(xj[m])) (18)
where i[m] and j[m] are the indices of sorted θl ·z̃ms and θl ·φ(xm) with respect tom, 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(LMlog(M)) operations.
Algorithm 1 Sliced-Wasserstein Auto-Encoder (SWAE) Require: 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 {z̃1, ..., z̃M} from qZ Sample {θ1, ..., θL} from SK−1 Sort θl · z̃M such that θl · z̃i[m] ≤ θl · z̃i[m+1] Sort θl · φ(xm) such that θl · φ(xj[m]) ≤ θl · φ(xj[m+1]) Update φ and ψ by descending: ∑M m=1 c(xm, ψ(φ(xm))) + λ ∑L l=1 ∑M m=1 c(θl · z̃i[m], θl · φ(xj[m]))
end while

heading: 4 EXPERIMENTS
text: 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 pX and pY and pZ and qZ . 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 pZ to closely follow qZ . 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 pZ and qZ as well as pX and pY 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 pZ to qZ , while the loss function for pX and pY 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 pX and pY ), it fails to provide a good match between pZ and qZ 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 P100 GPU, each batch iteration (batchsize=500) of WAEGAN took 0.2571± 0.0435(sec) while SWAE (with L = 50 projections) took 0.2437± 0.0391(sec).
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 qZ 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φ(I0) + (1 − t)φ(I1)) for t ∈ [0, 1], for all three datasets.

heading: 5 CONCLUSIONS
text: 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 pZ and qZ , 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 qZ comes from a source domain and the task is to encode the target domain pX in a latent space such that the distribution follows the distribution of the target domain.

heading: SUPPLEMENTARY MATERIAL

heading: COMPARISON OF DIFFERENT DISTANCES
text: 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 W1(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.

heading: LOG-LIKELIHOOD
text: To maximize (minimize) the similarity (dissimilarity) between pZ and qZ , we can write :
argmaxφ ∫ Z pZ(z)log(qZ(z))dz = ∫ Z ∫ X pX(x)δ(z − φ(x))log(qZ(z))dxdz
= ∫ X pX(x)log(qZ(φ(x)))dx
where we replaced pZ with equation 1. Furthermore, it is straightforward to show:
argmaxφ ∫ Z pZ(z)log(qZ(z))dz = argmaxφ ∫ Z pZ(z)log( qZ(z) pZ(z) )dz
= argminφDKL(pZ , qZ)

heading: PROOF OF EQUATION 3
text: The Wasserstein distance between the two probability measures ρX and ρY with respective densities pX and pY , can be measured via the Kantorovich formulation of the optimal mass transport problem:
Wc(pX , pY ) = infγ∈Γ ∫ X ∫ Y c(x, y)γ(x, y)dxdy
where Γ := {γ : X × Y → R+| ∫ Y γ(x, y)dy = pX(x), ∫ X γ(x, y)dx = pY (y)} is the set of all transportation plans (i.e., couplings or joint distributions) over pX and pY . Now, note that the two step process of encoding pX into the latent space Z and decoding it to pY , provides a unique decomposition of γ as γ0(x, y) = δ(y − ψ(φ(x)))pX(x) ∈ Γ.
Therefore we can write:
Wc(pX , pY ) = infγ∈Γ ∫ X ∫ Y c(x, y)γ(x, y)dxdy ≤
W ‡c (pX , pY ) := ∫ X ∫ Y c(x, y)γ0(x, y)dxdy = ∫ X c(x, ψ(φ(x)))pX(x)dx
which proves equation 3. Finally, taking the infimum of the two sides of the inequality, with respect to φ and ψ, we have:
infψ,φWc(pX , pY ) = infψ,φ infγ∈Γψ,φ ∫ X ∫ Y c(x, y)γ(x, y)dxdy ≤
infψ,φW ‡ c (pX , pY ) = infψ,φ ∫ X c(x, ψ(φ(x)))pX(x)dx
where Γψ,φ := {γ| ∫ Y γ(x, y)dy = pX(x), ∫ X γ(x, y)dx = ∫ X pX(x)δ(y − ψ(φ(x)))dx}. Figure 7 demonstrates a simple scenario were the Wasserstein distance, Wc(pX , pY ), is zero however, W ‡c (pX , pY ) is non-zero. Finally, we note that ψ(φ(·)) = id(·) is a global optima for both Wc(pX , pY ) and W ‡c (pX , pY ).

heading: SLICING EMPIRICAL DISTRIBUTIONS
text: Following equation 10 a distribution can be sliced via: RpX(t, θ) = ∫ X pX(x)δ(t− θ · x)dx
Figure 8 visualizes two sample slices for an example distribution pX . Here we calculate a Radon slice of the empirical distribution pX(x) = 1M ∑M m=1 δ(x− xm) with respect to θ ∈ Sd−1. Using the definition of the Radon transform in equation 10 and RVT in equation 1 we have:
RpX(t, θ) = 1
M M∑ m=1 ∫ X δ(x− xm)δ(t− θ · x)dx
= 1
M M∑ m=1 δ(t− θ · xm)

heading: DIMENSIONALITY OF THE LATENT SPACE
text: 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.

heading: CALCULATING THE SLICED WASSERSTEIN DISTANCE AS A MEASURE OF GOODNESS OF FIT
text: 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 ∼ qZ and x ∼ pX , 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:
argmaxθ θTST θ
θT (SW + αI)θ
s.t. ‖θ‖ = 1
where SW is the within class covariance matrix, ST is the data covariance matrix, I is the identity matrix, and α identifies the interpolation between PCA and LDA (i.e. α = 0 leads to LDA and α→∞ leads to PCA).

heading: ERROR ANALYSIS OF WASSERSTEIN DISTANCE
text: For p ≥ 1 we can use the triangle inequality and write
Wp(pX , pY ) ≤ Wp(pX , pX,M ) +Wp(pY , pX,M ) ≤ Wp(pX , pX,M ) +Wp(pY , pY,M ) +Wp(pX,M , pY,M )
which leads to
Wp(pX , pY )−Wp(pX,M , pY,M ) ≤ Wp(pX , pX,M ) +Wp(pY , pY,M )
Taking the expectation of both sides of the inequality and using the empirical convergence bounds of Wp (in this case W1) we have,
E(W1(pX , pY )−W1(pX,M , pY,M )) ≤ E(W1(pX , pX,M )) + E(W1(pY , pY,M ))
≤ C√ M
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).
Fi gu
re 11
:S am
pl e
co nv
er ge
nc e
be ha
vi or
fo ro
ur m
et ho
d co
m pa
re d
to th
e W
A E
-G A
N To
ls tik
hi n
et al
.( 20
17 ),
w he
re q Z
is se
tt o
th e
ri ng
di st
ri bu
tio n.
T he
co lu
m ns
re pr
es en
tb at
ch ite
ra tio
ns (b
at ch
si ze
= 50
0) .T
he to
p ha
lf of
th e
ta bl
e sh
ow s
re su
lts of ψ
(z )
fo rz ∼ q Z
,a nd
th e
bo tto
m ha
lf sh
ow s z ∼ q Z
an d φ
(x )
fo rx ∼ p X
.I t
ca n
be se
en th
at SW
A E
pr ov
id es
a su
pe ri
or m
at ch
be tw
ee n φ
(x )
(i .e
.p Z
)a nd q Z
w hi
le be
in g
le ss
co m
pu ta
tio na
lly ex
pe ns
iv e
an d
no tr
eq ui
ri ng
ad ve
rs ar
ia lt
ra in
in g.
Fi gu
re 12
:S am
pl e
co nv
er ge
nc e
be ha
vi or
fo ro
ur m
et ho
d co
m pa
re d
to th
e W
A E
-G A
N To
ls tik
hi n
et al
.( 20
17 ),
w he
re q Z
is se
tt o
th e
un if
or m
po la
rd is
tr ib
ut io
n. T he co lu m ns re pr es en tb at ch ite ra tio ns (b at ch si ze = 50 0) .T he to p ha lf of th e ta bl e sh ow s re su lts of ψ (z ) fo rz ∼ q Z ,a nd th e bo tto m ha lf sh ow s z ∼ q Z an d φ (x ) fo r x ∼ p X .It ca n be se en th at SW A E pr ov id es a su pe rio rm at ch be tw ee n φ (x ) (i. e. p Z )a nd q Z w hi le be in g le ss co m pu ta tio na lly ex pe ns iv e an d no tr eq ui rin g ad ve rs ar ia l tr ai ni ng .

Review summaries:
The authors propose a new autoencoding algorithm for the unsupervised generative modeling which they call Sliced Wasserstein Autoencoders (SWAE). SWAE minimizes a reconstruction cost (measured with respect to the non-negative cost function c(x,x') defined for pairs of input images x, x'), regularized by a penalty measuring a discrepancy between the prior distribution over the latent space qz and the push-forward pz of the unknown data distribution through the deterministic encoder. The authors present an extensive theoretical argument supporting the choice of this objective and a number of empirical results performed on MNIST, LSUN bedrooms, and Celeba. 

This paper presents an extension of Wasserstein autoencoder (WAE) by modifying the regularization term in learning objective of variational autoencoder. This term measures the divergence between the distribution of the encoded training samples and the samplable prior distribution. The modification is based on the sliced-Wasserstein distance where the distance between two distributions is measured through slicing or projecting the high-dimensional distributions into one-dimensional marginal distributions. As a result, a closed-form solution to the integral in Eq. (9) is obtained via a numerical method. The adversarial learning in WAE, designed to fulfill the calculation of high-dimensional distance, can be avoided. In general, this is an interesting work by introducing new idea of sliced-Wasserstein distance.