What is a VAE?
A Variational Autoencoder (VAE) is a neural network that learns to compress data into a smaller representation (latent space) and then reconstruct it. Unlike a regular autoencoder, a VAE learns a probability distribution over the latent space, which means you can generate new data by sampling from that distribution.
Training
Latent Space 2D
Click anywhere to sample. Drag points to morph. Colors show digit classes.
Latent Vector Manipulation
Morphing Animation
Reconstruction Grid
How It Works
The Variational Part
In a regular autoencoder, the encoder outputs a single latent vector. In a VAE, the encoder outputs two things: a mean (μ) and a log-variance (log σ²). We then sample from this distribution using the reparameterization trick:
z = μ + σ · ε, where ε ~ N(0, I)
This trick lets us backpropagate through the random sampling, which is essential for training.
KL Divergence
The KL divergence term pushes our learned distribution toward a standard normal distribution N(0, I). Without it, the VAE would just be a regular autoencoder. With it, nearby points in latent space produce similar outputs, and we can generate new images by sampling from N(0, I).
KL = -0.5 · Σ(1 + log(σ²) - μ² - σ²)
Loss Function
The total VAE loss has two parts:
- Reconstruction loss: How well does the decoder recreate the input? (binary cross-entropy)
- KL divergence: How close is the latent distribution to a standard normal?
Loss = ReconstructionLoss + β · KL
When β = 1, it's a standard VAE. When β < 1, reconstructions get sharper but the latent space becomes less organized. When β > 1, the latent space becomes more structured but reconstructions get blurrier.
Architecture Modes
Simple mode uses fully-connected layers only. It's fast to train and easy to understand, but the reconstructions aren't as sharp because it doesn't exploit the spatial structure of images.
Modern mode uses convolutional layers, which preserve spatial relationships in images. This produces much better reconstructions but takes longer to train and has more parameters.