Auto-Encoding Variational Bayes

“Auto-Encoding Variational Bayes” by Diederik P. Kingma and Max Welling of the University of Amsterdam was posted to arXiv on 20 December 2013 and accepted at the International Conference on Learning Representations (ICLR) 2014. It proposed a general way to train directed models with continuous latent variables and intractable posteriors: write samples from the approximate posterior as a differentiable function of independent noise, the reparameterization trick, so that the variational lower bound can be estimated without bias and optimized by stochastic gradient methods. With neural networks as encoder and decoder, the method gives what the paper named the variational auto-encoder. The variational autoencoder and evidence lower bound articles cover the model and the bound as understood today. In 2024 the paper received ICLR’s inaugural Test of Time Award.

Versions. The ICLR 2014 proceedings page links to the arXiv record, not a fixed version, so which version served as the camera-ready is not recorded. Ten versions appeared by 1 May 2014; v11 (2022) changes only an abstract typo and reference capitalization. This article follows v10, the last substantive version, whose section, figure and appendix numbers v11 keeps. The nine-page v1 had no appendices. v4 added appendices on the Gaussian KL, the marginal-likelihood estimator, Monte Carlo EM and full VB; v6 (21 January 2014) introduced the names SGVB and estimators A and B, the MLP appendix and the note on Rezende et al.; v10 restored the name AEVB, absent from v6 to v9, and rewrote the abstract and introduction.

Problem

The paper asks how to do inference and learning when the marginal likelihood $p_\theta(x) = \int p_\theta(z)\, p_\theta(x \mid z)\, dz$ and the posterior $p_\theta(z \mid x)$ are intractable, as they already are for a neural-network likelihood with one nonlinear hidden layer (§2.1). That rules out EM and mean-field variational Bayes, and Monte Carlo EM’s per-datapoint sampling loop is too slow for large datasets. The authors wanted generative parameters, fast inference of an image’s code $z$, and a model of $x$ usable as an image prior for denoising, inpainting and super-resolution.

Variational inference offered a lower bound to maximize (§2.2), but its gradient with respect to the variational parameters $\phi$ was the obstacle: the usual score-function estimator, which multiplies the integrand by $\nabla_\phi \log q_\phi(z)$, has in the paper’s words very high variance and is impractical.

Contribution

• SGVB (Stochastic Gradient Variational Bayes): a differentiable, unbiased estimator of the lower bound for almost any model with continuous latent variables (§2.3).
• AEVB (Auto-Encoding VB): for i.i.d. data with a latent variable per datapoint, a recognition model $q_\phi(z \mid x)$, or probabilistic encoder, trained jointly with the generative model by SGVB, so inference needs one encoder pass instead of per-datapoint MCMC (Algorithm 1).
• The variational auto-encoder as the worked example (§3), with experiments against wake-sleep and Monte Carlo EM on MNIST and Frey Face images (§5).

Method

Estimators A and B. Pick a differentiable transformation $z = g_\phi(\epsilon, x)$ of auxiliary noise $\epsilon \sim p(\epsilon)$ that produces samples from $q_\phi(z \mid x)$ (Eq. 4). Expectations under $q_\phi$ become expectations under $p(\epsilon)$, which does not depend on $\phi$, so their Monte Carlo estimates can be differentiated. Applied to the two forms of the bound, this gives the generic estimator A (Eq. 6) and estimator B (Eq. 7), which integrates the KL term analytically when possible:

$$\tilde{\mathcal{L}}^{A}(\theta, \phi; x) = \frac{1}{L} \sum_{l=1}^{L} \Big[ \log p_\theta(x, z^{(l)}) - \log q_\phi(z^{(l)} \mid x) \Big], \qquad \tilde{\mathcal{L}}^{B}(\theta, \phi; x) = -\mathrm{KL}\big(q_\phi(z \mid x) \,\Vert\, p_\theta(z)\big) + \frac{1}{L} \sum_{l=1}^{L} \log p_\theta(x \mid z^{(l)}),$$

with $z^{(l)} = g_\phi(\epsilon^{(l)}, x)$. The authors note that B typically has less variance than A, and read it as an autoencoder objective: a KL regularizer plus an expected negative reconstruction error. Algorithm 1 draws a minibatch of $M$ datapoints and noise, differentiates the minibatch estimate ($N/M$ times the sum, Eq. 8) with respect to $\theta$ and $\phi$ together, and updates with SGD or Adagrad. The experiments use $M = 100$ and $L = 1$.

The reparameterization trick (§2.4). The paper’s argument is a change of variables: because $q_\phi(z \mid x) \prod_i dz_i = p(\epsilon) \prod_i d\epsilon_i$, an integral of $f(z)$ against $q_\phi$ equals the integral of $f(g_\phi(\epsilon, x))$ against $p(\epsilon)$. For a univariate Gaussian, $z = \mu + \sigma \epsilon$ with $\epsilon \sim \mathcal{N}(0, 1)$. The paper lists three families that admit such a $g_\phi$:

1. Tractable inverse CDF: $\epsilon$ uniform and $g_\phi$ the inverse CDF (exponential, Cauchy, Weibull, Gumbel and others).
2. Location-scale families: the standard member as $\epsilon$ and $g = \text{location} + \text{scale} \cdot \epsilon$ (Laplace, Student’s t, uniform, Gaussian and others).
3. Composition of such variables: the log-normal as an exponentiated Gaussian, the Gamma as a sum of exponentials, and the Dirichlet, Beta, chi-squared and F distributions.

In v1 to v4, the section was titled “The deterministic parameterization trick” and related the idea to non-centred parameterizations used in Gibbs sampling; v5 dropped that remark, and v6 introduced the present title.

The VAE (§3, Appendices B and C). The prior is $\mathcal{N}(0, I)$. The encoder is an MLP with one tanh hidden layer that outputs the mean and $\log \sigma^2$ of a diagonal Gaussian (Appendix C.2); the decoder is a Bernoulli MLP for binary data (C.1) or a Gaussian MLP for real-valued data. With $z = \mu + \sigma \odot \epsilon$, estimator B uses the closed-form Gaussian KL of Appendix B, given in the ELBO article. The paper calls the Gaussian posterior a simplifying choice, not a limitation of the method.

Other appendices. Appendix D’s marginal-likelihood estimator draws posterior samples with Hybrid Monte Carlo, fits a density $q(z)$ to them, and averages $q(z) / p_\theta(x, z)$ over fresh posterior samples to estimate $1/p_\theta(x)$; it is reliable below about five latent dimensions. Appendix E specifies the Monte Carlo EM baseline, and Appendix F extends the estimator to a posterior over $\theta$ (Algorithm 2), without experiments.

Related Work as the Paper Saw It

The paper calls the wake-sleep algorithm (Hinton et al., 1995) the only other on-line method, to the authors’ knowledge, for the same class of models. Wake-sleep also trains a recognition model at the same cost per datapoint and handles discrete latents, but it optimizes two objectives at once that together do not correspond to a bound on the marginal likelihood (§4). v1 to v5 put the criticism differently: the recognition network’s updates minimize $\mathrm{KL}(p \Vert q)$ instead of the $\mathrm{KL}(q \Vert p)$ the bound calls for.

The paper also cites control-variate schemes for the score-function estimator and a similar reparameterization by Salimans and Knowles (2013) for exponential-family posteriors. Unlike regularized autoencoders (denoising, contractive, sparse), its regularizer comes from the bound and needs no extra hyperparameter; deep autoregressive networks (DARN, added in v7) learn a directed model with an auto-encoding structure, but only for binary latents.

The paper notes that Rezende, Mohamed and Wierstra (2014), posted to arXiv on 16 January 2014, independently linked autoencoders, directed models and stochastic variational inference through the same reparameterization. Their ICML paper, which calls the two works simultaneous and complementary (§3.2), treats deep latent Gaussian models with several stochastic layers and names the gradient rules stochastic backpropagation.

Results

Kingma and Welling trained models of MNIST digits and Frey Face images (§5), the latter with a Gaussian decoder whose means pass through a sigmoid. Encoder and decoder had 500 hidden units each for MNIST and 200 for the smaller Frey Face set; a small weight decay made training approximate MAP estimation, and Adagrad adapted the step sizes. (The section’s second sentence swaps the labels, calling the generative model the encoder.)

Lower bound against wake-sleep (Fig. 2). Lower-bound curves over up to $10^8$ training samples cover MNIST with 3, 5, 10, 20 and 200 latent dimensions and Frey Face with 2, 5, 10 and 20, with wake-sleep given the same encoder. The authors report that AEVB converged considerably faster and reached a better solution in every case, and that superfluous latent dimensions did not cause overfitting, which they attribute to the regularizing effect of the bound. Training took about 20 to 40 minutes per million samples on a CPU. The 200-dimensional run and the overfitting remark first appear in v5.

Marginal likelihood (Fig. 3). With 100 hidden units and 3 latent dimensions, low enough for the Appendix D estimator, AEVB, wake-sleep and Monte Carlo EM with a Hybrid Monte Carlo sampler were trained on 1,000 and on 50,000 MNIST images. Both figures give curves without tables, so this article reports only the authors’ textual findings. The caption notes that Monte Carlo EM is not on-line and cannot be applied efficiently to the full dataset.

Worked Example: Why the Reparameterized Estimator Matters

The paper motivates SGVB by the score-function estimator’s variance but does not measure it. The NumPy script below does, for $8 \times 8$ binary blob images whose centre a 2-D code $z$ sets. For four images and three widths $\sigma$, it fixes an encoder mean slightly off the true code, $\mu = z^* + 0.2$, and draws 20,000 single-sample ($L = 1$) gradients in $(\mu, \log \sigma)$ from the naive estimator, estimators A and B, and the naive estimator with an oracle baseline $f - \mathbb{E}_q[f]$, standing in for control variates. Gauss–Hermite quadrature gives the exact gradient; the hand-written gradients were checked against finite differences.

import sys
import numpy as np
from scipy.special import expit

# Decoder p(x|z): an 8x8 binary image of a blob whose centre is set by a 2-D latent z,
# prior p(z) = N(0, I). Each pixel is on with probability sigmoid(4 - 0.5 d^2), where d
# is the pixel's distance from the centre (3.5 + 1.5 z_1, 3.5 + 1.5 z_2).
py, px = np.mgrid[0:8, 0:8].reshape(2, -1).astype(float)
S, A, B = 1.5, 4.0, 0.5

def decoder(z):                                   # z: (n, 2) -> logits (n, 64), centres
    cy, cx = 3.5 + S * z[:, :1], 3.5 + S * z[:, 1:]
    return A - B * ((py - cy) ** 2 + (px - cx) ** 2), cy, cx

def log_lik_and_grad(x, z):                       # log p(x|z) and its gradient in z
    l, cy, cx = decoder(z)
    r = x - expit(l)                              # d log p / d logit
    g = np.stack([(r * 2 * B * (py - cy)).sum(1), (r * 2 * B * (px - cx)).sum(1)], 1) * S
    return (x * l - np.logaddexp(0, l)).sum(1), g

def estimators(x, mu, sigma, eps):
    """Single-sample (L = 1) gradients of the bound w.r.t. (mu, log sigma), one row per eps."""
    z = mu + sigma * eps
    ll, g_ll = log_lik_and_grad(x, z)
    log_prior = -0.5 * (z ** 2).sum(1) - np.log(2 * np.pi)
    log_q = (-0.5 * eps ** 2 - np.log(sigma) - 0.5 * np.log(2 * np.pi)).sum(1)
    f = ll + log_prior - log_q                    # integrand of the bound, eq. (2)
    score = np.hstack([eps / sigma, eps ** 2 - 1])            # grad of log q at fixed z
    naive = f[:, None] * score                                # paper's naive estimator, §2.2
    g_h = g_ll - z                                            # grad of log p(x|z) + log p(z)
    est_a = np.hstack([g_h, g_h * sigma * eps + 1])           # SGVB estimator A, eq. (6)
    est_b = np.hstack([g_ll - mu, g_ll * sigma * eps + 1 - sigma ** 2])   # estimator B, eq. (7)
    return {"naive (score function)": naive, "SGVB estimator A": est_a,
            "SGVB estimator B": est_b}, f, score

# Exact gradient by 2-D Gauss-Hermite quadrature over eps (reparameterized form).
t, w = np.polynomial.hermite_e.hermegauss(60)
eps_q = np.stack(np.meshgrid(t, t, indexing="ij"), -1).reshape(-1, 2)
w_q = np.outer(w, w).ravel() / (2 * np.pi)

FIRST = int(sys.argv[1]) if len(sys.argv) > 1 else 0
N = 20000                                          # gradient draws per image and sigma
print(f"{'sigma':>5s}  {'estimator':24s}{'var mu-grad':>12s}{'var log-sigma-grad':>19s}"
      f"{'L to match B':>13s}{'rms bias/s.e.':>14s}")
for sigma_val in [1.0, 0.3, 0.1]:
    tot = {}
    for seed in range(FIRST, FIRST + 4):           # four images, blob centres from the prior
        rng = np.random.default_rng(seed)
        z_true = rng.normal(size=(1, 2))
        x = (rng.random((1, 64)) < expit(decoder(z_true)[0])).astype(float)
        mu, sigma = z_true + 0.2, np.full((1, 2), sigma_val)     # an imperfect encoder output
        exact = (w_q[:, None] * estimators(x, mu, sigma, eps_q)[0]["SGVB estimator B"]).sum(0)
        ests, f, score = estimators(x, mu, sigma, rng.normal(size=(N, 2)))
        # naive estimator with an oracle constant baseline: f - E_q[f] (control variate)
        elbo = (w_q * estimators(x, mu, sigma, eps_q)[1]).sum()
        ests["naive, oracle baseline"] = (f - elbo)[:, None] * score
        for name, g in ests.items():
            v = g.var(0, ddof=1)
            bias = np.abs(g.mean(0) - exact) / np.sqrt(v / N)
            a = tot.setdefault(name, [0.0, 0.0, 0.0])     # variances summed over components,
            a[0] += v[:2].sum() / 4; a[1] += v[2:].sum() / 4  # averaged over the four images
            a[2] += (bias ** 2).sum() / 16
    vb = tot["SGVB estimator B"]
    for name, (vm, vs, b) in tot.items():
        print(f"{sigma_val:5.1f}  {name:24s}{vm:12.1f}{vs:19.1f}"
              f"{(vm + vs) / (vb[0] + vb[1]):13.2f}{np.sqrt(b):14.2f}")

Output (default argument 0, images from seeds 0 to 3; runs in under a second):

sigma  estimator                var mu-grad var log-sigma-grad L to match B rms bias/s.e.
  1.0  naive (score function)       32409.4           130358.2        14.01          0.83
  1.0  SGVB estimator A              3972.6             8180.1         1.05          0.97
  1.0  SGVB estimator B              3806.4             7809.3         1.00          0.97
  1.0  naive, oracle baseline       12631.7            66543.2         6.82          0.80
  0.3  naive (score function)       21416.7             4940.3        56.19          0.90
  0.3  SGVB estimator A               397.7               91.2         1.04          0.97
  0.3  SGVB estimator B               382.2               86.8         1.00          0.95
  0.3  naive, oracle baseline        1515.5              631.1         4.58          0.72
  0.1  naive (score function)      126233.5             2458.8      2709.71          0.93
  0.1  SGVB estimator A                45.9                3.5         1.04          0.96
  0.1  SGVB estimator B                44.1                3.3         1.00          0.92
  0.1  naive, oracle baseline         994.0               41.8        21.81          1.17

All four estimators are unbiased: their means lie within about one standard error of the exact gradient (last column). They differ in variance. “L to match B” (total variance relative to B) is roughly how many samples an estimator must average to equal one sample of B. The naive estimator needs 14 at $\sigma = 1$, 56 at $\sigma = 0.3$ and about 2,700 at $\sigma = 0.1$: as the encoder sharpens, the reparameterized variance falls, while the naive $\mu$-gradient, scaled by $\epsilon / \sigma$, grows noisier. The oracle baseline still leaves a factor of about 5 to 22, and A has 4 to 5% more variance than B, as the paper remarks. Arguments 4 and 8 (seeds 4 to 11) gave naive factors of 12.0 to 12.6, 41.8 to 62.5 and 1,849 to 3,438.

Impact

ICLR’s program chairs chose the paper from the ICLR 2013 and 2014 papers for the inaugural ICLR Test of Time Award in 2024. The award statement credits it with leading the integration of deep learning with scalable probabilistic inference, giving rise to the VAE and sparking many later probabilistic models, and names the concurrent paper of Rezende et al. With generative adversarial networks, published the same year, the VAE is one of the families compared in the generative model article.

Limitations

• Continuous latents only. Unlike wake-sleep, AEVB cannot handle discrete latents, as the paper concedes.
• Small models and datasets. The networks have one hidden layer, the datasets are MNIST digits and the small Frey Face set, and the paper defers convolutional and hierarchical encoders and decoders to future work (§7).
• Evaluation. Results are plotted, not tabulated, and the marginal-likelihood comparison uses only three latent dimensions.
• A diagonal Gaussian posterior. The paper does not measure how loose a bound this simplifying choice gives.

What Came After

Kingma, Rezende, Mohamed and Welling (2014) applied the framework to semi-supervised learning. Burda, Grosse and Salakhutdinov (2016) trained the same kind of model on a tighter, importance-weighted bound. Jang, Gu and Poole (2017) and Maddison, Mnih and Teh (2017) addressed the discrete-latent gap with continuous relaxations of categorical variables that can be reparameterized. In image generation, latent diffusion (Rombach et al., 2022) gives one of its two first-stage autoencoders a slight KL penalty toward a standard normal, similar to a VAE. Kingma and Welling’s (2019) monograph reviews the framework and its extensions.