[Paper Review] From optimal transport to generative modeling: the VEGAN cookbook

Olivier Bousquet, Sylvain Gelly+3 more|arXiv (Cornell University)|May 22, 2017

Advanced Numerical Methods in Computational MathematicsEngineering10 references99 citations

TL;DR

The paper reformulates unsupervised generative modeling as an optimal transport problem between true data distribution and a latent variable model, introduces the Penalized Optimal Transport (POT) objective, and links POT to AAEs, VAEs, and WGANs.

ABSTRACT

We study unsupervised generative modeling in terms of the optimal transport (OT) problem between true (but unknown) data distribution $P_X$ and the latent variable model distribution $P_G$. We show that the OT problem can be equivalently written in terms of probabilistic encoders, which are constrained to match the posterior and prior distributions over the latent space. When relaxed, this constrained optimization problem leads to a penalized optimal transport (POT) objective, which can be efficiently minimized using stochastic gradient descent by sampling from $P_X$ and $P_G$. We show that POT for the 2-Wasserstein distance coincides with the objective heuristically employed in adversarial auto-encoders (AAE) (Makhzani et al., 2016), which provides the first theoretical justification for AAEs known to the authors. We also compare POT to other popular techniques like variational auto-encoders (VAE) (Kingma and Welling, 2014). Our theoretical results include (a) a better understanding of the commonly observed blurriness of images generated by VAEs, and (b) establishing duality between Wasserstein GAN (Arjovsky and Bottou, 2017) and POT for the 1-Wasserstein distance.

Motivation & Objective

• Provide a principled OT-based formulation for unsupervised generative modeling.
• Derive an equivalent latent-space reparametrization of OT couplings.
• Relax constraints to obtain the Penalized Optimal Transport (POT) objective.
• Show connections between POT, AAEs, VAEs, AVB, and GAN variants.
• Discuss implications for training stability and output sharpness/blurriness in generated samples.

Proposed method

• Re-express the primal OT problem W_c(P_X,P_G) via a latent encoder Q(Z|X) and a generator P_G(Y|Z).
• Introduce the constrained optimization on encoders that match the aggregated posterior Q_Z to the prior P_Z.
• Relax constraints with a penalty to obtain the POT objective: D_POT = inf_Q(Z|X) E_{P_X,X} E_{Q(Z|X)}[c(X,G(Z))] + lambda D_GAN(Q_Z, P_Z).
• Show that for squared Euclidean cost and Gaussian decoders POT coincides with Adversarial Autoencoders (AAE).
• Relate POT with VAE/AVB in terms of blurriness and to WGAN when using the 1-Wasserstein cost.
• Outline the primal/dual equivalence for the 1-Wasserstein case and discuss gradient implications.

Experimental results

Research questions

• RQ1 How can the OT cost between true data distribution and a latent variable model be reformulated to explicit latent encoders?
• RQ2 What is the effect of relaxing OT constraints via a penalty on encoders, and how does this relate to existing generative modeling methods?
• RQ3 When does the penalized OT objective align with AAEs or VAEs, and what are the implications for sample quality and training stability?
• RQ4 How do primal and dual formulations behave under 1-Wasserstein and 2-Wasserstein costs in generative modeling?

Key findings

• The OT problem can be equivalently written in terms of probabilistic encoders constrained to match posterior and prior distributions.
• Relaxing these constraints yields the Penalized Optimal Transport (POT) objective that can be minimized by SGD by sampling from P_X and P_G.
• For squared Euclidean cost, POT coincides with the objective of Adversarial Auto-Encoders (AAE), providing a theoretical justification for AAEs as approximately minimizing W_2(P_X,P_G).
• POT and WGAN align under the Euclidean cost in the 1-Wasserstein setting, with dual/primal perspectives offering different training dynamics.
• VAEs and AVB minimize marginal log-likelihoods and tend to produce blurrier outputs, whereas POT/AAE can avoid this blurriness under certain conditions.
• When decoders are Gaussian, AAE corresponds to optimizing the POT objective with lambda tied to decoder variance.