D7.2 · LESSON 2 OF 6

Variational autoencoders

Explain the different training objectives of autoencoders, VAEs, GANs and diffusion models.

THE IDEA

Make the codes sampleable

A variational autoencoder (VAE) changes two things. The encoder does not output a single code but a distribution: a mean μ and a spread σ for each latent dimension. And the loss gets a second term that pulls every such distribution toward a fixed prior, usually N(0, 1).

loss = reconstruction + β · KL( encoder distribution ‖ prior )

For one dimension the KL term has a closed form: ½(σ² + μ² − 1 − ln σ²). It is exactly 0 when μ = 0 and σ = 1 (the prior itself) and grows as the encoder drifts away. Because all codes are kept near the prior, sampling a code from N(0, 1) and decoding it gives something sensible.

encoder mean μ1.5encoder spread σ0.5

Grey: the prior N(0, 1). Blue: what the encoder says about this input.

KL = ½(σ² + μ² − 1 − ln σ²) = ½(0.25 + 2.25 − 1 − -1.386) = 1.4431

weight β on the KL term0.5

loss = reconstruction + β·KL = 2 + 0.5·1.443 = 2.722

Move the encoder's mean and spread. Grey is the prior; the KL penalty is the price of straying from it.

Predict first. Does a VAE force every individual reconstruction to be exact?

Chapter context

Generative modeling tries to represent or sample from a data distribution. Different model families construct that distribution differently; a plausible-looking sample alone does not tell us how the model was trained or how well it covers the data.

TRY IT YOURSELF

Part 1. Do it hands-on

Compute KL values by hand and make the penalty vanish. The Python part implements the KL term, the full loss and a collapse check.

encoder mean μ1.5encoder spread σ0.5

Grey: the prior N(0, 1). Blue: what the encoder says about this input.

KL = ½(σ² + μ² − 1 − ln σ²) = ½(0.25 + 2.25 − 1 − -1.386) = 1.4431

weight β on the KL term0.5

loss = reconstruction + β·KL = 2 + 0.5·1.443 = 2.722

KL(μ=1, σ=1) KL(μ=0, σ=0.5) 2.0 + 0.5·0.5

To do: Set the encoder's mean and spread so the KL penalty is exactly zero
To do: KL for mean 1 and spread 1: ½(σ² + μ² − 1 − ln σ²). Type it
To do: KL for mean 0 and spread 0.5 (3 decimals)
To do: Total loss = reconstruction + β · KL with reconstruction 2.0, KL 0.5 and β = 0.5. Type it

Part 2. Do it in code

kl_gaussian(mu, sigma): KL( N(mu, sigma²) || N(0, 1) ) summed over all entries: 0.5 * sum(sigma² + mu² − 1 − log(sigma²)). mu and sigma are arrays. Raise ValueError if any sigma is not positive.
vae_loss(x, x_hat, mu, sigma, beta=1.0): reconstruction (sum of squared differences between x and x_hat) plus beta times kl_gaussian(mu, sigma). Return a float.
posterior_collapsed(mu, sigma, tol=1e-3): True if the encoder output is (within tol) the prior in every dimension, meaning it carries no information about the input: all |mu| <= tol and all |sigma − 1| <= tol.

import numpy as np

def kl_gaussian(mu, sigma):
    pass

def vae_loss(x, x_hat, mu, sigma, beta=1.0):
    pass

def posterior_collapsed(mu, sigma, tol=1e-3):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does a VAE force every individual reconstruction to be exact?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

A VAE uses $z=\mu+\sigma\odot\epsilon$ , with $\epsilon\sim\mathcal N(0,I)$ . A standard diffusion forward marginal is

x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon.

Here alpha-bar summarizes retained signal. This is a noising equation, not a complete sampling algorithm.

A VAE optimizes an evidence lower bound:

\operatorname{ELBO}=E_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)]-D_{KL}(q_\phi(z\mid x)\Vert p(z)).

For diagonal Gaussian q and standard normal p, its KL is $\frac12\sum_j(\mu_j^2+\sigma_j^2-1-\log\sigma_j^2)$ . The original GAN minimax objective is

\min_G\max_D E_{x\sim p_{data}}[\log D(x)]+E_{z\sim p(z)}[\log(1-D(G(z)))].

Practical generator losses often use a non-saturating variant. A common diffusion noise-prediction loss is $E\|\epsilon-\epsilon_\theta(x_t,t)\|^2$ under a specified noising schedule.

Worked example

For a scalar diffusion example, take x0=2, alpha-bar=0.64 and epsilon=−1. The signal term is 0.8×2=1.6 and the noise term is 0.6×(−1)=−0.6, giving xt=1.0. The clean value cannot be inferred from xt alone without additional assumptions or a learned model.

For q=N(1,1) against a standard normal prior, the one-dimensional VAE KL is 0.5×(1²+1−1−ln 1)=0.5 nats. This is only the regularization term, not the complete VAE objective.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import sqrt, isclose
x0, alpha, noise = 2,.64,-1
xt=sqrt(alpha)*x0+sqrt(1-alpha)*noise
assert isclose(xt,1)
assert 2+.5*(-1) == 1.5
from math import log
assert isclose(.5*(1**2+1**2-1-log(1**2)),.5)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

A VAE's approximate posterior is not identical to its prior. GAN loss is not a universal sample-quality score. A diffusion noising equation is not sufficient to implement the reverse model.