Mathematical Foundations of VAEs

Most VAE tutorials hand you a loss function: reconstruction error plus a KL term. The code works, the samples look like blurry digits, and you move on. Then you change the reconstruction loss from a sum to a mean, or swap binary cross-entropy for squared error, and the model stops generating anything useful. The reason is that the VAE loss is not a recipe someone tuned by hand. It is a lower bound on a log-likelihood, and every piece of it has a specific job.

This article derives that loss from the start: what a VAE is trying to model, why the obvious approach is intractable, how the evidence lower bound (ELBO) gets around it, and how the reparameterization trick makes the whole thing trainable with ordinary backpropagation. If you want the intuition first, read the introduction to variational autoencoders. This piece assumes you know what an encoder and decoder are and are comfortable with expectations and logarithms.

The problem a VAE solves

A generative model tries to learn the distribution your data came from, written p(x). If you had p(x), you could sample new data from it and score how likely a given input is. The generative vs. discriminative models article covers why that is a harder goal than predicting a label.

Images, audio and text are high-dimensional, but they are not random. A 28 by 28 handwritten digit has 784 pixels, yet the things that actually vary between digits are few: which digit, the slant, the stroke thickness. A latent variable model makes that explicit. It assumes each data point x was produced by first drawing a hidden code z from a simple distribution, then generating x from z:

z ~ p(z)            prior, usually a standard normal N(0, I)
x ~ p_θ(x | z)      decoder, a neural network with parameters θ

The decoder network outputs the parameters of a distribution over x, for example the probability of each pixel being on. The model’s distribution over data is what you get by averaging over every possible code:

p_θ(x) = ∫ p_θ(x | z) p(z) dz

Training means choosing θ to maximize the log-likelihood of the training data, the sum of log p_θ(x) over your examples. That is where the trouble starts.

Why the likelihood is intractable

The integral runs over every point in latent space, and the decoder is a nonlinear neural network, so there is no closed form. You could estimate it by sampling many z from the prior and averaging p_θ(x | z), but almost every random z produces something that looks nothing like a given x. The estimate is dominated by the rare samples that happen to land in the right region, and you would need an enormous number of them per training example.

The related quantity you would really like is the posterior, which codes are likely to have produced this particular x:

p_θ(z | x) = p_θ(x | z) p(z) / p_θ(x)

It has p_θ(x) in the denominator, so it is just as intractable. This is the core difficulty: to learn the decoder you need to know which codes explain each input, and to know that you need the quantity you are trying to learn.

Variational inference and the encoder

Variational inference replaces the true posterior with an approximation you can compute. You pick a family of simple distributions, and for each input you find the member closest to p_θ(z | x). In a VAE that family is a diagonal Gaussian, and instead of solving a separate optimization problem per data point, a second network predicts its parameters directly from x:

q_φ(z | x) = N(z; μ_φ(x), diag(σ²_φ(x)))

That network is the encoder, with parameters φ. Predicting posterior parameters with a shared network is called amortized inference: you pay the cost once during training, and afterwards inference for a new input is a single forward pass.

flowchart LR
    subgraph Inference["Inference model q_φ"]
        X1["Input x"] --> E["Encoder"]
        E --> MS["μ and σ"]
    end
    MS --> Z["Code z"]
    subgraph Generative["Generative model p_θ"]
        P["Prior N(0, I)"] -.->|"sampling new data"| Z
        Z --> D["Decoder"]
        D --> X2["Distribution over x"]
    end

During training, codes come from the encoder. When you generate new data, you skip the encoder and draw z from the prior. That second path only works if the codes the encoder produces look like samples from the prior, which is exactly what the KL term in the loss enforces.

Here is the notation used in the rest of the article:

Deriving the evidence lower bound

The goal is still to maximize log p_θ(x). The trick is to write it in terms of the encoder. Because q_φ(z | x) integrates to 1 and log p_θ(x) does not depend on z, you can take its expectation under q without changing it:

log p_θ(x) = E_q[ log p_θ(x) ]
           = E_q[ log ( p_θ(x, z) / p_θ(z | x) ) ]                         Bayes' rule
           = E_q[ log ( p_θ(x, z) / q_φ(z | x) ) ]
             + E_q[ log ( q_φ(z | x) / p_θ(z | x) ) ]                      multiply and divide by q
           = ELBO(θ, φ; x) + KL( q_φ(z | x) ‖ p_θ(z | x) )

Every step is exact. The log-likelihood splits into two parts:

• The ELBO, which contains only things you can evaluate: the joint p_θ(x, z) = p_θ(x | z) p(z) and the encoder q_φ.
• The KL divergence between the encoder and the true posterior, which you cannot evaluate, because it contains p_θ(z | x).

KL divergence is never negative, so the ELBO is always less than or equal to log p_θ(x). That is where the name comes from: it is a lower bound on the evidence. The gap between the two is exactly how wrong the encoder’s approximation is.

flowchart TD
    L["log p_θ(x): fixed for a given θ"] --> A["ELBO: computable"]
    L --> B["KL(q ‖ true posterior): gap, not computable"]
    A --> C["Raise ELBO with θ: better model"]
    A --> D["Raise ELBO with φ: smaller gap"]

This explains why a single objective trains both networks. For a fixed decoder, log p_θ(x) is a constant, so increasing the ELBO by changing φ can only shrink the gap: the encoder moves toward the true posterior. Changing θ increases the ELBO by making the data more likely under the model, and usually raises the log-likelihood itself. You never compute the gap, but maximizing the ELBO pushes on both fronts at once.

You will also see the bound derived in one line with Jensen’s inequality, since log is concave and the log of an average is at least the average of the log. That route gets the same result but hides what the gap is, which is why the identity above is more useful.

The two terms of the loss

Expanding the joint inside the ELBO gives the form every implementation uses:

ELBO = E_q[ log p_θ(x | z) ]  −  KL( q_φ(z | x) ‖ p(z) )
       reconstruction term       regularization term

The first term rewards codes that let the decoder reproduce x. On its own, it would turn the VAE into an ordinary autoencoder: the encoder would shrink σ toward zero and put each input at its own isolated point. The second term penalizes the encoder for moving away from the prior. It pulls every posterior toward N(0, I), so the codes for different inputs overlap and fill the space, and a random draw from the prior lands somewhere the decoder knows how to handle.

The training loss is the negative ELBO, averaged over a mini-batch. Minimizing it is maximizing a bound on the likelihood, not an ad hoc mix of two goals.

The KL term in closed form

With a diagonal Gaussian encoder and a standard normal prior, the KL term needs no sampling. For a single latent dimension with mean μ and standard deviation σ:

KL( N(μ, σ²) ‖ N(0, 1) ) = −log σ + (σ² + μ²) / 2 − 1/2
                         = ½ ( μ² + σ² − log σ² − 1 )

Because the dimensions of a diagonal Gaussian are independent, the KL for the whole code is the sum over dimensions:

KL = ½ Σ_j ( μ_j² + σ_j² − log σ_j² − 1 )

Each part has an interpretation. The μ² term charges for moving a code away from the origin. The σ² − log σ² − 1 part is zero at σ = 1 and grows in both directions, so it charges for being either too certain or too spread out. A dimension that matches the prior exactly, μ = 0 and σ = 1, costs nothing, and as you will see below, that is also the signature of a dimension the model has stopped using.

A worked example: one dimension with μ = 0.5 and σ = 0.8.

½ ( 0.25 + 0.64 − ln 0.64 − 1 ) = ½ ( 0.25 + 0.64 + 0.4463 − 1 ) ≈ 0.168 nats

Encoders usually output log σ² rather than σ. It can take any real value, so the network needs no activation to keep it positive, and σ² = exp(log σ²) is always valid. You can check the closed form against PyTorch’s own implementation and against a Monte Carlo estimate:

import torch
from torch.distributions import Normal, kl_divergence

mu = torch.tensor([0.5])
sigma = torch.tensor([0.8])
logvar = torch.log(sigma ** 2)

closed_form = 0.5 * (mu ** 2 + logvar.exp() - logvar - 1)

q = Normal(mu, sigma)
p = Normal(torch.zeros(1), torch.ones(1))
library = kl_divergence(q, p)

z = q.sample((1_000_000,))
monte_carlo = (q.log_prob(z) - p.log_prob(z)).mean()

print(closed_form.item(), library.item(), monte_carlo.item())
# about 0.1682, 0.1682, and a value close to 0.168

The reconstruction term depends on your decoder

The reconstruction term is E_q[log p_θ(x | z)], a log-likelihood, so its exact form depends on the distribution the decoder outputs. Picking a loss function is really picking that distribution.

The Gaussian row matters most in practice. If you use plain squared error, you have implicitly fixed σ_x² = 0.5, and the scale of the reconstruction term relative to the KL term is set by that choice. Use a smaller σ_x and reconstruction dominates: sharper outputs, codes that drift from the prior, poorer samples. Use a larger one and the KL term dominates: blurrier outputs and a smoother latent space. Much of the tuning people do on VAEs is this trade-off under a different name.

The same reasoning explains a common bug. The ELBO is a sum over every pixel and every latent dimension. If you average the reconstruction error over 784 pixels but sum the KL over 20 latent dimensions, you have shrunk the reconstruction term by a factor of 784 relative to the bound. The KL term wins, the encoder outputs the prior for every input, and the decoder produces the same average image every time. Sum over the elements of each example, then average over the batch.

The reparameterization trick

The ELBO is an expectation over z ~ q_φ(z | x), and you estimate it by sampling z. Gradients with respect to θ are straightforward: sample z, run the decoder, backpropagate. Gradients with respect to φ are the problem, because φ controls the distribution you are sampling from, and sampling is not a differentiable operation.

The fix is to move the randomness out of the path from φ to the loss. Sample noise from a fixed distribution, then build z from it with a deterministic function:

ε ~ N(0, I)
z = μ_φ(x) + σ_φ(x) ⊙ ε

Here ⊙ is element-wise multiplication. The resulting z has exactly the same distribution as before, but now z is a differentiable function of μ and σ, and ε is just an input that needs no gradient.

flowchart LR
    X["x"] --> ENC["Encoder φ"]
    ENC --> MU["μ"]
    ENC --> LV["log σ²"]
    EPS["ε ~ N(0, I)"] --> Z["z = μ + σ ⊙ ε"]
    MU --> Z
    LV --> Z
    Z --> DEC["Decoder θ"]
    DEC --> LOSS["Negative ELBO"]
    MU --> KL["KL term"]
    LV --> KL
    KL --> LOSS

The alternative, the score function estimator (also known as REINFORCE), can differentiate through sampling without this trick, but its gradient estimates have high variance and need many samples or careful variance reduction. The reparameterized estimator uses the gradient of the decoder itself, which carries much more information per sample. In the original paper, Auto-Encoding Variational Bayes by Kingma and Welling, a single sample of z per data point was enough, as long as the mini-batch was reasonably large.

The trick applies to any distribution you can write as a differentiable transformation of fixed noise. Gaussians are the easy case. Discrete latents are not, which is why discrete VAEs need extra machinery such as continuous relaxations.

Putting it together in PyTorch

Here is the full loss for a Bernoulli decoder, written to match the derivation term by term. The model’s forward pass returns the decoder logits plus the encoder’s μ and log σ²:

import torch
import torch.nn as nn
import torch.nn.functional as F


class VAE(nn.Module):
    def __init__(self, x_dim=784, h_dim=400, z_dim=20):
        super().__init__()
        self.enc = nn.Linear(x_dim, h_dim)
        self.mu = nn.Linear(h_dim, z_dim)
        self.logvar = nn.Linear(h_dim, z_dim)
        self.dec1 = nn.Linear(z_dim, h_dim)
        self.dec2 = nn.Linear(h_dim, x_dim)

    def encode(self, x):
        h = F.relu(self.enc(x))
        return self.mu(h), self.logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + std * eps

    def decode(self, z):
        return self.dec2(F.relu(self.dec1(z)))  # logits of p_θ(x | z)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


def negative_elbo(logits, x, mu, logvar, beta=1.0):
    # Reconstruction: -log p_θ(x | z), summed over pixels of each example
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="none").sum(dim=1)
    # KL(q_φ(z | x) || N(0, I)), summed over latent dimensions
    kl = 0.5 * (mu.pow(2) + logvar.exp() - logvar - 1).sum(dim=1)
    # Average over the batch
    return (recon + beta * kl).mean()

Inputs are flattened images with values in [0, 1]. Using binary_cross_entropy_with_logits instead of a sigmoid followed by binary_cross_entropy is more numerically stable, since it never computes log(0). The beta argument is 1 for a standard VAE; the next section covers why you might change it.

Sampling new data uses only the decoder and the prior:

model.eval()
with torch.no_grad():
    z = torch.randn(16, 20)              # 16 codes from N(0, I)
    images = torch.sigmoid(model.decode(z)).view(16, 28, 28)

For training loops, optimizers and learning rates, the deep learning basics article covers the general mechanics, and the next article in this series builds a complete VAE step by step.

Where the theory shows up in practice

Weighting the KL term

Multiplying the KL term by a factor β gives the objective used by β-VAE (Higgins et al., ICLR 2017). With β greater than 1, the model is pushed harder toward the prior, which squeezes how much information each latent dimension can carry. The authors found this encourages disentangled codes, where individual dimensions line up with separate factors such as position or rotation, at the cost of blurrier reconstructions. With β not equal to 1, the objective is no longer a bound on log p_θ(x), so you are trading likelihood for structure on purpose.

Posterior collapse

If the decoder is powerful enough to model x without any help from z, the cheapest way to maximize the ELBO is to ignore the code entirely: the encoder outputs the prior for every input, the KL term drops to zero, and the decoder models the data on its own. This is posterior collapse. Bowman et al. reported it in VAEs for text with recurrent decoders, where most straightforward training runs produced an encoder equal to the prior.

You can detect it from the closed-form KL: compute the average KL per latent dimension over a validation set. Dimensions near zero carry no information about the input. Common remedies:

• KL annealing: start β at 0 and ramp it to 1 over the first part of training, so the model learns to use z before the penalty applies in full.
• Free bits: stop penalizing each dimension’s KL below a small threshold, so using a little capacity is free.
• A weaker decoder: limit what the decoder can do alone, for example by dropping some of its inputs during training.

Tighter bounds

The gap between the ELBO and log p_θ(x) is the KL from the encoder to the true posterior. A diagonal Gaussian cannot represent posteriors that are correlated or multi-modal, so the gap never fully closes. Importance-weighted autoencoders average several samples inside the logarithm to get a tighter bound, and normalizing flows make q more flexible. Both cost more compute per step, and for many uses the standard bound is good enough.

Further reading

• Auto-Encoding Variational Bayes, Kingma and Welling: the original paper, including the derivation of the Gaussian KL in its appendix.
• PyTorch torch.distributions documentation: distributions, rsample for reparameterized sampling and kl_divergence.
• Overview of generative models and introduction to GANs, for how VAEs compare with other model families.

Key takeaways

• A VAE is a latent variable model: draw a code z from a simple prior, then generate x from z with a decoder network.
• The likelihood p_θ(x) and the true posterior p_θ(z | x) are intractable, so an encoder learns an approximate Gaussian posterior q_φ(z | x).
• log p_θ(x) equals the ELBO plus the KL divergence from q to the true posterior. Maximizing the ELBO improves the model and the approximation at the same time.
• The ELBO splits into a reconstruction log-likelihood and KL(q ‖ prior). For Gaussians the KL is ½ Σ (μ² + σ² − log σ² − 1), computed without sampling.
• The reconstruction loss is the negative log-likelihood of your chosen decoder distribution. Its scale relative to the KL term matters, so sum over elements before averaging over the batch.
• The reparameterization trick, z = μ + σ ⊙ ε, makes sampling differentiable so both networks train with ordinary backpropagation.
• β weighting, KL annealing and free bits all adjust the balance between the two terms, mainly to get disentangled codes or to avoid posterior collapse.