In a Nutshell…

VAEs are really neat way to encode information without ruining the “regularity” of the distribution the data lies on. Instead of simply learning a deterministic mapping that condenses the data, we also require that the distribution of the compressed space is as close to possible to a predictable distribution.

This shows up as a balancing of a reconstruction loss and KL regularization loss.

VAEs are extremely useful in modern generative modeling. They are usually seen alongside Diffusion or Flow Matching models (I wrote a whole other tutorial on those here) that rely on having a relatively slim latent space to generate into. Learning about them is a must for a modern computer scientist or anyone interested in understanding the algorithms around AI today.

So let’s get to it!

Training a VAE

Historically, VAEs were introduced first with the work of Kingma and Welling [1].

The important thing to understand here is the general setup of probabilistic modeling.

💡 Probabilistic Modeling with Latent Variables

Suppose we have some data $X$ that is sampled via a random process. The process is as follows. First, a latent variable $z$ is generated according to some distribution $p(z)$. Then, data point $x$ is sampled from some conditional distribution $p(x \mid z)$.

These latent variables can be thought of as an unobserved, internal representation of the dataset $X$. For instance, if $X$ consists of pictures of dogs, the latent variables could be thought of as living in a highly specialized submanifold that encodes the characteristics of dogs.

There are four important actors to consider in this regime:

  1. $p(z)$: the distribution of the latent variable.
  2. $p(x\mid z)$: the conditional likelihood - expresses the probability of seeing a dataset given a specific latent.
  3. $p(z \mid x)$: the posterior. It is a conditional distribution over latents given the data point.

  4. $p(x)$: the distribution of the data points that nature provides (e.g. a distribution over images of dogs). Our goal is ultimately to sample from this distribution! This marginal distribution is intractable (we don’t know it and cannot calculate it), though we can think of it like:

    $$ p_\theta(x) = \int p_\theta(z)p_\theta(x\mid z)dz $$

💡 Probability Distribution Proxy

The key idea is that because we don’t know the distributions involved with $X$ we approximate them via proxies. These will parameterized using deep neural networks:
$$ p(z\mid x)\approx q_\phi(z\mid x),\quad p(x\mid z)\approx p_\theta(x\mid z) $$ The parameters $\phi$ and $\theta$ will be learned eventually.

These two proxy distributions are associated with an encoder and a decoder respectively. $q_\phi(z\mid x)$ is created via a Neural network that compresses the data $x$ into a latent representation. $p_\theta(x \mid z)$ does the opposite.

To parameterize the distributions via neural networks, we only ask the networks to predict the parameters of the distributions. This will work as follows: given $x$ the encoder network predicts $(\mu_\phi(x),\sigma_\phi(x)^2)=\text{Enc}_\phi(x)$. Then:
$$ q_\phi(z\mid x) = \mathcal{N}(z\mid \mu_\phi(x),\sigma_\phi(x)^2\cdot I). $$ Similarly, if we fix some hyperparameter $\sigma_\text{dec}$, the decoder network gives:
$$ p_\theta(x\mid z)=\mathcal{N}(x\mid f_\theta(z),\sigma^2_\text{dec}I) $$

Training and Evidence Lower Bound (ELBO)

Training the encoder and decoder networks is done by simple maximum likelihood estimation. We try to set the parameters such that the likelihood $p_\theta(x)$ of observing data $x$ is maximized. The issue is that we cannot calculate $p_\theta(x)$ directly. To do so we’d have to consider all possible values of the latent variable $z$ and write:

$$ p_\theta(x) = \int p_\theta(x \mid z)\cdot p(z)dz $$

That is intractable. Instead, we come up with a different trick. First, we write:

$$ \log p_\theta(x) = \log \int q_\phi(z\mid x)\frac{p_\theta(x,z)}{q_\phi(z \mid x)}dz=\log\mathbb{E}_{z\sim q_\phi(z \mid x)} \frac{p_{\theta}(x,z)}{q_{\phi}(z\mid x)} $$

Estimating this directly is possible (in fact it gives a model called the Importance Weighted Autoencoder (IWAE)). However, it requires great care to avoid high variance. A simpler approach is to use a lower bound. Jensen’s inequality on the concave $\log$ function gives:

$$ \log p_\theta(x) \geq \mathbb{E}_{z \sim q_\phi(z\mid x)}\left[\log\frac{p_\theta(x,z)}{q_\phi(z\mid x)}\right]:=\text{ELBO}(x) = \mathcal{L}(x;\theta,\phi) $$

The quantity $\mathcal{L}(x;\theta,\phi)$ is called the Evidence Lower Bound (ELBO). It is a quantity we can easily compute and differentiate, meaning we can maximize it. It serves as a proxy to the true likelihood $p_\theta(x)$ which we cannot directly estimate.

A different way to look at the ELBO is by writing as follows (I’ll leave the few lines of algebra to you):

$$ \text{ELBO}(x) = \mathbb{E}_{z \sim q_\phi(z\mid x)}[\log p_\theta(x \mid z)]-\mathbb{D}_{\text{KL}}(q_\phi(z \mid x)\mid\mid p(z)) $$

where $\mathbb{D}_{\text{KL}}(p\mid\mid q):=\mathbb{E}_{z \sim p}\left[\log\frac{p(z)}{q(z)}\right]$ is the KL divergence between two distributions. For a sanity check, the KL divergence is always non-negative, so that confirms that ELBO is indeed a proper lower bound to the likelihood. Furthermore, it shows that our loss can be decomposed into two parts:

  • $\mathbb{E}_{z \sim q_\phi(z\mid x)} \log p_\theta(x \mid z)$: this expresses the quality of the decoder. Given a latent $z$ produced by the encoder, how well can we decode it back into a data point $x$? This is easy to compute over multiple samples of $x$: calculate $q_\phi(z \mid x)$, sample a point and compute $\log p_\theta(x \mid z)$ via the decoder network.
  • $\mathbb{D}_{\text{KL}}(q_\phi(z \mid x) \mid \mid p(z))$: This is a regularizer term. It measures how far the encoder proxy distribution is from the prior on the latents. Because we choose $p(z)$ as a standard normal, this becomes analytically easy to compute!
💡 Training a VAE

To train a VAE we have to solve the following optimization problem:
$$ (\phi,\theta) \in \arg\max_{\phi,\theta} \sum\limits_{x \in X} \text{ELBO}(x) $$

Gradient-based optimization has us solve this by gradient descent. We first need to compute the gradients:

$$ \nabla_{\phi,\theta}\text{ELBO}(x) = \nabla_{\phi,\theta} \left\{\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x,z) - \log q_\phi(z\mid x)]\right\} $$

  1. The gradient with respect to $\theta$ is easy to compute:

    $$ \nabla_\theta \text{ELBO}(x)=\nabla_\theta\left\{\int[\log p_\theta(x,z)-\log q_\phi(z\mid x)]q_\phi(z\mid x)dz\right\}\\ =\int\nabla_\theta\{\log p_\theta(x,z)-\log q_\phi(z \mid x)\}q_\phi(z\mid x)dz\\ = \mathbb{E}_{q_\phi(z\mid x)}[\nabla_\theta \log p_\theta(x,z)] $$

    We can approximate the last sum via Monte-Carlo estimation: $\frac{1}{L}\sum\limits_{i=1}^L \nabla_\theta \{\log p_\theta(x, z_i)\}.$ Note that $p_\theta(x,z_i)$ can be computed easily for $p(z) = \mathcal{N}(0,I)$ via Bayes’ Rule:

    $$ \log p_\theta(x,z) = \log p_\theta(x\mid z)+\log p(z) $$

  2. The gradient with respect to $\phi$ creates a different challenge:

    $$ \nabla_\phi \text{ELBO}(x)=\nabla_\phi\left\{\int[\log p_\theta(x,z)-\log q_\phi(z\mid x)]q_\phi(z\mid x)dz\right\} $$

    If we wrote the same expression as above, we wouldn’t be able to break up the integral into an expectation anymore! To solve this issue, Kingma and Welling [1] introduce the reparameterization trick.


    💡 The Reparameterization Trick



    We want $z$ to be a transformation of a random variable $\varepsilon$ whose distribution is independent of $\phi$ and $x$. To be more specific, let us write:


    $$ z=\mu_\phi(x)+\sigma_\phi(x)\cdot \varepsilon := g(\varepsilon, \phi,x) $$
    where $\varepsilon \sim \mathcal{N}(0, I)$. We can choose any transformation $g$ for $\varepsilon \sim p(\varepsilon)$ as long as the following condition holds:


    $$ q_\phi(z\mid x)\cdot |\det(\partial z/\partial \varepsilon)|=p(\varepsilon) $$
    This is indeed the case with the transformation $g$ we chose above (I’ll skip it, but it’s a few lines of algebra). Then, we can do a change of variables in the integral:


    $$ \mathbb{E}_{q_\phi(z\mid x)}[f(z)] = \int f(z)q_\phi(z \mid x)dz\\ =\int f(g(\varepsilon))\cdot q_\phi(g(\varepsilon)\mid x)dg(\varepsilon)\\ =\int f(g(\varepsilon))\cdot q_\phi(g(\varepsilon)\mid x)\cdot \left|\det\left(\frac{\partial g(\varepsilon)}{\partial \varepsilon}\right)\right|d\varepsilon\\ =\int f(z)\cdot p(\varepsilon)d\varepsilon\\ =\mathbb{E}_{p(\varepsilon)} f(z) $$
    Then we can differentiate:


    $$ \nabla_\phi\mathbb{E}_{q_\phi(z\mid x)}[f(z)]=\nabla_\phi\mathbb{E}_{p(\varepsilon)}[f(z)] = \nabla_\phi\left\{\int f(z)p(\varepsilon)d\varepsilon\right\}\\ =\int \nabla_\phi \{f(z)\cdot p(\varepsilon)\} d\varepsilon=\int \{\nabla_\phi f(z)\} \cdot p(\varepsilon) d\varepsilon\\ =\mathbb{E}_{p(\varepsilon)}[\nabla_\phi f(z)] $$
    Plugging in $f(z) = -\log q_\phi(z\mid x)$ we get:


    $$ \nabla_\phi \mathbb{E}_{q_\phi(z\mid x)}[-\log q_\phi(z\mid x)]=\mathbb{E}_{p(\varepsilon)}[-\nabla_\phi \log q_\phi(z\mid x)]\\ \approx -\frac{1}{L}\sum\limits_{\ell=1}^L \nabla_\phi \left[\log p(\varepsilon_\ell)-\log\left|\det \frac{\partial z_\ell}{\partial \varepsilon_\ell}\right]\right]=\frac{1}{L}\sum\limits_{\ell=1}^L \nabla_\phi\log\left|\det \frac{\partial z_\ell}{\partial \varepsilon_\ell}\right| $$
    So now this is tractable!

Putting it all together

Based on our discussion so far, the structure of our VAE is as follows:

  1. Encoder: $\mu_\phi(x),\sigma_\phi(x)^2$ → Latent distribution: $q_\phi(z\mid x) = \mathcal{N}(z\mid \mu, \sigma^2)$.

    1. Loss: KL divergence term in ELBO (measures quality of decoder). Because the distributions involved are normal, we can calculate this analytically:

      $$ \mathbb{D}_{\text{KL}}(q_\phi(z\mid x)\mid\mid p(z)) = \frac{1}{2}(\sigma_\phi(x)^2d - d + ||\mu_\phi(x)||^2 - 2d\log \sigma_\phi(x)) $$

    2. We generate $z \sim q_\phi$ using the reparameterization trick.

      $$ z = \mu+\sigma\cdot \varepsilon, \quad \varepsilon\sim \mathcal{N}(0,I) $$

      This allows us to backpropagate the loss. The loss only has gradient with respect to $\phi$ here, not $\theta$!

  2. Decoder: $f_\theta(z)$ → Output distribution: $p_\theta(x\mid z) = \mathcal{N}(x\mid f_\theta(z), \sigma_{\text{dec}}^2 I)$.

    1. Loss: 1st ELBO term. Plugging in the Gaussian distribution above (and doing a bit of algebra), we get:

      $$ \mathbb{E}_{q_\phi(z\mid x)} \log p_\theta(x \mid z)\approx-\frac{1}{M}\sum\limits_{m=1}^M \frac{||x-f_\theta(\mu_\phi(x)+\sigma_\phi(x)\varepsilon_m)||_2^2}{2\sigma_\text{dec}^2} $$

      • Notice something interesting about this loss? It’s sort of a reconstruction loss. Given the original image, we sample a latent from the conditional distribution and then try to recover the image again. In that sense, the encoder really does compress information!

      Note that given a single $x$ we take $M$ different samples $\varepsilon_m$, which correspond to $M$ different samples from $q_\phi(z\mid x)$ which we can only take using the encoder + reparameterization trick.

During training, we optimize $\phi$ and $\theta$ to maximize the ELBO, which is the sum of the terms we derived above.

Inference

Inference is quite easy! All we have to do is generate a latent vector $z\sim \mathcal{N}(0,I_d)$ and pass it through the decoder!

References

[1] Kingma, D. P., & Welling, M. (2013). Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114.

[2] Burda, Yuri, Roger Grosse, and Ruslan Salakhutdinov. "Importance weighted autoencoders." arXiv preprint arXiv:1509.00519 (2015).