1 K-L Divergence

Generative modeling relies heavily on metrics of similarities between two distributions, among which the most commonly used is called the K-L divergence, short for Kullback–Leibler divergence. It is defined below for two distributions with probability density functions $p_1\left(x\right)$ and $p_2\left(x\right)$:

K-L divergence has two important properties.

1. It is obvious that K-L divergence is not symmetric in terms of $p_1\left(x\right)$ and $p_2\left(x\right)$.
2. It is always non-negative, and it is 0 iff $p_1\left(x\right)$ and $p_2\left(x\right)$ are the same everywhere. To see why, we can break DL-divergence into two parts:

The second term in (2), with the negative sign, is $p_1$’s information theoretic entropy. The first term, also with the negative sign, is the cross entropy between $p_1$ and $p_2$. The first term is always no smaller than the second term per Gibb’s inequality.

2 Intuition

The concept of autoencoders predate the VAE. An autoencoder, shown in Figure 1, consists of an encoder $E_{\varphi }$ and a decoder $D_{𝜃}$. $E_{\varphi }$, a deep neural network parameterized by $\varphi$, takes a sample $x$ from population $𝕏$ and maps it to $z=E_{𝜃}\left(x\right)$ in $\mathbb{Z}$. $D_{𝜃}$, another deep neural network parameterized by $𝜃$, aiming to reconstruct $x$, takes $z$ as input and maps it to $\tilde{x}=D_{𝜃}\left(z\right)=D_{𝜃}\left(E_{\varphi }\right(x\left)\right)$. $\mathbb{Z}$ is usually of a lower dimension than $𝕏$ , and thus $E_{\varphi }$ is considerd to posess some compression capability and unsupervised feature extraction capability.

The training of an autoencoder minimizes the reconstruction loss: the expected $L_2$ distance between $x$ and $\tilde{x}$:

Once trained, the decoder $D_{𝜃}$, to some extent, is already a generative model in that it can create samples in $𝕏$ given a sample $z$. The distribution of $z$ or even the range of $z$, however, is unknown, which prevents its effective sampling. Ideally, we’d like $z$ to follow some simple distribution, such as $N(0,I)$, as it is easy to sample from. As summarized in Figure 2, VAE makes a few changes to the autoencoder architecture to make $D_{𝜃}$ able to take samples from $N(0,I)$ as input and map them to $𝕏$.

• Instead of giving out concrete samples in $\mathbb{Z}$, $E_{\varphi }$ outputs the parameters for the probability density function $p_{\varphi }\left(z\right|x)$.
• $p_{𝜃}\left(z\right|x)$ is required to be a multivariant normal distribution with independent components. That is, $p_{𝜃}\left(z\right|x)=N({\mu }_{\varphi }\left(x\right),{\sigma }_{\varphi }^2\left(x\right))$ where ${\sigma }_{\varphi }^2\left(x\right)$ is a diagonal matrix.
• ${\mu }_{\varphi }\left(x\right)$ is penalized for being different from $0$, and ${\sigma }_{\varphi }^2\left(x\right)$ for being different from $I$. With this penalty, $p_{𝜃}\left(z\right|x)$ approximately follows $N(0,I)$, so does $p\left(z\right)$ as $p\left(z\right)=\mathrm{\int \; <!--nolimits-->}<!--FUNCTION\; APPLICATION-->p\left(x\right)p\left(z\right|x)dx=\mathrm{\int \; <!--nolimits-->}<!--FUNCTION\; APPLICATION-->p(x\left)N\right(z;0,I)dx=N(z;0,I)$.
• A new sampler component $S$ is introduced which, given ${\mu }_{\varphi }\left(x\right)$ and ${\sigma }_{\varphi }^2\left(x\right)$, draws a sample $z\sim N\left({\mu }_{\varphi }\right(x),{\sigma }_{\varphi }^2(x\left)\right)$. $z$ is then fed to the decoder $D_{𝜃}$, just as in a regular autoencoder.

How exactly are ${\mu }_{\varphi }\left(x\right)$ and ${\sigma }_{\varphi }^2\left(x\right)$ penalized? Compute the K-L divergence between $N\left({\mu }_{\varphi }\right(x),{\sigma }_{\varphi }^2(x\left)\right)$ and $N(0,I)$ as below:

In (3), $d$ is the dimension of $\mathbb{Z}$. Removing the constants from (3) and estimating it with samples, our final K-L divergence loss is

The reconstruction loss for VAE, is also slightly different from that for a regular autoencoder. It can be estimated by the following equation, given a function $S(\mu ,{\sigma }^2)$ that returns a sample from $N(\mu ,{\sigma }^2)$.

This formulation has one big problem. $S(\cdot ,\cdot )$ is not differentiable, which makes the reconstruction loss not amenable to back-propagation based optimization. Luckily, it is easy to rewrite $S\left({\mu }_{\varphi }\right(x_i),{\sigma }_{\varphi }^2(x_i\left)\right)$ as ${\mu }_{\varphi }\left(x_i\right)+S(0,I)\odot {\sigma }_{\varphi }\left(x_i\right)$, where $\odot$ is the element-wise product and ${\sigma }_{\varphi }\left(x_i\right)$ is ${\sigma }_{\varphi }^2\left(x_i\right)$’s diagonals arranged in a vector form, by leveraging the reparameterization trick for normal distributions. The final formulation for the reconstruction loss therefore is

The total loss combines the K-L divergence loss in (4) and the reconstruction loss in (5) with a weight hyperparameter $\lambda$:

$\lambda$ controls the relative importance between reconstructing the original samples and making sure $z$ follows $N(0,I)$. It is likely that different data sets require different $\lambda$.

In practice, instead of outputting ${\sigma }_{\varphi }^2\left(x_i\right)$, $E_{\varphi }$ outputs $\log <!--FUNCTION\; APPLICATION-->{\sigma }_{\varphi }^2\left(x_i\right)$, but that is only a minor engineering detail.

3 Bayesian View

This section derives the total loss objective function through a Bayesian view.

The maximum likelihood method is often used to optimize a neural network that takes samples $x_i$ as input and produces $p_{𝜃}\left(x_i\right)$. Assuming each of these $x_i$ are i.i.d samples, the likelihood of observing all of them is $p(x_1,x_2,x_3,\cdots ,x_n)=\prod <!--FUNCTION\; APPLICATION-->p_{𝜃}\left(x_i\right)$. The training objective is to maximize $\prod <!--FUNCTION\; APPLICATION-->p_{𝜃}\left(x_i\right)$, which is equivalent to minimizing the expected negative log odds:

Considering for now only the decoder $D_{𝜃}$ part of VAE. It maps a sample $z$ to $\tilde{x}$, but it can also be viewed as spitting out parameters for $p_{𝜃}\left(x\right|z)$. More specifically it spits out ${\mu }_{𝜃}\left(z\right)$, the parameters in $N(x;{\mu }_{𝜃}(z),I)$. If the maximum likelihood method is to be used for finding the optimal $𝜃$, $p_{𝜃}\left(x\right)$ is needed which can be calculated this way: $p_{𝜃}\left(x\right)=\mathrm{\int \; <!--nolimits-->}<!--FUNCTION\; APPLICATION-->p_{𝜃}(x,z)dz=\mathrm{\int \; <!--nolimits-->}<!--FUNCTION\; APPLICATION-->p\left(z\right)p_{𝜃}\left(x\right|z)dz=E_{z\sim p\left(z\right)}{p}_{𝜃}(x\left|z\right)$. Estimating $p_{𝜃}\left(x\right)$ this way, however, is intractable due to the number of dimensions $z$ potentially has.

Assuming there is an effective way of sampling $z$ that follows a distribution $p_{\varphi }\left(z\right|x)$, which may or may not be equal to $p_{𝜃}\left(z\right|x)$, $p_{𝜃}\left(x\right)$ can be calcualted the following way:

Note that in the derivation above, the only requirement of $p_{\varphi }\left(z\right|x)$ is to be a valid probability density function. Is there such a $p_{\varphi }\left(z\right|x)$ that is easy to sample from? Yes, that is exactly the responsibility of VAE’s encoder $E_{\varphi }$ which takes in $x$ and spits out the parameters for normal distribution $p_{\varphi }\left(z\right|x)$: ${\mu }_{\varphi }\left(x\right)$ and ${\sigma }_{\varphi }^2(x$).

With $p_{𝜃}\left(x\right)$ estimated this way, $\log <!--FUNCTION\; APPLICATION-->p_{𝜃}\left(x\right)$ becomes:

3.1 Change of Optimization Objective

With all the derivation steps, it is still not clear how to calcuate $\log <!--FUNCTION\; APPLICATION-->p_{𝜃}\left(x\right)$ precisely.

Given $\log <!--FUNCTION\; APPLICATION-->(\cdot )$ is a concave function, Jensen’s inequality states that $\log <!--FUNCTION\; APPLICATION-->E\left(x\right)\ge E(\log <!--FUNCTION\; APPLICATION-->(x\left)\right)$. We thus have:

It is now possible to estimate the right hand side of the inequality, broadly known as the Evidence Lower BOund (ELBO), which can be rewritten further:

Since $p\left(z\right)=N(0,I)$, the second term in (7), as already calculated by (3) in Section 2, is:

At the beginning of this section, $p_{𝜃}\left(x\right|z)$ is already required to take the form of $N\left({\mu }_{𝜃}\right(z),I)$ which means:

If in the process of estimating $E_{z\sim p_{\varphi }\left(z\right|x)}\log <!--FUNCTION\; APPLICATION-->p_{𝜃}\left(x\right|z)$, only one single sample $z$ drawn which is equal to ${\mu }_{\varphi }\left(x_i\right)+S(0,I)\odot {\sigma }_{\varphi }\left(x_i\right)$, the first term in (7) gives

Putting the maximum likelihood and the two terms of the ELBO together, we’ve arrived at:

Removing the constants and $\frac{1}{2}$, our final optimization objective shown above is identical to (6) in Section 2, keeping in mind that

• ${\mu }_{𝜃}$ and $D_{𝜃}$ are the same funcion, and
• with this theoretical foundation, the need for a $\lambda$ is also eliminated.

3.2 The ELBO gap

Since the optimization objective is changed from the log likelihoods to the ELBO, it is helpful to understand the gap betwen the two.

The gap is 0 when $p_{\varphi }\left(z\right|x)$ and $p_{𝜃}\left(z\right|x)$ are identical.

4 Joint Distribution View

Section 3’s derivation starts from the maximum likelihood objective, and then switches to maximizing the ELBO. This section provides a simpler joint distribution approach to derive the ELBO objective directly, inspired by Jianlin Su at http://kexue.fm.

In VAE, once trained, the decoder $D_{𝜃}$ can be used as an independent generative model, without the encoder. The encoder can also be used without the decoder as a descriminative model. The training process is what links both components together. It is reasonable to require them to work on the same distribution about both $x$ and $z$. That is, our objective is to minimize the K-L divergence between $p_{𝜃}(x,z)$ and $p_{\varphi }(x,z)$:

The definition of the ELBO in Section 3 can be used to verify this.

5 Latent Space

In VAE’s training process, $p\left(x\right)$ and $p\left(z\right)=N(0,I)$ are given, VAE learns $p_{\varphi }\left(z\right|x)$ and $p_{𝜃}\left(x\right|z)$ simultaneously. Note however that $p_{𝜃}\left(x\right)$ is never directly optimized to match $p\left(x\right).$ This could be one major reason why VAE is not known to generate very realistic images.

VAE’s encoder, on the other hand, is a very reasonable feature extraction tool. Suppose there are a bunch of sample human face pictures labelled with whether the person has large eyes or not. Denote these samples by $(x,y)$ where $x$ is the image, and $y=1$ if the person has large eyes and 0 otherwise. A vector $e$ in $\mathbb{Z}$ calculated the following way probably captures the latent representation of large eyes.

Given any human face picture $x$, ${\mu }_{𝜃}(x+\lambda e)$ should generate a variation of $x$ that has big or small eyes as $\lambda$ varies.