Sinkhorn AutoEncoders

1 INTRODUCTION . 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 . 2 BACKGROUND . 2.1 WASSERSTEIN DISTANCE AND WASSERSTEIN AUTOENCODERS . 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 . 2.2 THE SINKHORN ALGORITHM . 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 .

Summary: 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.