Wasserstein Auto-Encoders

Ilya Tolstikhin; Olivier Bousquet; Sylvain Gelly; Bernhard Schoelkopf

Source-linked AI summary

Wasserstein Auto-Encoders

Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, Bernhard Schoelkopf

arXiv:1711.01558v4Nov 5, 2017stat.MLcs.LG

TL;DR

The paper addresses generative modeling through a new regularized auto-encoder formulation based on optimal transport. WAE matches encoded data to a prior using latent-space regularization, and experiments report better sample quality than VAEs while retaining several VAE properties.

Problem
The paper seeks to minimize the optimal transport distance between the unknown data distribution and a latent-variable model distribution.
Method
WAE combines reconstruction cost with a discrepancy penalty matching the encoded data distribution QZ to the prior PZ, using adversarial or MMD-based regularization.
Results
WAE preserves stable training, an encoder-decoder architecture, and a structured latent manifold while generating samples of better quality than VAEs and approaching GANs.
Takeaways & Limitations
WAE provides a family of optimal-transport-based generative-model algorithms, including an adversary-free MMD variant and a formulation related to adversarial auto-encoders.
Takeaways & Limitations
MMD-based training can provide no gradient for encoded outliers far from the prior's support, depending on kernel bandwidth.

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces Wasserstein Auto-Encoders, which approach generative modeling through optimal transport and regularized auto-encoding. WAE offers alternative latent-space regularizers and aims to retain VAE properties while improving sample quality.

Motivation: Optimal transport provides a weaker distributional topology than many f-divergences, which is useful for data supported on low-dimensional manifolds.The paper notes that stronger divergences can max out and provide no useful training gradients.
Proposed method: WAE is a family of regularized auto-encoders that minimizes the optimal transport cost Wc(PX, PG) for any cost function c.Its objective combines reconstruction cost with a latent-distribution discrepancy regularizer.
Experiments: Experiments on MNIST and CelebA evaluate WAE with squared reconstruction cost and compare its generated samples with established approaches.The paper reports that WAE preserves stable training, an encoder-decoder architecture, and a structured latent manifold while producing higher-quality samples approaching GANs.
Proposed method: WAE regularizes the encoded data distribution toward the prior using either adversarial training or maximum mean discrepancy.The MMD-based option yields a fully adversary-free min-min optimization problem.
Theory: Theoretical analysis shows that the primal optimal transport problem for generative models is equivalent to one involving a probabilistic encoder Q(Z|X).This equivalence provides the basis for the WAE objective.

2 Proposed method

WAE minimizes optimal transport between the data distribution and a latent-variable model through an auto-encoder objective combining reconstruction with a latent-distribution penalty. It supports adversarial and MMD regularizers, including a fully adversary-free variant.

2 Proposed method: WAE minimizes optimal transport Wc(PX, PG) using an encoder-decoder objective with reconstruction cost and a penalty matching QZ to the prior PZ.QZ is the encoded-data distribution, defined as EPX[Q(Z|X)].
2 Proposed method: The optimal-transport formulation replaces coupling optimization with optimization over probabilistic encoders whose aggregate latent distribution is constrained to match PZ.Relaxing this constraint with a divergence penalty yields the WAE objective.
2 Proposed method: The method can use arbitrary latent-space divergences and input-space cost functions, with λ controlling the penalty that relaxes exact prior matching.The squared-cost, GAN-penalty case coincides with adversarial auto-encoders.
2 Proposed method: WAE permits deterministic encoders that map each input directly to a latent code, unlike VAEs' required non-degenerate Gaussian encoders.The algorithms can sample from a deterministic encoder by returning its mean mapping.
2 Proposed method: WAE-GAN matches QZ and PZ with adversarial training in latent space, while WAE-MMD uses maximum mean discrepancy and supports fully adversary-free min-min optimization.MMD has an unbiased U-statistic estimator suitable for stochastic gradient descent and is expected to work well with Gaussian priors.

3 Related work

Related work positions WAE against unregularized auto-encoders, VAEs, GAN-based models, and other regularized auto-encoders. Its distinguishing features are aggregate latent matching, flexible penalties and costs, and a primal optimal-transport foundation.

3 Related work: Unregularized auto-encoders can scatter codes into disconnected regions, leaving decoder “holes” and making latent-space sampling difficult.Their training minimizes reconstruction cost without a regularizer.
3 Related work: Unlike VAE regularization, WAE matches the overall encoded distribution EPX[Q(Z|X)] to PZ rather than separately matching each conditional encoder distribution.The VAE regularizer can also be decomposed into aggregate KL divergence plus mutual information, which WAE drops.
3 Related work: WAE-GAN is equivalent to adversarial auto-encoders and generalizes them through arbitrary input costs and latent discrepancy measures, including MMD.The theory provides a Wasserstein interpretation and theoretical justification for AAEs.
3 Related work: Compared with WGAN, WAE uses the primal optimal-transport formulation, naturally includes an encoder, and applies to any cost function c rather than only W1.WGAN approaches the problem through the dual form and does not provide an encoder.
3 Related work: Earlier auto-encoder/adversarial approaches lacked an incentive for reciprocal encoders and decoders, whereas WAE uses a different penalty and has a theoretical foundation.WAE need not require a min-max game.

4 Experiments

Experiments on MNIST and CelebA evaluate WAE-GAN and WAE-MMD for reconstruction, latent-manifold geometry, and sample quality. Matching the encoded distribution QZ to the prior PZ is important for sample quality, with WAE-GAN sometimes outperforming WAE-MMD but training less stably.

Experimental setup: WAE-GAN and WAE-MMD were trained on MNIST and CelebA to assess reconstructions, latent geometry, and random-sample quality.The datasets contain 70k MNIST images and roughly 203k CelebA images.
Experimental setup: The experiments used deterministic encoder-decoder pairs with Euclidean latent spaces, Gaussian priors, squared reconstruction cost, Adam, and convolutional networks.Latent dimensions were 8 for MNIST and 64 for CelebA.
Evaluation: Random samples are produced by sampling from PZ and decoding the resulting latent vectors.The decoder maps sampled vectors z through Gθ(z).
Evaluation: Accurate matching between QZ and PZ is important because the decoder is trained on encoded data and may perform poorly on prior samples when the distributions differ.Even slight differences between QZ and PZ affected sample quality in the experiments.
Results: WAE-GAN sometimes produced better matching and samples than WAE-MMD, but WAE-GAN was less stable because of adversarial training.WAE-MMD showed very stable training, similar to VAE.
Results: WAE-MMD produced samples of slightly better quality than VAE, while WAE-GAN achieved the best overall results in the reported CelebA evaluation.Sample quality was assessed with FID and a heuristic sharpness measure.

5 Conclusion

The paper derives Wasserstein auto-encoders as a family of generative-model algorithms and evaluates WAE-GAN and WAE-MMD against VAEs. The experiments report better sample quality without compromising training stability or reconstruction quality.

Conclusion: Wasserstein auto-encoders are derived from the optimal transport cost as a new family of algorithms for generative modeling.The paper studies WAE-GAN and WAE-MMD as particular implementations.
Conclusion: Compared with VAEs, trained WAE models generated higher-quality images while retaining training stability and reconstruction quality.The conclusion identifies further work on matching QZ to PZ, adversarially training c, and dual formulations.
Evaluation: Sharpness evaluation converts images to grayscale, applies a Laplace edge detector, and averages activation variances across 1000 generated images.Blurrier images yield smaller variances because they contain fewer edges.

A Implicit generative models: a short tour of GANs and VAEs

Implicit generative models can be sampled and trained with SGD, but VAEs and GANs optimize different distributional objectives. Regularization shapes latent representations and sampling behavior, while Wasserstein distances address difficulties caused by strong divergences on disjoint manifolds.

Shared properties: GANs and VAEs can both be trained by sampling from PG without an analytical density and scaled with stochastic gradient descent.This permits flexible implicit models built by sampling a latent code and mapping it to data.
GANs and VAEs: GANs minimize a lower bound on the Jensen-Shannon divergence, while VAEs minimize an upper bound on negative log-likelihood or KL divergence.The VAE formulation uses a standard normal prior and Gaussian conditional encoders for tractability.
Wasserstein perspective: Strong divergences can saturate when data and model distributions lie on largely disjoint low-dimensional manifolds, producing unhelpful gradients.The 1-Wasserstein distance remains meaningful and provides stable gradients in such cases.
Regularization: Removing the VAE KL regularizer can scatter encoded training points across latent space, leaving holes where the decoder is untrained and making sampling difficult.The regularizer is therefore connected to the usefulness of the latent representation.
Adversarial auto-encoders: Adversarial auto-encoders replace the VAE KL term with a regularizer matching the aggregated posterior QZ to the prior PZ.The aggregated posterior is the marginal distribution obtained by sampling X from PX and then Z from Q(Z|X).

B Proof of Theorem 1 and further details

The proof represents couplings between data and generated variables through latent codes, yielding an optimal-transport upper bound under encoder and generator factorizations.

Optimal-transport couplings Γ(X, Y) can be viewed as conditional mappings Γ(Y|X) from data points to generated points.The marginal constraint permits factoring the joint distribution as Γ(Y|X)PX(X).
Theorem 1 factors this mapping through latent codes using an encoder Q(Z|X) and generator PG(Y|Z).
The admissible three-variable distributions enforce X∼PX, (Y,Z)∼PG,Z, and conditional independence Y ⊥⊥ X | Z.
The construction uses PX,Z = P(X∼PX, Z∼PZ), while the optimal-transport cost is defined by minimizing expected c(X,Y) over admissible couplings.
When PG(Y|Z) is deterministic, the factored coupling set equals the full set of couplings between PX and PG.For non-deterministic decoders, only set inclusion is guaranteed.

B.1 Random decoders PG(Y |Z)

For random decoders, the factorized coupling class generally provides an upper bound on optimal transport rather than an exact equality, with decoder variance contributing to the bound.

For non-deterministic decoders, PX,Y is only a subset of the full coupling set P(PX, PG), producing an upper bound on optimal transport.
The random-decoder analysis assumes X∈R^d and decoder conditional means G(z)∈R^d with coordinate-wise marginal variances σ1^2,...,σd^2.
The bound specializes to the squared cost c(x,y)=∥x−y∥2.
The appendix conjectures that the deterministic-decoder condition may also be necessary, but explicitly states that this necessity is not used in the paper.

C.1 MNIST

The MNIST experiments use convolutional encoder-decoder models, adversarial architecture for WAE-GAN, and training heuristics designed to align encoded data with the prior and randomize encoders.

MNIST models use mini-batches of 100 and are trained for 100 epochs with λ=10 and σ_z^2=1.
Both encoder and decoder use fully convolutional architectures with 4x4 filters.
WAE-GAN includes a separate adversary architecture for latent-space adversarial training.
The convolutional stack uses stride-2 convolutions with SAME padding, while Conv, fractional strided convolution, batch normalization, ReLU, and fully connected layers are explicitly specified.
Training first pretrains the encoder to align the encoded distribution's mean and covariance with the prior, then adds truncated pixel-wise Gaussian noise to encourage random encoders.
The VAE baseline uses cross-entropy loss with a Bernoulli decoder and otherwise matches the listed architectures and hyperparameters.

C.2 CelebA

The CelebA experiments compare WAE-MMD, WAE-GAN, and VAE models using convolutional architectures, specified optimization schedules, and a latent adversary informed by the known prior density.

CelebA images are center-cropped to 140x140, resized to 64x64, and trained with mini-batches of 100.
Reported WAE models are trained for 55 epochs, while the VAE is trained for 68 epochs.
WAE-MMD uses λ=100 and WAE-GAN uses λ=1, with σ_z^2=2 for both.
Both encoder and decoder use fully convolutional architectures with 5x5 convolutional filters.
WAE-GAN adds the analytically known log prior log p_Z(z) to the adversary output so it can learn the remaining Q_Z density term.
The reported VAE uses squared loss with Gaussian decoders, while Bernoulli decoders were observed to match Gaussian-decoder performance under the best σ_G^2 choice.

D Extended experimental setup (bigVAE and bigWAE)

The extended study trains over 3,000 WAE-GAN, WAE-MMD, and VAE models under matched computational budgets, broad hyperparameter sweeps, and varied architectures. Models are compared by selecting checkpoints with the lowest FID score from 10,000 generated samples.

Architectures and encoders: Encoder and decoder mappings used either DCGAN-style convolutional architectures or ResNet50-v2, while WAE encoders were deterministic or diagonal-covariance Gaussian.The larger-scale study specifically included ResNet50-v2 encoder-decoder mappings.
Experimental setup: Over 3,000 WAE-GAN, WAE-MMD, and VAE models were trained with identical computational budgets.Each model used 8 Google Cloud TPU-v2 accelerators for 100,000 mini-batch steps with Adam.
Hyperparameter sweep: The sweep varied latent dimensionality, mini-batch size, encoder-decoder learning rate, WAE-GAN adversary learning rate, and regularization strength.Latent dimensionality ranged from 16 to 512, batch size from 512 to 1024, and encoder-decoder learning rate from 10^-5 to 10^-2.
Hyperparameter sweep: WAE-MMD used inverse multiquadratics kernels across fixed scales, with the base scale sampled from [0.1, 16].The fixed scale set was S = {0.1, 0.2, 0.5, 1, 2, 5, 10}.
Evaluation: For each configuration, intermediate checkpoints were evaluated using FID on 10,000 samples, and the checkpoint with the lowest FID was selected.The procedure produced the statistics shown in the subsequent figures.