Mathematics in VAE for Engineers.(WIP)

---

Name	Notation	Description
Prior	$𝑝\left(𝑧\right)$	The overall distribution of the latent variable $𝑧$. For example, $𝑝\left(𝑧\right)=\mathrm{𝒩︀}\left(0,𝐼\right)$
likelihood	$𝑝\left(𝑥|𝑧\right)$	The probability of the observed variable $𝑥$ given the latent variable $𝑧$. It also describes the generative process, so it is also called probabilistic decoder.
posterior	$𝑝\left(𝑧|𝑥\right)$	The probability of the latent variable $𝑧$ given the observed variable $𝑥$.
Evidence	$𝑝\left(𝑥\right)$	You can think of Evidence as how likely a generative model is to generate a particular image or how reasonable the model considers this image to be. Also called Marginal Likelihood in original VAE paper

These four distributions are related by Bayes’ theorem:

$$𝑝\left(𝑧|𝑥\right)=\frac{𝑝\left(𝑥|𝑧\right)𝑝\left(𝑧\right)}{𝑝\left(𝑥\right)}$$

Since Encoder-Decoder architecture has become super common in deep learning, we won’t discuss the application of VAE in detail. Instead, we will focus on the math behind VAE and try to restore the history background of when VAE was proposed.

A common misconception is that VAE was originally motivated purely by the task of generating new data. In fact, one of its central motivations was a classical problem in probabilistic modeling: when the posterior over the latent variables is intractable, how can we efficiently infer the latent variables and learn the parameters of a generative model? You might ask, why is the posterior inference problem important? Generative models seem to be much more popular. In statistics, a classical scenario is that we observed some data $𝑥$, which can be huge, and we want to know what process produced them.¹

Since we are interested in the posterior inference problem, we can first assume that the decoder $𝑝\left(𝑥|𝑧\right)$ is a neural network, while its parameters remain unknown. We denote it as ${𝑝}_{𝜃}\left(𝑥|𝑧\right)$, where $𝜃$ represents the parameters of the neural network. And, $𝑝\left(𝑧\right)$ as the distribution over latent space, is determined by Encoder.² Since both encoder and decoder are something we are going to find, we can simply impose a prior distribution(e.g., a standard isotropic Gaussian) to act as a “constraint” on both models. Then according to Bayes’ theorem, we can compute the posterior $𝑝\left(𝑧|𝑥\right)=\frac{{𝑝}_{𝜃}\left(𝑥|𝑧\right)𝑝\left(𝑧\right)}{𝑝\left(𝑥\right)}$. The only problem is that the evidence $𝑝\left(𝑥\right)=\int {𝑝}_{𝜃}\left(𝑥|𝑧\right)𝑝\left(𝑧\right)𝑑𝑧$ is often intractable when the likelihood is parameterized by a nonlinear neural network, especially when the latent space is high-dimensional.

Wait, why can you just assume that the generative model ${𝑝}_{𝜃}\left(𝑥|𝑧\right)$ is known? How about we assume that the posterior $𝑝\left(𝑧|𝑥\right)$ is known? Just use a neural network to fit it, then the problem is solved! Of course you can do that, but then the Bayes’ theorem would become $𝑝\left(𝑥|𝑧\right)=\frac{𝑝\left(𝑧|𝑥\right)𝑝\left(𝑥\right)}{𝑝\left(𝑧\right)}$, then $𝑝\left(𝑥\right)=\int 𝑝\left(𝑧|𝑥\right)𝑝\left(𝑥\right)𝑑𝑧$ is still intractable. In other words, once you know side of the Bayes’ theorem, you already have a theoretical solution provided by it; we just cannot obtain its explicit form.

What if we use two neural networks to model the generative process and the posterior inference process simultaneously? This is the central idea behind VAE. We use a decoder ${𝑝}_{𝜃}\left(𝑥|𝑧\right)$ to model how data are generated from latent variables, and an encoder ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$ to approximate the intractable posterior (${𝑝}_{𝜃}\left(𝑧|𝑥\right)$). However, this introduces a new challenge: how can we jointly optimize the two networks when the true posterior and the marginal likelihood are intractable?

Remember that we now have a Decoder ${𝑝}_{𝜃}\left(𝑥|𝑧\right)$ and a theoretical Encoder ${𝑝}_{𝜃}\left(𝑧|𝑥\right)$ derived from our decoder via Bayes’ theorem. Since we can’t obtain the explicit form of this theoretical Encoder, we are now using another neural network, with its parameters denoted as $\mathrm{\Phi }$, to approximate it. This approximated Encoder is denoted as ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$.

For simplicity’s sake, I will use ${𝑞}_{\mathrm{\Phi }}$ to represent ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$ and ${𝑝}_{𝜃}$ to represent ${𝑝}_{𝜃}\left(𝑧|𝑥\right)$ in this section.

Since one of our goals is to make ${𝑞}_{\mathrm{\Phi }}$ to be as close to ${𝑝}_{𝜃}$ as possible, we can simply minimize its KL divergence.

$${𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}={\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\cdot log\left(\frac{{𝑞}_{\mathrm{\Phi }}}{{𝑝}_{𝜃}}\right)𝑑𝑧$$

And don’t forget that we still don’t know the Decoder’s parameters $𝜃$ at this point. So another goal is naturally optimize our Decoder, which means to maximize the probability of generating data in our datasets ${𝑝}_{𝜃}\left(𝑥\right)$. We will start with only KL divergence, in later steps you will see ${𝑝}_{𝜃}\left(𝑥\right)$ emerges by itself.

Because of the definition of expected value, we can rewrite this integral as

$$\begin{array}{cc}{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}& ={\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\cdot log\frac{{𝑞}_{\mathrm{\Phi }}}{{𝑝}_{𝜃}}𝑑𝑧\\ & ={𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log\frac{{𝑞}_{\mathrm{\Phi }}}{{𝑝}_{𝜃}}\right]\\ & ={𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑞}_{\mathrm{\Phi }}\right]-{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\right]\end{array}$$

We know that ${𝑝}_{𝜃}\left(𝑧|𝑥\right)=\frac{{𝑝}_{𝜃}\left(𝑧,𝑥\right)}{{𝑝}_{𝜃}\left(𝑥\right)}$, thus

$$\begin{array}{cc}{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}& ={𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑞}_{\mathrm{\Phi }}\right]-{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[\frac{log\left({𝑝}_{𝜃}\left(𝑧,𝑥\right)\right)}{{𝑝}_{𝜃}\left(𝑥\right)}\right]\\ & ={𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑞}_{\mathrm{\Phi }}\right]-{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)\right]+{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑥\right)\right]\end{array}$$

${𝑝}_{𝜃}\left(𝑥\right)$ is not related to $𝑧$, thus ${𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑥\right)\right]=log{𝑝}_{𝜃}\left(𝑥\right)$.

$${𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}=log{𝑝}_{𝜃}\left(𝑥\right)-{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]$$

Here we get both variables we want to optimize in a single equation. And notice that

$${𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]=log{𝑝}_{𝜃}\left(𝑥\right)-{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}$$

meaning by maximizing ${𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]$ we can minimize the KL divergence and maximize ${𝑝}_{𝜃}\left(𝑥\right)$ at the same time! We now call this useful term Evidence Lower Bound or ELBO for short

$$\text{ELBO}={𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]$$

, since

$$log{𝑝}_{𝜃}\left(𝑥\right)={𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}+{𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]$$

and ${𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\Vert {𝑝}_{𝜃}\right)}\ge 0$, thus ${𝔼}_{{𝑞}_{\mathrm{\Phi }}}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\right]$ is the lower bound of log evidence $log{𝑝}_{𝜃}\left(𝑥\right)$.

Now the issue becomes how do we maximize ELBO. A natural solution is to find its gradient, and then we can perform gradient descent to $-\text{ELBO}$. Let’s first try to find the gradient w.r.t⁴ $𝜃$.

$$\begin{array}{cc}{\nabla }_{𝜃}\text{ELBO}& ={\nabla }_{𝜃}\left({𝔼}_{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]\right)\\ & ={\nabla }_{𝜃}\left({\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]𝑑𝑧\right)\\ & ={\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right){\nabla }_{𝜃}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]𝑑𝑧\end{array}$$

Since the second term $log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$ is not related to $𝜃$, its gradient w.r.t $𝜃$ is just zero.

$$\begin{array}{cc}& {\nabla }_{𝜃}\text{ELBO}\\ & ={\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right){\nabla }_{𝜃}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)\right]𝑑𝑧\\ & ={𝔼}_{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}{\nabla }_{𝜃}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)\right]\end{array}$$

After The Reparameterization Trick, we can perform Monte Carlo sampling to estimate the value of this expectation. The basic idea is just sampling bunch of latents $𝑧$ out of the distribution ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$, which can be done after The Reparameterization Trick, then calculate the average value of the gradient given those z.

$$\begin{array}{cc}& {𝔼}_{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}{\nabla }_{𝜃}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)\right]\\ & \approx \frac{1}{𝐿}\sum _{𝑙=1}^{𝐿}{\nabla }_{𝜃}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)\right]\end{array}$$

Where $𝐿$ means the number of Monte Carlo samples.

To sum up, so far there is nothing blocking us to find the gradient of the ELBO w.r.t $𝜃$. But we won’t be that lucky when we try to find the gradient w.r.t $\mathrm{\Phi }$.

$$\begin{array}{cc}{\nabla }_{\mathrm{\Phi }}\text{ELBO}& ={\nabla }_{\mathrm{\Phi }}\left({𝔼}_{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]\right)\\ & ={\nabla }_{\mathrm{\Phi }}{\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]𝑑𝑧\\ & ={\int }_{𝑧}{\nabla }_{\mathrm{\Phi }}\left[{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]\right]𝑑𝑧\\ & ={\int }_{𝑧}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right){\nabla }_{\mathrm{\Phi }}\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]𝑑𝑧+{\int }_{𝑧}{\nabla }_{\mathrm{\Phi }}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\left[log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]𝑑𝑧\end{array}$$

Once again, we can do Monte Carlo sampling to the first term since it is an expectation. But we can’t do so to the second term. That’s exactly where The Reparameterization Trick comes in.

The origin of the problem is that ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$, which is the distribution we’re sampling from, is related to the value we are finding gradient with respect to ($\mathrm{\Phi }$). So the solution here is to replace ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$ with an unrelated distribution, while keeping the value of ELBO unchanged. That sounds impossible first, but a special design is applied here to make this possible.

The encoder network doesn’t directly output a latent variable $𝑧$, instead, it gives you mean and variance according to $𝑥$ deterministically. Later we sample an $𝑧$ from the Gaussian distribution given those parameters.

We denote the mean and variance as ${𝜇}_{\mathrm{\Phi }}\left(𝑥\right)$ and ${𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)$, thus we have

$${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)=\mathrm{𝒩︀}\left({𝜇}_{\mathrm{\Phi }}\left(𝑥\right),{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)\right)$$

This is also true for the decoder. (In some cases, the variance is fixed.)

$${𝑝}_{𝜃}\left(𝑥|𝑧\right)=\mathrm{𝒩︀}\left({𝜇}_{𝜃}\left(𝑧\right),{𝜎}_{𝜃}^2\left(𝑧\right)\right)$$

We will use this feature later on.

Now we know that $𝑧$ is sampled from an Gaussian distribution, thus

$$\begin{array}{cc}{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)& =\frac{1}{{𝜎}_{\mathrm{\Phi }}\left(𝑥\right)\sqrt{2𝜋}}exp\left(-\frac{{\left(𝑧-{𝜇}_{\mathrm{\Phi }}\left(𝑥\right)\right)}^2}{2{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)}\right)\\ \text{let}𝑓\left(𝑧\right)& \stackrel{\text{def}}{=}log{𝑝}_{𝜃}\left(𝑧,𝑥\right)-log{𝑞}_{\mathrm{\Phi }}\\ \text{ELBO}& =\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)𝑓\left(𝑧\right)𝑑𝑧\\ & =\int \frac{1}{{𝜎}_{\mathrm{\Phi }}\left(𝑥\right)\sqrt{2𝜋}}exp\left(-\frac{{\left(𝑧-{𝜇}_{\mathrm{\Phi }}\left(𝑥\right)\right)}^2}{2{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)}\right)𝑓\left(𝑧\right)𝑑𝑧\end{array}$$

With that in mind, now let’s define a new variable $𝜀$, sampled from Standard Normal distribution $\mathrm{𝒩︀}\left(0,1\right)$. Then there is the key step in The Reparameterization Trick: we can use linear transformation(denoted as ${𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)$) to convert $𝜀$ to $𝑧$.

$$\begin{array}{c}\text{If}𝜀\sim \mathrm{𝒩︀}\left(0,1\right)\text{and}𝑧\sim \mathrm{𝒩︀}\left({𝜇}_{\mathrm{\Phi }}\left(𝑥\right),{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)\right)\\ \text{Then}𝑧={𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)=𝜇+𝜎𝜀\\ 𝑑𝑧=𝜎𝑑𝜀\end{array}$$

Let’s now replace all the $𝑧$ with $𝜀$

$$\begin{array}{cc}\text{ELBO}& =\int \frac{1}{𝜎\sqrt{2𝜋}}exp\left(-\frac{{\left(𝜎𝜀\right)}^2}{2{𝜎}^2}\right)𝑓\left(𝜇+𝜎𝜀\right)𝜎𝑑𝜀\\ & =\int \frac{1}{\sqrt{2𝜋}}exp\left(-\frac{{𝜀}^2}{2}\right)𝑓\left(𝜇+𝜎𝜀\right)𝑑𝜀\\ & ={𝔼}_{𝜀\sim \mathrm{𝒩︀}\left(0,1\right)}\left[𝑓\left(𝜇+𝜎𝜀\right)\right]\end{array}$$

At this point, the distribution used for sampling the expectation has been changed from the problematic ${𝑞}_{\mathrm{\Phi }}$ to $𝜀$, which is independent of the optimizable parameters $\mathrm{\Phi }$. Thus, finding gradient is not a problem anymore:

$${\nabla }_{𝜃,\mathrm{\Phi }}\text{ELBO}\approx \frac{1}{𝐿}\sum _{𝑖=1}^{𝐿}{\nabla }_{𝜃,\mathrm{\Phi }}\left[log{𝑝}_{𝜃}\left(𝑥,{𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)\right)-log{𝑞}_{\mathrm{\Phi }}\left({𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)|𝑥\right)\right]$$

This optimization method is called Stochastic Gradient Variational Bayes or SGVB. To sum up, SGVB did 2 things:

1. Monte Carlo estimates the expectation.
2. Reparameterization separates the randomness from the optimizable parameters, making gradient can be propagated to $\mathrm{\Phi }$.

Having gone through the journey above, we are now ready to optimize the ELBO. But we can do even better.

For simplicity, I will use the letter $𝑧$ below to denote ${𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)$, but please keep in mind that this is the result after reparameterization.

$$\begin{array}{cc}\text{ELBO}& ={𝔼}_{𝜀\sim \mathrm{𝒩︀}\left(0,1\right)}\left[\underset{\begin{array}{cc}& =log\left({𝑝}_{𝜃}\left(𝑥|𝑧\right)\cdot {𝑝}_{𝜃}\left(𝑧\right)\right)\\ & =log{𝑝}_{𝜃}\left(𝑥|𝑧\right)+log{𝑝}_{𝜃}\left(𝑧\right)\end{array}}{\underbrace{log{𝑝}_{𝜃}\left(𝑥,𝑧\right)}}-log{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\right]\\ & ={𝔼}_{𝜀\sim \mathrm{𝒩︀}\left(0,1\right)}\left[log{𝑝}_{𝜃}\left(𝑥|𝑧\right)\right]-\underset{{𝐷}_{\text{KL}}\left({𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\Vert {𝑝}_{𝜃}\left(𝑧\right)\right)}{\underbrace{{𝔼}_{𝜀\sim \mathrm{𝒩︀}\left(0,1\right)}\left[log\frac{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}{{𝑝}_{𝜃}\left(𝑧\right)}\right]}}\\ & \approx \underset{\text{Reconstruction Term}}{\underbrace{\frac{1}{𝐿}\sum _{𝑖=1}^{𝐿}log{𝑝}_{𝜃}\left(𝑥|{𝑧}^{\left(𝑖\right)}\right)}}-\underset{\text{Regularization Term}}{\underbrace{{𝐷}_{\text{KL}}\left({𝑞}_{\mathrm{\Phi }}\left({𝑧}^{\left(𝑖\right)}|𝑥\right)\Vert {𝑝}_{𝜃}\left({𝑧}^{\left(𝑖\right)}\right)\right)}}\end{array}$$

This way, the Reconstruction Term can keep being estimated using Monte Carlo sampling, while the KL Divergence can be computed analytically when both distributions are Gaussian.(${𝑝}_{𝜃}\left(𝑧\right)=\mathrm{𝒩︀}\left(0,1\right)$ because we want our latent space to be Gaussian.)

As the same trick played in The Reparameterization Trick, ${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)$ can be represented as

$${𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)=\mathrm{𝒩︀}\left({𝜇}_{\mathrm{\Phi }}\left(𝑥\right),{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)\right)=\frac{1}{{𝜎}_{\mathrm{\Phi }}\left(𝑥\right)\sqrt{2𝜋}}exp\left(-\frac{{\left(𝑧-{𝜇}_{\mathrm{\Phi }}\left(𝑥\right)\right)}^2}{2{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)}\right)$$

Again, for simplicity, I will use the $𝜎$ and $𝜇$ below to denote ${𝜇}_{\mathrm{\Phi }}\left(𝑥\right)$ and ${𝜎}_{\mathrm{\Phi }}\left(𝑥\right)$. Then the KL Divergence can be expanded as follows:

$$\begin{array}{cc}{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\parallel {𝑝}_{𝜃}\left(𝑧\right)\right)}& =\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)log\left(\frac{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}{{𝑝}_{𝜃}\left(𝑧\right)}\right)d𝑧\\ & =\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)log\left(\frac{1}{𝜎}exp\left(-\frac{1}{2}\left(\frac{{\left(𝑧-𝜇\right)}^2}{{𝜎}^2}-{𝑧}^2\right)\right)\right)d𝑧\\ & =-\frac{1}{2}\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\left[log{𝜎}^2-{𝑧}^2+\frac{1}{{𝜎}^2}{\left(𝑧-𝜇\right)}^2\right]d𝑧\\ & =-\frac{1}{2}\left[log{𝜎}^2\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧-\int {𝑧}^2{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧+\frac{1}{{𝜎}^2}\int {\left(𝑧-𝜇\right)}^2{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧\right]\end{array}$$

Using the properties of the Gaussian distribution:

$$\begin{array}{cc}\int {𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧& =1\\ \int {𝑧}^2{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧& ={𝜇}^2+{𝜎}^2\\ \int {\left(𝑧-𝜇\right)}^2{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)d𝑧& ={𝜎}^2\end{array}$$

Substituting these results:

$$\begin{array}{cc}{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\parallel {𝑝}_{𝜃}\left(𝑧\right)\right)}& =-\frac{1}{2}\left[log{𝜎}^2-\left({𝜇}^2+{𝜎}^2\right)+\frac{1}{{𝜎}^2}{𝜎}^2\right]\\ & =-\frac{1}{2}\left[log{𝜎}^2-{𝜇}^2-{𝜎}^2+1\right]\\ & =\frac{1}{2}\left[{𝜇}^2+{𝜎}^2-1-log{𝜎}^2\right]\end{array}$$

Substituting the closed-form KL Divergence into the ELBO:

$$\begin{array}{cc}\text{ELBO}& ={𝔼}_{{𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)}\left[log{𝑝}_{𝜃}\left(𝑥|𝑧\right)\right]-{𝐷}_{\text{KL}\left({𝑞}_{\mathrm{\Phi }}\left(𝑧|𝑥\right)\parallel {𝑝}_{𝜃}\left(𝑧\right)\right)}\\ & \approx \underset{\text{Reconstruction Term}}{\underbrace{\frac{1}{𝐿}\sum _{𝑖=1}^{𝐿}log{𝑝}_{𝜃}\left(𝑥|{𝑧}^{\left(𝑖\right)}\right)}}+\underset{\text{Regularization Term}}{\underbrace{\frac{1}{2}\left[log{𝜎}^2-{𝜇}^2-{𝜎}^2+1\right]}}\end{array}$$

where ${𝑧}^{\left(𝑖\right)}=𝜇+𝜎{𝜀}^{\left(𝑖\right)}$ and ${𝜀}^{\left(𝑖\right)}\sim \mathrm{𝒩︀}\left(0,1\right)$.

As we mentioned earlier, the network gives you means and variances, thus ${𝑝}_{𝜃}\left(𝑥|𝑧\right)$ is also Gaussian.

$$\begin{array}{cc}{𝑝}_{𝜃}\left(𝑥|𝑧\right)& =\mathrm{𝒩︀}\left({𝜇}_{𝜃}\left(𝑧\right),{𝜎}_{𝜃}^2\left(𝑧\right)\right)\\ & =\frac{1}{{𝜎}_{𝜃}\left(𝑧\right)\sqrt{2𝜋}}exp\left(-\frac{{\left(𝑥-{𝜇}_{𝜃}\left(𝑧\right)\right)}^2}{2{𝜎}_{𝜃}^2\left(𝑧\right)}\right)\end{array}$$ $$\begin{array}{cc}{\mathrm{ℒ︀}}_{\text{recon}}& =-\frac{1}{𝐿}\sum _{𝑖=1}^{𝐿}log{𝑝}_{𝜃}\left(𝑥|{𝑧}^{\left(𝑖\right)}\right)\\ & =\frac{1}{𝐿}\sum _{𝑖=1}^{𝐿}\left(log{𝜎}_{𝜃}\left(𝑧\right)+\frac{1}{2}log\left(2𝜋\right)+\frac{{\left(𝑥-{𝜇}_{𝜃}\left(𝑧\right)\right)}^2}{2{𝜎}_{𝜃}^2\left(𝑧\right)}\right)\end{array}$$

In the paper’s experiment they found that the number of samples $𝐿$ per data point can be set to 1 as long as the minibatch size was large enough. Thus ${\mathrm{ℒ︀}}_{\text{recon}}$ can be further simplified as

$${\mathrm{ℒ︀}}_{\text{recon}}=log{𝜎}_{𝜃}\left(𝑧\right)+\frac{1}{2}log\left(2𝜋\right)+\frac{{\left(𝑥-{𝜇}_{𝜃}\left(𝑧\right)\right)}^2}{2{𝜎}_{𝜃}^2\left(𝑧\right)}$$

If we now substitute $𝑧={𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)$ back and don’t consider $𝜎$ for now, ${\mathrm{ℒ︀}}_{\text{recon}}$ is proportional to ${\left(𝑥-{𝜇}_{𝜃}\left({𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)\right)\right)}^2$, which is exactly MSE loss.

For the Gaussian Decoder with a fixed unit variance, the training Loss for each data point in the end is:

$${\mathrm{ℒ︀}}_{\text{VAE}}=\frac{1}{2}\Vert 𝑥-{𝜇}_{𝜃}\left({𝑔}_{\mathrm{\Phi }}\left(𝜀,𝑥\right)\right){\Vert }^2-\frac{1}{2}\left[log{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)-{𝜇}_{\mathrm{\Phi }}^2\left(𝑥\right)-{𝜎}_{\mathrm{\Phi }}^2\left(𝑥\right)+1\right]$$

1. ¹For example, we observed a lot of symptoms of your body $𝑥$, then you might want to know what factors $𝑧$ cause these symptoms, such as disease.
2. ²In inference perspective, if you integrate $𝑝\left(𝑧|𝑥\right)𝑝\left(𝑥\right)$ over all the data $𝑥$, you do get the distribution of $𝑝\left(𝑧\right)$.
3. ³Strictly speaking, Monte Carlo sampling is used to estimate the expectation of a distribution.
4. ⁴with respect to