S LICED -W ASSERSTEIN A UTO -E NCODERS
A BSTRACT
In this paper we use the geometric properties of the optimal transport (OT) problem and the Wasserstein distances to deﬁne 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 speciﬁed.
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 efﬁcient numerical solution that provides similar capabilities to
Wasserstein Auto-Encoders (WAE) and Variational Auto-Encoders (VAE), while beneﬁting 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.
; Murez et al.
, super-resolution Ledig et al.
Kolouri & Rohde , 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.
Radford et al.
Arjovsky et al.
; Berthelot et al.
and Variational auto-encoders (VAE) Kingma & Welling (2013);
Mescheder et al.
; 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.
; Gulrajani et al.
; Bousquet et al.
(2017); Arjovsky et al.
; Kolouri et al.
(2017) due to their exquisite geometric characteristics
Santambrogio (2015).
See the supplementary material for an intuitive example showing the beneﬁt 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 speciﬁcally the follow-up work by Tolstikhin et al.
However, our approach avoids the need to perform adversarial training in the encoding space and is not restricted to closed-form distributions, while still beneﬁting 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 ¸Sim¸sekli 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.
¸Sim¸sekli et al.
(2018), on the other hand, take an interesting but different approach of parameter-free generative modeling via sliced-Wasserstein ﬂows.
N OTATION AND P RELIMINARIES
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 deﬁned on X .
We deﬁne the probability density function p X ( x ) for input data x to be:
[Equation]
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 simpliﬁes 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 brieﬂy review the existing dissimilarity measures for these distributions.
1. M INIMIZING D ISSIMILARITY BETWEEN p X AND p Y
We ﬁrst 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, JensenShannon, 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.
; Makhzani et al.
; Mescheder et al.
Others have utilized the rich mathematical foundation of the OT problem and Wasserstein distances
Arjovsky et al.
Gulrajani et al.
; Bousquet et al.
; Tolstikhin et al.
(2017) to deﬁne 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):
[Equation]
Furthermore, the r.h.s.
of equation 3 supports a simple implementation where for i.i.d samples of the input distribution, { x n } Nn =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. M INIMIZING D ISSIMILARITY 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:
[Equation]
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.
; Kolouri et al.
(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 } Nj =1 from q Z and Z = { z n = φ ( x n ) } Nn =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 ˜
z
j s and z n s.
In their inﬂuential 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) deﬁnes a discriminator network, D Z ( p Z , q Z ) , to classify ˜
z
j 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-deﬁnite 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.
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.
; Kolouri et al.
(2016b)
; Carriere et al.
; 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 efﬁcient, 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 } Mn =1 and { y n = ψ ( φ ( x n )) }
Mn =1 , and the problem simpliﬁes to matching the lower-dimensional point clouds { φ ( x n ) }
Mn =1 and { ˜ z
j ∼ q Z } 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.
P ROPOSED METHOD
In what follows we ﬁrst 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. W ASSERSTEIN 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 deﬁned 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]
[Table]
[Equation] [Table] [Equation] where det ( D · ) is the determinant of the Jacobian, then the Wasserstein distance could be deﬁned based on the Monge formulation of the problem (see Villani (2008) and Kolouri et al.
as:
[Equation]
[Equation]
where MP is the set of all diffeomorphisms that satisfy equation 7.
As can be seen from equationand equation 8, obtaining the Wasserstein distance requires solving an optimization problem.
We note that various efﬁcient optimization techniques have been proposed in the past (e.g. Cuturi (2013);
Solomon et al.
; 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]
[Equation]
This closed-form solution motivates the deﬁnition of sliced-Wasserstein distances.
2. S LICED -W ASSERSTEIN 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 deﬁne positive deﬁnite kernels for distributions and Carriere et al.
(2017) to deﬁne a kernel for persistence diagrams.
Sliced-Wasserstein was recently used for learning Gaussian mixture models in Kolouri et al.
and it was also used as a measure of goodness of ﬁt for
GANs in Karras et al.
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 ﬁeld of
Integral Geometry and speciﬁcally 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 deﬁned as:
[Equation]
where S d − 1 is the d-dimensional unit sphere.
Note that for any ﬁxed θ ∈ S d − 1 , R p X ( · ; θ ) is a one-dimensional slice of distribution p X .
In other words, R p 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 deﬁned as:
[Equation]
Given that R p X ( · ; θ ) and R p 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 (cid:96) 22 , c ( x, y ) = (cid:107) x − y (cid:107) 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 ) = (cid:107) x − y (cid:107) pp for p ≥ 2 , the following upper bound hold for the SW distance:
[Equation]
12) where, α d,p = 1 d (cid:82) S d − 1 (cid:107) θ (cid:107) pp dθ ≤ 1 .
Chapter 5 in Bonnotte (2013) proves this inequality.
In our paper, we are interested in p = 2 , for which α p,d = 1 d , and we have:
[Equation]
In the Numerical Implementation Section, we provide a numerical experiment to compare W 2 and
SW 2 , that conﬁrms the above equation.
2. S LICED -W ASSERSTEIN A UTO -E NCODER (SWAE)
Our proposed formulation for the SWAE is as follows:
[Equation]
[Equation]
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 predeﬁned 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 ˜ z
i 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 .
N UMERICAL IMPLEMENTATION
We now describe the numerical details of our approach.
3. N UMERICAL IMPLEMENTATION OF THE W ASSERSTEIN 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 (cid:80) Mm =1 a m , where a m =
c ( P − 1 X ( τ m ) , P − 1 Y ( τ m )) and τ m = 2 m − 1 2 M (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 (cid:80) Mm
=
δ x m and p Y ≈ p Y,M = 1 M (cid:80)
Mm =1 δ y m , where δ x m 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 (cid:80) Mm =1 u ( t − x m ) where u ( . ) is the step function ( P Y,M ( t ) is deﬁned 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:
Prob (sup t | P X,M ( t ) − P X ( t ) | > (cid:15) ) ≤ 2 exp ( − 2 M(cid:15) 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 ﬁrst sorting x m s and y m s and then calculating:
[Equation]
The problem of calculating the Wasserstein distance between samples from one-dimensional densities simpliﬁes to solving two sorting problems (solved in O ( M ) /
O ( Mlog ( M )) best/worst case).
We need to address one ﬁnal question here.
How well does equation 15 approximate the Wasserstein distance, W c ( p X , p Y ) ?
We ﬁrst 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.
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 ( W 1 ( p X,M , p X ) ≤ C √ M where C is an absolute constant.
Similar results are achieved for
E ( W p ( p X,M , p X )) and ( E ( W pp ( p X,M , p X )))
1 p , although under more strict assumptions on p X (i.e., slightly stronger assumptions than having a ﬁnite 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. S LICING 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 (cid:80) Mm =1 δ x m .
Following equation 10 it is straightforward to show that the marginal densities (i.e. slices) are obtained from:
[Equation]
see the supplementary material for a proof.
The Dvoretzky–Kiefer–Wolfowitz upper bound holds for
R p X ( t, θ ) and R p X,M ( t, θ ) .
3. M INIMIZING SLICED -W ASSERSTEIN 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 [Equation] -sample [Equation]
[Equation]
Such Monte Carlo estimation was used in Rabin & Peyré (2011), and later used in Bonneel et al.
Kolouri et al.
(2018); ¸Sim¸sekli et al.
; Deshpande et al.
Moreover, the global minimum for SW c ( p Z , q Z ) is also a global minimum for each
[Equation]
Note that
[Equation]
A ﬁne 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 ﬁnite 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]
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 ( SW 2 ( p X,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. P UTTING 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
Mm
=
and { ˜ z m ∼ q Z } Mm =1 be i.i.d random samples from the input data and the predeﬁned distribution, q Z , correspondingly.
Let { θ l } Ll = 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 efﬁciently), 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.
Require: Regularization coefﬁcient λ , and number of random projections, L .
Initialize the parameters of the encoder, φ , and decoder, ψ while φ and ψ have not converged do
Sample { x 1 , ..., x M } from training set (i.e. p X ) Sample { ˜ z 1 , ..., ˜ z M } from q Z
Sample { θ 1 , ..., θ L } from S K −
Sort θ l · ˜ z M such that θ l · ˜ z i [ m ] ≤ θ l · ˜ z
i [ m +1] Sort θ l · φ ( x m ) such that θ l · φ ( x j [ m ] ) ≤ θ l · φ ( x j [ m +1] )
Update φ and ψ by descending: (cid:80) Mm =1 c ( x m , ψ ( φ ( x m )))
+ λ (cid:80)
Ll =1 (cid:80)
Mm =1 c ( θ l · ˜ z
i [ m ] , θ l · φ ( x j [ m ] )) end while
E XPERIMENTS
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 rectiﬁed 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.
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) (speciﬁcally 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 signiﬁcant 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 WAEGAN took 0 .
2571 ± 0 .
0435 (sec) while SWAE (with L = 50 projections) took 0 .
2437 ± 0 .
0391 (sec).
[Table] [Table]
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 sufﬁce 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.
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.
C ONCLUSIONS
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 speciﬁed.
In addition, we provided a simple and efﬁcient 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.
Note that the far right visualization demonstrates the decoding of a 25 × 25 grid in [ − 1 , 1] 2 .
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.
S UPPLEMENTARY M ATERIAL C OMPARISON OF DIFFERENT DISTANCES
Following the example by Arjovsky et al.
(2017) and later Kolouri et al.
here we show a simple example comparing the Jensen-Shannon divergence with the Wasserstein distance.
First note that the Jensen-Shannon divergence is deﬁned as,
[Equation]
where KL ( p, q ) = (cid:82) 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.
L OG - LIKELIHOOD
To maximize (minimize) the similarity (dissimilarity) between p Z and q Z , we can write :
[Equation]
where we replaced p Z with equation 1.
Furthermore, it is straightforward to show:
[Equation]
P ROOF OF EQUATION
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]
where Γ := { γ : X × Y → R + | (cid:82) Y γ ( x, y ) dy =
p X ( x ) , (cid:82) X γ ( x, y ) dx
= p Y ( 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:
[Equation]
which proves equation 3.
Finally, taking the inﬁmum of the two sides of the inequality, with respect to φ and ψ , we have:
[Equation]
where Γ ψ,φ := { γ | (cid:82) Y γ ( x, y ) dy =
p X ( x ) , (cid:82) X γ ( x, y ) dx =
(cid:82) X p X
( x ) δ ( y − ψ
( φ ( x ))) dx } .
Figure 7 demonstrates a simple scenario were the Wasserstein distance, W c ( p X , p Y ) , is zero however, 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 ) .
S LICING EMPIRICAL DISTRIBUTIONS
Following equation 10 a distribution can be sliced via:
[Equation]
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 (cid:80) Mm =1 δ ( x − x m ) with respect to θ ∈ S d − 1 .
Using the deﬁnition of the Radon transform in equation 10 and RVT in equation 1 we have:
[Equation]
D IMENSIONALITY 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.
C ALCULATING THE SLICED W ASSERSTEIN DISTANCE AS A MEASURE OF GOODNESS OF FIT
In this paper we also used the sliced Wasserstein distance as a measure of goodness of ﬁt (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:
[Equation]
The parameter t ∈ [0 , 1] indicates the interpolation parameter.
where S W is the within class covariance matrix, S T is the data covariance matrix, I is the identity matrix, and α identiﬁes the interpolation between PCA and LDA (i.e. α = 0 leads to LDA and α → ∞ leads to PCA).
E RROR ANALYSIS OF W ASSERSTEIN DISTANCE
For p ≥ 1 we can use the triangle inequality and write
[Equation]
which leads to
[Equation]
Taking the expectation of both sides of the inequality and using the empirical convergence bounds of
W p (in this case W 1 ) we have,
[Equation]
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).
F i gu r e 11 :
S a m p l e c onv e r g e n ce b e h a v
i o r f o r ou r m e t hod c o m p a r e d t o t h e
W A
EGAN
T o l s ti kh i n e t a l .
( ) , w h e r e q
Z
i s s e t t o t h e r
i ng
d
i s t
r
i bu ti on .
T h e c o l u m n s r e p r e s e n t b a t c h
i t e r a ti on s ( b a t c h s i ze = 500 ) .
T h e t op h a l f o f t h e t a b l e s ho w s r e s u lt s o f ψ ( z ) f
o r z ∼ q Z , a nd t
h e b o tt o m h a l f s ho
w s z ∼ q Z a nd φ ( x ) f
o r x ∼ p X .
I t ca n b e s ee n t h a t
S W
A E p r ov i
d
e s a s up e r
i o r m a t c h b e t w ee n φ ( x ) ( i .
e .
p Z ) a nd q Z w h il e b e i ng l e ss c o m pu t a ti on a ll y e xp e n s i v e a nd no t r e qu i r
i ng a dv e r s a r i a l t r a i n i ng .
F i gu r e 12 :
S a m p l e c onv e r g e n ce b e h a v
i o r f o r ou r m e t hod c o m p a r e d t o t h e
W A
EGAN
T o l s ti kh i n e t a l .
( ) , w h e r e q
Z
i s s e t t o t h e un
i f
o r m po l a r d i s t
r
i bu ti on .
T h e c o l u m n s r e p r e s e n t b a t c h it e r a ti on s ( b a t
c h s i ze = 500 ) .
T h e t op h a l f o f t h e t a b l e s ho w s r e s u lt s o f ψ ( z ) f
o r z ∼ q Z , a nd t h e bo tt o m h a l f s ho
w s z ∼ q Z a nd φ ( x ) f
o r x ∼ p X .
I t ca n b e s ee n t h a t
S W
A E p r ov i
d
e s a s up e r
i o r m a t c h b e t w ee n φ ( x ) ( i .
e .
p Z ) a nd q Z w h il e b e i ng l e ss c o m p u t a ti on a ll y e xp
e n s i v e a nd no t r e qu
i r i ng a dv e r s a r i a l t r a i n i ng .