Clarification of variational autoencoders

Question

I've been spending almost a month trying to understand VAE. I was reading a bunch of tutorials, and at first, it made sense and seemed straightforward. Then I was experimenting with it, it produced weird results, had no idea what was wrong. While trying to figure it out realised that I don't understand it at all.

This question will consist of multiple small questions that I probably should ask separately, but it would lose the context, so I will try to separate them.

My first impression was that the only difference between AE and VAE is:

1. the encoder, instead of producing a single point in N-dimensional space, is producing an N-dimensional range where the point can be;
2. the KL divergence loss - added to the reconstruction loss - is somehow magically forcing the encoder to produce the ranges in the way that they will overlap ("touch") each other, and be centred, so the ranges will try to be close to 0 as possible, but at the same time ranges (distributions) with different classes will push each other.

When I just simply looked at the equation (of KL loss), it sort of made sense (however, didn't understand why it did) why it does centring the distributions, and why it made the distributions overlap, but didn't (and doesn't) make sense at all how it forces the distributions of same classes to "group" with each other. Simply because the KL loss equation (seemingly) has no idea about which class is which, and the reconstruction loss equation doesn't care how far away the same classes are from each other (and has no direct connection with the latent distribution of the given observation) as long as the observations are close to the expected output.

---

Then I realised that classes in VAE don't even exist, it is just in my diagram rendering the means given the MNIST dataset I used for training, and colored the points with the corresponding classes.

So the encoder is trying to organise the inputs to put similar ones close to each other, and slowly transform them into different inputs. But still not understanding how.

I started to learn about KL divergence, and how it works. KL divergence (or relative entropy) compares probability distributions and measures how different they are from each other. The question I was trying to find the answer for is in terms of VAE, what are we comparing with what?

After reading, seeing examples, and experimenting, I thought we're comparing the distributions on the batch axis, that would make sense. It would try to make the distributions very similar to each other, but the reconstruction loss wouldn't let them, and the tension between them would produce the beautiful transition of the different outputs.

---

But then I realised that this isn't the case, as the equation doesn't directly make the different distributions to contact with each other. I found the brief explanation of the equation we're using in Wikipedia:

A special case, and a common quantity in variational inference, is the relative entropy between a diagonal multivariate normal, and a standard normal distribution (with zero mean and unit variance):

So I realised that we're comparing our distributions with the distribution where we sampled $\epsilon$ randomly, which is a normal distribution with mean zero and standard deviation equal to one.

Does this mean that we're just forcing our distributions to have 0 means, and standard deviations 1? Or if there are other goals, how do we earn them with this equation?

---

I see the elements of VAE mentioned as the terms of Bayes's theorem. For example, a very detailed, and nice tutorial mentions:

The weight of KL divergence in the loss function is a hyperparameter we shouldn’t ignore at all: it adjusts the “distance” between the prior and posterior distribution of z, and plays a decisive role in the performance of the model.

So as I read about the Bayes theorem, in short, the posterior shows the probability that our hypothesis is true given the evidence is true, by multiplying the likelihood (probability of the evidence is true given the hypothesis is true) with the prior (independent probability of the hypothesis is true), divided with the independent probability of evidence is true.

Given the examples that made me sort of understand the Bayes theorem, it is not too clear in VAEs, which part is the prior, posterior, likelihood, hypothesis, and evidence?

Answer 1

From a probabilistic perspective, a variational autoencoder is a latent variable model e.g. a generative model for $X$ using latent variables $Z$. By standard rules of probability, we can write the true distribution over $X$ as $$ p(X) = \int p(X,Z) dZ$$ Generally we the goal is to describe high dimensional data $X$ as e.g. images with low dimensional latent variables $Z$. From a Bayesian point of view it is clear how we get this if the generative model $p(X,Z)$ is given, just using Bayes theorem $p(Z|X) = \frac{p(X,Z)}{p(X)} = \frac{p(X|Z)p(Z)}{p(X)}$. However sometimes it is hard to write down $p(X,Z)$ analytically, instead one may use deep generative models to capture the latent structure e.g. we can use $$ p_\theta (X) = \int p_\theta (X,Z) dZ = \int p_\theta (X|Z)p(Z) dZ$$ where the likelihood $p_\theta (X|Z)$ is parameterized by a neural network with parameters $\theta$. The whole goal now is the find the parameters $\theta$ that minimize the KL-divergence between $p_\theta (x)$ and the observed data distribution $p_D(x)$ i.e. $$KL(p_D(x) || p_\theta(x)) = E[\log p_D(x)] - E[\log p_\theta (x)]$$ Note that the first expectation is constant w.r.t $\theta$ so we do not need it for minimization. Hence this problem is equivalent to maximize the expected log evidence $E[\log p_\theta(x)]$. This, unfortunately, is intractable. However, we can maximize its lower bound the ELBO. Todo so we have to introduce a variation distribution $q_\phi (z|x) \approx p(z|x)$ parameterized by $\phi$ approximating the posterior over the latent variables. Note that this will be our "Encoder", hence typically also implemented by a neural network. We can write $$ E[\log p_\theta (x)] = \mathcal{L} (x,\theta, \phi) + KL(q_\phi (z|x) || p_\theta (z|x)) \geq \mathcal{L} (x,\theta, \phi)$$

Which uses the fact that the KL divergence is always positive. The ELBO is now given by

$$ \mathcal{L} (x,\theta, \phi) = \sum_n E_q[\log p_\theta (x_n|z_n)] - KL(q_\phi (z_n|x_n) || p(z_n)) $$

So for Variational autoencoders (VAE) the goal is now to find both the parameters for the generative model $\theta$ a.k.a the decoder and the parameters for the variational distribution $\phi$ a.k.a the encoder! This is done by maximizing the ELBO. As we can see this will either maximize $E_q[\log p_\theta (x_n|z_n)]$ (likelihood, decoder performance) or minimize $KL(q_\phi (z_n|x_n) || p(z_n))$ (keeps the encoder as similar to the prior). This may be a confusing step because it does not include maximizing the "encoder" performance. But notice that the ELBO is an lower bound to $E[\log p_\theta(x)] = \mathcal{L} (x,\theta, \phi) + KL(q_\phi (z,x) || p_\theta (z|x))$ , hence maximizing $\mathcal{L}$ also minimizes $KL(q_\phi (z|x) || p_\theta (z|x))$ (the encoder is near the exact posterior)!

For the "standard" VAE you typically choose the prior $p(z) = \mathcal{N}(0,1)$ and $q_\phi(z|x) = \mathcal{N}(\mu(x), \sigma(x))$ where $\mu(x), \sigma(x)$ are given by a neural network. Thus is because the the KL divergence becomes analytical as you notices $$KL(\mathcal{N}(\mu,\sigma)||\mathcal{N}(0,1)) = \frac{1}{2}(\sigma + \mu^2 - 1 - \log \sigma )$$ Hence we can easily calculate the ELBO. What now typically is done is to pass $x$ through the encoder and obtain $\mu, \sigma$. Notice the ELBO still has intractable expectation which however can be approximated using MC methods by sampling $z = \mu + \sigma\epsilon$ with $\epsilon \sim N(0,1)$. Here I also used the reparameterization trick which allows us backpropagation of gradient for $\mu, \sigma$. Then we pass the latent variables through the decoder and compute gradient/update parameters...

Hope this helps and there are not too many typos :)