Wasserstein Auto-Encoders

TLDR

The paper proposes the Wasserstein Auto‑Encoder (WAE), a new generative model algorithm. WAE minimizes a penalized Wasserstein distance to match the model to the target distribution, using a regularizer distinct from the VAE’s and generalizing adversarial auto‑encoders. Experiments demonstrate that WAE’s regularizer aligns the encoded distribution with the prior, generalizes adversarial auto‑encoders, and yields samples of higher quality than VAEs, as indicated by improved FID scores.

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, as measured by the FID score.