Variational Autoencoders: ELBO, Reparameterization, and Latent Space

A variational autoencoder combines a probabilistic decoder with an encoder that approximates inference about latent variables. It specifies a prior from which to generate new codes and trains with a lower bound on data log-likelihood. An ordinary autoencoder can also support generation if we fit a distribution to its codes; its reconstruction objective alone does not supply that distribution. Here we derive the VAE objective, implement it, and train a small binary-pattern model.

The generative view

Assume data is produced by sampling a latent vector \(z\sim p(z)=\mathcal N(0,I)\), then an observation \(x\sim p_\theta(x\mid z)\). The decoder’s parameters are \(\theta\); I means independent unit-variance latent coordinates. The decoder outputs the parameters of an observation distribution, such as a Bernoulli probability for each binary pixel. The probability of x (or its density for a continuous observation model) averages over possible latent causes:

\[p_\theta(x)=\int p_\theta(x\mid z)p(z)\,dz\]

For a nonlinear neural decoder, this integral usually lacks a tractable analytic form. The difficulty depends on the model: some linear-Gaussian cases can be integrated exactly. Prior sampling gives a Monte Carlo estimate, but it can have high relative variance if the latent regions explaining a particular observation have little prior mass.

The true posterior \(p_\theta(z\mid x)\) describes latent causes conditional on the observed x. An encoder with parameters \(\phi\) supplies a tractable approximation \(q_\phi(z\mid x)\). In this article it is a diagonal Gaussian with an input-dependent mean and variance. The same encoder is used for every observation, rather than solving a separate inference optimization for each one.

The reconstruction objective this builds on is covered in Autoencoders and Representation Learning.

The ELBO

ELBO stands for evidence lower bound, where the evidence is \(p_\theta(x)\). Write \(p_\theta(x,z)=p_\theta(x\mid z)p(z)\), multiply and divide by q, and apply the concavity of log. Assuming q covers the relevant support and the expectations are finite:

\[\log p_\theta(x)=\log\mathbb E_q\!\left[\frac{p_\theta(x,z)}{q_\phi(z\mid x)}\right]\ge\mathbb E_q[\log p_\theta(x,z)-\log q_\phi(z\mid x)]\]

Log of an average is at least the average of the logs. Splitting the joint log probability into a likelihood term and a prior term gives:

\[\log p_\theta(x)\ge\mathbb{E}_{q_\phi(z\mid x)}\left[\log p_\theta(x\mid z)\right]-D_{KL}\left(q_\phi(z\mid x)\,\Vert\,p(z)\right)\]

The expected log-likelihood rewards assigning probability to the observed x when z is drawn from the encoder. The KL term compares that encoder distribution with the prior. Training minimizes the negative ELBO: a reconstruction negative log-likelihood plus KL. A squared-error reconstruction term follows only after choosing an appropriate Gaussian observation model; it is not part of the definition of a VAE.

The exact gap is \(\log p_\theta(x)-\mathrm{ELBO}=D_{KL}(q_\phi(z\mid x)\Vert p_\theta(z\mid x))\). This KL uses the true posterior, whereas the loss above uses the prior. At fixed decoder parameters, improving the encoder’s ELBO reduces the gap. When both parameter sets change, an increase in ELBO does not guarantee that log-likelihood rises and the gap shrinks separately on every step.

The prior already defines a sampling distribution before training. The KL penalty encourages encoder distributions to stay close to it, helping align training codes with generation codes. It neither creates mathematical continuity nor guarantees that every sampled or interpolated code decodes well. The continuous decoder used below is continuous by construction, even without this penalty.

The reparameterization trick

To optimize the expected reconstruction term, we need gradients with respect to the distribution parameters. PyTorch’s Normal.sample() returns samples without an autograd path to those parameters. A naive backward pass through those samples therefore cannot train the encoder through this term. This API behavior does not mean every sampling-based gradient estimator is impossible.

For a diagonal Gaussian, sample parameter-independent noise \(\varepsilon\sim\mathcal N(0,I)\) and express z as a differentiable transformation:

\[z=\mu_\phi(x)+\sigma_\phi(x)\odot\varepsilon\]

For each sampled epsilon, z is differentiable in the encoder outputs. Coordinatewise, \(\partial z/\partial\mu=1\) and \(\partial z/\partial\log\sigma^2=\tfrac12\sigma\varepsilon\). Gradients of a sampled loss then estimate gradients of its expectation under the usual differentiability and integrability conditions. PyTorch’s Normal.rsample() implements this pathwise construction. Score-function estimators provide another route, including for discrete latent variables; the appropriate estimator depends on the distribution and objective.

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

class VAE(nn.Module):
    def __init__(self, in_dim=784, hidden=400, latent=20):
        super().__init__()
        if min(in_dim, hidden, latent) < 1:
            raise ValueError("model dimensions must be positive")
        self.in_dim, self.latent_dim = in_dim, latent
        self.enc = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU())
        self.fc_mu = nn.Linear(hidden, latent)
        self.fc_logvar = nn.Linear(hidden, latent)
        self.dec = nn.Sequential(
            nn.Linear(latent, hidden), nn.ReLU(),
            nn.Linear(hidden, in_dim))

    def encode(self, x):
        h = self.enc(x)
        return self.fc_mu(h), self.fc_logvar(h)

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

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.dec(z), mu, logvar        # decoder returns logits

The encoder emits \(\mu\) and \(\ell=\log\sigma^2\), each of shape (N,d) for batch size N and latent width d. The formula \(\sigma=\exp(\ell/2)\) makes the standard deviation positive mathematically, but very large or small log variances can overflow or underflow numerically. Monitor finite losses and gradients; a softplus parameterization is another positive-scale choice. Inputs are flattened (N,D) tensors, and the decoder returns (N,D) logits. Run the code blocks in order. Calling eval() does not disable the explicit random draw in forward().

The loss in closed form

For the diagonal encoder Gaussian and standard normal prior used here, the per-example KL is analytic. Summing over its d latent coordinates gives:

\[D_{KL}=-\tfrac{1}{2}\sum_{j=1}^{d}\left(1+\log\sigma_j^2-\mu_j^2-\sigma_j^2\right)\]

def vae_loss(recon_logits, x, mu, logvar, beta=1.0):
    if x.ndim != 2 or x.size(0) == 0 or recon_logits.shape != x.shape:
        raise ValueError("inputs and logits must have the same nonempty (N,D) shape")
    if mu.ndim != 2 or mu.shape != logvar.shape or mu.size(0) != x.size(0):
        raise ValueError("mu and logvar must share shape (N,d)")
    if not 0 <= beta < float("inf"):
        raise ValueError("beta must be finite and nonnegative")
    if not torch.isfinite(x).all() or torch.any((x < 0) | (x > 1)):
        raise ValueError("BCE targets must be finite and in [0,1]")
    recon = F.binary_cross_entropy_with_logits(
        recon_logits, x, reduction="none").sum(dim=1).mean()
    kl_per_dim = 0.5 * (mu.square() + logvar.exp() - 1 - logvar)
    kld = kl_per_dim.sum(dim=1).mean()
    return recon + beta * kld, recon.item(), kld.item()

torch.manual_seed(0)
model = VAE()
x = torch.bernoulli(torch.full((16, 784), 0.5))
logits, mu, logvar = model(x)
total, r, k = vae_loss(logits, x, mu, logvar)
print("reconstruction", round(r, 2), "KL", round(k, 4))
# reconstruction 548.56 KL 0.3779

This decoder models conditionally independent Bernoulli coordinates, so the reconstruction term sums binary log losses over all D coordinates. It is a likelihood for binary targets; using continuous targets in [0,1] makes BCE a surrogate, not a normalized continuous-data likelihood. KL sums over d coordinates, and both terms average over the batch. Changing just the KL sum to a mean divides its weight by d; changing just reconstruction to a mean multiplies the effective relative KL weight by D. At D=784 and d=20 these are different rescalings.

With mu=0 and variance=1, a coordinate’s KL is zero; with mu=1 and variance=1, it is 0.5 nats. Natural logs give nats; divide by log 2 for bits. The untrained model’s reconstruction term is around 549, compared with 784 log 2 ≈ 543 for uniform binary predictions. A large total is expected when summing hundreds of coordinates.

Posterior collapse

Posterior collapse occurs when the encoder becomes uninformative about the input, often approaching the prior. Exact q(z|x)=p(z) for every x gives zero KL and no input information in z. For the factorized Bernoulli decoder here, an input-independent optimum predicts each coordinate’s marginal frequency. An autoregressive decoder can ignore z while still modeling complex dependencies and producing varied outputs; collapse does not universally imply a dataset-mean image or visible blur.

A decoder that models the data well without z may gain little reconstruction benefit from paying a KL cost to use it. Optimization and the variational family also matter. Possible interventions include:

• KL annealing. Start with a reduced KL weight and raise it during training. This can help the model begin using z, but does not ensure that it keeps doing so.
• Free bits. For example, floor the batch-averaged KL of each dimension: \(\sum_j\max(\lambda,\overline{KL}_j)\). Below the floor this penalty has no gradient pushing that KL lower. It does not force a dimension to carry at least lambda information; grouping and averaging conventions must be specified.
• Weaken the decoder. Restricting a powerful decoder can make z more useful, but excessive restriction can also cause underfitting.

Track per-dimension KL averaged over representative held-out inputs, reconstruction loss, and sensitivity to replacing or shuffling z. Some near-zero dimensions can be unused capacity in an otherwise useful representation. Positive KL alone does not establish input information: q(z|x)=N(2,1) for every x has KL 2 nats from N(0,1), yet is independent of x. Encoder-mean variation and downstream probes provide complementary evidence.

Beta-VAE and the rate–distortion trade-off

The objective \(\mathcal D+\beta R\) weighs reconstruction negative log-likelihood \(\mathcal D\) against average KL R. At beta=1 it is the negative ELBO; other weights modify the optimization target, and beta below 1 does not generally retain a log-likelihood lower-bound interpretation. Increasing beta strengthens pressure toward the prior, which can reduce input information or worsen reconstruction. Sharper images, smooth interpolation, and disentangled factors are not guaranteed at any particular beta. At beta=0 this implementation remains stochastic unless its encoder variances become negligible.

To interpret the average KL as rate, let \(q(z)=\mathbb E_{p_{data}(x)}q(z\mid x)\) be the aggregate encoder distribution. Then \(R=\mathbb E_xD_{KL}(q(z\mid x)\Vert p(z))=I_q(X;Z)+D_{KL}(q(z)\Vert p(z))\). Here \(I_q(X;Z)\) is mutual information, measuring dependence between input and code under the encoder’s joint distribution. Rate includes both this information and mismatch to the prior; it is an upper bound on that mutual information, not exactly bits transmitted. Divide nats by log 2 to express the same quantities in bits. Beta changes the trade-off, whose practical result must be measured for the chosen model and optimization.

Sampling and interpolation

@torch.no_grad()
def sample(model, n):
    if not isinstance(n, int) or n < 1:
        raise ValueError("n must be a positive integer")
    param = next(model.parameters())
    was_training = model.training
    model.eval()
    try:
        z = torch.randn(n, model.latent_dim, device=param.device, dtype=param.dtype)
        probabilities = torch.sigmoid(model.dec(z))
        observations = torch.bernoulli(probabilities)
        return observations, probabilities
    finally:
        model.train(was_training)


@torch.no_grad()
def interpolate(model, x1, x2, steps=8):
    if not isinstance(steps, int) or steps < 2:
        raise ValueError("use at least two interpolation steps")
    if x1.shape not in ((model.in_dim,), (1, model.in_dim)) or x2.shape != x1.shape:
        raise ValueError("provide two single inputs with matching feature shape")
    param = next(model.parameters())
    was_training = model.training
    model.eval()
    try:
        x1 = x1.reshape(1, -1).to(device=param.device, dtype=param.dtype)
        x2 = x2.reshape(1, -1).to(device=param.device, dtype=param.dtype)
        mu1, _ = model.encode(x1)
        mu2, _ = model.encode(x2)
        ts = torch.linspace(0, 1, steps, device=param.device, dtype=param.dtype)[:, None]
        z = (1 - ts) * mu1 + ts * mu2
        return torch.sigmoid(model.dec(z))
    finally:
        model.train(was_training)

sample() draws both a latent code and a binary observation, and also returns the decoder probabilities for inspection. A probability image is a conditional mean at that sampled z, not a sampled binary image. interpolate() decodes a line between two encoder means; its endpoints use those point estimates, not an average over the entire posterior. Smooth numerical changes follow from the decoder’s continuity, but semantic quality needs separate evaluation. A straight path can leave regions commonly visited during training; even for a standard normal prior, points near the origin are not where most probability mass lies in high dimensions.

A different approach to generation is developed in Diffusion Models Explained: DDPM, Samplers, and Guidance.

A small train-to-sample example

These eight-bit patterns let us run the full model on a CPU without downloading data. Training and validation are independent draws from the same two-pattern distribution; repeated patterns are expected here, and this is not a test of generalization to unseen image types. Both observation outcomes are binary, matching the likelihood. The evaluation averages 32 independent latent draws per input with beta=1. It estimates the negative ELBO; a finite Monte Carlo estimate is not itself guaranteed to bound the true negative log-likelihood.

torch.manual_seed(5)
patterns = torch.tensor([[1., 1., 1., 1., 0., 0., 0., 0.],
                         [0., 0., 0., 0., 1., 1., 1., 1.]])
train_x = patterns[torch.randint(2, (256,))]
valid_x = patterns[torch.randint(2, (64,))]
toy = VAE(in_dim=8, hidden=32, latent=2)
opt = torch.optim.Adam(toy.parameters(), lr=0.005)

@torch.no_grad()
def evaluate(model, x, draws=32):
    if draws < 1:
        raise ValueError("draws must be positive")
    was_training = model.training
    model.eval()
    try:
        records = []
        for _ in range(draws):
            logits, mu, logvar = model(x)
            loss, recon, kl = vae_loss(logits, x, mu, logvar)
            records.append([loss.item(), recon, kl])
        return torch.tensor(records).mean(0)
    finally:
        model.train(was_training)

before = evaluate(toy, valid_x)
for step in range(300):
    toy.train()
    xb = train_x[torch.randint(len(train_x), (32,))]
    logits, mu, logvar = toy(xb)
    loss, _, _ = vae_loss(logits, xb, mu, logvar, beta=1.0)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
after = evaluate(toy, valid_x)
print("negative ELBO estimate before/after", round(before[0].item(), 3),
      round(after[0].item(), 3))
print("after reconstruction/KL", round(after[1].item(), 3), round(after[2].item(), 3))
observations, probabilities = sample(toy, 4)
print("sample and probability shapes", tuple(observations.shape), tuple(probabilities.shape))
print("first sample", observations[0].int().tolist())
print("its probabilities", [round(v, 2) for v in probabilities[0].tolist()])
print("interpolation shape", tuple(interpolate(toy, patterns[0], patterns[1]).shape))
# negative ELBO estimate before/after 5.97 1.565
# after reconstruction/KL 0.069 1.496
# sample and probability shapes (4, 8) (4, 8)
# first sample [1, 1, 1, 1, 0, 0, 0, 0]
# its probabilities [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]
# interpolation shape (8, 8)

In this CPU run, the estimated negative ELBO falls from 5.970 to 1.565. Its final reconstruction and KL terms are about 0.069 and 1.496; the low reconstruction term alone would omit most of this objective. Exact values vary with the random draws and numerical environment. The displayed probabilities are rounded to two decimals, so 0.0 and 1.0 need not be exact zeros or ones. For real data, split before fitting preprocessing, use a likelihood appropriate to the observations, and reserve held-out data for evaluating the selected model. The helper expects its input tensor on the model’s device.

Likelihood choices and latent image models

For a Gaussian decoder with fixed variance s², negative log-likelihood is \(\lVert x-m_\theta(z)\rVert^2/(2s^2)\) plus a constant. The optimal mean averages ambiguous targets, so displaying that mean can blur details. Sampling from the Gaussian adds observation noise; it does not recover which ambiguous structure was intended. Our Bernoulli decoder also returns conditional means as probabilities, but uses binary cross-entropy rather than MSE. Likelihood family, latent information, architecture, and optimization all affect visual quality, so this averaging argument is not a universal explanation of differences from GANs or diffusion models.

Latent diffusion uses a trained image autoencoder to reduce spatial resolution before diffusion modeling. Such autoencoders may use KL or other regularization together with perceptual and adversarial losses, so they need not be trained with the simple Bernoulli ELBO above. The diffusion model learns a distribution of the resulting image codes; dimensionality reduction lowers its computational cost without guaranteeing any particular sample quality.

How to choose an objective for a given problem is the subject of Loss Function Design: Choosing an Objective That Matches the Problem.

Exercises

1. The closed-form KL. Compare the analytic Gaussian KL with Monte Carlo estimates at three sample sizes. Repeat each estimate 200 times and report the mean and standard deviation across those repeats. Distinguish the precision of one estimate from the precision of their average.

With independent samples and finite variance, the standard deviation of one Monte Carlo mean scales as 1/√n.

import torch
torch.manual_seed(0)
mu, logvar = torch.randn(1, 8), torch.randn(1, 8)
std = (0.5 * logvar).exp()

kl_closed = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum()
print("closed form", round(kl_closed.item(), 4))

q = torch.distributions.Normal(mu, std)
p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))

import statistics as st
for n in (10, 100, 10_000):
    est = []
    for _ in range(200):                       # 200 independent estimates
        z = q.sample((n,))
        est.append((q.log_prob(z) - p.log_prob(z)).sum(-1).mean().item())
    print(f"n={n:6d} mean {st.mean(est):.4f}  sd {st.stdev(est):.4f}")
# closed form 6.7896
# n=    10 mean 6.9506  sd 1.2161
# n=   100 mean 6.7607  sd 0.3821
# n= 10000 mean 6.7910  sd 0.0368

2. A detached sample versus a pathwise gradient. Estimate \(\partial\mathbb E[z^2]/\partial\mu\) for \(z\sim\mathcal N(\mu,\sigma^2)\). Compare Normal.sample() with a reparameterized sample and report whether the first path supplies an autograd gradient.

A missing gradient is different from a gradient whose value is zero. The pathwise estimate should be near 2μ, with sampling noise.

torch.manual_seed(0)
mu = torch.tensor(1.5, requires_grad=True)
sigma = torch.tensor(0.5)
z = torch.distributions.Normal(mu, sigma).sample((100_000,))
print("direct grad_fn", z.grad_fn)
try:
    z.square().mean().backward()
except RuntimeError:
    print("direct backward: no autograd path")
print("mu.grad", mu.grad)

mu2 = torch.tensor(1.5, requires_grad=True)
eps = torch.randn(100_000)
z2 = mu2 + sigma * eps
z2.square().mean().backward()
print("reparam grad", round(mu2.grad.item(), 4), "analytic", 3.0)
# direct grad_fn None
# direct backward: no autograd path
# mu.grad None
# reparam grad 2.9951 analytic 3.0

3. Interpreting per-dimension KL. Treat the following illustrative values as KL averages over evaluation inputs. Count coordinates above a chosen threshold and calculate the share in the top three. Does that establish how much input information each coordinate carries?

The threshold summarizes KL allocation. It does not identify mutual information or decoder use by itself.

import torch
kl_per_dim = torch.tensor([2.41, 1.87, 0.93, 0.02, 0.01, 0.00,
                           0.00, 0.00, 0.00, 0.00])
above_threshold = (kl_per_dim > 0.01).sum().item()
print("above 0.01 nats", above_threshold, "of", len(kl_per_dim))
print("total KL   ", round(kl_per_dim.sum().item(), 3))
print("share in top 3",
      round((kl_per_dim[:3].sum() / kl_per_dim.sum()).item(), 4))
# above 0.01 nats 4 of 10
# total KL    5.24
# share in top 3 0.9943

References

• Kingma and Welling (2014). Auto-Encoding Variational Bayes.
• Rezende, Mohamed, and Wierstra (2014). Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML.
• Bowman et al. (2016). Generating Sentences from a Continuous Space. CoNLL.
• Kingma et al. (2016). Improving Variational Inference with Inverse Autoregressive Flow.
• PyTorch distributions: sample(), rsample(), and gradient estimators.
• Higgins et al. (2017). beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR.
• Alemi et al. (2018). Fixing a Broken ELBO. ICML.
• Rombach, Blattmann, Lorenz, Esser, and Ommer (2022). High-Resolution Image Synthesis with Latent Diffusion Models. CVPR.