Year: 2018

Abstract:
Optimal Transport offers an alternative to maximum likelihood for learning generative autoencoding models. We show how this principle dictates the minimization of the Wasserstein distance between the encoder aggregated posterior and the prior, plus a reconstruction error. We prove that in the non-parametric limit the autoencoder generates the data distribution if and only if the two distributions match exactly, and that the optimum can be obtained by deterministic autoencoders. We then introduce the Sinkhorn AutoEncoder (SAE), which casts the problem into Optimal Transport on the latent space. The resulting Wasserstein distance is minimized by backpropagating through the Sinkhorn algorithm. SAE models the aggregated posterior as an implicit distribution and therefore does not need a reparameterization trick for gradients estimation. Moreover, it requires virtually no adaptation to different prior distributions. We demonstrate its flexibility by considering models with hyperspherical and Dirichlet priors, as well as a simple case of probabilistic programming. SAE matches or outperforms other autoencoding models in visual quality and FID scores.

heading: 1 INTRODUCTION
text: Unsupervised learning aims to find the underlying rules that govern a given data distribution. It can be approached by learning to mimic the data generation process, or by finding an adequate representation of the data. Generative Adversarial Networks (GAN) (Goodfellow et al., 2014) belong to the former class, by learning to transform noise into a distribution that matches the given one. AutoEncoders (AE) (Hinton & Salakhutdinov, 2006) are of the latter type, by learning a representation that maximizes the mutual information between the data and its reconstruction, subject to an information bottleneck. Variational AutoEncoders (VAE) (Kingma & Welling, 2013; Rezende et al., 2014), provide both a generative model — i.e. a prior distribution on the latent space with a decoder that models the conditional likelihood — and an encoder — approximating the posterior distribution of the generative model. Optimizing the exact marginal likelihood is intractable in latent variable models such as VAE’s. Instead one maximizes the Evidence Lower BOund (ELBO) as a surrogate. This objective trades off a reconstruction error of the input and a regularization term that aims at minimizing the Kullback-Leibler (KL) divergence from the approximate posterior to the prior.
An alternative principle for learning generative autoencoders is proposed by Tolstikhin et al. (2018). The theory of Optimal Transport (OT) (Villani, 2008) prescribes a different regularizer: one that matches the prior with the aggregated posterior — the average (approximate) posterior over the training data. In Wasserstein AutoEncoders (WAE) (Tolstikhin et al., 2018), this is enforced by the choice of either the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012)), or by adversarial training on the latent space. WAE empirically improves upon VAE. More recently, a family of Wasserstein divergences has been used by Ambrogioni et al. (2018) in the context of variational inference. The particular choice of Wasserstein distances may be crucial for convergence, due to the induced weaker topology as compared to other divergences, such as the KL (Arjovsky et al., 2017).
We contribute to the formal analysis of autoencoders with OT. First, we prove that in order to minimize the Wasserstein distance between the generative model and the data distribution, we can miniminize the usual reconstruction-plus-regularizer cost, where the regularizer is the Wasserstein distance between the encoder aggregated posterior and the prior. Second, in the non-parametric limit, the model learns the data distribution if and only if the aggregated posterior matches the prior exactly. Third, as a consequence of the Monge-Kontorovich equivalence (Villani, 2008), the functional space of this learning problem can be limited to that of deterministic autoencoders.
The theory supports practical innovations. We learn deterministic autoencoders by minimizing a reconstruction error and the Wasserstein distance on the latent space between samples of the aggregated posterior and the prior. The latter is known to be costly, but a fast approximate solution is provided by the Sinkhorn algorithm (Cuturi, 2013). We follow Frogner et al. (2015) and Genevay et al. (2018), by exploiting the differentiability of the Sinkhorn iterations, and unroll it for backpropagation. Altogether, we call our method the Sinkhorn AutoEncoder (SAE).
The Sinkhorn AutoEncoder is agnostic to the analytical form of the prior, as it optimizes a samplebased cost function which is aware of the geometry of the latent space. Furthermore, as a byproduct of using deterministic networks, it models the aggregated posterior as an implicit distribution (Mohamed & Lakshminarayanan, 2016) with no need of the reparametrization trick for learning the encoder (Kingma & Welling, 2013). Therefore, with essentially no change in the algorithm, we can learn models with Normally distributed priors and aggregated posteriors, as well as distributions living on manifolds such as hyperspheres (Davidson et al., 2018) and probability simplices.
We start our experiments by studying unsupervised representation learning by training an encoder in isolation. Our results demonstrate the capability of the Sinkhorn algorithm to produce embeddings that conserve the local geometry of the data, echoing results from Bojanowski & Joulin (2017). Next we move to the autoencoder. In an ablation study, we compare with the exact Hungarian algorithm in place of the Sinkhorn and show that our method performs equally well, while converging faster. We then compare against prior work on autoencoders with Normal and spherical priors on MNIST, CIFAR10 and CelebA. SAE with a spherical prior produces visually more appealing interpolations, crisper samples and comparable or lower FID (Heusel et al., 2017). Finally, we further show the flexibility of SAE with qualitative results by using a Dirichlet prior, which defines the latent space on a probability simplex, as well as with a simple probabilistic programming task.

heading: 2 BACKGROUND

heading: 2.1 WASSERSTEIN DISTANCE AND WASSERSTEIN AUTOENCODERS
text: We follow Tolstikhin et al. (2018) and denote with X ,Y,Z the sample spaces and with X,Y, Z and PX , PY , PZ the corresponding random variables and distributions. Given a map F : X → Y we denote by F# the push-forward map acting on a distribution P as P ◦ F−1. If F (Y |X) is non-deterministic we define the push-forward of a distribution P as the induced marginal of the joint distribution F (Y |X)PX (denoted by F (Y |X)#PX ). For any measurable non-negative cost c : X × Y → R+ ∪ {∞}, one can define the following OT-cost between marginal distributions PX and PY via:
Wc(PX , PY ) = inf Γ∈Π(PX ,PY ) E(X,Y )∼Γ[c(X,Y )], (1)
where Π(PX , PY ) is the set of all joint distributions that have as marginals the given PX and PY . The elements from Π(PX , PY ) are called couplings from PX to PY . From now on we will assume that X = Y and c(x, y) is a distance. In this case Wc(PX , PY ) is the Wasserstein distance w.r.t the cost c. If c(x, y) = ‖x− y‖pp for p ≥ 1 then Wp = p √ Wc is called the p-th Wasserstein distance.
Let PX denote the true data distribution on X . We define a latent variable model given as follows: we fix a latent space Z and a prior distribution PZ on Z and consider the conditional distribution G(X|Z) (the decoder) parameterized by a neural network G. Together they specify a generative model as G(X|Z)PZ . The induced marginal will be denoted by PG. Learning PG to approximate the true PX is then defined as:
min G Wc(PX , PG).
Because of the infimum over Π(PX , PG) inside Wc, this is intractable. To rewrite this objective we consider the posterior distribution Q(Z|X) (the encoder) and its aggregated posterior QZ :
QZ = Q(Z|X)#PX = EX∼PXQ(Z|X), (2)
the induced marginal of the joint distribution Q(Z|X)PX . Tolstikhin et al. (2018) show that, if the decoder G(X|Z) is deterministic, i.e. PG = G#PZ , or in other words, if all stochasticity of the generative model is captured by Z, then:
Wc(PX , PG) = inf Q(Z|X): QZ=PZ EX∼PXEZ∼Q(Z|X)[c(X,G(Z))]. (3)
Learning the generative model G with the Wasserstein AutoEncoder amounts to: min G min Q(Z|X) EX∼PXEZ∼Q(Z|X)[c(X,G(Z))] + β ·DZ(QZ , PZ), (4)
where β > 0 is a Lagrange multiplier andDZ is any divergence measure on probability distributions on Z , which choice is left open. WAE uses either MMD or a discriminator trained adversarially for DZ . As discussed in Bousquet et al. (2017), Equation 4 is a lower bound of Equation 3 for any value of β > 0. Minimizing this lower bound does not ensure a minimization of the original objective of Equation 3.

heading: 2.2 THE SINKHORN ALGORITHM
text: In place of any choice of DZ , in Section 3 we formally support the minimization of a Wasserstein distance on latent space. The distance is notoriously hard to compute, which is the reason why the rewriting of Equation 3 is of practical interest. When restricting to discrete distributions, the problem becomes more amenable and efficient approximations exist. To motivate this direction, recall that we can always see samples of a continuous distribution as Dirac deltas, whose expectation defines a discrete distribution. Let two discrete distributions with support on M points be P̂ = 1 M ∑M i=1 δzi , Q̂ = 1 M ∑M i=1 δz′i . Given a cost c ′, their (empirical) Wasserstein distance is:
Wc′(Q̂, P̂ ) = min R∈SM
1 M 〈R,C ′〉F , (5)
where C ′ij = c ′(z′i, zj) is the matrix associated to the cost c ′, R is a doubly stochastic matrix as defined in SM = {R ∈ RM×M≥0 | R1 = 1, RT 1 = 1}, and 〈·, ·〉F denotes the Frobenius inner product; 1 is the vector of ones. Eq. (5) is known to converge to the Wasserstein distance between the continuous distributions as M tends to infinity (Weed & Bach, 2017). This linear program has solutions on the vertices of SM , which is the set of permutation matrices (Peyré & Cuturi, 2018). The Hungarian algorithm finds an optimal solution in O(M3) time (Kuhn, 1955).
An entropy-regularized version of problem (5) can be solved more efficiently. Let the entropy of R be H(R) = − ∑M i,j=1Ri,j logRi,j . For ε > 0, Cuturi (2013) defines the Sinkhorn distance Sc′ :
R∗ = arg min R∈SM 1 M 〈R,C ′〉F − εH(R), Sc′(Q̂, P̂ ) = 〈R∗, C ′〉F , (6)
and shows that the (Sinkhorn, 1964)’s algorithm returns its regularized optimum — that is also unique due to strong convexity of the entropy. The Sinkhorn is a fixed point algorithm that runs nearly in M2 time (Altschuler et al., 2017) and can be efficiently implemented with matrix multiplications; see Algorithm 1. Its convergence to the Wasserstein distance is studied by Weed (2018).
The smaller the ε, the smaller the entropy and the better the approximation of the Wasserstein distance. At the same time, a larger number of steps O(L) is needed to converge. Conversely, high entropy encourages the solution to lie far from a permutation matrix. Note that all Sinkhorn operations are differentiable. So when the distance is used as a cost function, we can unroll O(L) iterations and backpropagate (Genevay et al., 2018). In conclusion, we obtain a differentiable surrogate for Wasserstein distances between empirical distributions; the approximation arises from sampling, entropy regularization and the finite amount of steps in place of convergence.

heading: 2.3 NOISE AS TARGETS
text: Bojanowski & Joulin (2017) introduce Noise As Targets (NAT), an algorithm for unsupervised representation learning. The method learns a neural network fθ by embedding images into a uniform hypersphere. A sample z is drawn from the sphere for each training image and fixed. The goal is to learn θ such that 1-to-1 matching between images and samples is improved: matching is coded with a permutation matrix R, and updated with the Hungarian algorithm. The objective is:
max θ max R∈PM
Tr(RZfθ(X) >), (7)
where Tr(·) is the trace operator, Z and X are respectively prior samples and images stacked in a matrix and PM ⊂ SM is the set of M -dimensional permutations. NAT learns by alternating SGD and the Hungarian. One can interpret this problem as supervised learning, where the samples are targets (sampled only once) but their assignment is learned; notice that freely learnable Z would make the problem ill-defined. The authors relate NAT to OT, a link that we make formal below.

heading: 3 PRINCIPLES OF WASSERSTEIN AUTOENCODING
text: With Equation 3, Tolstikhin et al. (2018) reformulate the Wasserstein distance in image space in terms of autoencoders. The hard constraint QZ = PZ is in practice replaced with a soft constraint by adding a penalty in the form of a divergence DZ(QZ , PZ). The resulting objective (4) is a lower bound to the Wasserstein difference and the choice of a divergence DZ is left open. In contrast, we show that one should opt for minimizing a Wasserstein distance in latent space and that this leads to an equality with — not a bound for — the original Wasserstein distance in image space.
More precisely, Theorem 3.1 first proves that the Wasserstein distance between the generative model and data distribution is bounded from above by a quantity consisting of the reconstruction error and the Wasserstein distance between PZ and QZ . Theorem 3.2 shows that we can restrict learning to the class of deterministic (auto)encoders. Put together, Corollary 3.3 provides a principled learning objective in the framework of Optimal Transport by rewriting the Wasserstein distance in image space into an equivalent tractable form. We start with the following bound:
Theorem 3.1. If G(X|Z) is deterministic and γ-Lipschitz then:
Wp(PX , PG) ≤Wp(PX , G#QZ) + γ ·Wp(QZ , PZ).
The proof (A.2) exploits the triangle inequality of the Wasserstein distance and its behaviour under composition with Lipschitz maps – a property not shared with divergences such as the KL. To effectively minimize the right-hand side in Theorem 3.1 over a class of encoders we need to further upper bound the reconstruction term with the following1:
Wp(PX , G#QZ) ≤ p √ EX∼PXEZ∼Q(Z|X)EX′∼G(X|Z)[‖X −X ′‖ p p], (8)
which reduces to the p-th root of EX∼PX [‖X−G(Q(X))‖pp] if bothG andQ are deterministic. The tightness of this bound and its use as an objective function for learning are discussed below.
We now improve the characterization of Equation 3, which is formulated in terms of stochastic encoders Q(Z|X) and deterministic decoders G(X|Z). In fact, it is possible to restrict the learning class to that of deterministic autoencoders:
Theorem 3.2. Let PX be not atomic2 and G(X|Z) deterministic. Then for every continuous cost c:
Wc(PX , PG) = inf Q(Z|X) deterministic: QZ=PZ EX∼PX [c(X,G(Q(X)))].
Using the cost c(x, y) = ‖x− y‖pp, the equation holds with W pp (PX , PG) in place of Wc(PX , PG).
The statement is a direct consequence of the equivalence between the Kantorovich and Monge formulations of OT (Villani, 2008); see the proof in A.3. We remark that this result is stronger than, and can be used to deduce Equation 3; see A.4 for a proof. Combining the two previous results, we are now in position to prove that the bound in Theorem 3.1 is tight for deterministic (auto)encoders:
Wp(PX , PG) Th.3.1 ≤ inf
Q det.
p √ EX∼PX [‖X −G(Q(X))‖ p p] + γ ·Wp(QZ , PZ) (9)
≤ inf Q det.,QZ=PZ p
√ EX∼PX [‖X −G(Q(X))‖
p p] + γ ·Wp(QZ , PZ)︸ ︷︷ ︸
=0
(10)
Th.3.2 = Wp(PX , PG). (11)
Inequality in Step 10 holds because we restrict the domain of the infimum, which in turns implies Wp(QZ , PZ) = 0. As a consequence we obtain the following Corollary, which provides us with an objective for learning generative autoencoders:
1This is from the fact that (idX , G)#Q(Z|X)PX ∈ Π(PX , G#QZ). 2A probability measure is non-atomic if every point in its support has zero measure. It is important to distin-
guish between the empirical data distribution P̂X , which is always atomic, and the underlying true distribution PX , only which we need to assume to be non-atomic.
Corollary 3.3. Let PX be non-atomic and G(X|Z) be deterministic and γ-Lipschitz. Then we have the equality:
Wp(PX , PG) = inf Q(Z|X) deterministic
p √ EX∼PX [‖X −G(Q(X))‖ p p] + γ ·Wp(QZ , PZ). (12)
More precisely, we can now formulate our learning problem as the minimization of the right-hand side of Equation 12 over deterministic decoders:
min G min Q
p √ EX∼PX [‖X −G(Q(X))‖ p p] + γ ·Wp(QZ , PZ) (13)
When the encoder is a neural network of limited capacity, enforcingQZ ≈ PZ might not be feasible in the general case of dimension mismatch (Rubenstein et al., 2018). In fact, since the class of deterministic neural networks is much smaller than the class of deterministic measurable maps, one might consider adding noise to the output, i.e. use stochastic networks instead. Nonetheless, neural networks can approximate any measurable map up to arbitrarily small error (Hornik, 1991), and we prove a related bound for the Wasserstein distance in A.5. It follows that learning deterministic autoencoders is sufficient to approach the theoretical upper bound and thus it will be our empirical choice.
Finally, Theorems 3.1 and 3.2 strengthen the relevance of matching aggregated posterior and prior, which we show to be a sufficient and necessary condition for generative autoencoding. Justified by the previous results, we state it for deterministic autoencoders (proof in A.6).
Theorem 3.4 (Sufficiency and necessity for generative autoencoding). Suppose perfect reconstruction, that is, PX = (G ◦Q)#PX . Then:
i) PZ = QZ =⇒ PX = PG, ii) PZ 6= QZ =⇒ PX 6= PG. (14)
In particular, Theorem 3.4 ii) certifies that, under perfect reconstruction, failing to match aggregated posterior and prior makes learning the data distribution impossible. Matching in latent space should be seen as fundamental as minimizing the reconstruction error, a fact known about the performance of VAE (Hoffman & Johnson, 2016; Higgins et al., 2017; Alemi et al., 2018; Rosca et al., 2018).

heading: 4 SINKHORN AUTOENCODERS
text: In light of our theory, we minimize the Wasserstein distance between the aggregated posterior and the prior, and we do so by running the Sinkhorn on their empirical samples. Let {xi}Mi=1 be the data input to the deterministic encoder Q(z′i|xi) = δz′i and {zi} M i=1 the samples from the prior PZ . The
empirical distributions are Q̂Z = 1M ∑M i=1 δz′i and P̂Z = 1 M ∑M i=1 δzi . With C ′ ij = c(z ′ i, zj), the Sinkhorn distance is Sc′(Q̂Z , P̂Z) as defined in Equation 6.
We compute the Sinkhorn distance in two steps: first obtain the optimal regularized coupling R∗ and then multiply it with the cost, i.e. set ε = 0:
R∗ = arg min R∈SM 1 M 〈R,C ′〉F − εH(R)
Sc′(Q̂Z , P̂Z) = 1 M 〈R ∗, C ′〉F . (15)
See Algorithm 1. Note that we do not sacrifice differentiability: we stackO(L) Sinkhorn operations on top of the encoder, without additional learnable parameters, and run auto-differentiation.
Algorithm 1 SINKHORN Input: {zi}mi=1 ∼ PZ , {z′i}mi=1 ∼ QZ , ε, L ∀i, j, Cij = c(zi, z′j) K = e−C/ε, u← 1 # elem-wise exp repeat L times: v ← 1/(K>u) # elem-wise division u← 1/(Kv)
R∗ ← Diag(u)KDiag(v) Output: 1
M 〈R∗, C〉F
With a deterministic decoder G and encoder Q, we arrive at the objective for the Sinkhorn AutoEncoder (SAE):
min G min Q EX∼P̂X [‖X −G(Q(X))‖ p p] + β · Sc′(Q̂Z , P̂Z). (16)
In practice, we drop the p-th root and tune a β > 0 hyper-parameter to mix the two terms. Small ε and hence large L worsen the numerical stability of the Sinkhorn; thus it is more convenient to scale Sc′ by β to explore the trade off, as in the WAE. In most experiments, both c and c′ will be ‖ · ‖22. This objective is minimized by mini-batch SGD, which requires the re-calculation of an optimal regularized coupling R∗ at each iteration. Experimentally we found that this is not a significant overhead, unless a large L is needed for convergence due to a small ε. In practice, Algorithm 1 loops for L iterations but can exit earlier if the updates of u reach a fixed point.
We have not specified our distribution PZ yet. In fact, SAE can work in principle with arbitrary priors. The only requirement coming from the Sinkhorn is the ability to generate samples. The choice should be motivated by the desired geometric properties of the latent space; Theorem 3.4 stresses the importance of such choice for the generative model. For quantitative comparison with prior work, we focus primarily on hyperspheres, as in the Hyperspherical VAE (HVAE) (Davidson et al., 2018). Moreover, considering the Wasserstein distance (ε = 0) from a uniform hyperspherical prior with squared Euclidean cost, we recover the NAT objective as a special case of ours (see Appendix A.7); yet, our method enjoys lower complexity and differentiability. The remarkable performance of NAT on representation learning on ImageNet confirms the value of the spherical prior. Other distributions are also considered in the paper, in particular the Dirichlet prior — with a tunable bias towards the simplex vertices — as a choice for controlling latent space clustering.
Deterministic encoders model implicit distributions. Distributions are said to be implicit when their probability density may be intractable or even unknown, but it is possible to obtain samples and gradients for their parameters; GANs are examples of models with implicit distributions. Implicit distributions can give more flexibility as they are not limited by families of distributions with tractable density (Mohamed & Lakshminarayanan, 2016; Huszár, 2017). Moreover, by encoding with deterministic neural networks, we bypass the use of reparametrization tricks for gradient estimation.

heading: 5 RELATED WORK
text: The normal prior is common in VAE for the reason of tractability. In fact, changing the prior and/or the approximate posterior distributions requires the use of tractable densities and the appropriate reparametrization trick. A hyperspherical prior is used by Davidson et al. (2018) with improved experimental performance; the algorithm models a Von Mises-Fisher posterior, with a non-trivial posterior sampling procedure and a reparametrization trick based on rejection sampling. Our implicit encoder distribution sidesteps these difficulties; recent advances on variable reparametrization can also simplify these requirements (Figurnov et al., 2018). We are not aware of methods embedding on probability simplices, except the use of Dirichlet priors by the same Figurnov et al. (2018).
Hoffman & Johnson (2016) showed that the objective of a VAE does not force the aggregated posterior and prior to match, and that the mutual information of input and codes may be minimized instead. Just like the WAE, SAE avoids this effect by construction. Makhzani et al. (2015) and WAE improve latent matching by GAN/MMD. With the same goal, Alemi et al. (2017), Tomczak & Welling (2017) introduce learnable priors in the form of a mixture of approximate posteriors, which can be used in SAE as well.
The Sinkhorn (1964) algorithm gained interest after Cuturi (2013) showed its application for fast computation of Wasserstein distances. The algorithm has been applied to ranking (Adams & Zemel, 2011), domain adaptation (Courty et al., 2014), multi-label classification (Frogner et al., 2015), metric learning (Huang et al., 2016) and ecological inference (Muzellec et al., 2017). Santa Cruz et al. (2017); Linderman et al. (2018) used it for supervised combinatorial losses. Our use of the Sinkhorn for generative modeling is akin to that of Genevay et al. (2018), which matches data and model samples with adversarial training, and to Ambrogioni et al. (2018), which matches samples from the model joint distribution and a variational joint approximation. WAE and WGAN objectives are linked respectively to primal and dual formulations of OT (Tolstikhin et al., 2018).
Our approach for training the encoder alone qualifies as self-supervised representation learning (Donahue et al., 2017; Noroozi & Favaro, 2016; Noroozi et al., 2017). As in NAT (Bojanowski & Joulin, 2017) and in constrast to most other methods, we can sample pseudo labels (from the prior) independently from the input. In Appendix A.7 we show a formal connection with NAT.

heading: 6 EXPERIMENTS
text: We start our empirical analysis with a qualitative assessment of the representation learned with the Sinkhorn algorithm. In the rest we focus on the autoencoder. We compare with NAT and confirm the Sinkhorn to be a better choice than the Hungarian. We display interpolations and samples of SAE and compare numerically with AE, (β)-VAE, HVAE and WAE-MMD. We further show the flexibility of SAE by using a Dirichlet prior and on a toy probabilistic programming task.
We experiment on MNIST, CIFAR10 (Krizhevsky & Hinton, 2009) and CelebA (Liu et al., 2015). MNIST is dynamically binarized and the reconstruction error is the binary cross-entropy (although not a distance, it is a commonly used divergence for binary data). For CIFAR10 and CelebA the reconstruction is the squared Euclidean distance; in every experiment, the latent cost is also squared Euclidean. We train fully connected neural networks for MNIST and the convolutional architectures from Tolstikhin et al. (2018) for the rest; the latent space dimensions are respectively 10, 64, 64. We run Adam (Kingma & Ba, 2014) with mini-batches of 128. Hyperspherical embedding is hardcoded in the architectures by L2 normalization of the encoder output as in Bojanowski & Joulin (2017). The Sinkhorn runs with = 0.1, L = 50, except when otherwise stated. FID scores for CIFAR10 and CelebA are calculated as in Heusel et al. (2017), while for MNIST we train a 2-layer convolutional network to extract features for the Fréchet distance, similarly to Odena et al. (2018). Notice that the FID score is a Wasserstein-2 distance and hence our theory applies directly.

heading: 6.1 REPRESENTATION LEARNING WITH SINKHORN ENCODERS
text: We demonstrate qualitatively that the Sinkhorn distance is a valid objective for unsupervised feature learning by showing that we can learn the encoder in isolation. The task is to embed the input distribution in a lower dimensional space, preserving the local data geometry, by solving Problem 14 with no reconstruction cost. We display the representation of a 3D Swiss Roll and MNIST. For the Swiss Roll we set ε = 10−3, while for MNIST it is set to 0.5, and L is picked to ensure convergence. For the Swiss roll (Figure 1a), we use a 50-50 fully connected network with ReLUs.
Figures 1b, 1c show that the local geometry of the Swiss Roll is conserved in the new representational spaces — a square and a sphere. While the global shape is not necessarily more unfolded than the original, it looks qualitatively more amenable for further computation. Figure 1d shows the t-SNE visualization (Maaten & Hinton, 2008) of the learned representation of the MNIST test set. With neither labels nor reconstruction error, we learn an embedding that is aware of class-wise clusters. Minimization of the Sinkhorn distance achieves this by encoding onto a d-dimensional uniform sphere, such that points are encouraged to map far apart; in particular, in high dimension we can prove (see A.8) that the collapse probability decreases with d:
Proposition 6.1. Let z, z′ be two uniform samples from a d-dimensional sphere. In the high dimensional regime, for any δ <
√ 2 we have P (‖z − z′‖2 > δ) ≥ 1− 14d(√2−δ)2 .
Other than this repulsive effect — the uniform distribution has max-entropy on any compact space —, a contractive force is present due to the inductive prior of neural networks, which are known to be Lipschitz functions (Balan et al., 2017). On the one hand, points in the latent space disperse in order to fill up the sphere; on the other hand, points close on image space cannot be mapped too far from
each other. As a result, local distances are conserved while the overall distribution is spread. When the encoder is combined with a decoder G — the topic of the experiments below —, the contractive force strenghtens: they collaborate in learning a latent space which makes reconstruction possible despite finite capacity and hence favours the conservation of local similarities; see Figure 1e.

heading: 6.2 AUTOENCODING WITH THE SINKHORN DISTANCE AND NAT
text: We investigate the advantages of the Sinkhorn with respect to NAT in training autoencoders; this is an ablation study for our method. First, Sinkhorn has a lower complexity than the Hungarian. In both cases, the complexity can be reduced by mini-batch optimization. Yet, training with large minibatches (> 200) becomes quickly impractical with the Hungarian. Second, the differentiability of the Sinkhorn allows us to avoid the alternating minimization and instead backpropagate on the joint parameter space of encoder and doubly stochastic matrices. Third, the Sinkhorn approximates the empirical Wasserstein distance, while the Hungarian is optimal. Last, NAT draws samples once and uses them as targets throughout learning; their assignment to training points is updated by optimizing a permutation matrix over mini-batches and storing the local optimal result. We term NAT in this context Hungarian-targets and our method Sinkhorn-sample. We can design two hybrid methods. Hungarian-sample: a permutation R can be used to compute the cost 〈R,C ′〉F and backpropagate. Sinkhorn-targets: a doubly stochastic matrix R solution of the Sinkhorn can be used for sampling a permutation3 and targets can be re-assigned. We test the impact of these choices experimentally by test set reconstruction error and FID score on MNIST and CIFAR10; we measure latent space mismatch by the MMD with Gaussian kernel over the test set.
Table 1 shows the results. From the FID scores, we conclude that there is no significant difference in generative performance between either Sinkhorn vs. Hungarian, or samples vs. targets. The parameter β trading off reconstruction and latent space cost is more influential than any of these choices. On MNIST, MMD is often lower with fixed targets; this is a sign that the FID does not fully account for all model qualities. Due to the additional overhead of the Hungarian and the targets updating, our algorithm implements the Sinkhorn with mini-batch sampling. In the rest, we also fix β for MNIST and CIFAR as the best found here.

heading: 6.3 COMPARISON WITH OTHER AUTOENCODERS
text: We compare with AE, (β-)VAE, HVAE4 and WAE. Figures 2 shows interpolations and samples of SAE and VAE from CIFAR10 and CelebA. SAE interpolations are defined on geodesics connecting points on the hypersphere. SAE tends to produce crisper images, with higher contrast, and avoids averaging effects as particularly evident in the CelebA interpolations. The CelebA samples are also interesting: while SAE generally maintains a crisper look than VAE’s, faces appear more often malformed. Table 2 reports a quantitative comparison. Each baseline model has a version with normal and spherical prior. FID scores of SAE are on par or superior to that of VAE and consistently better than WAE. The spherical prior appears to reduce FID scores in several cases.

heading: 6.4 DIRICHLET PRIORS
text: We further demonstrate the flexibility of SAE by using Dirichlet priors on MNIST. The prior draws samples on the probability simplex; hence, here we constrain the encoder by a final softmax layer. We use priors that concentrate on the vertices, by the intuition that digits would naturally cluster around them. A 10-dimensional Dir(1/2) prior (Figure 3a) results in an embedding qualitatively similar to the uniform sphere (1e). With a more skewed prior Dir(1/5), we could expect an organization in latent space where each digit is mapped to a vertex, as little mass lies in the center. We
3Obtaining the closest permutation to a double stochastic matrix is costly. We use a stochastic heuristic due to Fogel et al. (2013) that reduces to sorting. We select permutation minimizing 〈C, ·〉F out of 10 draws.
4Comparing with Davidson et al. (2018) in high dimension was unfeasible. The HVAE likelihood requires evaluating the Bessel function, which is computed on CPU. Note that SAE is oblivious to likelihood functions.
found that in dimension 10 this is seldom the case, as multiple vertices can be taken by the same digit to model different styles, while other digits share the same vertex.
We thus experiment with a 16-dimensional Dir(1/5), which yields more disconnected clusters (3b); the effect is evident when showing the prior and the aggregated posterior that tries to cover it (3c). Figure 3d (leftmost and rightmost columns) shows that every digit 0−9 is indeed represented on one of the 16 vertices, while some digits are present with multiple styles, e.g. the 7. The central samples in the Figure are the interpolations obtained by sampling on edges connecting vertices – no real data is autoencoded. Samples from the vertices appear much crisper than other prior samples (3e), a sign of mismatch between prior and aggregated posterior on areas with lower probability mass. Finally, we point out that we could even learn the Dirichlet hyperparameter(s) with a reparametrization trick (Figurnov et al., 2018) and let the data inform the model on the best prior.

heading: 6.5 TOY PROBABILISTIC PROGRAMMING
text: We run a final experiment to showcase that SAE can handle more complex implicit distributions, on a toy example of probabilistic programming. The goal is to learn a generative model for MNIST digits positioned on a larger canvas; the data is corrupted with salt noise that we do not model explicitly and which are model is thus required to ignore. The generative model samples from a factored prior distribution for zwhat — the digit appearance — from a 10-dimensional sphere and for zwhere — the location and scale — from a 3-dimensional Normal. A decoder network is fed with zwhat and generates the digit; the digit is then positioned on the black canvas on the coordinates given by a spatial transformer (Jaderberg et al., 2015) which is fed with zwhere. The inference model produces zwhat, zwhere from the canvas, by using a spatial transformer and a encoder mirroring the generator.
Our autoencoder is fully deterministic. The cost in latent space amounts to the sum of the Sinkhorn distances in the two prior components, Normal and hyperspherical. Figure 4 compares qualitatively with a simplified version of AIR (Eslami et al., 2016), that is built on variational inference with an explicit modelling of the approximate posterior distribution for this program. SAE is able to replicate the behaviour of AIR by locating the digit on the canvas, ignoring the noise in reconstruction and generating realistic samples.

heading: 7 CONCLUSIONS
text: We introduced a new generative model built on the principles of Optimal Transport. Working with empirical Wasserstein distances and deterministic networks provides us with a flexible likelihoodfree framework for latent variable modeling. Besides, the theory suggests improving matching in latent space which could be achieved by the use of parametric implicit prior distributions.

heading: A APPENDIX

heading: A.1 LEMMA
text: As a useful helper Lemma, we prove a Lipschitz property for the Wasserstein distance Wp.
Lemma A.1. For every PX , PY distributions on a sample space S and a Lipschitz map F we have that Wp(F#PX , F#PY ) ≤ γ ·Wp(PX , PY ) , where γ is the Lipschitz constant of F .
Proof. Recall that
Wp(F#PX , F#PY ) p = inf
Γ∈Π(F#PX ,F#PY ) ∫ S×S ‖x− y‖pp dΓ(x, y).
Notice then that for every Γ ∈ Π(PX , PY ) we have that (F × F )#Γ ∈ Π(F#PX , F#PY ). Hence
{(F × F )#Γ : Γ ∈ Π(PX , PY )} ⊂ Π(F#PX , F#PY ). (17)
From (16) we deduce that
Wp(F#PX , F#PY ) p ≤ inf
Γ∈Π(PX ,PY ) ∫ S×S ‖x− y‖pp d(F × F )#Γ
= inf Γ∈Π(PX ,PY ) ∫ S×S ‖F (x)− F (y)‖pp dΓ ≤ γp · (Wp(PX , PY ))p .
Taking the p-root on both sides we conclude.

heading: A.2 PROOF OF THEOREM 3.1
text: Proof. Using the triangle inequality of the Wasserstein distance we obtain
Wp(PX , G#PZ) ≤Wp(PX , G#QZ) +Wp(G#QZ , G#PZ) ≤Wp(PX , G#QZ) + γ ·Wp(QZ , PZ), (18)
where in line (17) we have used Lemma A.1.
In case G is not deterministic, defining
γ = sup P,Q Wp(G(X|Z)#P, G(X|Z)#Q) Wp(P,Q)
exp we can still formulate a bound. The result follows directly from the first line of (17).

heading: A.3 PROOF OF THEOREM 3.2
text: The basic tool to prove Theorem 3.2 is the equivalence between Monge and Kantorovich formulation of optimal transport. For convenience we formulate its statement and we refer to Villani (2008) for a more detailed explanation.
Theorem A.2 (Monge-Kontorovich equivalence). Given PX and PY probability distributions on X such that PX is not atomic, c : X × X → R continuous, we have
Wc(PX , PY ) = inf T :X→X : T#PX=PY
∫ X c(x, T (x)) dPX(x). (19)
We are now in position to prove Theorem 3.2. We will prove it for a general continuous cost c.
Proof. Notice that as the encoder Q(Z|X) is deterministic there exists Q : X → Z such that QZ = Q#PX and Q(Z|X) = δ{Q(x)=z}. Hence
EX∼PXEZ∼Q(Z|X)[c(X,G(Z))] = ∫ X×Z c(x,G(z)) dPX(x)dδ{Q(x)=z}(z)
= ∫ X dPX(x) ∫ Z c(x,G(z))dδ{Q(x)=z}(z)
= ∫ X c(x,G(Q(x))) dPX(x) .
Therefore
inf Q(Z|X) deterministic: QZ=PZ EX∼PXEZ∼Q(Z|X)[c(X,G(Z))] = inf Q:X→Z QZ=PZ
∫ X c(x,G(Q(x))) dPX .
(20)
We now want to prove that
{G ◦Q : Q#PX = PZ} = {T : X → X : T#PX = PG} . (21)
For the first inclusion ⊂ notice that for every Q : X → Z such that QZ = PZ we have that G ◦Q : X → X and
(G ◦Q)#PX = G#Q#PX = G#PZ .
For the other inclusion ⊃ consider T : X → X such that T#PX = PG = G#PZ . We want first to prove that there exists a set A ⊂ X with PX(A) = 1 such that G : Z → T (A) is surjective. Indeed if it does not hold there exists B ⊂ X with PX(B) > 0 and G−1(T (B)) = ∅. Hence
0 = G#PZ(T (B)) = T#PX(T (B)) = PX(B) > 0
that is a contraddiction. Therefore by standard set theory the mapG : Z → T (A) has a right inverse that we denote by G̃. Then define Q = G̃ ◦ T . Notice that G ◦Q = G ◦ G̃ ◦ T = T almost surely in PX and also
(G̃ ◦ T )#PX = PZ .
Indeed for any A ⊂ Z Borel we have
(G̃ ◦ T )#PX(A) = (G̃ ◦G)#PZ(A) = PZ(G̃−1(G−1(A)) = PZ(A) .
This concludes the proof of the claim in (20). Now we have
inf Q:X→Z
Q#(PX)=PZ
∫ X c(x,G(Q(x))) dPX(x) = inf
T :X→X T#(PX)=PG
∫ X c(x, T (x)) dPX(x) .
Notice that this is exactly the Monge formulation of optimal transport. Therefore by Theorem A.2 we conclude that
inf Q(Z|X) deterministic: QZ=PZ EX∼PXEZ∼Q(Z|X)[c(X,G(Z))] = inf Γ∈Π(PX ,PG) E(X,Y )∼Γ[c(X,Y )]
as we aimed.

heading: A.4 TOLSTIKHIN ET AL. (2018)’S THEOREM AS A CONSEQUENCE
text: Proof. Thanks to Theorem 3.2 we have that
Wc(PX , PG) = inf Q(Z|X) deterministic: QZ=PZ EX∼PXEZ∼Q(Z|X)[c(X,G(Z))]
≥ inf Q(Z|X): QZ=PZ EX∼PXEZ∼Q(Z|X)[c(X,G(Z))] .
For the opposite inequality given Q(Z|X) such that PZ = ∫ Q(Z|X)dPX define Q(X,Y ) = PX× [G#Q(Z|X)]. It is a distribution on X ×X and it is easy to check that π1#Q(X,Y ) = PX and π2#Q(X,Y ) = G#PZ , where π
1 and π2 are the projection on the first and the second component. Therefore
{Q(X,Z) : Q(Z|X) such that QZ = PZ} ⊂ Π(PX , PG)
and so
Wc(PX , PG) ≤ inf Q(Z|X):QZ=PZ ∫ X×X c(x, y) dQ(x, y)
= ∫ X [∫ X c(x, y) dG#Q(Z|X)(y) ] dPX
= ∫ X [∫ X c(x,G(z)) dQ(Z|X)(z) ] dPX .

heading: A.5 BOUNDS FOR NEURAL NETWORKS
text: Theorem A.3. Let PX be non-atomic and G(X|Z) be deterministic and γ-Lipschitz. If QNN is a neural network approximating a near optimal deterministic encoder up to an error of ε ≥ 0 in Lp-norm then we have the inequality:
0 ≤ { p √ EX∼PX [‖X −G(QNN(X))‖ p p] + γ ·Wp(QNNZ , PZ) } −Wp(PX , PG) ≤ 3γε.
Proof. Let Q∗ be an optimal measurable deterministic encoder that optimizes the right-hand side of Theorem 3.1 among measurable deterministic encoder (or at least δ ≤ γε close to it) and QNN a
neural network approximation of Q∗ such that p √
EX∼PX [‖Q∗(X)−QNN(X))‖ p p] ≤ ε (existence
by Hornik (1991)). Then we get:
p √ EX∼PX [‖X −G(QNN(X))‖ p p] + γ ·Wp(QNNZ , PZ)
triangle ineq. ≤ p √ EX∼PX [‖X −G(Q∗(X))‖ p p]
+ p √
EX∼PX [‖G(Q∗(X))−G(QNN(X))‖ p p]︸ ︷︷ ︸
≤γ·ε
+ γ ·Wp(QNNZ , Q∗Z)︸ ︷︷ ︸ ≤ε +γ ·Wp(Q∗Z , PZ)
≤ p √ EX∼PX [‖X −G(Q∗(X))‖ p p] + γ ·Wp(Q∗Z , PZ) + 2γε
Def. ofQ∗
≤ inf Q
{ p √ EX∼PX [‖X −G(Q(X))‖ p p] + γ ·Wp(QZ , PZ) } + δ + 2γε
3.2 ≤ Wp(PX , PG) + 3γε,
where in the last inequality we use additionally that the upper bound in Theorem 3.1 is sharp.
Finally, we can also formulate a version of Corollary 3.3 restricted to deterministic neural networks as follows:
Theorem A.4. Let PX be non-atomic and G(X|Z) be deterministic and γ-Lipschitz. Then we have the equality:
Wp(PX , PG) = inf Q NN
p √ EX∼PX [‖X −G(Q(X))‖ p p] + γ ·Wp(QZ , PZ),
where Q runs through all deterministic neural network encoders (or any other class of universal approximators).
Proof. This directly follows from A.3 in combination with Hornik (1991).

heading: A.6 PROOF OF THEOREM 3.4
text: Proof. Statement i) follows directly from the definition of push-forward of a measure.
For ii) notice that if PZ 6= QZ then there exists A ⊂ Z a Borel set such that PZ(A) 6= Q#PX(A). Then
G#PZ(G(A)) = PZ((G −1 ◦G)(A)) = PZ(A) 6= Q#PX(A)
= Q#PX(A)((G −1 ◦G)(A)) = (G ◦Q)#PX(G(A)).
Hence as PX = (G ◦Q)#PX by hypothesis, we immediately deduce that PX 6= PG.

heading: A.7 COMPARISON WITH BOJANOWSKI & JOULIN (2017)
text: We prove that the cost function of NAT is equivalent to ours when the encoder output is L2 normalized, c′ is squared Euclidean and the Sinkhorn distance is considered with ε = 0:
arg max θ max R∈PM
Tr(RZfθ(X) >) (22)
= arg max θ max R∈PM
〈RZ, fθ(X)〉F (23)
= arg min θ min R∈PM
2− 2〈RZ, fθ(X)〉F (24)
= arg min θ min R∈PM
‖RZ‖2F + ‖fθ(X)‖2F − 2〈RZ, fθ(X)〉F (25)
= arg min θ min R∈PM
‖RZ − fθ(X)‖2F (26)
= arg min θ min R∈PM ∑ i,j Ri,j‖zi − fθ(xj)‖22 (27)
= arg min θ min R∈PM
〈R,C〉F (28)
⊆ arg min θ min R∈SM 1 M 〈R,C〉F − 0 ·H(R) . (29)
Step 24 holds because both R and fθ(X) are row normalized. Step 25 exploits R being a permutation matrix. The inclusion in Step 28 extend to degenerate solutions of the linear program that may not lie on vertices. We have discussed several differences between our Sinkhorn encoder and NAT. There are other minor ones with Bojanowski & Joulin (2017): ImageNet inputs are first converted to grey and passed through Sobel filters and the permutations are updated with the Hungarian only every 3 epochs. Preliminary experiments ruled our any clear gain of those choices in our setting.

heading: A.8 PROOF OF PROPOSITION 6.1
text: Proof. Let z, z′ two points sampled uniformrly from a d-dimensional sphere. Let α be the Euclidean distance between the two points. α has an analytical form (Wu et al., 2017) :
p(‖z − z′‖2) = p(α) = αd−2
c(d)
[ 1− 1 4 α2 ] d−3 2 , where c(d) = √ π Γ ( d−1 2 ) Γ ( d 2
) . For high dimension, it approaches a Gaussian: p(α) ≈ N ( √ 2, 12d ) as d→ +∞. By the Chebischev inequality, for every t > 0
P (|α− √
2| ≥ t) ≤ 1 2dt2 .
Choosing t = −δ+ √ 2 for δ < √
2 and using the symmetry of the Gaussian around the expectation we obtain
1
2d( √ 2− δ)2 ≥ P (|α−
√ 2| ≥ −δ + √ 2)
= 2P (α ≤ √ 2 + δ − √ 2)
= 2 (1− P (α ≥ δ)) .
Hence
P (α ≥ δ) ≥ 1− 1 4d( √ 2− δ)2 .

Review summaries:
This paper builds upon the recent Wasserstein Auto-Encoders; the main innovation is the use of the Sinkhorn distance between the prior and the aggregate posterior. This distance, is used as a differentiable and more tractable surrogate for the wasserstein distance, which is proposed as an alternative to heuristics (e.g. MMD) in Wasserstein-auto encoders.