We propose the Wasserstein Auto-Encoder (WAE)---a new algorithm for building a generative model of the data distribution. WAE minimizes a penalized form of the Wasserstein distance between the model distribution and the target distribution, which leads to a different regularizer than the one used by the Variational Auto-Encoder (VAE).
 This regularizer encourages the encoded training distribution to match the prior. We compare our algorithm with several other techniques and show that it is a generalization of adversarial auto-encoders (AAE). Our experiments show that WAE shares many of the properties of VAEs (stable training, encoder-decoder architecture, nice latent manifold structure) while generating samples of better quality. The field of representation learning was initially driven by supervised approaches, with impressive results using large labelled datasets. Unsupervised generative modeling, in contrast, used to be a domain governed by probabilistic approaches focusing on low-dimensional data. Recent years have seen a convergence of those two approaches. In the new field that formed at the intersection, variational auto-encoders (VAEs) BID16 constitute one well-established approach, theoretically elegant yet with the drawback that they tend to generate blurry samples when applied to natural images. In contrast, generative adversarial networks (GANs) BID9 turned out to be more impressive in terms of the visual quality of images sampled from the model, but come without an encoder, have been reported harder to train, and suffer from the "mode collapse" problem where the resulting model is unable to capture all the variability in the true data distribution. There has been a flurry of activity in assaying numerous configurations of GANs as well as combinations of VAEs and GANs. A unifying framework combining the best of GANs and VAEs in a principled way is yet to be discovered.This work builds up on the theoretical analysis presented in BID3 . Following ; BID3 , we approach generative modeling from the optimal transport (OT) point of view. The OT cost (Villani, 2003) is a way to measure a distance between probability distributions and provides a much weaker topology than many others, including f -divergences associated with the original GAN algorithms BID25 . This is particularly important in applications, where data is usually supported on low dimensional manifolds in the input space X . As a result, stronger notions of distances (such as f -divergences, which capture the density ratio between distributions) often max out, providing no useful gradients for training. In contrast, OT was claimed to have a nicer behaviour BID11 although it requires, in its GAN-like implementation, the addition of a constraint or a regularization term into the objective. : Both VAE and WAE minimize two terms: the reconstruction cost and the regularizer penalizing discrepancy between P Z and distribution induced by the encoder Q. VAE forces Q(Z|X = x) to match P Z for all the different input examples x drawn from P X . This is illustrated on picture (a), where every single red ball is forced to match P Z depicted as the white shape. Red balls start intersecting, which leads to problems with reconstruction. In contrast, WAE forces the continuous mixture Q Z := Q(Z|X)dP X to match P Z , as depicted with the green ball in picture (b). As a result latent codes of different examples get a chance to stay far away from each other, promoting a better reconstruction.In this work we aim at minimizing OT W c (P X , P G ) between the true (but unknown) data distribution P X and a latent variable model P G specified by the prior distribution P Z of latent codes Z ∈ Z and the generative model P G (X|Z) of the data points X ∈ X given Z. Our main contributions are listed below (cf. also FIG0 ):• A new family of regularized auto-encoders (Algorithms 1, 2 and Eq. 4), which we call Wasserstein Auto-Encoders (WAE), that minimize the optimal transport W c (P X , P G ) for any cost function c. Similarly to VAE, the objective of WAE is composed of two terms: the c-reconstruction cost and a regularizer D Z (P Z , Q Z ) penalizing a discrepancy between two distributions in Z: P Z and a distribution of encoded data points, i.e. DISPLAYFORM0 When c is the squared cost and D Z is the GAN objective, WAE coincides with adversarial auto-encoders of BID23 .• Empirical evaluation of WAE on MNIST and CelebA datasets with squared cost c(x, y) = x − y 2 2 . Our experiments show that WAE keeps the good properties of VAEs (stable training, encoder-decoder architecture, and a nice latent manifold structure) while generating samples of better quality, approaching those of GANs.• We propose and examine two different regularizers D Z (P Z , Q Z ). One is based on GANs and adversarial training in the latent space Z. The other uses the maximum mean discrepancy, which is known to perform well when matching high-dimensional standard normal distributions P Z BID10 . Importantly , the second option leads to a fully adversary-free min-min optimization problem.• Finally, the theoretical considerations presented in BID3 and used here to derive the WAE objective might be interesting in their own right. In particular , Theorem 1 shows that in the case of generative models, the primal form of W c (P X , P G ) is equivalent to a problem involving the optimization of a probabilistic encoder Q(Z|X) .The paper is structured as follows. In Section 2 we review a novel auto-encoder formulation for OT between P X and the latent variable model P G derived in BID3 . Relaxing the resulting constrained optimization problem we arrive at an objective of Wasserstein auto-encoders. We propose two different regularizers, leading to WAE-GAN and WAE-MMD algorithms. Section 3 discusses the related work. We present the experimental results in Section 4 and conclude by pointing out some promising directions for future work. Using the optimal transport cost, we have derived Wasserstein auto-encoders-a new family of algorithms for building generative models. We discussed their relations to other probabilistic modeling techniques. We conducted experiments using two particular implementations of the proposed method, showing that in comparison to VAEs, the images sampled from the trained WAE models are of better quality, without compromising the stability of training and the quality of reconstruction. Future work will include further exploration of the criteria for matching the encoded distribution Q Z to the prior distribution P Z , assaying the possibility of adversarially training the cost function c in the input space X , and a theoretical analysis of the dual formulations for WAE-GAN and WAE-MMD.
We propose a new auto-encoder based on the Wasserstein distance, which improves on the sampling properties of VAE.
{"src": "scitldr-AIC_train", "paper_id": "HkL7n1-0b"}