Variational Autoencoder

A variational autoencoder (VAE) is a latent-variable generative model built in the shape of an autoencoder. Its encoder maps an image to a distribution over latent codes rather than a single code, its decoder maps a code to a distribution over images, and both are trained by maximizing the evidence lower bound (ELBO) on the log-likelihood of the training images. A KL term pulls every image’s code distribution toward a standard normal prior, so codes drawn from that prior decode to new images, which a plain autoencoder cannot guarantee. The VAE was introduced in 2014 by Kingma and Welling and, independently, by Rezende, Mohamed and Wierstra; the Auto-Encoding Variational Bayes article covers that publication.

Core Idea

The model generates an image $x$ by drawing a code $z$ from the prior $p(z) = \mathcal{N}(0, I)$ and then $x$ from the decoder $p_\theta(x \mid z)$. Maximum likelihood would need $p_\theta(x) = \int p_\theta(x \mid z)\, p(z)\, dz$, which has no closed form when the decoder is a neural network. The VAE adds an encoder $q_\phi(z \mid x)$ that proposes the codes likely to have produced $x$, and maximizes the ELBO instead.

The encoder is one network shared by all images, so inference takes one forward pass instead of a per-image optimization. Rezende et al. (2014, §4.1) call this amortised inference.

Building Blocks

Encoder. A network that outputs a mean $\mu_\phi(x)$ and a log-variance $\log \sigma^2_\phi(x)$, defining a diagonal Gaussian $q_\phi(z \mid x) = \mathcal{N}(\mu_\phi(x), \operatorname{diag}(\sigma^2_\phi(x)))$. Kingma and Welling (2014, Appendix C.2, arXiv v11) output $\log \sigma^2$, which keeps the variance positive without a constraint.

Reparameterized sampling. A code is drawn as

$$z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I),$$

so gradients flow through $\mu$ and $\sigma$ into the encoder by ordinary backpropagation; the ELBO article explains why this trick is needed.

Decoder. A network, often built with transposed convolutions or upsampling layers, that maps $z$ to the parameters of $p_\theta(x \mid z)$, usually a distribution that factorizes over pixels. Kingma and Welling (2014, Appendix C) used two:

• Bernoulli, for binary or $[0, 1]$ images: the decoder outputs a probability $y_i$ per pixel, and $\log p_\theta(x \mid z) = \sum_i x_i \log y_i + (1 - x_i) \log (1 - y_i)$, the negative binary cross-entropy.
• Gaussian: their decoder also outputs a variance; with a mean $\mu_\theta(z)$ and the variance fixed at $\sigma^2$ over $D$ pixels,

$$\log p_\theta(x \mid z) = -\frac{1}{2\sigma^2} \lVert x - \mu_\theta(z) \rVert^2 - \frac{D}{2} \log (2\pi\sigma^2),$$

so the reconstruction term is a scaled squared error.

Mathematical Formulation

Training maximizes the ELBO averaged over training images, in the notation of the ELBO article, which derives it:

$$\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}\big[\log p_\theta(x \mid z)\big] - \mathrm{KL}\big(q_\phi(z \mid x) \,\Vert\, p(z)\big) \le \log p_\theta(x).$$

The first term is estimated with one reparameterized sample; the KL term has a closed form (given in the ELBO article) with one term per latent dimension. With a fixed-variance Gaussian decoder, multiplying the objective by $2\sigma^2$ gives the squared error plus $2\sigma^2$ times the KL term (our derivation), so the decoder variance is a hidden weight on the KL term, and a plain squared-error loss implicitly chooses it.

Training Considerations

Posterior collapse. Bowman et al. (2016) trained a VAE whose decoder was an LSTM language model and found that, except in rare cases, training set $q(z \mid x)$ equal to the prior and drove the KL term to zero. The decoder learns to explain the data without $z$, after which little gradient passes between encoder and decoder, a stable equilibrium (§3.1). The failure is now called posterior collapse, as in van den Oord, Vinyals and Kavukcuoglu (2017). Image VAEs show a related effect, inactive latent dimensions: Sønderby et al. (2016, §2.3) saw units whose posterior matched the prior early in training and never recovered.

KL annealing. Bowman et al. (2016, §3.1) multiply the KL term by a weight that rises from zero, where training is a plain autoencoder, to one, the true bound, on a sigmoid schedule. Sønderby et al. (2016, Eq. 20) raised the weight linearly over the first epochs (“warm-up”) and found warm-up and batch normalization crucial for VAEs with several stochastic layers.

Free bits. Kingma et al. (2016, Appendix C.8, arXiv v2) replace each latent group’s KL term with $\max(\lambda, \mathbb{E}[\mathrm{KL}_j])$, so using fewer than $\lambda$ nats per group brings no benefit.

Example

The PyTorch example below (PyTorch 2.14.1, CPU, one thread, about 13 s) trains a VAE with an 8-dimensional latent and a Bernoulli decoder on $16 \times 16$ binary images of a disk or square generated from four factors: class, size and the two centre coordinates. It prints the ELBO terms, describes decoded images by statistics (area; bounding-box fill, which separates squares from disks; the fraction of in-between pixel probabilities), interpolates between two codes, and retrains with the KL term weighted by $\beta$.

import sys, time
import torch
from torch import nn
import torch.nn.functional as F

torch.set_num_threads(1)
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 0
t0 = time.time()

yy, xx = torch.meshgrid(torch.arange(16.), torch.arange(16.), indexing="ij")
def render(is_square, r, cy, cx):
    """Binary 16x16 images of a filled disk or square with 'radius' r centred at (cy, cx)."""
    dy, dx = (yy - cy[:, None, None]).abs(), (xx - cx[:, None, None]).abs()
    r = r[:, None, None]
    disk = dy**2 + dx**2 <= r**2
    square = torch.maximum(dy, dx) <= 0.85 * r      # similar area to the disk
    return torch.where(is_square[:, None, None], square, disk).float().view(-1, 256)

def shapes(n, g):
    """Four generative factors: shape class, size, and the two coordinates of the centre."""
    is_square = torch.rand(n, generator=g) < 0.5
    r = 2.5 + 3.5 * torch.rand(n, generator=g)
    cy, cx = (r + (15 - 2 * r) * torch.rand(2, n, generator=g)).unbind(0)
    return render(is_square, r, cy, cx), is_square

def describe(p):
    """Statistics of decoded probability maps p (n x 256): area and bounding-box fill
    of the thresholded shape (a square fills its box, a disk about 80% of it), and the
    fraction of the shape's pixels whose probability is in-between (0.1 < p < 0.9)."""
    on = (p > 0.5).view(-1, 16, 16)
    area = on.sum((1, 2)).float()
    rows, cols = on.any(2).float(), on.any(1).float()
    box = rows.sum(1) * cols.sum(1)
    fill = area / box.clamp(min=1)
    grey = ((p > 0.1) & (p < 0.9)).sum(1) / ((p > 0.1).sum(1)).clamp(min=1)
    return area, fill, grey

J = 8                                                   # latent dimensions
class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 2 * J))
        self.dec = nn.Sequential(nn.Linear(J, 256), nn.ReLU(), nn.Linear(256, 256))

    def forward(self, x):
        mu, logvar = self.enc(x).chunk(2, dim=1)        # q(z|x) = N(mu, diag(exp(logvar)))
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)   # reparameterization
        logits = self.dec(z)                            # Bernoulli decoder p(x|z)
        rec = F.binary_cross_entropy_with_logits(logits, x, reduction="none").sum(1)
        kl = 0.5 * (mu**2 + logvar.exp() - 1 - logvar)  # per-dimension KL to N(0, I)
        return rec, kl, mu

g = torch.Generator().manual_seed(SEED)
x_train, _ = shapes(20000, g)
gt = torch.Generator().manual_seed(1234)                # same test set for every seed
x_test, sq_test = shapes(2000, gt)
area, fill, _ = describe(x_test)
print(f"PyTorch {torch.__version__}; rule 'fill > 0.9 means square' is right on "
      f"{((fill > 0.9) == sq_test).float().mean():.1%} of real test shapes")

def train(beta, steps=5000, log=False):
    torch.manual_seed(SEED)
    vae = VAE()
    opt = torch.optim.Adam(vae.parameters(), lr=1e-3)
    for step in range(1, steps + 1):
        x = x_train[torch.randint(0, len(x_train), (128,), generator=g)]
        rec, kl, _ = vae(x)
        loss = (rec + beta * kl.sum(1)).mean()          # negative (beta-weighted) ELBO
        opt.zero_grad(); loss.backward(); opt.step()
        if log and step in (1, 100, 500, 2000, 5000):
            with torch.no_grad(), torch.random.fork_rng():
                torch.manual_seed(SEED + 100)           # same noise for every evaluation
                rec, kl, _ = vae(x_test)
            r, k = rec.mean().item(), kl.sum(1).mean().item()
            print(f"{step:6d}{r:11.2f}{k:9.2f}{-r - k:11.2f}")
    return vae

print(f"{'step':>6s}{'recon':>11s}{'KL':>9s}{'ELBO':>11s}   (test set, nats per image, beta = 1)")
vae = train(1.0, log=True)

torch.manual_seed(SEED + 100)
with torch.no_grad():
    rec, kl, mu = vae(x_test)
    print("KL per latent dimension:", " ".join(f"{v:.2f}" for v in kl.mean(0)))
    p_rec = torch.sigmoid(vae.dec(mu))                  # decoder mean at the encoder mean
    p_prior = torch.sigmoid(vae.dec(torch.randn(2000, J)))
for name, p in [("test shapes", x_test), ("reconstructions", p_rec), ("prior samples", p_prior)]:
    a, f, gr = describe(p)
    print(f"{name:16s} square {(f > 0.9).float().mean():5.1%}  area {a.mean():5.1f} px"
          f"  in-between pixels {gr.mean():5.1%}  empty {(a == 0).float().mean():4.1%}")

# Interpolate between a small square and a large disk, both centred.
ends = render(torch.tensor([True, False]), torch.tensor([3.0, 6.0]),
              torch.tensor([7.5, 7.5]), torch.tensor([7.5, 7.5]))
with torch.no_grad():
    m = vae.enc(ends)[:, :J]
    t = torch.linspace(0, 1, 9)[:, None]
    p_lat = torch.sigmoid(vae.dec((1 - t) * m[0] + t * m[1]))
    p_pix = (1 - t) * ends[0] + t * ends[1]             # cross-fade in pixel space
print(f"{'t':>5s}| latent path: {'area':>4s} {'fill':>5s} {'centre (r, c)':>14s} | pixel cross-fade: in-between")
for i in range(9):
    a, f, _ = describe(p_lat[i:i + 1])
    on = (p_lat[i] > 0.5).view(16, 16).float()
    cy = (on.sum(1) * torch.arange(16.)).sum() / on.sum().clamp(min=1)
    cx = (on.sum(0) * torch.arange(16.)).sum() / on.sum().clamp(min=1)
    _, _, gp = describe(p_pix[i:i + 1])
    print(f"{t[i].item():5.3f}| {'':13s}{a.item():4.0f} {f.item():5.2f}   ({cy:4.1f}, {cx:4.1f}) |"
          f" {gp.item():18.0%}")

print(f"{'beta':>5s}{'recon':>9s}{'KL':>8s}{'active dims':>13s}")
for beta in [0.25, 1.0, 4.0, 16.0]:
    v = vae if beta == 1.0 else train(beta)
    torch.manual_seed(SEED + 100)
    with torch.no_grad():
        rec, kl, _ = v(x_test)
    print(f"{beta:5.2f}{rec.mean():9.2f}{kl.sum(1).mean():8.2f}{int((kl.mean(0) > 0.1).sum()):13d}")
print(f"time {time.time() - t0:.0f} s", file=sys.stderr)

Output (seed 0; two runs gave identical output):

PyTorch 2.14.1; rule 'fill > 0.9 means square' is right on 99.3% of real test shapes
  step      recon       KL       ELBO   (test set, nats per image, beta = 1)
     1     174.61     0.07    -174.68
   100      31.88    11.98     -43.86
   500      14.85    10.98     -25.83
  2000       7.76    11.16     -18.92
  5000       5.57    10.90     -16.47
KL per latent dimension: 0.85 1.83 0.00 0.70 0.02 2.48 2.51 2.52
test shapes      square 51.2%  area  57.5 px  in-between pixels  0.0%  empty 0.0%
reconstructions  square 51.7%  area  57.6 px  in-between pixels 13.7%  empty 0.0%
prior samples    square 25.2%  area  53.7 px  in-between pixels 28.7%  empty 3.2%
    t| latent path: area  fill  centre (r, c) | pixel cross-fade: in-between
0.000|                36  1.00   ( 7.5,  7.5) |                 0%
0.125|                43  0.88   ( 7.7,  7.9) |                68%
0.250|                61  0.95   ( 7.6,  7.6) |                68%
0.375|                64  1.00   ( 7.5,  7.5) |                68%
0.500|                85  0.85   ( 7.4,  7.5) |                68%
0.625|                92  0.92   ( 7.4,  7.5) |                68%
0.750|               104  0.72   ( 7.5,  7.5) |                68%
0.875|               112  0.78   ( 7.5,  7.5) |                68%
1.000|               112  0.78   ( 7.5,  7.5) |                 0%
 beta    recon      KL  active dims
 0.25     2.34   19.61            8
 1.00     5.57   10.90            6
 4.00    14.85    5.90            4
16.00    50.97    1.77            3

The reconstruction term falls from 175 to 5.6 nats while the KL term settles near 11, and two of eight latent dimensions are inactive (KL 0.00 and 0.02). Reconstructions keep the data’s class balance (51.7% squares), but only 25.2% of prior samples read as squares and 3.2% decode to nothing. In our reading, the prior puts probability on codes that no training image’s encoder distribution covers well, so the decoder was never trained on them.

Along the latent path the shape stays centred and its area grows steadily from 36 to 112 pixels; it reads as a square up to $t = 0.375$ and a disk at $t = 0.5$ and from $t = 0.75$, with less clear-cut shapes at $t = 0.125$ and $0.625$. The pixel cross-fade superimposes the two shapes, leaving 68% of covered pixels in-between. From $\beta = 0.25$ to $16$ the KL falls from 19.6 to 1.8 nats, reconstruction error rises from 2.3 to 51 nats, and active dimensions (KL above 0.1) fall from 8 to 3. At $\beta = 4$ four stay active, as many as the data’s factors; we did not check whether each tracks one factor.

Over seeds 0 to 5, passed as the script’s command-line argument, the final ELBO ranged from $-16.33$ to $-17.15$ nats, prior samples read as squares from 18.9% to 25.4%, and active dimensions at $\beta = 1$ from 5 to 7; every seed kept 8, 4 and 3 at $\beta = 0.25$, 4 and 16, and no path’s area ever shrank, each reading square at $t = 0.25$ and $0.375$ and disk at $t = 0.5$, $0.75$ and $0.875$.

Variants

β-VAE. Higgins et al. (2017) multiply the KL term by a hyperparameter $\beta$:

$$\mathcal{L}_\beta(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}\big[\log p_\theta(x \mid z)\big] - \beta\, \mathrm{KL}\big(q_\phi(z \mid x) \,\Vert\, p(z)\big).$$

As summarized by Burgess et al. (2018, §3), well-chosen values, usually above 1, encourage disentangled codes, in which single latent units respond to single generative factors, at the price of worse reconstructions. Since KL is non-negative, $\beta > 1$ loosens the bound and $\beta < 1$ removes the guarantee (our note).

Conditional VAE. Sohn, Yan and Lee (2015) model an output $y$, such as a segmentation map, given an input $x$, with a prior network $p_\theta(z \mid x)$ and a recognition network $q_\phi(z \mid x, y)$, maximizing

$$\log p_\theta(y \mid x) \ge \mathbb{E}_{q_\phi(z \mid x, y)}\big[\log p_\theta(y \mid x, z)\big] - \mathrm{KL}\big(q_\phi(z \mid x, y) \,\Vert\, p_\theta(z \mid x)\big).$$

Different draws of $z$ give different plausible outputs for one input; they tested it on completing MNIST digits from some of their quadrants and on segmenting Caltech-UCSD Birds and a subset of Labeled Faces in the Wild.

Hierarchical VAEs. Layers of latents $z_1, \dots, z_L$ are generated top-down, each conditioned on the layers above. The Ladder VAE (Sønderby et al., 2016) combines the encoder’s bottom-up estimate with the top-down prior at each layer by precision weighting. NVAE (Vahdat & Kautz, 2020) uses depthwise separable convolutions, batch normalization, residual Normal distributions and spectral regularization; it reached 2.91 bits per dimension on CIFAR-10 (Table 1, with flows), and its authors call it, to their knowledge, the first successful VAE on natural images as large as $256 \times 256$, sampled at a reduced prior temperature (0.6 on CelebA HQ, Fig. 4). Child (2021, Table 2) reached 2.87 bits per dimension on CIFAR-10 with 45 stochastic layers (mean of four seeds), better than PixelCNN++ (2.92) with one network evaluation per sample; with nearest-neighbour upsampling he saw no posterior collapse and needed neither free bits nor warm-up (§4.1).

VQ-VAE (related). The VQ-VAE (van den Oord et al., 2017) snaps each encoder output vector to its nearest entry in a learned codebook of $K$ vectors. Its posterior is one-hot and deterministic, so under a uniform prior the KL term is the constant $\log K$ and is dropped (§3.1–3.2). Training uses a reconstruction term, a codebook term and a commitment term weighted by 0.25, with straight-through gradients through the quantizer (§3.2, Eq. 3). A PixelCNN prior over the codes is fitted after training (§3.3), and the authors describe the design as avoiding posterior collapse.

Use in Computer Vision

First stage of latent diffusion. Latent diffusion models (Rombach et al., 2022) train a diffusion model on the latent $z = E(x)$ of a pretrained autoencoder. One of their two autoencoder variants, KL-reg., imposes a slight KL penalty toward a standard normal, similar to a VAE (§3.1): the encoder outputs a mean and standard deviation, codes are reparameterized samples, the KL weight is about $10^{-6}$, and the latent can be rescaled to unit standard deviation by a single factor estimated from the first batch of data (arXiv v2, Appendix G; the text calls it a component-wise variance, but its formula averages over all components). A perceptual loss and a patch-based adversarial loss replace pixel-space $L^2$ or $L^1$ losses, whose blur the authors avoid. With so small a weight, in our reading, the KL term mainly bounds the latent’s scale, and the diffusion model, not the prior, supplies the distribution of codes.

Other uses. Conditional VAEs predict diverse segmentations and hierarchical VAEs generate faces (see Variants). NVAE’s authors argue that VAEs, unlike GANs, do not suffer from mode collapse (impact statement); the generative model article compares the families.

Limitations

Blurry samples. VAE samples of complex natural images tend to be blurry, an effect commonly associated with the squared-error loss (Zhao, Song & Ermon, 2017, §3.1), that is, a fixed-variance Gaussian decoder. Larsen et al. (2016) argue that such element-wise measures suit images poorly: a small translation causes a large pixel error that a human barely notices. Zhao et al. (2017, Proposition 4) explain the blur for that decoder: the optimal output for a code $z$ is the average of the images the encoder maps to $z$, and unless encoding is lossless, different images share codes and their average is fuzzy. They argue that, under some conditions, the cause is the inference distribution rather than maximum likelihood itself. The example is consistent with this averaging for a Bernoulli decoder (our reading): training images have no in-between pixels, but 13.7% of the covered pixels in reconstructions and 28.7% in prior samples have probabilities between 0.1 and 0.9. Sampling $x$ from the factorized decoder instead of showing its mean would add independent pixel noise, not sharp detail (our reasoning). Remedies include feature-space or adversarial reconstruction losses (Larsen et al., 2016), as in latent diffusion, and more expressive inference, such as Zhao et al.’s sequential VAE, which produced sharp LSUN samples with a pixel-wise loss.

Other limits. The ELBO is a bound whose gap depends on the encoder; collapsed or inactive latents waste capacity; and the prior can put probability on codes that decode poorly, so samples are worse than reconstructions.

Important Models

Model	Year	Key change
VAE (Kingma & Welling; Rezende et al.)	2014	Amortized Gaussian encoder, reparameterized ELBO training
Conditional VAE (Sohn et al.)	2015	Conditioning on an input for structured prediction
Ladder VAE (Sønderby et al.)	2016	Top-down hierarchical inference, KL warm-up
IAF-VAE (Kingma et al.)	2016	Flow-based posterior, free bits
β-VAE (Higgins et al.)	2017	Weighted KL term for disentangling
VQ-VAE (van den Oord et al.)	2017	Discrete codebook latents, learned prior after training
NVAE (Vahdat & Kautz)	2020	Deep hierarchical VAE for images up to $256 \times 256$
VDVAE (Child)	2021	Up to 78 stochastic layers, gradient skipping
Latent diffusion autoencoder, KL-reg. (Rombach et al.)	2022	Very weak KL penalty, latent space for diffusion