[Paper Review] Wasserstein Auto-Encoders

Ilya Tolstikhin, Olivier Bousquet+2 more|arXiv (Cornell University)|Nov 5, 2017

Generative Adversarial Networks and Image SynthesisComputer Science32 references192 citations

TL;DR

Wasserstein Auto-Encoders (WAE) minimize the Wasserstein distance between data and model distributions using an encoder–decoder framework with either GAN-based or MMD-based penalties to align the encoded latent distribution with a prior, yielding stable training and improved sample quality over VAEs.

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.

Motivation & Objective

• Motivate a generative modeling framework rooted in optimal transport that yields stable training and interpretable latent manifolds.
• Develop a new auto-encoder objective that minimizes the Wasserstein distance between true data and model distributions.
• Propose two regularization strategies to align the encoded latent distribution with a prior: GAN-based (WAE-GAN) and MMD-based (WAE-MMD).
• Demonstrate that WAEs retain VAE-like benefits while producing higher-quality samples on MNIST and CelebA.

Proposed method

• Formulate the Wasserstein auto-encoder objective as the minimum over probabilistic encoders Q(Z|X) of the expected reconstruction cost plus a penalty that enforces QZ to match the prior PZ (D_Z(QZ, PZ)).
• Use the primal form of the OT cost; derive that the latent marginal QZ must equal the prior PZ for optimal transport coupling.
• Provide two concrete penalties: (a) GAN-based D_Z with adversarial training in latent space (WAE-GAN), and (b) maximum mean discrepancy (MMD) with a characteristic kernel (WAE-MMD).
• Opt for either deterministic or probabilistic encoders, and use deep networks to parameterize Q and G.
• Empirically evaluate on MNIST and CelebA using squared cost c(x,y)=||x−y||^2, comparing to VAEs and showing improved sample quality.

Experimental results

Research questions

• RQ1 Can we construct a regularized auto-encoder that minimizes the Wasserstein distance between data and generated distributions in a tractable way?
• RQ2 How can we regularize the latent encoding to match a specified prior, and what are the trade-offs of GAN-based versus MMD-based penalties?
• RQ3 Do WAEs produce stable training and latent manifolds while improving sample quality compared to VAEs?
• RQ4 How do WAEs perform on standard benchmarks (MNIST, CelebA) in terms of reconstruction, latent structure, and sample realism?

Key findings

• WAEs generalize adversarial auto-encoders and allow any cost c and any latent-space discrepancy D_Z between QZ and PZ.
• WAE-GAN and WAE-MMD achieve stable training and latent-space regularization, producing higher-quality samples than VAEs on MNIST and CelebA.
• WAE samples approach GAN quality while preserving the encoder–decoder architecture and latent manifold structure characteristic of VAEs.
• WAE-MMD exhibits stable training and competitive sample quality, with WAE-GAN occasionally outperforming in sample realism on CelebA.
• Quantitative results on CelebA show WAE-GAN achieving the best Fréchet Inception Distance (FID) among reported variants, with WAE-MMD also outperforming VAE.