Variational autoencoder (VAE)

Prerequisites

• Autoencoders.
• Normal distributions: mean $\mu$, standard deviation $\sigma$.
• Probabilistic models (helpful).

The analogy

How it works

Loss: $L = \underbrace{\lVert x - \hat x \rVert^2}_{\text{reconstruction}} + \underbrace{\tfrac12 \big(\sigma^2 + \mu^2 - 1 - \ln \sigma^2\big)}_{\text{KL}(\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(0, 1))}$ per latent dimension; $z = \mu + \sigma \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, 1)$.

Start with one VAE step by hand

1-D code: the encoder gives $\mu = 0.5$, $\sigma = 0.8$. The random draw is $\varepsilon = 0.5$. The input is $x = [1, 0]$ and the decoder outputs $\hat x = [0.8, 0.1]$. Predict, then open "Check".

Finish the picture: the two forces

The two loss terms pull in different directions.

• Reconstruction alone. Predict what the encoder would do to $\sigma$.
  Check

  Shrink $\sigma$ towards 0 and spread the means apart, to make each code precise: the VAE would turn back into a plain autoencoder with gaps.

• KL alone. Predict the codes.
  Check

  Every input would get $\mu = 0$, $\sigma = 1$: a perfectly regular code space that carries no information about the input. The sum balances precision against regularity.

Full worked calculation

A VAE with a 2-D latent space. For one input the encoder outputs $\mu = [0.5, -1]$ and $\sigma = [0.8, 0.5]$. We sample a code with the reparameterization trick, compute the KL term per dimension, and assemble the loss.

1. 1The encoder outputs a distribution

   $$q(z \mid x) = \mathcal N\big([0.5, -1],\; \text{diag}(0.8^2, 0.5^2)\big)$$

   Instead of one code, the encoder describes a small cloud of plausible codes: a centre $\mu$ and a spread $\sigma$ per dimension.

2. 2Reparameterization: sample

   $$\varepsilon \sim \mathcal N(0, 1) = [0.5, -1] \;\Rightarrow\; z = [0.5 + 0.8 \cdot 0.5,\; -1 + 0.5 \cdot (-1)] = [0.9,\; -1.5]$$

   The randomness is drawn separately, then shifted and scaled. Because $z$ is now a simple formula in $\mu$ and $\sigma$, gradients can flow through the sampling into the encoder.

3. 3KL term, dimension 1

   $$\tfrac12\big(\mu^2 + \sigma^2 - 1 - \ln\sigma^2\big) = \tfrac12\big(0.25 + 0.64 - 1 + 0.446\big) = 0.168$$

   This measures how different this dimension's cloud is from the standard normal $\mathcal N(0, 1)$. It is 0 only when $\mu = 0$ and $\sigma = 1$.

4. 4KL term, dimension 2

   $$\tfrac12\big(1 + 0.25 - 1 + 1.386\big) = 0.818$$

   Dimension 2 is penalised much more: its centre is far from 0 ($\mu = -1$) and its cloud is narrow ($\sigma = 0.5$, so $-\ln\sigma^2 = 1.386$). The KL pushes clouds to stay wide and centred.

5. 5Total loss

   $$L = \underbrace{\lVert x - \hat x \rVert^2}_{\text{reconstruction}} + \underbrace{0.168 + 0.818}_{\text{KL} = 0.986}$$

   The decoder turns $z = [0.9, -1.5]$ into $\hat x$; its error is the first part. With a reconstruction error of 0.05 the total would be 1.036. The two parts pull in opposite directions: precise codes vs. well-organised latent space.

6. 6Generating new data

   $$z \sim \mathcal N(0, I), \quad \text{e.g. } z = [0.3, -0.2] \;\Rightarrow\; \hat x = \text{decoder}(z)$$

   Because the KL kept every cloud close to $\mathcal N(0, 1)$, codes drawn from it land in regions the decoder knows. That is what a plain autoencoder cannot do.

The rule in plain words

Encode $x$ into $\mu$ and $\sigma$; sample $z = \mu + \sigma \varepsilon$; decode to $\hat x$. Minimise reconstruction error + KL divergence from $\mathcal{N}(0, 1)$.

To generate, sample $z \sim \mathcal{N}(0, 1)$ and decode. Nearby codes give similar outputs, so you can interpolate between examples.

The code

Read this code after you have computed the loss by hand. It will make more sense.

double reparam(double mu, double sigma, double eps) { return mu + sigma * eps; }

double kl_term(double mu, double sigma) {               /* KL(N(mu, sigma^2) || N(0, 1)) */
    return 0.5 * (sigma * sigma + mu * mu - 1 - log(sigma * sigma));
}

double vae_loss(double x[], double xh[], int d, double mu, double sigma) {
    double rec = 0;
    for (int j = 0; j < d; j++) rec += (x[j] - xh[j]) * (x[j] - xh[j]);   /* reconstruction */
    return rec + kl_term(mu, sigma);
}

flowchart LR
  X["x"] --> ENC["encoder"]
  ENC --> MU["μ"]
  ENC --> SG["σ"]
  EPS["ε ~ N(0, 1)"] --> Z["z = μ + σ ε"]
  MU --> Z
  SG --> Z
  Z --> DEC["decoder"]
  DEC --> XH["x̂"]
          

Putting numbers on it

The two loss terms for our example:

Common mistakes

• Sampling $z$ directly from $\mathcal{N}(\mu, \sigma^2)$ in code, which blocks gradients to $\mu$ and $\sigma$.
• Forgetting the KL term, which turns the VAE into a plain autoencoder with a gappy code space.
• Having the encoder output $\sigma$ directly; in practice it outputs $\ln \sigma^2$ so that $\sigma$ stays positive.

Try it yourself

Did you know?

Key points to remember

If you remember one thing from each line, remember the because.

• The encoder outputs a distribution because that fills the code space instead of leaving isolated points.
• The KL term keeps codes near $\mathcal{N}(0, 1)$ because random samples from $\mathcal{N}(0, 1)$ must decode to sensible data.
• The reparameterization trick enables training because it moves the randomness into $\varepsilon$, outside the path of the gradients.
• The two loss terms compete because precise codes help reconstruction but cost KL.
• VAEs can generate because decoding any sample from the prior lands in a region the decoder learned.

Self-check questions

Answer these in order, out loud. They follow the page, so your answers become the script of your video.

1. How is a VAE like filing photos as fuzzy regions on a shelf?
   Sample answer

   Each photo covers a region, and a rule packs the regions around the middle, so any spot you pick decodes to something sensible.

2. Compute $z$, the reconstruction error, the KL term and the total loss for the example.
   Sample answer

   $z = 0.9$; reconstruction 0.05; KL ≈ 0.168; total ≈ 0.218.

3. Why does reparam take eps as an input rather than sampling inside?

   A good answer mentions gradients.

   Sample answer

   Keeping the random draw $\varepsilon$ outside makes $z$ a simple deterministic function of $\mu$ and $\sigma$, so backpropagation can compute gradients for them.

4. So what is a VAE, in one sentence?
   Sample answer

   An autoencoder whose encoder outputs a distribution, trained with reconstruction error plus a KL penalty towards $\mathcal{N}(0, 1)$, so it can generate new data.

5. Why does a plain autoencoder generate poorly while a VAE generates well?

   A good answer mentions gaps in the code space.

   Sample answer

   A plain autoencoder's codes are isolated points with gaps that decode to garbage; the VAE's KL term and sampling fill the space around $\mathcal{N}(0, 1)$ with decodable codes.

6. What is the KL term for one latent dimension?
   Sample answer

   $\tfrac12(\sigma^2 + \mu^2 - 1 - \ln \sigma^2)$.

7. How do you generate a new sample with a trained VAE?
   Sample answer

   Sample $z$ from $\mathcal{N}(0, 1)$ and pass it through the decoder only.

8. If you had to teach this to a friend in 2 minutes, what would you leave out, and why?
   Sample answer

   I'd leave out the KL formula and the blurriness issue. I'd keep the fuzzy-shelf analogy, $z = \mu + \sigma\varepsilon = 0.9$, and "loss = reconstruction + KL; sample from $\mathcal{N}(0, 1)$ to generate".

External resources

• Kingma and Welling: Auto-Encoding Variational Bayes (arXiv): the original VAE paper.
• Wikipedia: Variational autoencoder: the model, loss and reparameterization trick.
• Jeremy Jordan: Variational autoencoders: an illustrated, intuitive explanation.

Videos

• YouTube search: variational autoencoders explained: pick a visual explanation of the latent space.