Biadversarial Variational Autoencoder

TL;DR

This paper introduces a Biadversarial Variational Autoencoder that replaces Gaussian assumptions with adversarial networks, enabling better modeling of multimodal distributions and improving image quality.

Contribution

It proposes a novel VAE framework using adversarial networks to avoid Gaussian assumptions, enhancing the ability to model complex, multimodal data distributions.

Findings

Avoids Gaussian assumptions in VAE

Improves modeling of multimodal distributions

Produces sharper, higher-quality images

Abstract

In the original version of the Variational Autoencoder, Kingma et al. assume Gaussian distributions for the approximate posterior during the inference and for the output during the generative process. This assumptions are good for computational reasons, e.g. we can easily optimize the parameters of a neural network using the reparametrization trick and the KL divergence between two Gaussians can be computed in closed form. However it results in blurry images due to its difficulty to represent multimodal distributions. We show that using two adversarial networks, we can optimize the parameters without any Gaussian assumptions.

Equations

$$\begin{array}{rl}\log p_{\theta }(x)=& \log \int p_{\theta }(x,z)dz\\ =& \log \int \frac{p_{\theta }(x,z)q_{\varphi }(z|x)}{q_{\varphi }(z|x)}dz\\ \ge & {\mathbb{E}}_{q_{\varphi }(z|x)}(\log p_{\theta }(x|z)+\log p(z)-\log q_{\varphi }(z|x))\equiv \text{ELBO}\end{array}$$

$$\begin{array}{r}{q}_{\varphi }(z|x)=\mathcal{N}(z,{\mu }_z,{\sigma }_z^2\mathbf{I})\end{array}$$

$$\begin{array}{r}\text{ELBO}\equiv {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(\log p_{\theta }(x|z))+{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(\log p(z)-\log q_{\varphi }(z|x))\end{array}$$

$$\begin{array}{r}\underset{\theta }{\max }\underset{\varphi }{\max }{\mathbb{E}}_{x\sim {\hat{p}}_{data}}\text{ELBO}\end{array}$$

$$\begin{array}{rl}& \underset{\theta }{\min }\underset{\varphi }{\min }{\mathbb{E}}_{x\sim {\hat{p}}_{data}}{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(\log q_{\varphi }(z|x)-\log p(z))\\ =& \underset{\varphi }{\min }{\mathbb{E}}_{x\sim {\hat{p}}_{data}}KL(q_{\varphi }(z|x)||p(z))\end{array}$$

$$\begin{array}{r}\underset{{\mathcal{D}}_{\varphi }}{\max }V({\mathcal{D}}_{\varphi },\varphi )\equiv {\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(1-{\mathcal{D}}_{\varphi }(x,z))-{\mathbb{E}}_{z\sim p(z)}(\exp (-{\mathcal{D}}_{\varphi }(x,z))))\end{array}$$

$$\begin{array}{rl}& {\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(1-{\mathcal{D}}_{\varphi }(x,z))-{\mathbb{E}}_{z\sim p(z)}(\exp (-{\mathcal{D}}_{\varphi }(x,z))))=\\ & \int {\hat{p}}_{data}(x)(q_{\varphi }(z|x)(1-{\mathcal{D}}_{\varphi }(x,z))+p(z)\exp (-{\mathcal{D}}_{\varphi }(x,z)))dzdx\end{array}$$

$$\begin{array}{rl}& \forall x,z,{\mathcal{D}}_{\varphi }^{\ast }(x,z)=\log (\frac{p(z)}{q_{\varphi }(z|x)})\end{array}$$

$$\begin{array}{rl}V({\mathcal{D}}_{\theta ,\theta }^{\ast },\theta ,\theta )=& {\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(1-\log (\frac{p(z)}{q_{\varphi }(z|x)}))\\ & -{\mathbb{E}}_{z\sim p(z)}(\exp (-\log (\frac{p(z)}{q_{\varphi }(z|x)}))))\\ =& {\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(\log (\frac{q_{\varphi }(z|x)}{p(z)}))\\ =& {\mathbb{E}}_{x\sim {\hat{p}}_{data}}KL(q_{\varphi }(z|x)||p(z))\end{array}$$

$$\begin{array}{r}\underset{\varphi }{\min }\underset{{\mathcal{D}}_{\varphi }}{\max }V({\mathcal{D}}_{\varphi },\varphi )\end{array}$$

$$\begin{array}{rl}p(x|z)=& \mathcal{N}(x|\mu ,{\sigma }^2\mathbf{I})\\ =& \frac{1}{(2\pi {)}^{n/2}\sigma }\exp (-\frac{||x-\mu |{|}_2^2}{2{\sigma }^2})\end{array}$$

$$\begin{array}{rl}-\log p(x|z)=& \log ((2\pi {)}^{n/2}\sigma )+\frac{||x-\mu |{|}_2^2}{2{\sigma }^2}\end{array}$$

$$\begin{array}{rl}& \underset{\theta }{\mathrm{argmax}}\underset{\varphi }{\mathrm{argmax}}{\mathbb{E}}_{x\sim {\hat{p}}_{data}}{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}(\log p_{\theta }(x|z))\\ =& \underset{\theta }{\mathrm{argmin}}\underset{\varphi }{\mathrm{argmin}}{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}{\mathbb{E}}_{x\sim {\hat{p}}_{data}}(\log {\hat{p}}_{data}(x)-\log p_{\theta }(x|z))\\ =& \underset{\theta }{\mathrm{argmin}}\underset{\varphi }{\mathrm{argmin}}{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}KL({\hat{p}}_{data}(x)||p_{\theta }(x|z))\end{array}$$

$$\begin{array}{r}\underset{{\mathcal{D}}_{\theta ,\varphi }}{\max }V({\mathcal{D}}_{\theta ,\varphi },\theta ,\varphi )\equiv {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}({\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathcal{D}}_{\theta ,\varphi }(x,z))-{\mathbb{E}}_{x\sim p_{\theta }(x|z)}(\exp ({\mathcal{D}}_{\theta ,\varphi }(x,z)-1)))\end{array}$$

$$\begin{array}{rl}& {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}({\mathbb{E}}_{x\sim {\hat{p}}_{data}}({\mathcal{D}}_{\theta ,\varphi }(x,z))-{\mathbb{E}}_{x\sim p_{\theta }(x|z)}(\exp ({\mathcal{D}}_{\theta ,\varphi }(x,z)-1)))=\\ & \int q_{\varphi }(z|x)({\hat{p}}_{data}(x)({\mathcal{D}}_{\theta ,\varphi }(x,z))+p_{\theta }(x|z)\exp ({\mathcal{D}}_{\theta ,\varphi }(x,z)-1))dxdz\end{array}$$

$$\begin{array}{rl}& \forall x,z,{\mathcal{D}}_{\theta ,\varphi }^{\ast }(x,z)=1+\log (\frac{{\hat{p}}_{data}(x)}{p_{\theta }(x|z)})\end{array}$$

$$\begin{array}{rl}V({\mathcal{D}}_{\theta ,\varphi }^{\ast },\theta ,\varphi )=& {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}({\mathbb{E}}_{x\sim {\hat{p}}_{data}}(1+\log (\frac{{\hat{p}}_{data}(x)}{p_{\theta }(x|z)}))-\\ & {\mathbb{E}}_{x\sim p_{\theta }(x|z)}(\exp (1+\log (\frac{{\hat{p}}_{data}(x)}{p_{\theta }(x|z)})-1)))\\ =& {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}({\mathbb{E}}_{x\sim {\hat{p}}_{data}}(\log (\frac{{\hat{p}}_{data}(x)}{p_{\theta }(x|z)}))\\ =& {\mathbb{E}}_{z\sim q_{\varphi }}KL({\hat{p}}_{data}(x)||p_{\theta }(x|z))\end{array}$$

$$\begin{array}{r}\underset{\theta }{\min }\underset{\varphi }{\min }\underset{{\mathcal{D}}_{\theta ,\varphi }}{\max }V({\mathcal{D}}_{\theta ,\varphi },\theta ,\varphi )\end{array}$$

$$\begin{array}{r}\underset{\theta }{\min }\underset{\varphi }{\min }(\underset{{\mathcal{D}}_{\theta ,\varphi }}{\max }V({\mathcal{D}}_{\theta ,\varphi },\theta ,\varphi )+\underset{{\mathcal{D}}_{\varphi }}{\max }V({\mathcal{D}}_{\varphi },\varphi ))\end{array}$$

Full text

1 Introduction

We want to maximize the evidence lower bound (ELBO) of the marginal likelihood $p_{\theta}(x)$. We can derive the ELBO with the Jensen inequality by marginalizing out the latent variable $z$ and introducing the approximate posterior $q_{\phi}(z|x)$:

[TABLE]

2 Inference

Many works on variational autoencoder assume a Gaussian distribution for the approximate posterior distribution:

[TABLE]

where $\mu_{z}$ and $\sigma_{z}$ are neural network functions.

This is convenient for computation but very restrictive for $z$. We introduce an adversarial network that will optimize the parameters of the encoder without the need of any restrictive assumption. To do that, rearrange eq. (1):

[TABLE]

The objective being:

[TABLE]

where $\phi$ denotes the parameter of the encoder and $\theta$ denotes the parameters of the decoder.

Rewrite the second term of the ELBO in eq. (3) to bring out a Kullback-Leibler (KL) divergence:

[TABLE]

This term corresponds to the KL divergence between the approximate posterior $q_{\phi}(z|x)$ and the prior $p(z)$. Note that it is the reverse KL divergence, ie. the difference between both distributions is bounded by the approximate posterior, which is a better option to learn real modes in case of a multimodal distribution. Inspired by 2016arXiv160600709N , we define a network with an objective that differs slightly from the original adversarial network 2014arXiv1406.2661G so the associated generator learns to minimize the reverse KL divergence instead of the Jensen-Shannon divergence. In so doing we are able to optimize the parameters without doing any parametric assumption on the posterior. Introduce the network $\mathcal{D}_{\phi}:X\times Z\longrightarrow\mathbb{R}$ with the following objective:

[TABLE]

where the parameters $\phi$ is fixed.

Inspired by 2014arXiv1406.2661G , write the second term as an integral to find the optimal value of $\mathcal{D}_{\phi}$:

[TABLE]

Given a pair $(a,b)$ in $\mathbb{R}^{2}$, the function $d\in\mathbb{R}\mapsto a(1-d)-b\exp(-d)$ reaches its maximum at $d^{*}=\log(\frac{b}{a})$. Hence the maximum of the integral is reached if:

[TABLE]

By replacing eq. (8) in eq. (6), we show that the optimal value function $V(\mathcal{D}^{*}_{\theta,\theta},\theta,\theta)$ reached by the discriminator, the generator being fixed, is the KL divergence in eq. (5):

[TABLE]

In so doing we can optimize the second term of the ELBO with a minimax game with value function $V(\mathcal{D}_{\phi},\phi)$:

[TABLE]

3 Generative process

Many works on variational autoencoder assume also a Gaussian distribution for the output distribution:

[TABLE]

where $\mu$ is a neural network function and $\sigma^{2}$ is a hyperparameter.

The negative log likelihood of this distribution is an affine function of the L2 norm, hence we often encounter a L2 reconstruction term in works on variational autoencoders :

[TABLE]

Rearrange the first term of the objective in eq. (3) to bring out a direct KL divergence:

[TABLE]

This time we choose an adversarial objective so that the associated generator learns to minimize the direct KL divergence. Introduce the network $\mathcal{D}_{\theta,\phi}:X\times Z\longrightarrow\mathbb{R}$ with the following objective:

[TABLE]

where the parameters $\theta$ and $\phi$ are fixed.

Write the second term as an integral to find the optimal value of $\mathcal{D}_{\theta,\phi}$:

[TABLE]

Given a pair $(a,b)$ in $\mathbb{R}^{2}$, the function $d\in\mathbb{R}\mapsto ad-b\exp(d-1)$ reaches its maximum at $d^{*}=1+\log(\frac{a}{b})$. Hence the maximum of the integral is reached if:

[TABLE]

By replacing eq. (16) in eq. (14), we show that the optimal value function $V(\mathcal{D}^{*}_{\theta,\phi},\theta,\phi)$ reached by the discriminator, the generator being fixed, is the direct KL divergence in eq. (13):

[TABLE]

In so doing we can optimize the second term of the ELBO with a minimax game with value function $V(\mathcal{D}_{\theta,\phi},\theta,\phi)$:

[TABLE]

Finally we have transformed the optimization of the ELBO into a minimax game involving two discriminators:

[TABLE]

4 Implementation

The model is implemented with PyTorch. The implementation is available here:

https://github.com/ArnaudFickinger/BAVAE.