Variational Autoencoder (VAE)

This post is an English translation of my chapter on VAE from the Machine Learning Handbook by Yandex School of Data Analysis (originally in Russian). The text follows the latest revised version of the chapter; the figures are the original English-language versions.

Introduction

There is a fairly broad area of machine learning devoted to training generative models. Their goal is to learn the distribution from which the objects of the training set could have been sampled.

Once trained, a generative model can sample new objects from the learned distribution — objects that do not belong to the original data. Most often this is associated with generating new images: from pictures of handwritten digits to face swapping in videos.

The model this post is about is called the variational autoencoder, or VAE. It belongs to the family of generative models. Here is a quick roadmap of what lies ahead.

• The sections “Problem setup” and “Training a VAE” describe how a VAE is built and trained in its classical form. These two sections are enough to get a general understanding of VAE.

• The section “Key papers” is not required for a first understanding, but may be interesting to those who want to learn about the foundational works that grew out of the VAE idea.

Problem setup

Imagine that we need to draw a horse. How would we go about it?

We would probably first sketch the overall silhouette of the horse, its size and pose, and then start adding details: the mane, the tail, the hooves, picking a coat color, and so on. It seems that while learning to draw, we learn to identify a basic set of factors that matter most for generating a new image — overall silhouette, size, color and the like — and during drawing we simply plug in particular values of those factors.

At the same time, identical combinations of the same factors can lead to different pictures — after all, you most likely cannot draw something exactly the same way twice.

Let us try to formalize the process described above. Suppose we have a dataset $D$ living in a high-dimensional space of raw data $X^N$ — the objects we wish to generate — and a lower-dimensional space $Z^M$ of hidden (latent) variables, which encode the hidden factors in the data. The generative process then consists of two consecutive stages (see the picture below):

1. Sample $z\in{Z^M}$ from the distribution $p(z)$

2. Sample $x\in{X^N}$ from the distribution $p(x\mid z)$

Image source

That is, thinking in terms of drawing pictures of horses, we first mentally sample some $z$ (size, shape, color), then fill in all the necessary details — that is, we sample from the distribution $p(x\mid z)$ — and in the end we hope that the result will resemble a horse.

Thus, building a generative model in our case means being able to sample, via the two-stage process described above, objects that are close to the objects of the training set $D$.

More formally, we would like our model to maximize the likelihood $p(x)$ of the elements of the training set $D$ under the described generation procedure:

Assume that the joint distribution $p(x,z)$ is parameterized by some parameter $\theta\in{\Theta}$ and is expressed by a function continuous in $\theta$ for every fixed $x$ and $z$:

Then

and we can write down the following optimization problem:

Solving it, we obtain our generative model.

Remark 1. After the analogy with learning to draw, one might mistakenly conclude that latent variables always carry some nicely interpretable meaning. In practice this does not have to be the case: the latent variables we end up finding may or may not have a simple interpretation. The explanations above were primarily meant to illustrate the notion of “latent variables.”

Remark 2. It might seem that $p(x)$ is somehow already known to us, and then it is unclear why we need all these complications with latent variables and integrals. Indeed, we can actually build a statistical estimate $\hat{p}(x)$ from the data $D$ and even try to generate new data with such models (as done, for example, here). But statistical methods come with various limitations, the most serious of which appears to be the curse of dimensionality: the more dimensions your data has, the more diverse examples you need to build an adequate estimate $\hat{p}(x)$. We will talk about the curse of dimensionality in a bit more detail later.

Remark 3. Another natural question: why introduce latent variables at all, model the joint distribution $p(x,z)$, and define the target distribution $p(x)$ as the marginalization of $p(x,z)$ over $z$? Why should this approach work in the first place? The answer is that even with relatively simple expressions for $p(z)$ and $p(x \mid z)$, one can describe a rather complex distribution $p(x)$ — which is illustrated quite vividly in the example below.

Training a VAE

Before trying to solve optimization problem $(1)$, let us think about how we could even compute such an integral. The first thing that comes to mind is to approximate it with the Monte Carlo method:

where in the last step we use samples $z_k \sim p_\theta(z)$. However, if $z \in Z^M$ and $M$ is large enough, we run into the curse of dimensionality — the number of samples needed to cover $Z^M$ well grows exponentially with $M$:

Image source

Is there a way to somehow reduce the number of samples needed to compute $(1)$? As it happens, it is often the case that far from all possible $z$ map into elements of $D$, and the contribution of most $z$ to the estimate of $p_\theta(x \mid z)$ is practically zero. This suggests that for each $x$ we could benefit from knowing the distribution $q(z \mid x)$ of those $z$ that are preimages of $x$. We may assume that the distribution $q$ is parameterized by some family of parameters $\Phi$:

Knowing the distribution $q_\phi(z \mid x)$, we could sample only from it rather than from the whole $p_\theta(z)$, and if the distribution $q$ turns out to be good enough, the number of required samples will drop significantly.

We will discuss how to build $q_\phi$ later. For now, note that the processes of sampling from the distributions $q_\phi(z \mid x)$ and $p_\theta(x \mid z)$ are mutually inverse: the first maps dataset elements into a subset of the latent space $Z^M$, i.e. acts as an encoder, while the second maps latent variables into a subset of $X^N$, i.e. acts as a decoder:

Image source

Since both of these distributions will take part in training the VAE, an analogy arises between VAE and autoencoder models, which have a similar structure.

Deriving the loss function

We now have everything ready to write down the general form of the loss function for training a variational autoencoder. Recall that we train the model by maximizing the likelihood $p_\theta(x)$ with respect to $\theta$. For convenience, we switch to the log-likelihood:

Optimizing this expression directly is hard because of the curse of dimensionality discussed in the previous section. To defeat it, we would like to replace sampling from the prior distribution $p_\theta(z)$ with sampling from $q_\phi(z \mid x)$, which requires a certain trick. For any $q_\phi(z \mid x)$ that is nonzero for all $z \in Z^M$, we can write out the following chain of equalities:

The second term in the last equality is the $KL$ divergence between $q_\phi(z \mid x)$ and $p_\theta(z \mid x)$, which, as is well known, is nonnegative:

And the first term is the evidence lower bound (ELBO):

The first term in the last step is called the reconstruction loss, since it measures how well the decoder reconstructs the object $x$ from its latent representation $z$. The second one plays the role of a regularization term and pushes the distribution produced by the encoder to be closer to the prior distribution.

Since the $KL$ divergence is nonnegative, the ELBO is a lower bound on the log-likelihood of the data:

Let us take a closer look at the equalities we have written down.

• The function $\mathcal{L}_{\theta,\phi}$ can be optimized by gradient descent (SGD), once we choose a convenient form for $p_\theta(x \mid z)$, $q_\phi(z \mid x)$ and $p_\theta(z)$. By maximizing $\mathcal{L}_{\theta,\phi}$, we increase $\log p_\theta(x)$, thereby improving our generative model. We will discuss ELBO optimization with SGD in detail in the next section.

• By maximizing $\mathcal{L}_{\theta,\phi}$, we simultaneously minimize $D_{KL}(q_\phi(z \mid x) \| p_\theta(z \mid x))$. The distribution $p_\theta(z \mid x)$ estimates which $z$ the object $x$ could have been generated from, and it is not known to us in advance. But if we pick a sufficiently large model for $q_\phi(z \mid x)$, then during optimization $q_\phi(z \mid x)$ can get very close to $p_\theta(z \mid x)$, and then we will be directly optimizing $\log p_\theta(x)$. As a pleasant bonus, to estimate the distribution of preimages of $x$ we can use $q_\phi(z \mid x)$ instead of the intractable $p_\theta(z \mid x)$. That is, $q_\phi$, which we introduced in the derivation as an arbitrary distribution, will indeed play the role of the model’s encoder.

Training a VAE with gradient descent

An important property of the ELBO is that it can be optimized by gradient descent with respect to the parameters $\phi$ and $\theta$. If the objects of the dataset $D$ are independent and identically distributed, then $\mathcal{L}_{\theta,\phi}(\mathcal{D})$ is written as a sum (or mean) of the values $\mathcal{L}_{\theta,\phi}(x)$ over the objects $x \in D$:

The values $\mathcal{L}_{\theta,\phi}(x)$ and their gradients $\nabla \mathcal{L}_{\theta,\phi}(x)$ cannot be computed exactly in the general case, but we can obtain unbiased estimates of them, which will let us use stochastic gradient descent.

An estimate for the gradient with respect to the parameters $\theta$ is easy to get:

where in the last line $z_k \sim q_\phi(z \mid x)$. However, an estimate for the gradient with respect to the parameters $\phi$ is harder to obtain, because they also take part in the sampling:

In the general case this problem is unsolvable. However, some distributions admit the reparameterization trick: representing the variable $z$ as an invertible differentiable function of random noise, the parameters $\phi$, and the variable $x \in D$:

Here the distribution $\varepsilon \sim p_\varepsilon$ does not depend on $\phi$ or $x$. For example, let $\varepsilon \sim \mathcal{N}(0, I)$. Then $g$ can have the following form:

After this substitution we can obtain an estimate of the gradient with respect to $\phi$:

where in the last line $\varepsilon_k \sim p_\varepsilon$. The reparameterization trick is well illustrated by the following picture:

Image source

Here $f$ is the loss function. The values of $f$ are the same in both diagrams, but in the left one the gradients with respect to $\phi$ cannot be computed, since we cannot differentiate through the random variable $z$.

In the right diagram, however, the source of randomness moves into the input data thanks to reparameterization, and the gradients are computed with respect to deterministic variables. We thus arrive at a setup typical for optimization with SGD: there we approximate the gradient of the loss function over random batches of input data, and here the role of random batches is played jointly by batches of the variables $x$ and of the random variables $\varepsilon$.

Besides the normal distribution, there are quite a few examples of distributions that admit reparameterization. They can be found in the VAE paper, in the section “The reparameterization trick”. Most VAE implementations, however, use the normal distribution.

In the end, the rough algorithm for training a VAE looks like this:

dataset = np.array(...)
epsilon = RandomDistribution(...)

# Encoder q_phi(z|x) — a neural network with parameters phi
encoder = Encoder()

# Decoder p_theta(x|z) — a neural network with parameters theta
decoder = Decoder()

for step in range(max_steps):
     # Sample a batch of input data and of random noise
     batch_x = sample_batch(dataset)
     batch_noise = sample_batch(epsilon)

     # Compute the parameters of q(z | x) with the encoder
     latent_distribution_parameters = encoder(batch_x)

     # Apply reparameterization (sample from q(z | x))
     z = reparameterize(latent_distribution_parameters, batch_noise)

     # The decoder outputs the parameters of the output distribution
     output_distribution_parameters = decoder(z)

     # Compute the ELBO and update the model parameters
     L = -ELBO(
        latent_distribution_parameters, 
        output_distribution_parameters, 
        batch_x
     )
     L.backward()

It is worth emphasizing that the decoder outputs precisely the parameters of the output distribution, not a particular sample from it. For example, if you model output images with a normal distribution $\mathcal{N}(\mu(z), \sigma^2(z))$, the decoder will predict some $\hat{\mu}(z)$ and $\hat{\sigma}(z)$, which, together with the parameters of the latent distribution (the encoder’s output), are fed into the ELBO.

To generate an actual picture at inference time, you either honestly sample from $\mathcal{N}(\hat{\mu}(z), \hat{\sigma}^2(z))$, or, as is often done, simply take the mean $\hat{\mu}(z)$ as the output image. In general, the exact way inference is carried out depends on the type of the output distribution used.

Choosing the distributions

It is time to give examples of concrete $p_\theta(x \mid z)$, $q_\phi(z \mid x)$ and $p_\theta(z)$ with which a VAE can be built. To begin with, assume that $p_\theta(z)$ can be set to the standard normal distribution:

Note that in this case the prior distribution of $z$ has no dependence on the parameters $\theta$.

The distribution $p_\theta(x \mid z)$ depends on the distribution your data comes from. If your data has a continuous distribution, then $p_\theta(x \mid z)$ can be set, for example, to a Gaussian:

The mean vector in this example is given by a function $f$ of the variables $\theta$ and $z$, and the covariance matrix is a constant diagonal matrix. The function $f$ can be defined by a neural network with parameters $\theta$. If desired, the covariance matrix can also be given by some function rather than restricted to constant matrices. If your data is discrete, a categorical distribution may be suitable:

in which the probability vector $f_\theta(z) = (p_1, \dots, p_n)$ is the network output after applying $\text{softmax}$. If you have binary data, you can use a Bernoulli distribution:

where $f_\theta(z) = p$ is the neural network output after applying a sigmoid.

The distribution $q_\phi(z \mid x)$ can in principle be anything, but in the simplest case it is a Gaussian with a diagonal covariance matrix:

Such a distribution, in particular, admits the reparameterization trick discussed above. If we choose $z$ to be two-dimensional, the distributions defined by $q$ can be nicely visualized:

Image source

Now recall how the ELBO is defined:

Let us compute it for the distributions given above.

We start with $D_{KL}(q_\phi(z \mid x) \| p_\theta(z))$. The $KL$ divergence between the distributions $\mathcal{N}(\mu, \Sigma)$ and $\mathcal{N}(0, I)$ equals:

where $M$ is the dimension of these distributions. A derivation of this relation can be found here. In our case $\mu_\phi(x) = (\mu_1, \dots, \mu_M)$, $\sigma_\phi^2(x) = \operatorname{diag}(\sigma_1^2, \dots, \sigma_M^2)$ and

Then the ELBO is computed as:

where $z_k \sim \mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x))$. As mentioned in the paper by the authors of VAE, in section 2.3, the number of samples $K$ can be set to one, provided the batch size is large enough (e.g. 100).

If you choose a Bernoulli $p_\theta(x \mid z)$, then

If a Gaussian $\mathcal{N}(f_\theta(z), \sigma^2)$, then

An example implementation of training and using a VAE on the MNIST dataset can be found in Keras and in PyTorch.

Inference with a trained model

Once we have trained a VAE, we can generate new samples simply by feeding $z \sim \mathcal{N}(0, I)$ into the decoder:

Image source

The encoder is not needed for generating new samples. However, we may want to estimate $p(x) = \int p(x \mid z) p(z) dz$ for $x$ from a test set, to understand how likely the model is to generate $x$. To estimate the integral we need to sample some number of $z$, and if we take samples from $z \sim \mathcal{N}(0, I)$, the estimate may converge poorly. But we can again use the ELBO as a lower bound on $\log p(x)$ and estimate it instead, sampling from the distribution $q_\phi(z \mid x)$. Such an estimate converges faster and gives a rough idea of how well the model handles a particular example $x$.

It is also interesting to look at how the codes of the training examples are distributed in the latent space. This is what the distribution of latent codes of MNIST digits may look like for a trained VAE with a two-dimensional latent space:

Image source

Different types of digits are shown in different colors (the correspondence between digits and colors is shown on the scale at the side). One can see that the model distinguishes zeros and ones best of all, and eights and threes worst of all. It is worth noting, of course, that the latent space was chosen to be two-dimensional for visualization purposes; with a higher dimension the model could learn to distinguish the digits better.

For a two-dimensional latent space there is another interesting way to visualize the structure of the manifold learned by the VAE. One can take a uniform grid on the unit square and map it into the latent space by applying the inverse CDF of the normal distribution.

The resulting samples can be fed into the decoder to see which pictures correspond to the grid nodes:

Left: the learned Frey Face manifold; right: the learned MNIST manifold. Image source

Shown here are examples generated for the Frey Face and MNIST datasets (both available here). This visualization lets us see the smooth transition of latent codes of some objects into the codes of others, as well as the relative arrangement of the latent codes.

For MNIST we again see, in particular, that the model has placed the codes of zeros and ones far apart, while the codes of threes and eights are very close. It is also fun to observe the smooth transition from sixes to zeros and from sevens to ones. For Frey Face, the happy faces sit far from the sad ones, and along the main diagonal of the square one can trace a gradual transition from a serious face to a smiling one.

It is also interesting to look at how the quality of the generated digits changes depending on the dimension of the latent space (the pictures show random samples from the model):

Latent space dimension, left to right: 2, 5, 10, 20. Image source

A noticeable transition is visible between dimensions 2 and 5; further increase of the dimension has almost no significant effect.

Conditional VAE (CVAE)

Sometimes we may want to generate not just an arbitrary object from the dataset, but one belonging to a particular group or class. Earlier we wrote out an equation for $\log p_\theta(x)$:

We can make all the distributions participating in this equation conditional on a variable $y$:

The variable $y$ can be the label of the object $x$, or an entirely arbitrary tensor characterizing $x$ in some way. Instead of a single $p_\theta(z)$ shared by all $x$ in the training set, there is now a separate prior distribution $p_\theta(z \mid y)$ for each value of $y$.

The variable $y$ can take both discrete and continuous values. It can even be, for example, half of an image that the model is asked to complete. Just in case, let us stress that training a CVAE is not the same as training several independent VAEs, since the CVAE weights are shared across all classes.

At the implementation level this is quite simple: you just concatenate the inputs of the encoder and decoder with the tensor corresponding to $y$. If $y$ takes categorical values, it is often useful to first encode them as one-hot vectors. The algorithm looks roughly like this:

dataset, labels = np.array(...), np.array(...)
epsilon = RandomDistribution(...)

# Encoder q_phi(z|x) — a neural network with parameters phi
encoder = Encoder()

# Decoder p_theta(x|z) — a neural network with parameters theta
decoder = Decoder()

for step in range(max_steps):
     # Sample a batch of input data, labels and random noise
     batch_x = sample_batch(dataset)
     batch_y = sample_batch(labels)
     batch_noise = sample_batch(epsilon)

     # Feed the encoder the concatenation of inputs and labels
     encoder_input = concatenate([batch_x, batch_y])

     # Compute the parameters of the distribution of z with the encoder
     latent_distribution_parameters = encoder(encoder_input)
     # Apply reparameterization
     z = reparameterize(latent_distribution_parameters, batch_noise)

     # Concatenate the resulting random vector with the labels
     decoder_input = concatenate([z, batch_y])

     # The decoder outputs the output image
     output_distribution_parameters = decoder(decoder_input)

     # Compute the ELBO and update the parameters
     L = -ELBO(
        latent_distribution_parameters, 
        output_distribution_parameters, 
        batch_x
     )
     L.backward()

A CVAE implementation in PyTorch and TensorFlow can be found, for example, here.

If we visualize the distribution of latent codes for MNIST digits obtained after conditioning the model on the digit class, we see something like this:

Image source

We see an unintelligible mixture of points instead of the distinct clusters that the plain VAE produced. The point is that instead of trying to place all digits in a single space $p(z) \sim \mathcal{N}(0, I)$, the model uses a separate latent space $p(z \mid y) \sim \mathcal{N}(0, I)$ for each digit:

Image source

Image source

In each picture, the right part shows the prior distributions for the digits 6 and 7, and the left part visualizes the structure of the learned manifolds for these digits, built the same way as the analogous visualization for the VAE. The quality of the images of each individual digit improves noticeably:

Image source

The variability of the generated digits has also grown noticeably, and the model can imitate digits written in different handwriting styles.

Key papers

Beyond the standard description of how VAE works, we will discuss several works building on the VAE idea. Although these papers came out in 2017–2021, we cover them not as “recent results” but as the foundation on which modern generative models are built: the idea of discrete latent codes from VQ-VAE became the basis for tokenizing non-textual data — images, audio and video, — and DALL-E was the first demonstration that a large language model can be trained on top of such tokens.

VQ-VAE and VQ-VAE-2

The models VQ-VAE and VQ-VAE-2 are interesting in that they employ discrete distributions as priors. In which situations can discrete distributions be more suitable than continuous ones? For example, when we deal with tokens in NLP tasks or phonemes in speech processing. Images, too, could be encoded by a set of integers: one number could encode the object type, another its color, a third the background color, and so on:

Image source

Moreover, there exist quite powerful algorithms (for example, the Transformer) designed to work with discrete data. Learning good discrete representations makes it possible to use such algorithms effectively for, say, image generation.

Looking ahead: this very idea — compressing a continuous signal into a sequence of discrete codes — turned out to be one of the most fruitful in modern deep learning. The neural audio codecs SoundStream and EnCodec are VQ-VAEs with multi-level residual quantization; the voice cloning model VALL-E operates on top of EnCodec tokens. Visual tokenizers (VQGAN, MAGVIT) encode images and videos as discrete tokens for autoregressive generative models, and multimodal LLMs that generate images (for example, Meta’s Chameleon) treat pictures as sequences of the same kind of tokens as text. The idea of structuring latent representations is developing in other directions as well — for example, the nested, “matryoshka-like” variable-size embeddings of Matryoshka Representation Learning.

VQ-VAE

The authors of VQ-VAE introduce a discrete latent space in the form of $K$ real-valued vectors $e_1, \dots, e_K$ of dimension $D$. The vectors of this space are called code vectors, or codes. The figure below shows a rough scheme of training the proposed model.

Image source

The encoder takes an image $x$ as input and outputs a tensor $z_e(x)$. In the figure this tensor has shape $M \times M \times D$: the last dimension coincides with the length of the code vectors, and $M \times M$ is the spatial dimension of the CNN output (for simplicity we do not write out the batch dimension explicitly).

Each of the $M \times M$ vectors of $z_e(x)$ is mapped to the code vector nearest to it in $L_2$ distance. After this procedure the tensor $z_e(x)$ turns into a tensor $z_q(x)$ consisting of $M \times M$ code vectors. The decoder receives the tensor $z_q(x)$ as input and maps it back to the original image. For speech and text the authors used a two-dimensional tensor $z_e(x)$ instead of a three-dimensional one.

The output distribution of the encoder $q(z \mid x)$ is defined here as follows:

During training, a uniform distribution $p(z)=\frac 1K$ is used as the prior over the latent space, so the term $D_{KL}(q(z \mid x) \| p(z))$ turns out to be constant and equal to $\log K$:

At the points where $q(z = k \mid x) = 0$, the next-to-last expression is extended by zero by continuity. Thus, the ELBO for such distributions takes the form

where $\theta$ are the decoder parameters. During optimization the $\log K$ term can be ignored. The mapping of the encoder output to code vectors is not differentiable, so the following trick is used during training: on the backward pass, the gradient is copied directly from the decoder to the encoder, skipping the layer that maps encoder outputs to code vectors.

This trick is very close to the technique known as the straight-through estimator: on the backward pass the gradient flows through a non-differentiable operation as if it were not there. This technique was first proposed in this paper (and a simple explanation of it can be found here). Using the straight-through estimator, however, does not allow training the code vectors themselves, since no gradients are computed for them. Therefore the loss function for training the model consists of three components:

Here $\operatorname{sg}[\cdot]$ denotes the stop-gradient operator: no gradients flow through its argument.

In the paper the loss is written somewhat differently:

This notation seems somewhat confusing, for two reasons:

1. The letter $e$ in the subscript of $z_e(x)$ is only meant to indicate that this is the encoder output, not the existence of a connection between the code vectors $e$ and the encoder parameters. But the latter is quite easy to assume by mistake.

2. Subtracting $e$ means subtracting not all elements of the codebook from the corresponding position of the tensor $z_e(x)$, but only the nearest neighbor of the element of $z_e(x)$ at that position. That is, in effect, subtracting $e$ in this notation is equivalent to subtracting $z_q(x)$. This is not clarified in the paper, but can be seen in the official implementation.

The first term is the ELBO up to a constant. The second term is responsible for moving the code vectors toward the encoder outputs. To prevent a situation where the encoder outputs keep dragging the code vectors around via the second loss component while themselves producing, at every iteration, vectors far from the current code vectors, a third term is added. It makes the encoder strive to output vectors close to the code vectors, and its importance is regulated by the coefficient $\beta$.

However, during training we lost the regularization term $D_{KL}(q(z \mid x) \| p(z))$, because of which the encoder distribution was under no obligation to approximate the prior and remained a narrow subset of it. As a result, when sampling from the uniform categorical distribution, we will most likely get plain noise instead of nice pictures:

Left to right: test data, their reconstructions, samples from the uniform prior. Image source

To fix this problem, the authors propose to learn, with an additional model, the prior distribution $p(z)$ of those latent variables that the model learned to generate during training. Since any code representation can be flattened into a sequence, and the number of codes is finite and fixed in advance, this task is close to training a language model.

Indeed, there we must predict the next word from the available vocabulary given the sequence of preceding words of a sentence, and in our case — predict the next latent code given the input sequence of discrete latent codes.

For images, the authors proposed to model the prior distribution of latent codes with PixelCNN. The details of this model’s architecture and training can be found in the original paper; here we describe only the general idea.

PixelCNN generates the pixels of an image sequentially, moving from the top-left corner to the bottom-right. It traverses all rows one by one from top to bottom, and within each row moves left to right:

Image source

For color images, the channels (R, G, B) are also modeled sequentially: when generating, channel B depends on R and G, and G only on R. When predicting the value of each next pixel, the model uses the values of already-generated neighbors from some surrounding square. To prevent the model from reading pixels that come after the currently predicted one, a special mask is used, an example of which is shown in the right part of the figure.

In the case of VQ-VAE, PixelCNN is trained not on pixels but on latent codes. Sampling from the learned prior looks much better than attempts to sample from the uniform one:

Left: samples from the uniform prior; right: samples from the prior learned by PixelCNN. Image source

For audio, the authors use WaveNet instead of PixelCNN. When training the prior models, one can also feed in class labels, so that later one can sample from those classes (the same principle as for CVAE).

The results of reconstructing ImageNet images with VQ-VAE look quite good (by reconstruction we mean the output of the full model consisting of the encoder and decoder):

Left: original ImageNet images (128×128); right: their VQ-VAE reconstructions with a 32×32 latent space and a codebook of 512 code vectors. Image source

And this is what sampling from a VQ-VAE with a PixelCNN-learned prior looks like:

128×128 samples from a VQ-VAE with a PixelCNN prior trained on ImageNet. Left to right: kit fox, gray whale, brown bear, admiral butterfly, coral reef, alp, microwave, pickup. Image source

VQ-VAE-2

The VQ-VAE-2 model is not used today by itself, but it is interesting as an example of how the VQ-VAE idea evolved — above all, through a hierarchy of discrete latent spaces.

DALL-E

Another work that largely shaped the development of generative models is DALL-E by OpenAI. They trained a model with 12 billion parameters that generates pictures from their text descriptions. For training, the authors collected a dataset of 250 million image–caption pairs. Here are some examples of this model at work:

Prompt: “an armchair in the shape of an avocado”. Image source

Prompt: “a store front that has the word ‘openai’ written on it”. Image source

More examples of generations for various text descriptions can be found in OpenAI’s blog post about DALL-E.

Prompt: “a painting of a capybara sitting in a field at sunrise” in different styles: painting, pop art, cubism, surrealism, Van Gogh, Monet, pencil drawing, charcoal, crayons and chalk. Image source

Conceptually, DALL-E builds on the results of VQ-VAE: first, code vectors for pictures are learned, and then a Transformer is trained to model the joint prior distribution of texts and code vectors. (I’m preparing a separate series of posts about transformers — stay tuned.)

In essence, DALL-E is the first demonstration that the “textual” pretraining recipe (autoregressive next-token prediction) also works for non-textual data, provided it is first discretized with a VQ-VAE-style encoder. This recipe quickly spread beyond image generation: for example, the voice cloning system Tortoise TTS is modeled after DALL-E — its author explicitly described the project as “DALL-E for speech synthesis”.

DALL-E employs an architecture based on the decoder part of the original Transformer, so it is also worth reading up on the GPT-2 model, which works in a similar way.

Training proceeds in two stages:

• First, a discretized VAE (dVAE) is trained, with an encoder compressing $256 \times 256$ RGB images into a tensor of $32 \times 32 = 1024$ code vectors. This training stage strongly resembles VQ-VAE, but instead of adding extra loss terms for the code vectors, the DALL-E authors use the Gumbel relaxation — a trick that makes honest differentiation with respect to the encoder parameters possible. We will discuss dVAE training in more detail below.

• Then a Transformer is trained (more precisely, only the decoder part of the original Transformer architecture), whose task is to learn the joint distribution of pictures and their text descriptions. It takes as input the concatenation of the embeddings of the text tokens and the code vectors of the pictures, and learns to predict the continuation of each input sequence. Some details of the Transformer training will also be covered below.

Inference with the trained model goes like this: the embeddings of a picture’s text description are fed into the Transformer, which autoregressively predicts the code vectors of a picture matching that description; the resulting code vectors are then passed through the dVAE decoder.

dVAE

The dVAE is trained by maximizing the ELBO for pictures $x$ and their discrete latent representations $z$:

where $\phi$ and $\theta$ are the parameters of the encoder and decoder of the discretized VAE, and $p(z)$ is the uniform categorical distribution over the code vectors. Notice the extra coefficient $\beta$, which in a standard VAE always equals 1. The DALL-E authors introduced the additional parameter $\beta$ following the results of the paper on $\beta$-VAE. Unlike the original paper, though, in their experiments the value of $\beta$ is gradually decreased during training.

The dVAE encoder maps $256 \times 256$ images into a tensor $z_e(x)$ of shape $32 \times 32 \times 8192$, where $8192$ is the number of code vectors. That is, to each of the $32 \times 32$ positions the encoder assigns a categorical distribution over the $8192$ code vectors, parameterized by the output logits.

To obtain the tensor $z_q(x)$ of code vectors, one could first apply $\text{softmax}$ to the distributions at each of the $32 \times 32$ positions, and then assign to each position the code vector whose index has the maximum probability (take the $\arg\max$ for that position).

However, the $\arg\max$ operation is not differentiable — and besides, in the VAE framework the decoder should receive a sample from the distribution predicted by the encoder, and taking the $\arg\max$ at each position is not sampling from the predicted distribution.

We will therefore need a couple of tricks that will let us simultaneously:

• approximate sampling from the $\text{softmax}$

• make the sampling differentiable

Gumbel-Max Trick and Gumbel-Softmax

The first trick is known as the Gumbel-Max Trick. Imagine we have network output logits $x_1, \dots, x_k$, and we want to use them to obtain a sample from the categorical distribution — that is, to predict a class stochastically. For this we usually apply $\text{softmax}$ to the logits to obtain probabilities $\pi_i$:

and then sample a class from the resulting categorical distribution $\{\pi_1, \dots, \pi_k\}$. It turns out that these two steps are equivalent to the following procedure:

• sample numbers $g_1, \dots, g_k$ from the standard Gumbel distribution,

• add the sample $g_i$ to each logit $x_i$,

• pick the class $j$ such that $j = \arg\max_i (x_i + g_i)$.

Why this is indeed the case can be read here. But the Gumbel-Max Trick alone will not help us — the operation still is not differentiable. So we need one more trick, proposed almost simultaneously in two papers (first and second) and named Gumbel-Softmax in one of them.

To describe this trick, note that the result of the $\arg\max$ operation is the index of some class $j$. Such an index can be described by one-hot encoding — a vector of length $k$ in which all elements are zero except the $j$-th, which equals one.

Gumbel-Softmax consists of doing the following instead of taking the $\arg\max$ at the last step of the Gumbel-Max Trick:

• compute $y_i = \frac{\exp((x_i + g_i)/\tau)}{\sum_{j=1}^k \exp((x_j + g_j)/\tau)}$, $i = 1, \dots, k$, — an approximation of the one-hot vector via the $\text{softmax}$ activation with temperature;

• sum the code vectors $e_i$ with weights $y_i$: $z = \sum_i y_i e_i$;

• output the vector $z$ as the latent vector for the given position

Strictly speaking, the DALL-E authors did not specify how the output vector $z$ is aggregated from the code vectors and the $y_i$, but this is the approach taken in the PyTorch implementation of DALL-E.

As $\tau \to 0$, sampling from the distribution $\frac{\exp((x_i + g_i)/\tau)}{\sum_{j=1}^k \exp((x_j + g_j)/\tau)}$ tends to the $\arg\max$, and during dVAE training the authors gradually decreased the value of $\tau$. In the following picture, the left part shows the plain Gumbel-Max Trick, and the right part its differentiable variant:

Image source

Thus, training the code vectors of the dVAE requires no extra loss terms on top of the ELBO, nor copying gradients from the decoder to the encoder (as was the case in VQ-VAE).

Moreover, it is worth noting that $D_{KL}(q_\phi(z \mid x) \| p(z))$ in this case does not degenerate into a constant, but genuinely acts as a regularizer. Let us spell out how it works: the encoder outputs a $32 \times 32$ grid of positions and at each of them predicts its own categorical distribution over the $8192$ code vectors, while the prior $p(z)$ takes the code at each position to be equally probable among all options and independent of the other positions. The regularizer pulls the encoder’s prediction at each position separately toward this uniform distribution.

Importantly, such regularization says nothing about the joint structure of all $1024$ codes of a picture: a real image corresponds to a strongly correlated set of codes, not at all to an independent set of equally probable ones. So one cannot sample pictures directly from the dVAE (drawing the code at each position independently) — the joint distribution of the codes has to be modeled separately, and that is exactly what the Transformer does in the second stage. Essentially, this is the same role that PixelCNN played in modeling the prior over codes for VQ-VAE.

Finally, the non-constant $KL$ has one more practical advantage — it encourages high codebook utilization. By pulling the distribution at each position toward the uniform one, the regularizer prevents the model from collapsing to the use of just a handful of codes (a problem known as codebook collapse). In hard VQ-VAE there is no such regularizer, so codebook collapse is fought there with separate techniques — for example, updating the code vectors with an exponential moving average (EMA; this technique is covered above in the section on VQ-VAE-2) or re-initializing “dead” codes that have stopped being selected.

The Logit-Laplace distribution

One more trick in dVAE training concerns the output distribution $p_\theta(x \mid z)$. The DALL-E authors noticed a problem arising with the commonly chosen Laplace and Gaussian distributions for $p_\theta(x \mid z)$: both are defined on the entire real line, whereas pixels take values from a bounded interval. Thus, part of the density is “lost” during modeling, ending up outside the feasible range of pixel values.

To fix this problem, the authors propose to use a distribution they called “Logit-Laplace”. Its density is defined on the interval $(0,1)$ and is expressed by the following formula:

This density corresponds to a random variable obtained by applying the sigmoid to a Laplace-distributed random variable. The expression for the Logit-Laplace distribution can be derived from the standard formula for the density of a random variable obtained by applying a monotone differentiable function to another random variable (see the formula, for example, here). The logarithm of this density is substituted into the ELBO in place of $\ln p_\theta(x \mid z)$.

The decoder outputs 6 tensors: the first three correspond to $\mu$ for the RGB channels, the remaining three correspond to $\ln b$, and these 6 tensors are used to compute the loss. Before being fed to the encoder, the image values are normalized by the function $\phi: [0,255] \to (\varepsilon, 1-\varepsilon)$:

This way the authors ensure that the decoder models values from $(\varepsilon, 1-\varepsilon)$, which mitigates the computational problems associated with dividing by $x(1-x)$ in the density formula. At inference time, the reconstruction $\hat x$ of a picture $x$ is computed by the formula:

where $\mu$ is the first three tensors of the decoder output. The outputs corresponding to $\ln b$ are not used here.

The prior over texts and images

In the second stage, the authors freeze the parameters $\phi$ and $\theta$ and model the joint distribution of pictures and their text descriptions with a Sparse Transformer with 12 billion parameters. As input it receives the concatenation of a picture’s text description and its code vectors. A picture is represented by 1024 code vectors obtained from the encoder $q_\phi$, and when sampling code sequences the plain $\arg\max$ is used, without adding noise from the Gumbel distribution.

The text description is tokenized with the BPE procedure (see the section on BPE here), and each token is assigned a vector of real numbers representing it (an embedding). At most 256 tokens are used to represent the text, and the vocabulary size is 16,384 tokens.

The Transformer’s task during training is to predict, for each initial segment of the input sequence, the token that follows it. This can be either a text token or an image code vector. Since the code vectors of a picture always come after the text tokens, when generating code vectors the attention mechanism also attends to all the preceding text tokens.

Furthermore, the attention mask for the code vectors takes into account that they are originally arranged not linearly one after another, but on a rectangular grid. The paper presents several variants of geometric patterns used for the attention mask over code vectors.

The loss is a weighted sum of the cross-entropy for the text tokens and the cross-entropy for the picture code vectors, with weights $\frac 18$ and $\frac 78$ respectively (image generation is given higher priority, hence the larger weight for its loss).

Of course, training a huge Transformer is anything but easy, and a substantial part of the paper is devoted to the tricks the authors applied to train such a large model.

Inference

At inference time, the tokens of a picture’s text description are fed into the model, and based on them the model autoregressively predicts the code vectors:

Image source

The picture’s code vectors are fed into the dVAE decoder, which maps them into the final picture:

Image source

To improve prediction quality, the authors first generate 512 pictures for each text description, and then pick the best picture among the predictions. Different sets of code vectors for the same text can be obtained, for example, by randomly picking a code vector at each generation step according to the distribution predicted by the Transformer. The ranking of the resulting 512 pictures is done with CLIP — a large neural network trained without supervision on a large amount of data to model the joint distribution of pictures and texts.

Conclusion

So, in this post we discussed how the VAE works in its classical form — with a continuous distribution of latent variables — and covered the works based on the idea of using discrete distributions in VAEs.

Of course, the various modifications of VAE are not limited to swapping continuous latent variables for discrete ones. There are many other possible directions for improving the model: hierarchical latent distributions (which we saw, by the way, in the context of VQ-VAE-2), loss functions other than the ELBO, various shapes of latent spaces, adversarial training, and much more.

A good list of papers on VAE modifications can be found here. Among the works developing hierarchical distributions, NVAE is worth noting — there is a good video review of it by Yannic Kilcher. It deserves a separate mention that the ideas of VAE underlie latent diffusion models (such as Stable Diffusion): in them, diffusion happens not in pixel space but in the latent space of a trained autoencoder.

This concludes our story about VAE. Hopefully, it gave you a general picture both of the original ideas from which the VAE model grew and of the most interesting results connected with it.

Self-check questions

1. Why introduce hidden (latent) variables when building a generative model? Why not just estimate $p(x)$ from the data directly?

2. Why can’t the likelihood $p_\theta(x) = \mathbb{E}_{z \sim p_\theta(z)}[p_\theta(x \mid z)]$ simply be estimated by Monte Carlo, sampling $z$ from the prior?

3. What is the ELBO, and why is maximizing it a reasonable substitute for maximizing $\log p_\theta(x)$?

4. Which two terms does the ELBO consist of, and what is each of them responsible for?

5. Why is the reparameterization trick needed when training a VAE? Why can an unbiased estimate of the ELBO gradient with respect to $\theta$ be obtained without it, but not with respect to $\phi$?

6. Which distributions are usually chosen for $p_\theta(z)$, $q_\phi(z \mid x)$ and $p_\theta(x \mid z)$ in a classical VAE?

7. How do you generate a new object with a trained VAE? Is the encoder needed for this?

8. How does CVAE differ from a plain VAE? Why is it not the same as training a separate VAE per class?

9. How is the latent space of VQ-VAE organized? What happens to the $KL$ term of the ELBO during training, and how do gradients pass through the quantization operation?

10. Why does sampling from the uniform prior of a trained VQ-VAE yield noise instead of meaningful pictures, and how do the authors solve this problem?

11. In a plain (continuous) VAE one can sample directly from the prior $\mathcal N(0,I)$, while VQ-VAE required a separate prior model (PixelCNN). Why?

12. When deriving the ELBO we assumed that the components of an object are conditionally independent given the latent representation: $p_\theta(x \mid z) = \prod_j p_\theta(x_j \mid z)$. But in VQ-VAE and dVAE the codes at different positions are, on the contrary, strongly correlated. Doesn’t this violate the original assumption and render the ELBO incorrect?

13. What do the Gumbel-Max Trick and Gumbel-Softmax do, and why are they needed in the dVAE (DALL-E)?

14. Gumbel-Softmax allows training a dVAE by honest ELBO optimization, but in practice discrete tokenizers more often use hard VQ (straight-through + commitment loss). Does this mean that careful regularization of the token prior is not that important?

15. Describe the two training stages of DALL-E and the inference procedure.

Practice

The hands-on notebook accompanying this post lives right in this blog’s repository: notebooks/vae.ipynb. You can open it in Colab — no GPU needed, everything runs on a CPU in a few minutes.

In the notebook we build a VAE on MNIST from scratch and reproduce with our own hands all the key pictures of this chapter:

• the encoder, reparameterization, a Bernoulli decoder and the ELBO as the loss function;
• reconstructions and samples from the prior $\mathcal{N}(0, I)$;
• the latent space map for $M = 2$ — those very clusters of digits;
• the learned manifold: a uniform grid passed through $\Phi^{-1}$ and the decoder;
• the effect of the latent space dimension on sample quality;
• CVAE: conditioning on the class, generating a chosen digit, and per-class manifolds.

The exercises at the end of the notebook are worth doing separately. The most important one is to train the model with zero weight on the $KL$ term: the reconstructions get better, while the samples from the prior turn into garbage. This is exactly the situation we analyzed for VQ-VAE, and seeing it with your own eyes is more useful than reading about it.