GAN and VAE from an Optimal Transport Point of View

Minimum Kantorovitch Estimators

where $P_{1}(x,y)=x$ and $P_{2}(x,y)=y$, and $P_{1\sharp}$ and $P_{2\sharp}$ are marginalization operators that return for a given coupling $\gamma$ its first and second marginal, respectively.

The notations $P_{1\sharp}$ and $P_{2\sharp}$ above agree with the more general notion of pushforward measures: Given a measurable map $g:\mathcal{Z}\rightarrow\mathcal{X}$, which can be interpreted as a function “moving” points from a measurable space to another, one can naturally extend $g$ to become a more general map $g_{\sharp}$ that can now “move” an entire probability measure on $\mathcal{Z}$ towards a new probability measure on $\mathcal{X}$. The operator $g_{\sharp}$ “pushes forward” each elementary mass of a measure $\zeta$ in $\mathcal{P}(\mathcal{Z})$ by applying the map $g$ to obtain then a mass in $\mathcal{X}$, to build on aggregate a new measure in $\mathcal{P}(\mathcal{X})$ written $g_{\sharp}\zeta$. More rigorously, the pushforward measure of a measure $\zeta\in\mathcal{P}(\mathcal{Z})$ by a map $g:\mathcal{Z}\rightarrow\mathcal{X}$ is the measure denoted as $g_{\sharp}\zeta$ in $\mathcal{P}(\mathcal{X})$ such that for any set $B\subset\mathcal{X}$, $(g_{\sharp}\zeta)(B)\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\zeta(g^{-1}(B))=\zeta(\left\{z\in\mathcal{Z}\;;\;g(z)\in B\right\})$.

MKE-GM.

The map $g_{\theta}$ should be therefore thought as a “decoding” map from a low dimensional space to a high dimensional space. In such a setting, the maximum likelihood estimator is in general undefined or difficult to compute (because the support of the measures $\mu_{\theta}$ are singular) while MKEs are attractive because they are always well defined.

Dual Formulation and GAN

Because (1) is a linear program, it has a dual formulation, known as the Kantorovich problem [13, Thm. 5.9]:

In the dual formulation (2), $\theta$ does not appear anymore in the constraints. Therefore, the gradient of $E$ can be computed as

which is, indeed, given a candidate potential $h$ for the first variable, the best possible potential that can be paired with $h$ that satisfies the constraints of (2) (see [13, Thm. 5.9]). For this reason, one can parameterize problem (2) as depending on one potential function only.

As a side-note, and as previously commented in the literature, there is at this point no empirical evidence that supports the idea that using discriminative deep networks that way can result in accurate approximations of Wasserstein distances. These alternative formulations provide instead a very useful proxy for a quantity directly related to the Wasserstein distance.

Primal Formulation and VAE

Following , in the special case of a generative model $\mu_{\theta}=g_{\theta\sharp}\zeta$, formula (1) can be conveniently re-written as

This is advantageous because now $\pi$ is defined over $\mathcal{Z}\times\mathcal{X}$, which is lower-dimensional than $\mathcal{X}\times\mathcal{X}$, and also because, as in Equation (2), $\theta$ does not appear in the constraints either. This provides an alternative formula for the gradient of $E$:

suggests to look for couplings $\pi$ with a parametric form. A simple way to achieve this is to restrict couplings $\pi$ to those of the form

where $f_{\xi}:\mathcal{X}\rightarrow\mathcal{Z}$ is a parametric “encoding” map (typically a deep network), see Figure 1, right. This $\pi_{\xi}$ satisfies by construction the marginal constraint $P_{2\sharp}\pi=\nu$, but in general it cannot satisfy the other constraint $P_{1\sharp}\pi=\zeta$ (because $P_{1\sharp}\pi_{\xi}$ is discrete while $\zeta$ is not). So following , it makes sense to consider a relaxed “unbalanced” formulation (in the sense of ) of the form

Plugging the ansatz $\pi=\pi_{\xi}$ in (6), one obtains the Wasserstein-VAE formulation

Such a cost is usually associated with the Monge formulation of optimal transport , whose original motivation was to find an optimal map under that cost that would be able to push forward a given measure onto another[12, §1.1].

Conclusions

The WGAN and WVAE formulations are very different, and are in some sense dual one of each other. For GAN, the couple $(g_{\theta},h_{\xi})$ should be thought as a (primal, dual) pair (often referred to as adversarial pair, which is reminiscent of game theory saddle points). For VAE, the couple $(f_{\xi},g_{\theta})$ is rather an (encoding, decoding) pair, and both have the flavour of transportation maps.

In sharp contrast to the primal gradient formula (5) which only requires integrating against an optimal coupling $\pi^{\star}$, the dual gradient formula (3) involves the integration of the gradient of an optimal potential $h^{\star}$. The latter tends to be more unstable and thus necessitates accurate optimization sub-iterations to obtain an optimal dual potential $h^{\star}$ or an approximation $h_{\xi}^{\star}$ within a restricted parametric class . This is somehow inline with the empirical observation that training VAE is more stable than training GAN. One should however bear in mind that, although both formulations can be motivated by the same minimum Kantorovitch estimation problem (MKE-GM), they define quite different estimators. In particular, GAN is often credited for producing less blurry outputs when used for image generation.

Denoting $\theta_{\text{\tiny MKE}},\theta_{\text{\tiny WGAN}}$ and $\theta_{\text{\tiny WVAE}}$ the solutions of (MKE-GM), (WGAN) and (WVAE), one has in the limit $\lambda\rightarrow+\infty$ (to cancel the bias due to the marginal constraint relaxation),

furthermore mentions that in the “non-parametric limit” (i.e. when the number of parameters appearing in $\xi$ tends to $+\infty$, and also letting $\lambda\rightarrow+\infty$), the gap between the estimators should vanish. Indeed, $h_{\xi}$ and $f_{\xi}$ should capture the desired optimal map in the limit and one thus recovers the true solution to (MKE-GM). While it would be interesting from a theoretical perspective to prove and quantify such a claim, it is unclear wether it would be useful for the practitioner. Indeed, the convergence rate might be slow, so that in practice one can be quite far from this non-parametric limit. One could even argue that this limit may give poor estimators for complicated datasets, so that parameterizing the maps and using non-convex optimization solvers lead instead to a beneficial and implicit regularization of these estimators.

References