Variational Autoencoders (VAEs) Explained: From Theory to ELBO and the Reparameterization Trick

Introduction

In addition, we saw in the article about autoencoders that they can be used to solve other computer vision tasks, such as image denoising, object removal, or image inpainting.

Not familiar with autoencoders yet? I made a video that walks through them:

While vanilla autoencoders might seem powerful, they still have several limitations that can be improved. Given that, the researchers developed a more advanced version called variational autoencoders. We are going to study them in this article.

Motivation

In the previous article, we mostly discussed data compression as the main application of autoencoders. In reality, there are other interesting aspects we can look at.

Similarity

If you conduct experiments and try to calculate the distance between objects in latent space, you will find that, in many cases, the similarity properties between objects are not preserved to the same extent as in other common embedding algorithms. This happens because, during training, the model focuses primarily on reconstructing the original images rather than on the internal representations within latent space.

For example, if you take two similar images and pass them into a learned autoencoder, they are likely to be less close to each other in the latent space than with other algorithms. The same would go for dissimilar objects, which can be close to each other. As a result, similarity preservation is a weakness of vanilla autoencoders.

Image generation

Decoders are very good at reconstructing original images. As a result, a question that has attracted the attention of many researchers is whether, with the autoencoder architecture, it would be possible to generate high-quality new images.

Given that the decoder generates images from latent space, a natural way to generate new images is to sample new points from latent space and pass them as input to the decoder. Unfortunately, the quality of the generated images is usually poor.

Even if there is an encoded object in the latent space and a new point is sampled in its neighborhood, it seems logical that a generated image from that point would look very similar to the original. But again, this is not always true.

An analogous scenario would occur if you took a point midway between two embedded objects in the latent space. For instance, let’s imagine that object A is an image of a smiling person, while object B is an image of a person with a disappointed expression. If we sample a new point C midway between the two, in the ideal case, we would expect the decoder to generate an image of someone with a neutral facial expression. An autoencoder will likely fail to do this.

Given all of that, the main motivations for a new autoencoder version are related to how the latent space is constructed. We want the latent space to preserve similarity properties among embedded objects. This will allow us to perform geometric tricks with object distances and to sample new points in the space to generate new images.

Methodology

Let’s introduce formal notation to help us derive mathematical results.

We already have two distinct spaces:

The real space where the data distribution p(x) contains real images.
The latent space that contains all possible latent vectors.

Those two spaces are separate, and we would need functions to connect images from both spaces. To that end, we are going to introduce two distributions:

Posterior distribution p(z | x) returns the probability that a latent vector z was generated from image x.
Likelihood distribution p(x | z) that, given a vector z, reconstructs an image x from it.

The overall idea is simple: if we can sample vectors z from the posterior distribution p(z | x), we could use them to reconstruct and generate realistic images from the original distribution using the decoder (or the likelihood distribution) p(x | z).

The problem is that we do not know the form of the posterior distribution p(z | x). In other words, it is intractable, so we cannot compute or sample from it directly.

Approximation idea

To solve the problem above, VAE assumes that the latent distribution is normal, N(0, 1). This allows estimation of p(x | z).

In addition, VAE proposes to approximate the posterior distribution p(z | x) with a normal distribution q(z | x) = N(μ, σ) with learned parameters μ and σ, using an encoder.

To train the whole model, we would need to define a loss function. For that, we are going to use ELBO estimation.

Training objective

First, we set the objective to maximize the log-likelihood of p(x). Why?

Let’s suppose that xi is a random image from the training dataset. Then we are trying to maximize the following expression:

log p(x 1) + log p(x2) + … + log p(xn)

If a model learns to assign high probabilities to xi, indicating that they are realistic images, then it automatically learns how the real data look. It is the most logical training objective we can use. As a consequence, maximizing log-likelihood will make unrealistic images outside the training dataset have lower probabilities.

Given that, our goal is now to derive the loss function. To do so, we will use the ELBO.

ELBO*

ELBO stands for Evidence Lower Bound. It is a very useful lower bound on the log-likelihood. While it might seem math-heavy at first, we will go step by step to make the explanation easier.

Given log p(x), on line 1, we multiply it by 1, which we can write as ∫q(z | x) dz, since the integral of any probability distribution is 1.

On line 2, we place the expression log p(x) inside the integral – we can do this because log p(x) does not depend on the variable z.

The line 2 can be rewritten as an expected value with respect to the distribution q(z | x).

We then apply the conditional probability formula to the expression on line 3, rewriting p(x) as a quotient of two probabilities (line 4).

Then both the numerator and denominator are multiplied by q(z | x). We can then decompose the expression inside the logarithm into a sum of two logarithms.

Finally, from line 6, we can see that the second term in the sum can be rewritten as a Kullback-Leibler divergence.

On line 7, we obtain the final expression for log p(x). As we can see, it consists of two terms, where the first is called the ELBO. The second term represents the KL divergence, which is always non-negative.

To simplify the expression, let’s use the following trick: since KL divergence is always ≥ 0, we can turn the equation on line 7 into an inequality by writing that log p(x) is greater than or equal to the ELBO (line 8).

The training objective will remain coherent: to maximize log p(x), we need to maximize the ELBO.

Then, as before, we use the conditional probability formula (line 9), and on line 10 we decompose a logarithm into a sum of two terms.

The second term in the sum on line 10 is the negative KL divergence.

The final inequality is shown on line 11. We can see that it is expressed as a difference of two terms. Since our initial goal was to maximize ELBO, it is clear that we need:

to maximize the first term;
to minimize the second term.

Let’s now understand what these two terms semantically mean:

Term E q(z | x)[log p(x | z)]

This term makes the model maximize the image reconstruction quality of a sampled image x given its latent vector z.

p(x | z) is a Gaussian centered at the decoder’s output x̂. Given that the Gaussian formula includes a term exp(-(x – x̂)2 / 2σ2) and that our term contains a logarithm, the logarithm would remove the exponent and leave us with -(x – x̂)2 / 2σ2. The value 2σ2 is simply a constant that can be removed, so we are left only with -(x – x̂)2. Therefore, maximizing log p(x | z) is equivalent to minimizing MSE.

The term Eq(z | x) in the expected value simply means that an input image x is passed to the encoder, which returns z, and then the decoder is applied to z.

Term KL(q(z | x) || p(z))

KL-divergence measures how close two distributions are to each other. In our scenario, we want to minimize the KL-divergence, which means bringing both distributions closer.

We already know the form of the prior distribution over the latent space, p(z), which is a multivariate Gaussian distribution N(0, 1). At the same time, the posterior distribution q(z | x) is learned by the encoder, which, given an input image x from real data, returns a z vector. Therefore, the presence of the KL divergence term forces the encoder to return vectors that are essentially Gaussian N(0, 1).

Two terms – two different roles

So we found out that both terms in the final expression play different roles:

Similarly to the vanilla autoencoder, the term Efocuses on the decoder’s reconstruction quality.q(z | x)[log p(x | z)]
Unlike autoencoders, the term KL(q(z | x) || p(z))is also known as a regularization term and makes the encoder output vectors follow a Gaussian distribution.- Firstly, it lets us easily determine which distribution the latent vectors should be sampled from during inference for image generation.
By returning a distribution of points instead of a single isolated point, the latent space becomes filled more smoothly, with fewer gaps than in a vanilla autoencoder. This gives the decoder much broader exposure to data points, allowing it to successfully generate new images.

Reparametrization trick

We understood that during training, for a given input image, the encoder maps it to a normal distribution with predicted mean μ and variance σ2 in the latent space. A random image from that distribution is then sampled and passed to the decoder, which generates the image. The output loss is computed based on the generated image.

But there is a slight problem with this approach.

After computing the loss, we need to update the model’s weights by performing backpropagation. From a mathematical perspective, it is not clear how backpropagation can be applied through the sampling operator. This is where the reparametrization trick comes into play.

The idea of the trick is very simple. Instead of directly sampling a random image from the latent space, we introduce a random value ε sampled from a normal distribution N(0, 1). Then we scale it by the predicted variance σ2 and shift it by the predicted mean μ.

This process is essentially equivalent to the image sampling described above, except that the final result is now differentiable with respect to the mean μ and variance σ2. This smart move allows full backpropagation through the neural network!

Conclusion

We have studied in detail how the VAE is trained. Here is the summary:

An input image is passed to the encoder, which predicts the mean and variance of the image in the latent space. Based on that, the KL divergence loss is calculated to estimate how close the latent image distribution is to a normal distribution.
A sampled image from the latent space is passed to the decoder, which tries to reconstruct the original image. The reconstruction loss is computed by estimating how close the reconstructed image is in comparison to the original image.
Both loss values are summed, and backpropagation is performed.

For inference, it is enough to sample an image from the latent space and pass it to the decoder, which will generate a new image.

A great aspect of VAEs is that they can also be conditioned during training by concatenating embeddings that represent information about the respective classes, labels, or descriptions to generate specific images instead of random ones.

Resources

Auto-Encoding Variational Bayes | Diederik P. Kingma, Max Welling
Reparametrization trick | Wikipedia

All images unless otherwise noted are by the author.