Variational Autoencoders (VAEs)

Variational autoencoders for generation

About

Variational Autoencoders (VAEs) are models that can be used to

generate

new images similar to ones seen in a training dataset.

While VAEs themselves aren’t directly used in modern generative models, they form a basis for many of them including diffusion models.

Latent Variables

Since latent variables are at the heart of VAEs, it is important to define what they actually are.

If you have heard about latent variables before, you might have heard them described as hidden variables that determine the observed variables. I had a hard time wrapping my head around this so I will give a few examples:

Say you have a bunch of students who took an exam for the course MATH 101. Trivially, the observed variables are the exam scores. However, a latent variable would be the amount of MATH 101 that was actually learned by each student.
Plato’s Allegory of the Cave
. Lets say that you were born and raised inside a cave. You have never left this cave your entire life. Every day at dusk, there are people who walk by the outside of the cave. Due to the sun, you see shadows of these people cast into the cave. You never know what the people who are casting the shadows look like, instead you can only
observe
the shadows. Here the latent variables are the humans on the outside of the cave casting/controlling the shadows.
Le
t’s say you have 2 variables
${x,y}$
. These 2 variables are used to produce 5 dimensional vectors that are functions of the 2 variables.
For example:
${[x^2,x+2,y^3x,x^3+y^2,\sin(x)+e^y]}$
. We only
observe
the 5 dimensional vectors for different values of
${x}$
and
${y}$
. You can think of
${x}$
and
${y}$
as controlling the changes in the 5 dimensional vectors. Therefore,
${x}$
and
${y}$
are the latent variables.

High Level Overview

In VAEs, we essentially train and “organize” a latent space (vector space of latent variables).

We do this by using a neural network to take a high dimensional input vector

{x}

to a lower dimensional vector

{z}

in the latent space. We then use another neural network to take

{z}

back into a higher dimensional vector

{\hat{x}}

which is the same dimension as

{x}

. Our network is trained to both “organize” (I will expand on what “organize” means later) the latent space and reduce the difference between

{x}

and

{\hat x}

To generate new data, we simply take a random sample from the latent space and use the second neural network to produce a vector in the input space.

Overview of VAE architecture. Source: Lilian Weng

https://lilianweng.github.io/posts/2018-08-12-vae/

Autoencoders

Before we talk about VAEs, we should talk about autoencoders.

Illustration of an autoencoder. Source: Lilian Weng

https://lilianweng.github.io/posts/2018-08-12-vae/

Autoencoders are models that try to learn low dimensional representations of input vectors. The low dimensional representations should contain the “most important” features of the input vector. This allows you to (approximately) reconstruct the original input vector from the low dimensional vector.

Observe how this is very similar to

PCA

. Recall that PCA uses a linear transformation (matrix) to go from a high dimensional vector to a low dimensional vector and then a second linear transformation (matrix) to go from the low dimensional vector to a high dimensional vector. PCA then tries to minimize the difference (reconstruction loss) between the 2 high dimensional vectors. Autoencoders do the same, but instead of going from different dimensions using linear transformations, it uses neural networks. One neural network (encoder) decreases the dimension and another neural network (decoder) increases it. This is why a lot of people refer to autoencoders as a kind of non-linear PCA.

In autoencoders, the space of low dimensional vectors can be thought of as the space of latent variables i.e. the latent space.

Autoencoders for Generation

You may now be thinking that to generate new content, we can simply take a random sample of the latent space and use the decoder neural network to get a vector that is the same dimension as our input vectors.

However, this won’t work because the latent space isn’t “organized”. Meaning a random sample is likely to produce gibberish.

Latent space model. Source:

https://towardsdatascience.com/understanding-variational-autoencoders-vaes-f70510919f73

The above image tries to model a latent space. The gradient colors represent regions in the latent space that, if sampled and passed through the decoder neural network, will produce some meaningful reconstruction (like the triangle, square, and circle). However, notice how taking a random sample of this space is likely not going to fall in one of these gradients. Decoding a sample from the white space is going to produce gibberish, like the purple squiggly line.

Irregular latent space prevents use from using autoencoder for new content generation. Source:

https://towardsdatascience.com/understanding-variational-autoencoders-vaes-f70510919f73

We need some way to “organize” this space so that the majority of samples fall in these gradients.

VAEs

Definition

A variational autoencoder can be described as a model that not only tries to learn a good latent space, but also “organize” it such a way that random samples of the space can produce meaningful results.

💡

VAEs achieve this by trying to learn the probability distribution of the latent space.

VAEs are similar to autoencoders in that they use

probabilistic

encoders/decoders instead of normal encoders/decoders. This means that instead of using neural networks to get a vector in the latent space given a vector in the input space (and vice-versa), it uses neural networks to get a probability distribution of the latent space given an vector in the input space (and vice-versa).

Math

Objective

For all observed data

{x}

, we assume they are i.i.d. and come from some underlying probability distribution

{p(x)}

. We then try to learn a model that generates a probability distribution

{p_\theta}

, for model parameters

{\theta}

, such that the likelihood

{p(x|\theta)}

is maximized for all observed data points. This approach is called

likelihood maximization

We can imagine the latent variables

{z}

of the observed data as modeled by a joint probability distribution

{p_\theta(x,z)}

. By chain rule of probability, we have

{p_\theta(x)=\frac{p_\theta(x,z)}{p_\theta(z|x)}}

. We also have

{p_\theta(x)=\int p_\theta(x,z)dz}

by marginalization. Directly computing

{p_\theta(x)}

and trying to maximize it is difficult because you would have to either have access to the ground truth distribution of the latent encoder,

{p_\theta(z|x)}

, or integrate all latent variables

{z}

out (intractable for high dimensional

{z}

Instead we try to approximate

{p_\theta(z|x)}

Let us first formally define our objective. As with many other probabilistic models, we try to maximize the log-likelihood of the observed data.

\argmax_\theta \log p_\theta(x)\\=\argmax_\theta \log \frac{p_\theta(x,z)}{p_\theta(z|x)}

Since we don’t have access to

{p_\theta (z|x)}

, we try to estimate it with a variational distribution

{q_\phi(z|x)}

for parameters

{\phi}

. More specifically,

\argmax_\theta \log p_\theta(x)\\=\argmax_\theta \log\int p_\theta (x,z)dz\\=\argmax_{\theta,\phi} \log\int p_\theta(x,z)\frac{q_\phi(z|x)}{q_\phi(z|x)}dz\\=\argmax_{\theta,\phi} \log \mathbb{E}_{q_\phi(z|x)}[\frac{p_\theta(x,z)}{q_\phi(z|x)}]\text{~(expectation definition)} \\\ge\argmax_{\theta,\phi} \mathbb{E}_{q_\phi(z|x)}[\log \frac{p_\theta(x,z)}{q_\phi(z|x)}] \text{~(Jensen's Inequality)}

The last step follows from

Jensen’s Inequality

. The last term,

{\mathbb{E}_{q_\phi(z|x)}[\log \frac{p_\theta(x,z)}{q_\phi(z|x)}]}

, is formally known as the evidence based lower bound (ELBO). Since it is a lower bound, maximizing this will also maximize the log-likelihood:

{\log p_\theta (x)}

ELBO can be simplified even further (and makes the objective function of VAEs more apparent)

\mathbb{E}_{q_\phi(z|x)}[\log \frac{p_\theta(x,z)}{q_\phi(z|x)}]=\mathbb{E}_{q_\phi(z|x)}[\log \frac{p_\theta(x|z)p_\theta((z)}{q_\phi(z|x)}]\\=\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]-\mathbb{E}_{q_\phi(z|x)}[\frac{p_\theta(z)}{q_\phi(z|x)}]\\= \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]-D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z)) \text{~(KL-divergence defintion)}

Here

{D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z))}

is known as the

KL-divergence

which is a measure of how close 2 probability distributions. Note that the KL-divergence is a non-symmetric measurement, but that is not too important right now.

We try to match

{q_\phi(z|x)}

{p_\theta(z|x)}

by reducing the difference (KL-divergence) between

{q_\phi(z|x)}

and

{p_\theta(z)}

. This works because

{p_\theta(z)}

is an upper bound on

{p_\theta(z|x)}

Note that

{q_\phi(z|x)}

gives a probability distribution of latent variables

{z}

given an observed variable

{x}

. Therefore,

{q_\phi(z|x)}

can be thought of as a

probabilistic encoder

. Similarly,

{p_\theta(x|z)}

gives a probability distribution of observed variables

{x}

given a latent variable

{z}

. Therefore

{p_\theta(x|z)}

can be thought of as a

probabilistic decoder

Probabilistic encoder and decoders. Source:

https://arxiv.org/pdf/2208.11970

Once again, our objective is to maximize ELBO over parameters

{\theta,\phi}

where ELBO is:

\text{ELBO}=\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]-D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z))

Optimizing

We first look at the second term in ELBO which is KL-divergence.

In VAEs, when trying to maximize ELBO, we assume

{p(z)}

to be the standard multivariate gaussian.

p(z)=N(0,I)

We also assume

{q_\phi(z|x)}

to be some multivariate gaussian with diagonal covariance.

q_\phi(z|x)=N(\mu_\phi(x), \sigma^2_\phi(x)I)

This means we learn models

{\mu_\phi(x)}

and

{\sigma_\phi(x)}

, which are smaller neural networks, in order to create a mean vector and a standard deviation vector given a datapoint

{x}

Since we are trying to reduce the KL-divergence and bring

{q_\phi(z|x)}

as close to

{p(z)}

as possible and

{p(z) }

is the standard normal distribution, we are essentially trying to bring the distribution of all latent variables to the standard normal distribution. This is how we “organize” the latent space.

The returned distributions of VAEs have to be “organized” to obtain a latent space with good properties. Source:

https://towardsdatascience.com/understanding-variational-autoencoders-vaes-f70510919f73

We now look at the first term in ELBO:

{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

. This term tells the model to maximize the log-likelihood for seeing the input data

{x}

given latent representation

{z}

Since we don’t have access to all possible latent variables

{z}

, we instead estimate this expectation using a

Monte-Carlo estimate

. Specifically, we do:

\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]\approx \frac{1}{L}\sum^L_{l=1}\log p_\theta (x|z_l)

This means for some relatively small

{L}

, we sample

{z_1,...,z_L\sim q_\phi(z|x)}

and compute

{\frac{1}{L}\sum^L_{l=1}\log p_\theta (x|z_l)}

. In practice, if the dataset is large, we can use

{L=1}

There is still one small problem though. Since sampling is a non-differentiable procedure, we can’t just simply generate samples from

{z_i\sim q_\phi(z|x)=N(\mu_\phi(x), \sigma^2_\phi(x)I)}

. This is because we want to be able to differentiate

{\mu_\phi}

and

{\sigma_\phi}

during back propagation.

Instead, we employ a method called the reparameterization trick. This trick rewrites

{q_\phi(z|x)}

as a deterministic function and offloads the randomness/sampling to a random variable

{\zeta\sim N(0,I)}

. Note that

{\zeta}

doesn’t need to be differentiated because it is not learned.

Therefore instead of sampling

{z\sim q_\phi(z|x)}

, we first sample

{\zeta\sim N(0,I)}

and then say

z=\mu_\phi(x)+\sigma_\phi(x)\cdot\zeta

Illustration of the reparametrisation trick. Source:

https://towardsdatascience.com/understanding-variational-autoencoders-vaes-f70510919f73

That’s it! We now have a well defined objective function to optimize over. Maximizing the

{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

term of ELBO tells our model to look for latent representations that allows for the reconstruction to be as close to the initial input and the

{D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z)) }

term of ELBO tells our model to “organize” the latent space for good sampling during generation.

In practice, however, we don’t use the probabilistic decoder

{p_\theta(x|z)}

to once again model the distribution and compute

{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

. Instead, we directly compute

{\hat{x}}

(reconstruction of

{x}

) from

{z}

using the probabilistic decoder neural network and try to make it as close to

{x}

as possible. That is we try to minimize

{l(x,\hat{x})}

for some loss function

{l}

instead of maximizing

{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

Overview of VAE architecture. Source: Lilian Weng

https://lilianweng.github.io/posts/2018-08-12-vae/

Code

The following shows how we can create a VAE on the

MNIST dataset

We use the following libraries and starter code:

import torch import torch.nn as nn import math import torch.optim as optim import torchvision.datasets as datasets import matplotlib.pyplot as plt import torchvision.transforms as transforms device = "cpu" if torch.cuda.is_available(): device = "cuda" elif torch.backends.mps.is_available(): device = "mps"

Loading the data

# download the MNIST dataset train = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True) test = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True) # binarize the data train.data = torch.where(train.data > 128, 1, 0) test.data = torch.where(test.data > 128, 1, 0) # plot samples sample_train = train.data[:2] sample_test = test.data[:2] for sample in sample_train: plt.figure() plt.imshow(sample.numpy(), cmap='gray') for sample in sample_test: plt.figure() plt.imshow(sample.numpy(), cmap='gray')

Loss term

Recall that we are trying to maximize the ELBO which is formally:

\argmax_{\theta,\phi}\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]-D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z))

Let’s look at the first term which is

{\argmax_{\theta,\phi}\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

. Recall that in practice, we don’t use the decoder neural network to model the distribution

{p_\theta(x|z)}

, but instead just use it to directly compute

{\hat{x}}

from

{z}

which is a reconstruction of

{x}

. Therefore, maximizing this is equivalent to minimizing the reconstruction loss between

{\hat{x} }

and

{x}

. Therefore, we use the

binary cross entropy loss

which is a good loss to measure reconstruction. Note that other losses can be used, such as

MSE loss

. So instead of

{\argmax_{\theta,\phi}\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}

we do:

\argmin_{\theta,\phi} \text{BCE-Loss}(\hat{x},x)

Now let’s look at the second term which is

{D_{\text{KL}}(q_\phi(z|x)~||~p_\theta(z)) }

. Instead of maximizing the negative of the KL-divergence, we simply just minimize it. Recall that we assume

{p_\theta(z)}

to be the standard normal distribution

{N(0,I)}

and

{q_\phi(z|x)}

to be the normal distribution

{N(\mu_\phi(x),\sigma^2_\phi(x)I)}

. Using the formula for KL-divergence on

{k}

dimensional multivariate gaussians and simplifying based on the fact that

{p_\theta(z)=N(0,I)}

, we get that the KL-divergence is:

\argmin_\phi\frac{1}{2}[\mu_\phi(x)^T\mu_\phi(x) + tr(\sigma^2_\phi(x)I)-k-\sum_{i=1}^k\log(\sigma^2_\phi(x)_i)]

Note that the

{-k}

can be omitted because it is a constant.

Putting this into code, we get the following:

def elbo_loss(x_hat, x, mu, std): bce_loss = nn.BCELoss(reduction='sum') kl_divergence = 0.5 * torch.sum(mu ** 2 + std ** 2 - torch.log(std ** 2)) return cross_entropy_loss(gen_x, x) + 5 * kl_divergence

Approximated ELBO loss function

Note that we chose to multiply the KL-divergence by a positive constant in order to tell the model to put more weight on organizing the latent space. This allows for slightly more meaningful generation.

VAE Model

The VAE model is essentially just a bunch of neural networks.

More specifically, for a latent space dimension

{d}

, the probabilistic encoder is a neural network that takes an input vector and produces a

{d}

dimensional mean vector (representing

{\mu_\phi(x)}

) and a

{d}

dimensional standard deviation vector (representing

{\sigma_\phi(x)}

The probabilistic decoder is another neural network that takes a

{d}

dimensional vector in the latent space and outputs a vector in the input space.

For the model, we chose to use a mix of non-linearities. The most important part is that we have

{tanh}

non-linearities at the end of our probabilistic encoder network layers (in order to constrain the outputs between -1 and 1) and a

{sigmoid}

non-linearity at the end of the decoder layer (to allow for good inputs for the

{\log}

term in the loss function).

class VAE(nn.Module): def __init__(self, input_dim): super(VAE, self).__init__() self.input_dim = input_dim self.latent_dim = 12 self.encoder = nn.Sequential( nn.Linear(self.input_dim, 512), nn.ReLU(), nn.Linear(512, 64), nn.ReLU(), ) self.mu = nn.Sequential( nn.Linear(64, self.latent_dim), nn.Tanh() ) self.std = nn.Sequential( nn.Linear(64, self.latent_dim), nn.Tanh() ) self.decoder = nn.Sequential( nn.Linear(self.latent_dim, 64), nn.ReLU(), nn.Linear(64, 512), nn.ReLU(), nn.Linear(512, self.input_dim,), nn.Sigmoid() ) def get_latent_distribution(self, x): mu = self.mu(self.encoder(x)) std = self.std(self.encoder(x)) return mu, std def sample_latent_distribution(self, mu, std): # Samples a vector z using the reparameterization trick # Here we are also using a monte-carlo estimate on 1 sample. # This is sufficient since our dataset is large. return std * torch.randn(std.shape[0]).unsqueeze(-1).to(device) + mu def encode(self, x): mu, std = self.get_latent_distribution(x) return self.sample_latent_distribution(mu, std), mu, std def decode(self, z): return self.decoder(z) def forward(self, x): x_hat, mu, std = self.encode(x) return self.decode(x_hat), mu, std def generate(self): # Take a random sample z with the assumption the latent space # is the standard gaussian z = self.sample_latent_distribution(torch.zeros(self.latent_dim).to(device), torch.ones(self.latent_dim).to(device)) size = int(math.sqrt(self.input_dim)) return self.decode(z).reshape(z.shape[0], size, size)

VAE model

Training Loop

The training loop is defined as usual:

vae = VAE(28**2).to(device) optimizer = optim.Adam(vae.parameters(), lr=1e-3) batch_size = 32 train_loader = torch.utils.data.DataLoader(train.data, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(test.data, batch_size=100, shuffle=True) epochs = 15 for epoch in range(epochs): epoch_loss = 0 for batch_idx, x in enumerate(train_loader): optimizer.zero_grad() x = torch.flatten(x.float(), start_dim=1).to(device) x_hat, mu, std = vae(x) loss = elbo(x_hat, x, mu, std) loss.backward() optimizer.step() epoch_loss += loss.item() print('Epoch: {}, Loss: {}'.format(epoch + 1, epoch_loss / len(train_loader)))

Results

The results and all the code can be found at:

GitHub - varunchitturi/variational-auto-encoders: Generative model using variation auto encoders (VAEs)

Generative model using variation auto encoders (VAEs) - varunchitturi/variational-auto-encoders

https://github.com/varunchitturi/variational-auto-encoders

Decoding

The following images are reconstructed from a single pass through the model. The image on the left of each figure is the original image from the dataset and the image on the right of each figure is the reconstruction.