Wasserstein Auto-Encoders

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

Reviews and author responses:

Let me clarify the markov chain point.

In the case Q(Z|X) is stochastic, the encode/decode chain X->Z->X' is stochastic. Namely, P(X'|X) is not a deterministic function, it is a distribution. A markov chain can be constructed if we sample X from P_X and use P(X'|X) as the transition probability.

By optimizing the Wasserstein distance between P(X') and P_X, we hope to get the parameter such that P(X') == P_X. The reconstruction term in this paper requires that X' == X, which is stronger than P(X') == P_X.

Thank you for the question.

Unfortunately, we did not quite get the point of your Markov chain example. But we would like to make it clear that the paper does not assume anything specific about the encoder Q(Z|X). As long as the aggregated posterior Qz matches the prior Pz, the encoder can be either deterministic or random. The same holds true for the WAE algorithm. We will try to emphasize it better in the updated version of the paper.

The decoder is indeed a different story: for Theorem 1 we need it to be deterministic, but a very similar result holds also for the random decoder (Supplementary B).

Thanks for the great work. It's nice to see there is theoretical support for the (auto-encoder + constraint on Z) objective.

It seems to me the expectation over X could not be moved out in theorem 1, as this breaks the independence of Z and X.
Consider the case Q(Z|X) is not deterministic, we can have a markov chain X_{t+1} ~ \int_Z [ P_G(X'|Z)Q(Z|X_t) ], which has a stationary distribution same as P_X.  The algorithm in this paper gives a special case where Q(Z|X) is deterministic and X_{t+1} = X_{t}.

In supplementary B, the case where the decoder is random is discussed. It would be nice to also discuss the cases where Q(Z|X) is random vs deterministic.

Do correct me if I'm wrong, thanks.

This paper satisfies the following necessary conditions for
acceptance. The writing is clear and I was able to understand the
presented method (and its motivation) despite not being too familiar
with the relevant literature. Explicitly writing the auto-encoder(s)
as pseudo-code algorithms was particular helpful. I found no technical
errors. The problem addressed is one worth solving - building a
generative model of observed data. There is some empirical testing
which show the presented method in a good light.

The authors are careful to relate the presented method with existing
ones, most notably VAE and AAE. I suppose one could argue that the
close connection to existing methods means that this paper is not
innovative enough. I think that would be unfair - most new methods
have close relations with existing ones - it is just that sometimes
the authors do not flag this up as they should.

WAE is a bit oversold. The authors state that WAE generates "samples
of better quality" (than VAE) without any condition being put on when
it does this. There is no proof that it is always better, and I can't
see how there could be. Any method of inferring a generative model
from data must make some 'inductive' assumptions. Surely one could
devise situations where VAE outperforms WAE. I think this issue should
have been examined in more depth.

I found no typo or grammatical errors which is unusual - good careful
job!



This very well written paper covers the span between W-GAN and VAE. For a reviewer who is not an expert in the domain, it reads very well, and would have been of tutorial quality if space had allowed for more detailed explanations. The appendix are very useful, and tutorial paper material (especially A). 

While I am not sure description would be enough to reproduce and no code is provided, every aspect of the architecture, if not described, if referred as similar to some previous work. There are also some notation shortcuts (not explained) in the proof of theorems that can lead to initial confusion, but they turn out to be non-ambiguous. One that could be improved is P(P_X, P_G) where one loses the fact that the second random variable is Y.


This work contains plenty of novel material, which is clearly compared to previous work:
- The main consequence of the use of Wasserstein distance is the surprisingly simple and useful Theorem 1. I could not verify its novelty, but this seems to be a great contribution.
- Blending GAN and auto-encoders has been tried in the past, but the authors claim better theoretical foundations that lead to solutions that do not rquire min-max
- The use of MMD in the context of GANs has also been tried. The authors claim that their use in the latent space makes it more practival

The experiments are very convincing, both numerically and visually.

Source of confusion: in algorithm 1 and 2, \tilde{z} is "sampled" from Q_TH(Z|xi), some one is lead to believe that this is the sampling process as in VAEs, while in reality Q_TH(Z|xi) is deterministic in the experiments.

This paper provides a reasonably comprehensive generalization to VAEs and Adversarial Auto-encoders through the lens of the Wasserstein metric. By posing the auto-encoder design as a dual formulation of optimal transport, the proposed work supports the use of both deterministic and random decoders under a common framework. In my opinion, this is one of the crucial contributions of this paper. While the existing properties of auto-encoders are preserved, stability characteristics of W-GANs are also observed in the proposed architecture. The results from MNIST and CelebA datasets look convincing, though could include additional evaluation to compare the adversarial loss with the straightforward MMD metric and potentially discuss their pros and cons. In some sense, given the challenges in evaluating and comparing closely related auto-encoder solutions, the authors could design demonstrative experiments for cases where Wassersterin distance helps and may be  its potential limitations.

The closest work to this paper is the adversarial variational bayes framework by Mescheder et.al. which also attempts at unifying VAEs and GANs. While the authors describe the conceptual differences and advantages over that approach, it will be beneficial to actually include some comparisons in the results section.

Dear Mathieu,

thank you for the suggestion. We will update the paper accordingly.

Congratulations on this nice paper. 

The ability to remove one of the two marginal constraints in Theorem 1 relies on the assumption that P_G(Y|Z=z) is a Dirac. I know that you stated in the intro that you focus on deterministic maps but it would be nice to repeat the assumptions made in Theorem 1.

We thank the reviewer for the positive feedback and the kind words regarding the overview part of the paper.

We will make sure to make notations clearer and include all the details of architectures used in experiments in the updated version of the paper. Of course we will also open source the code.

We are pleased that the reviewer found the paper well written. 

We tried to be modest in our claims, in particular we never implied that WAEs produce better samples for *all data distributions*. As noticed by the reviewer this would be indeed impossible to prove, especially because the question of how to evaluate and compare sample qualities of unsupervised generative models is still open. We will double-check that there are no bold and unsupported statements in the final version of the paper.

We thank the reviewer for the positive feedback. 

Comparing properties of WAE-MMD and WAE-GAN is indeed an intriguing direction and we intend to look into the details in our future research. In this paper we only report initial empirical observations, which can be concluded by saying that WAE-MMD enjoys a stable training but does not match Pz and Qz perfectly, while the training of WAE-GAN is not so stable but leads to much better matches once succeeded. 

In this paper we decided that comparing to VAE was sufficient for our purposes: both VAE and AVB follow the same objective of maximizing the marginal log likelihood in contrast to the minimization of the optimal transport studied in our work. However, we do agree that in future it would be interesting to compute the FID scores of the AVB samples. 

You state in the paper that the variational auto-encoder objective is composed of reconstruction cost plus a KL divergence term that captures how distinct the image by the encoder of each training example is from the prior p(z), and then go on to say that this KL term is not guaranteeing that the overall encoded distribution matches the prior p(z).

However, as shown in the paper "ELBO surgery: yet another way to carve up the variational evidence lower bound" by Hoffman and Johnson, the KL term in the VAE objective can be decomposed into exactly this KL(q(z)||p(z)) between the average encoder distribution and the prior plus a mutual information term, and that the former is a heavy contributor towards the overall KL term. This means that VAE does indeed try to match the overall encoder distribution of q to the prior, but also includes a regularizing term that aims to minimize the mutual information between the hidden code z and the index of the observation x that encourages the VAE to have the encoder produce the same codes z for different observations.

In conclusion, it would be more accurate to state that in comparison to VAEs you simply exclude the mutual information regularisation term from the objective as formulated in the ELBO surgery paper.

Meta-review: This paper proposes a new generative model that has the stability of variational autoencoders (VAE) while producing better samples. The authors clearly compare their work to previous efforts that combine VAEs and Generative Adversarial Networks with similar goals.  Authors show that the proposed algorithm is a generalization of Adversarial Autoencoder (AAE) and minimizes Wasserstein distance between model and target distribution. The paper is well written with convincing results. Reviewers agree that the algorithm is novel and practical; and close connections of the algorithm to related approaches are clearly discussed with useful insights.  Overall, the paper is strong and I recommend acceptance.

Decision: accepted-oral-papers