Inferential Wasserstein Generative Adversarial Networks

One of the goals of generative modeling is to match the model distribution $P_{\theta}(x)$ with parameters $\theta$ to the true data distribution $P_{X}$ for a random variable $X\in{\cal X}$. For latent variable models, the data point $X$ is generated from a latent variable $Z\in{\cal Z}$ through a conditional distribution $p(x|z)$. Here ${\cal X}$ denotes the support for $P_{X}$ and ${\cal Z}$ denotes the support for $P_{Z}$. In this paper, we consider models with $Z\sim{\cal N}(0,I)$. There has been a surge of research on deep generative networks in recent years and the literature is too vast to summarize here (Kingma and Welling, 2014; Goodfellow et al., 2014; Li et al., 2015; Gao et al., 2020; Qiu and Wang, 2020). These models have provided a powerful framework for modeling complex high dimensional datasets.

We start introducing two main approaches for generative modeling. The first one is called variational auto-encoders (VAEs) (Kingma and Welling, 2014), which use variational inference (Blei et al., 2017) to learn a model by maximizing the lower bound of the likelihood function. Specifically, let the latent variable $Z$ be drawn from a prior $p(z)$ and the data $X$ have a likelihood $p(x|z)$ that is conditioned on $Z=z$. Unfortunately obtaining the marginal distribution of $X$ requires computing an intractable integral $p(x)=\int p(x|z)p(z)dz$. Variational inference approximates the posterior $p(z|x)$ by a family of distribution $q_{\eta}(z|x)$ with the parameter $\eta$. The objective is to maximize the lower bound of the log-likelihood function known as the evidence lower bound (ELBO). The ELBO is given by

Note that the difference between the log-likelihood $\log p(x)$ and the ELBO is the Kullback-Leibler divergence between $q_{\eta}(z|x)$ and the true posterior $p(z|x)$. Usually $q_{\eta}(z|x)$ is a normal density where both the conditional mean and the conditional covariance are modeled by deep neural networks (DNNs), so that the first term of the ELBO can be approximated efficiently by Monte Carlo methods and the second term can be calculated explicitly. Therefore, the ELBO allows us to do approximate posterior inference with tractable computation. VAEs have elegant theoretical foundations but the drawback is that they tend to produce blurry images. The second approach is called generative adversarial networks (GANs) (Goodfellow et al., 2014), which learn a model by using a powerful discriminator to distinguish between real data samples and generative data samples. Specifically, we define a generator $G:{\cal Z}\rightarrow{\cal X}$ and a discriminator $D:{\cal X}\rightarrow\mathbb{[}0,1]$. The generator and discriminator play a two-player minimax game by being alternatively updated, such that the generator tries to produce real-looking images and the discriminator tries to distinguish between generated images and observed images. The GAN objective can be written as $\min_{G}\max_{D}V(G,D)$, where

GANs produce more visually realistic images but suffer from the unstable training and the mode collapse problem. Although there are many variants of generative models trying to take advantages of both VAEs and GANs (Tolstikhin et al., 2018; Rosca et al., 2017), to the best of our knowledge, the model which provides a unifying framework combining the best of VAEs and GANs in a principled way is yet to be discovered.

The autoencoder generative model consists of two parts: an encoder $Q$ and a generator $G$. The encoder $Q$ maps a data sample $x\in{\cal X}$ to a latent variable $z\in{\cal Z}$, and the generator $G$ takes a latent variable $z\in{\cal Z}$ to produce a sample $G(z)$. In general, the autoencoder generative model should satisfy the following three conditions simultaneously: (a) The generator can generate images which have a similar distribution with observed images, i.e., the distribution of $G(Z)$ is similar to that of $P_{X}$; (b) The encoder can produce meaningful encodings in the latent space, i.e., $Q(X)$ has a similar distribution with $Z$; (c) The reconstruction errors of this model based on these meaningful encodings are small, i.e., the difference between $X$ and $G(Q(X))$ is small.

We emphasize that the benefit of using an autoencoder is to encourage the model to better represent all the data it is trained with, so that it discourages mode collapse. We first show that, for any distribution residing on a compact smooth Riemannian manifold ¹¹1A smooth manifold ${\cal X}$ is a manifold with a $C^{\infty}$ atlas on ${\cal X}$. A $C^{\infty}$ atlas is a collection of charts $\{\varphi_{\alpha}:U_{\alpha}\rightarrow\mathbb{R}^{d}\}$ such that $\{U_{\alpha}\}$ covers ${\cal X}$, and for all $\alpha$ and $\beta$, the transition map $\varphi_{\alpha}\cdot\varphi_{\beta}^{-1}$ is a $C^{\infty}$ map. Here $U_{\alpha}$ is an open subset of ${\cal X}$. For any point $p\in{\cal X}$, let $T_{p}{\cal X}$ be the tangent space of ${\cal X}$ at $p$. A Riemannian metric assigns to each $p$ a positive definite inner product $g_{p}:T_{p}{\cal X}\times T_{p}{\cal X}\rightarrow\mathbb{R}$, along with which comes a norm $\|\cdot\|_{p}:T_{p}{\cal X}\rightarrow\mathbb{R}$ defined by $\|v\|_{p}=\sqrt{g_{p}(v,v)}$. The smooth manifold ${\cal X}$ endowed with this metric $g$ is called a smooth Reimannian manifold., there always exists an encoder $Q^{*}:{\cal X}\rightarrow{\cal Z}$ which guarantees meaningful encodings and exists a generator $G^{*}:{\cal Z}\rightarrow{\cal X}$ which generates samples with the same distribution as data points by using these meaningful codes.

Theorem 1 is a natural consequence of the Nash embedding theorem (Nash, 1956; Günther, 1991) and the probability integral transformation (Rosenblatt, 1952). In Theorem 1, we have proved the existence of $Q^{*}$ and $G^{*}$, however, learning $Q^{*}$ and $G^{*}$ from the data points is still a challenging task. Consider a general $f$-GAN model (Nowozin et al., 2016). Let $h:\mathbb{R}\rightarrow(-\infty,\infty]$ be a convex function with $h(1)=0$. The $f$-GAN defines the $f$-divergence between the data distribution $P_{X}$ and the generative model distribution $P_{G(Z)}$ for the generator $G$ as:

where $h^{*}(x)=\sup_{y}\{x\cdot y-h(y)\}$ is the convex conjugate of $h$ and ${\cal F}=\{f|f:\cal X\rightarrow\mathbb{R}\}$ is a class of functions whose output range is the domain of $h^{*}$. When $f$ is approximated by a DNN, its output range can be achieved by choosing an appropriate activation function specific to the f-divergence used. For example, if $h(x)=x\log(x)-(x+1)\log(x+1)$, then the corresponding convex conjugate $h^{*}(x)=-\log(1-\exp(x))$. To satisfy the above condition, we select the output activation function of the DNN $f$ to be $\sigma(v)=-\log(1+\exp(-v))$ such that the $f$-GAN can recover the original vanilla GAN (Goodfellow et al., 2014). If $h(x)=0$ when $x=1$ and $h(x)=\infty$ otherwise, we have $h^{*}(x)=x$. With the property that $\cal F$ is 1-Lipschitz function class, the $f$-GAN turns to be the WGAN.

For ease of presentation, we illustrate our methodology by mainly focusing on the Wasserstein distance and the inferential WGAN (iWGAN) model. The extension to general inferential f-GANs (ifGANs) is straightforward and will be presented in Section 2.3.

Suppose that we observe $n$ samples $\{x_{1},\ldots,x_{n}\}$. In practice, we minimize the empirical version, denoted by $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})$, of $\overline{W}_{1}(P_{X},P_{G(Z)})$ to learn both the encoder and the generator, where,

Here $\hat{\mathbb{E}}_{obs}\{\cdot\}$ denotes the empirical average on the observed data $\{x_{i}\}$ and $\hat{\mathbb{E}}_{z}$ denotes the empirical average on a random sample of standard normal random variables. Before we present the details of the algorithm, we first establish the generalization error bound for the iWGAN in this section.

In the context of supervised learning, generalization error is defined as the gap between the empirical risk and the expected risk. The empirical risk is corresponding to the training error, and the expected risk is corresponding to the testing error. Mathematically, the difference between the expected risk and the empirical risk, i.e. the generalization error, is a measure of how accurately an algorithm is able to predict outcome values for previously unseen data. However, in the context of GANs, neither the training error nor the test error is well defined. But we can define the generalization error in a similar way. Explicitly, we define the “training error” as $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})$ in (7), which is minimized based on observed samples. Define the “test error” as $W_{1}(P_{X},P_{G(Z)})$ in (1), which is the true 1-Wasserstein distance between $P_{X}$ and $P_{G(Z)}$. The generalization error for the iWGAN is defined as the gap between these two “errors”. In other words, for an iWGAN model with the parameter $(G,Q,f)$, the generalization error is defined as $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})-W_{1}(P_{X},P_{G(Z)})$. For discussions of generalization performance of classical GANs, see Arora et al. (2017) and Jiang et al. (2019).

For a fixed generator $G$, Theorem 3 holds uniformly for any discriminator $f\in\mathcal{F}$. It indicates that the 1-Wasserstein distance between $P_{X}$ and $P_{G(Z)}$ can be dominantly upper bounded by the empirical $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})$ and Rademacher complexity of $\mathcal{F}$. Since $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})\leq\widehat{W}_{1}(P_{X},P_{G(Q(X))})+\widehat{W}_{1}(P_{G(Q(X))},P_{G(Z)})$ for any $Q\in{\cal Q}$, the capacity of $\mathcal{Q}$ determines the value of $\widehat{\overline{W}}_{1}(P_{X},P_{G(Z)})$. In learning theory, Rademacher complexity, named after Hans Rademacher, measures richness of a class of real-valued functions with respect to a probability distribution. There are several existing results on the empirical Rademacher complexity of neural networks. For example, when $\mathcal{F}$ is a set of 1-Lipschitz neural networks, we can apply the conclusion from Bartlett et al. (2017) to $\widehat{\mathfrak{R}}_{n}(\mathcal{F})$, which produces an upper bound scaling as $\mathcal{O}(B\sqrt{L^{3}/n})$. Here $L$ denotes the depth of network $f\in\mathcal{F}$. Similar upper bound with an order of $\mathcal{O}(B\sqrt{Ld^{2}/n})$ can be obtained by utilizing the results from Li et al. (2019), where $d$ is the width of the network.

Next, we introduce the details of the algorithm. Our target is to solve the following optimization problem:

where $J_{1}(f)$ and $J_{2}(Q)$ are regularization terms for $f$ and $Q$ respectively. We approximate $G,Q,f$ by three neural networks with pre-specified architectures.

Since $f$ is assumed to be 1-Lipschitz, we adopt the gradient penalty defined as $J_{1}(f)=\mathbb{E}_{\hat{x}}\big{\{}(\mathbb{\|}\nabla_{\hat{x}}f(\hat{x})\|_{2}-1)^{2}\big{\}}$ in (Gulrajani et al., 2017) to enforce the 1-Lipschitz constraint on $f\in{\cal F}$. Furthermore, since we expect $Q(X)$ follows approximately standard normal, we use the maximum mean discrepancy (MMD) penalty (Gretton et al., 2012), denoted by $J_{2}(Q)=\mbox{MMD}_{k}(P_{Q(X)},P_{Z})$, to enforce $Q(X)$ to converge to $P_{Z}$. In particular,

where $k$ is set to be the Gaussian radial kernel function $k(x,y)=\exp(\frac{-\|x-y\|^{2}}{2})$.

We have adopted the stochastic gradient descent algorithm called the ADAM (Kingma and Ba, 2015) to estimate the unknown parameters in neural networks. The ADAM is an algorithm for first-order gradient-based optimization of stochastic objection functions, based on adaptive estimates of lower-order moments. Given the current tuple $(G^{i},Q^{i},f^{i})$ at the $i$th iteration, we sample a batch of observations $\{x^{i}_{k}\}_{k=1}^{n}\sim{P}_{X}$, latent variable $\{z^{i}_{k}\}_{k=1}^{n}\sim{P}_{Z}$, and $\{\epsilon_{k}\}_{k=1}^{n}\sim U[0,1]$. Then we construct $\hat{x}^{i}_{k}\leftarrow\epsilon_{k}x^{i}_{k}+(1-\epsilon_{k})G^{i}(z^{i}_{k})$, $i=1,\ldots,n$, for computing the gradient penalty. Let $L^{i}=L(G^{i},Q^{i},f^{i})$ and $J_{1}(f^{i})=(\|\nabla_{\hat{x}^{i}}f^{i}(\hat{x}^{i})\|_{2}-1)^{2}$. We can evaluate the gradient with respect to the parameters in $f$, which is denoted by

Then we can update $f^{i}$ by the ADAM using this gradient. Similarly, we can evaluate the gradient with respect to the parameters in $G$ and $Q$, which is denoted by

Then we can update $(G^{i},Q^{i})$ by the ADAM using this gradient. The stopping criteria are both the DualGap$(G^{i},Q^{i},f^{i})$ in (4) and the objective function $L(G^{i},Q^{i},f^{i})$ are less than pre-specified error tolerances $\epsilon_{1}$ and $\epsilon_{2}$, respectively. Specifically, based on the definition of the duality gap in (4), we approximate DualGap$(G^{i},Q^{i},f^{i})$ by the difference between $L(G^{i},Q^{i},f^{i+1})$ and The optimization (9) consists of two tuning parameters $\lambda_{1}$ and $\lambda_{2}$. We pre-specify some values for $\lambda_{1}$ and $\lambda_{2}$ and select the optimal tuning parameters by grid search using cross validation. The details of the algorithm are presented in Algorithm 1.

The iWGAN has proposed an efficient framework to stably and automatically estimate both the encoder and the generator. In this section, we provide a probabilistic interpretation of the iWGAN under the framework of maximum likelihood estimation.

Maximum likelihood estimator (MLE) is a fundamental statistical framework for learning models from data. However, for complex models, MLE can be computationally prohibitive due to the intractable normalization constant. MCMC has been used to approximate the intractable likelihood function but do not work efficiently in practice, since running MCMC till convergence to obtain a sample can be computationally expensive. For example, to reduce the computational complexity, Hinton (2002) proposed a simple and fast algorithm, called the contrastive divergence (CD). The basic idea of CD is to truncate MCMC at the $k$-th step, where $k$ is a fixed integer as small as one. The simplicity and computational efficiency of CD makes it widely used in many popular energy-based models. However, the success of CD also raised a lot of questions regarding its convergence property. Both theoretical and empirical results show that CD in general does not converge to a local minimum of the likelihood function (Carreira-Perpinan and Hinton, 2005; Qiu et al., 2020), and diverges even in some simple models (Schulz et al., 2010; Fischer and Igel, 2010). The iWGAN can be treated as an adaptive method for the MLE training, which not only provides computational advantages but also allows us to generate more realistic-looking images. Furthermore, this probabilistic interpretation enables other novel applications such as image quality checking and outlier detection.

Let $X$ denote the image. Define the density of $X$ by an energy-based model based on an autoencoder (Gu and Zhu, 2001; Zhao et al., 2017; Berthelot et al., 2017):

where

and $\theta\in\Theta$ is the unknown parameter and $V(\theta)$ is the log normalization constant. The major difficulty for the likelihood inference is due to the intractable function $V(\theta)$. Suppose that we have the observed data $\{x_{i}:i=1,\ldots,n\}$. The log-likelihood function of $\theta\in\Theta$ is $\ell(\theta)=n^{-1}\sum_{i=1}^{n}\log~{}p(x_{i}|\theta)$, whose gradient is