A Variational Autoencoder (VAE), introduced by Kingma and Welling in 2013, is an autoencoder whose latent code is a probability distribution rather than a fixed point. Training balances reconstruction quality against keeping the latent distribution close to a simple prior, which yields a smooth, continuous latent space you can sample from to generate new data.
Parent topics: Generative Models and Autoencoders.
The encoder outputs the parameters of a Gaussian qφ(z|x) = N(μ, diag(σ2)). Sampling is not differentiable, so the VAE rewrites the sample as z = μ + σ ⊙ ε with ε ~ N(0, I). The randomness now sits in an input, and gradients flow through μ and σ to the encoder. The network predicts log σ2 for numerical stability.
The log-likelihood log p(x) is intractable, so the VAE maximises a lower bound: ELBO = Eq(z|x)[log p(x|z)] − KL(q(z|x) ‖ p(z)). The first term is reconstruction (MSE for Gaussian decoders, BCE for Bernoulli). The KL term regularises each posterior toward the prior N(0, I) and has a closed form for Gaussians. In code we minimise the negative ELBO.
Because every input maps to a region (not a point) and all regions are pulled toward the same prior, neighbouring latent points decode to similar outputs. That enables sampling new data, interpolation between examples, latent-space arithmetic, and semi-supervised learning. The encoder's μ is also a good low-dimensional embedding, and the ELBO can serve as an anomaly score.
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import datasets, transforms
class VAE(nn.Module):
def __init__(self, d_in=784, h=400, d_z=20):
super().__init__()
self.fc1 = nn.Linear(d_in, h)
self.mu = nn.Linear(h, d_z)
self.logvar = nn.Linear(h, d_z)
self.dec = nn.Sequential(nn.Linear(d_z, h), nn.ReLU(), nn.Linear(h, d_in))
def forward(self, x):
h = F.relu(self.fc1(x))
mu, logvar = self.mu(h), self.logvar(h)
z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu) # reparameterization
return self.dec(z), mu, logvar
def neg_elbo(logits, x, mu, logvar, beta=1.0):
rec = F.binary_cross_entropy_with_logits(logits, x, reduction="sum")
kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return (rec + beta * kl) / x.size(0)
ds = datasets.MNIST("data", train=True, download=True, transform=transforms.ToTensor())
dl = torch.utils.data.DataLoader(ds, batch_size=128, shuffle=True)
model = VAE(); opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
for x, _ in dl:
x = x.view(x.size(0), -1)
logits, mu, logvar = model(x)
loss = neg_elbo(logits, x, mu, logvar)
opt.zero_grad(); loss.backward(); opt.step()
print(epoch, round(loss.item(), 1))
# generate: sample from the prior and decode
with torch.no_grad():
samples = torch.sigmoid(model.dec(torch.randn(16, 20))).view(16, 1, 28, 28)z and KL goes to zero. Use KL annealing (warm up β from 0), free bits, or a weaker decoder.μ at inference for deterministic features.A VAE is a latent-variable model: draw z from a prior p(z) = N(0, I), then draw x from a decoder pθ(x|z). The data likelihood p(x) = ∫ p(x|z) p(z) dz is an integral over every latent value and has no closed form when the decoder is a neural network. Variational inference introduces an approximate posterior qφ(z|x), produced by the encoder, and uses the identity
log p(x) = ELBO(x) + KL(q(z|x) ‖ p(z|x)).
The second KL is always non-negative, so the ELBO is a lower bound on log p(x), and maximizing it both improves the model and pulls q toward the true posterior. For a diagonal Gaussian encoder and standard normal prior, the regularizer has a closed form per latent dimension j: KL = ½ Σj (μj2 + σj2 − log σj2 − 1). A dimension with μ = 0 and σ = 1 costs nothing; that is exactly what an unused, "collapsed" dimension looks like.
Worked example. Suppose for one input a latent dimension has μ = 1.0 and log σ2 = −1 (so σ2 ≈ 0.368). Its KL contribution is ½(1.0 + 0.368 + 1 − 1) ≈ 0.68 nats. Moving μ toward 0 or widening σ toward 1 lowers the cost; encoding precise, distinct information raises it. The decoder's reconstruction term pushes the other way, and the balance between the two is the whole design space of the model.
Each minibatch passes through the encoder to get μ and log σ2, draws one ε per example (one Monte Carlo sample is usually enough), decodes, and sums reconstruction and weighted KL. Gradients reach the encoder because z is a deterministic function of μ, σ and the external noise. At generation time the encoder is not used: sample z from N(0, I) and decode. For reconstruction or embedding, use μ directly rather than a noisy sample.
| Setting | Effect and guidance |
|---|---|
| Latent size | Small (2 for visualization, 16 to 128 for images) keeps the space smooth; too large leaves many dimensions collapsed to the prior, which is harmless but wasteful. |
| β (KL weight) | β = 1 is the true ELBO. β > 1 gives smoother, more disentangled but blurrier results; β < 1 sharpens reconstructions at the cost of sample quality. Latent diffusion autoencoders use a very small KL weight. |
| KL annealing | Ramp β from 0 to its target over the first few epochs or thousands of steps to prevent early posterior collapse, especially with autoregressive text decoders. |
| Reconstruction likelihood | Bernoulli (BCE) for binarized images, Gaussian (MSE) for continuous data. The implied variance of the Gaussian silently acts like a β, so keep the scaling consistent. |
| Loss reduction | Sum over pixels and latent dimensions, then average over the batch. Mixing sum and mean is the most common bug in VAE code. |
Report the negative ELBO on held-out data (in nats, or bits per dimension for images); it is an upper bound on the true negative log-likelihood, which importance-weighted estimates (Burda et al., 2015) tighten. Track the reconstruction and KL terms separately: a KL near zero signals collapse. For samples, visual inspection, Fréchet Inception Distance on image data, and latent interpolations between real examples reveal whether the latent space is smooth. For a representation, measure downstream accuracy of a simple classifier trained on μ.
| Model | Likelihood | Encoder | Sample quality | Training stability |
|---|---|---|---|---|
| VAE | Lower bound | Yes | Often blurry | Stable |
| GAN | None | No (by default) | Sharp | Can be unstable, mode collapse |
| Diffusion | Bound | Fixed noising process | High | Stable, slow sampling |
| Autoregressive | Exact | No | High | Stable, sequential sampling |
The models are increasingly combined rather than chosen between: a VAE or VQ-VAE compresses data, and a diffusion or autoregressive prior models the compressed codes.
Further reading. Kingma and Welling (2013), "Auto-Encoding Variational Bayes", arXiv:1312.6114; Rezende, Mohamed and Wierstra (2014), "Stochastic Backpropagation and Approximate Inference in Deep Generative Models"; Higgins et al. (2017), "β-VAE", ICLR; Bowman et al. (2016), "Generating Sentences from a Continuous Space" (KL annealing); Burda, Grosse and Salakhutdinov (2015), "Importance Weighted Autoencoders"; Rombach et al. (2022), "High-Resolution Image Synthesis with Latent Diffusion Models"; Kingma and Welling (2019), "An Introduction to Variational Autoencoders".