S INKHORN
A UTO E NCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬂexibility 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.
I NTRODUCTION
Unsupervised learning aims to ﬁnd the underlying rules that govern a given data distribution.
It can be approached by learning to mimic the data generation process, or by ﬁnding 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.
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 ﬂexibility of SAE with qualitative results by using a Dirichlet prior, which deﬁnes the latent space on a probability simplex, as well as with a simple probabilistic programming task.
B ACKGROUND 2. W ASSERSTEIN DISTANCE AND W ASSERSTEIN A UTO E NCODERS
We follow Tolstikhin et al.
(2018) and denote with X , Y
, Z the sample spaces and with X, Y, Z and P X , P Y , P Z 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 deﬁne the push-forward of a distribution P as the induced marginal of the joint distribution F ( Y | X ) P X (denoted by F ( Y | X ) # P X ).
For any measurable non-negative cost c : X × Y → R + ∪ {∞} , one can deﬁne the following OT-cost between marginal distributions P X and P Y via:
[Equation]
where Π( P X , P Y ) is the set of all joint distributions that have as marginals the given P X and P Y .
The elements from Π( P X , P Y ) are called couplings from P X to P Y .
From now on we will assume that X = Y and c ( x, y ) is a distance.
In this case W c ( P X , P Y ) is the Wasserstein distance w.r.t the cost c .
If c ( x, y ) = (cid:107) x − y (cid:107) pp for p ≥ 1 then W p = p √ W c is called the p -th Wasserstein distance .
Let P X denote the true data distribution on X .
We deﬁne a latent variable model given as follows: we ﬁx a latent space Z and a prior distribution P Z 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 ) P Z .
The induced marginal will be denoted by P G .
Learning P G to approximate the true P X is then deﬁned as:
[Equation]
Because of the inﬁmum over Π( P X , P G ) inside W c , this is intractable.
To rewrite this objective we consider the posterior distribution Q ( Z | X ) (the encoder) and its aggregated posterior Q Z :
[Equation]
the induced marginal of the joint distribution Q ( Z | X ) P X .
Tolstikhin et al.
(2018) show that, if the decoder G ( X | Z ) is deterministic, i.e. P G =
G # P Z , or in other words, if all stochasticity of the generative model is captured by Z , then:
[Equation]
Learning the generative model G with the Wasserstein AutoEncoder amounts to:
[Equation]
where β > 0 is a Lagrange multiplier and D Z is any divergence measure on probability distributions on Z , which choice is left open.
WAE uses either MMD or a discriminator trained adversarially for
D Z .
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.
2. T HE S INKHORN ALGORITHM
In place of any choice of D Z , 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 efﬁcient approximations exist.
To motivate this direction, recall that we can always see samples of a continuous distribution as Dirac deltas, whose expectation deﬁnes a discrete distribution.
Let two discrete distributions with support on M points be ˆ P =
M (cid:80)
[Equation]
i ,
[Equation]
i .
Given a cost c (cid:48) , their (empirical) Wasserstein distance is:
[Equation]
where C (cid:48) ij = c (cid:48)
( z (cid:48)
i , z j ) is the matrix associated to the cost c (cid:48) , R is a doubly stochastic matrix as deﬁned in S M =
{ R ∈ R M × M ≥
[Equation]
= 1 , R T 1 = 1 } , and (cid:104)· , ·(cid:105) 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 inﬁnity (Weed & Bach, 2017).
This linear program has solutions on the vertices of S M , which is the set of permutation matrices (Peyr ´ e & Cuturi, 2018).
The Hungarian algorithm ﬁnds an optimal solution in O ( M 3 ) time (Kuhn, 1955).
An entropy-regularized version of problem (5) can be solved more efﬁciently.
Let the entropy of R be H ( R )
= − (cid:80)
Mi,j =1 R i,j log R i,j .
For ε > 0 , Cuturi (2013) deﬁnes the Sinkhorn distance S c (cid:48) :
[Equation]
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 ﬁxed point algorithm that runs nearly in M 2 time (Altschuler et al., 2017) and can be efﬁciently 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 ﬁnite amount of steps in place of convergence.
2. N OISE A S T ARGETS
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 ﬁxed.
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:
[Equation]
where Tr( · ) is the trace operator, Z and X are respectively prior samples and images stacked in a matrix and P M ⊂ S M 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-deﬁned.
The authors relate NAT to OT, a link that we make formal below.
P RINCIPLES OF W ASSERSTEIN AUTOENCODING
With Equation 3, Tolstikhin et al.
(2018) reformulate the Wasserstein distance in image space in terms of autoencoders.
The hard constraint Q Z = P Z is in practice replaced with a soft constraint by adding a penalty in the form of a divergence D Z ( Q Z , P Z ) .
The resulting objective (4) is a lower bound to the Wasserstein difference and the choice of a divergence D Z 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 ﬁrst 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 P Z and Q Z .
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:
[Equation]
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 following 1 :
[Equation]
which reduces to the p -th root of E X ∼ P
X
[ (cid:107) X − G ( Q ( X ))
(cid:107) pp
if both G and Q 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 P X be not atomic
and G ( X | Z ) deterministic.
Then for every continuous cost c :
[Equation]
Using the cost c ( x, y )
= (cid:107) x − y (cid:107)
pp , the equation holds with W pp ( P X , P G ) in place of W c ( P X , P G ) .
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:
[Equation] [Equation]
Inequality in Step 10 holds because we restrict the domain of the inﬁmum , which in turns implies
W p ( Q Z , P Z ) = 0 .
As a consequence we obtain the following Corollary, which provides us with an objective for learning generative autoencoders:
Corollary 3.3.
Let P X be non-atomic and G ( X | Z ) be deterministic and γ -Lipschitz.
Then we have the equality:
[Equation]
More precisely, we can now formulate our learning problem as the minimization of the right-hand side of Equation 12 over deterministic decoders: [Equation]
When the encoder is a neural network of limited capacity, enforcing Q Z ≈ P Z 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 sufﬁcient 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 sufﬁcient and necessary condition for generative autoencoding.
Justiﬁed by the previous results, we state it for deterministic autoencoders (proof in A.6).
Theorem 3.4 (Sufﬁciency and necessity for generative autoencoding) .
Suppose perfect reconstruction, that is, P X
=
( G ◦ Q )
# P X .
Then:
i )
[Equation]
In particular, Theorem 3.4 ii) certiﬁes 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).
S INKHORN A UTO E NCODERS
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 { x i } Mi =1 be the data input to the deterministic encoder Q ( z (cid:48)
i | x i ) =
δ z (cid:48)
i and { z i } Mi =1 the samples from the prior P Z .
The empirical distributions are ˆ Q Z = 1 M (cid:80) Mi
=
δ z (cid:48)
i and ˆ P Z = 1 M (cid:80) Mi
=
δ z i .
With C (cid:48) ij =
c ( z (cid:48) i , z j ) , the
Sinkhorn distance is S c (cid:48)
( ˆ Q Z , ˆ P Z ) as deﬁned in Equation 6.
We compute the Sinkhorn distance in two steps: ﬁrst obtain the optimal regularized coupling R ∗ and then multiply it with the cost, i.e. set ε = 0 : [Equation]
[Equation]
See Algorithm 1.
Note that we do not sacriﬁce differentiability: we stack O
( L )
Sinkhorn operations on top of the encoder, without additional learnable parameters, and run auto-differentiation.
[Equation]
With a deterministic decoder G and encoder Q , we arrive at the objective for the Sinkhorn AutoEncoder (SAE):
[Equation]
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
S c (cid:48) by β to explore the trade off, as in the WAE.
In most experiments, both c and c (cid:48) will be (cid:107) · (cid:107) 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 signiﬁcant overhead, unless a large L is needed for convergence due to a small ε .
In practice, Algorithmloops for L iterations but can exit earlier if the updates of u reach a ﬁxed point.
We have not speciﬁed our distribution P Z 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.
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 conﬁrms 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 ﬂexibility as they are not limited by families of distributions with tractable density (Mohamed & Lakshminarayanan, 2016; Husz´ar, 2017).
Moreover, by encoding with deterministic neural networks, we bypass the use of reparametrization tricks for gradient estimation.
R ELATED WORK
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 difﬁculties; 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.
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.
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 classiﬁcation (Frogner et al., 2015), metric learning (Huang et al., 2016) and ecological inference (Muzellec et al., 2017).
Santa Cruz et al.
; 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 qualiﬁes 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. [Figure]
E XPERIMENTS
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 conﬁrm 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 ﬂexibility 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 CIFAR
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 L 2 normalization of the encoder output as in Bojanowski & Joulin (2017).
The Sinkhorn runs with (cid:15)
= 0 .
1 , L = 50 , except when otherwise stated.
FID scores for CIFAR 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 Fr´echet distance, similarly to Odena et al.
Notice that the FID score is a Wasserstein- 2 distance and hence our theory applies directly.
6. R EPRESENTATION LEARNING WITH S INKHORN ENCODERS
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 (cid:48) be two uniform samples from a d -dimensional sphere.
In the high dimensional regime, for any δ < √ 2 we have P ( (cid:107) z − z (cid:48)
(cid:107) 2 > δ ) ≥ 1 −
[Equation]
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 ﬁll up the sphere; on the other hand, points close on image space cannot be mapped too far from
[Table] [Table]
56.5 WAE normal
MMD 0.
61.6 AE sphere † - - 4.
83.9 HVAE sphere
KL 0.
- - - - - - - - WAE sphere
MMD 0.
61.3 SAE sphere
Sinkhorn 0.
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 ﬁnite capacity and hence favours the conservation of local similarities; see Figure 1e.
6. A UTOENCODING WITH THE S INKHORN DISTANCE AND NAT
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 (cid:104) R, C (cid:48) (cid:105) F and backpropagate.
Sinkhorn-targets : a doubly stochastic matrix R solution of the Sinkhorn can be used for sampling a permutation 3 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 signiﬁcant difference in generative performance between either Sinkhorn vs. Hungarian, or samples vs. targets.
The parameter β trading off reconstruction and latent space cost is more inﬂuential than any of these choices.
On MNIST, MMD is often lower with ﬁxed 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 ﬁx β for MNIST and CIFAR as the best found here.
[Figure]
6. C OMPARISON WITH OTHER AUTOENCODERS
We compare with AE, ( β -)VAE, HVAE 4 and WAE.
Figures 2 shows interpolations and samples of
SAE and VAE from CIFAR10 and CelebA. SAE interpolations are deﬁned 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.
6. D IRICHLET PRIORS
We further demonstrate the ﬂexibility of SAE by using Dirichlet priors on MNIST.
The prior draws samples on the probability simplex; hence, here we constrain the encoder by a ﬁnal 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 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.
6. T OY PROBABILISTIC PROGRAMMING
We run a ﬁnal 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 z what — the digit appearance — from a 10 -dimensional sphere and for z where — the location and scale — from a 3 -dimensional Normal.
A decoder network is fed with z what 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 z where .
The inference model produces z what , z where 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 simpliﬁed 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.
C ONCLUSIONS
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 ﬂexible 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.
A A PPENDIX A. L EMMA
As a useful helper Lemma we prove a Lipschitz property for the Wasserstein distance W p .
Lemma A.1.
For every P X , P Y distributions on a sample space S and a Lipschitz map F we have that
[Equation]
where γ is the Lipschitz constant of F .
Proof.
Recall that
[Equation]
Notice then that for every Γ ∈ Π( P X , P Y ) we have that ( F × F )
# Γ ∈ Π
( F # P X , F # P Y ) .
Hence {
[Equation]
F × F
[Equation]
P
[Equation]
From (16) we deduce that
[Equation]
Taking the p -root on both sides we conclude.
A. P ROOF OF T HEOREM 3.
Proof.
Using the triangle inequality of the Wasserstein distance we obtain
[Equation]
where in line (17) we have used Lemma A.1.
In case G is not deterministic, deﬁning
[Equation]
exp we can still formulate a bound.
The result follows directly from the ﬁrst line of (17).
A. P ROOF OF T HEOREM 3.
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 P X and P Y probability distributions on X such that P X is not atomic, c : X × X
→ R continuous, we have
[Equation]
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
Q Z = Q # P X and Q ( Z | X )
= δ { Q ( x )=
z } .
Hence
[Equation]
Therefore inf Q ( Z | X ) deterministic : Q Z
[Equation]
P Z
[Equation]
We now want to prove that
[Equation]
For the ﬁrst inclusion ⊂ notice that for every Q : X → Z such that Q Z
=
P
Z we have that
G ◦ Q : X → X and
[Equation]
For the other inclusion ⊃ consider T : X → X such that T # P X
= P G = G # P Z .
We want ﬁrst to prove that there exists a set A ⊂ X with P X ( A )
= 1 such that G : Z → T ( A ) is surjective.
Indeed if it does not hold there exists B ⊂ X with P X ( B ) > 0 and G − 1 ( T ( B )) = ∅ .
Hence
[Equation]
that is a contraddiction.
Therefore by standard set theory the map G : Z → T ( A ) has a right inverse that we denote by (cid:101) G .
Then deﬁne Q = (cid:101) G ◦ T .
Notice that G ◦ Q = G ◦ (cid:101) G ◦ T
= T almost surely in P X and also ( (cid:101) G ◦ T )
# P X
= P Z .
Indeed for any A ⊂ Z Borel we have
[Equation]
This concludes the proof of the claim in (20).
Now we have
[Equation]
Notice that this is exactly the Monge formulation of optimal transport.
Therefore by Theorem A. we conclude that
[Equation]
as we aimed.
A. T OLSTIKHIN ET AL . (2018)’ S T HEOREM AS A CONSEQUENCE
Proof.
Thanks to Theorem 3.2 we have that
[Equation]
X, G ( Z ))] .
For the opposite inequality given Q ( Z | X ) such that P Z
= (cid:82) Q ( Z | X
) dP X deﬁne Q ( X, Y ) =
P X × [ G # Q ( Z | X )] .
It is a distribution on X ×X and it is easy to check that π 1# Q ( X, Y ) =
P X and π 2# Q ( X, Y )
=
G # P Z , where π 1 and π 2 are the projection on the ﬁrst and the second component.
Therefore { Q ( X, Z ) : Q ( Z | X ) such that Q Z = P Z } ⊂ Π( P X , P G ) and so
[Equation]
A. B OUNDS FOR NEURAL NETWORKS
Theorem A.3.
Let P X be non-atomic and G ( X | Z ) be deterministic and γ -Lipschitz.
If Q NN is a neural network approximating a near optimal deterministic encoder up to an error of ε ≥ 0 in
L p -norm then we have the inequality:
[Equation]
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 Q NN a neural network approximation of Q ∗ such that p (cid:112) E X ∼ P X
[ (cid:107) Q ∗ ( X ) − Q NN ( X ))
(cid:107) pp ] ≤ ε (existence by Hornik (1991)).
Then we get:
[Equation]
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 P X be non-atomic and G ( X | Z ) be deterministic and γ -Lipschitz.
Then we have the equality:
[Equation]
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).
A. P ROOF OF T HEOREM 3.
Proof.
Statement i ) follows directly from the deﬁnition of push-forward of a measure.
For ii ) notice that if P Z (cid:54) =
Q Z then there exists A ⊂ Z a Borel set such that P Z ( A ) (cid:54) =
Q # P X ( A ) .
Then
[Equation]
Hence as P X = ( G ◦ Q ) # P X by hypothesis, we immediately deduce that P X (cid:54) =
P G .
A. C OMPARISON WITH B OJANOWSKI & J OULIN (2017)
We prove that the cost function of NAT is equivalent to ours when the encoder output is L 2 normalized, c (cid:48) is squared Euclidean and the Sinkhorn distance is considered with ε = 0 :
[Equation]
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 ﬁrst converted to grey and passed through Sobel ﬁlters 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.
A. P ROOF OF P ROPOSITION 6.
Proof.
Let z, z (cid:48) 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) :
[Equation]
For high dimension, it approaches a Gaussian: p ( α ) ≈ N ( √ 2 , 12 d ) as d → + ∞ .
By the Chebischev inequality, for every t >
[Equation]
Choosing t = − δ + √ 2 for δ < √ 2 and using the symmetry of the Gaussian around the expectation we obtain
[Equation]
Hence
[Equation]