Auto-Encoding Variational Bayes

How to backpropagate through a random sample.

A latent-variable model can dream up new data, but the likelihood you would train it on is an integral nobody can compute. The variational autoencoder gets around that with one trick: it rewrites a random sample so the gradient can flow straight through it, and a generative model you can train with plain backprop falls out.

What does it take to turn an autoencoder, a thing that just copies its input, into a model that invents new data you have never seen?

Here is the dream that organizes a lot of generative modeling. You would like a machine that can produce new examples of something. New handwritten digits, new faces, new molecules. The cleanest way to set that up is a latent-variable model. Pick a hidden code $\mathbf{z}$ from a simple distribution, push it through a network, and read out a datapoint $\mathbf{x}$. The network learns to turn easy random codes into hard, structured data. To make a new sample you draw a fresh code and decode it. That is the whole generative story, and it is simple.

The trouble is training it. To fit the network you want to maximize the probability the model assigns to your real data. Writing that probability down is easy. Computing it is not, because it hides an integral over every possible hidden code, and that integral has no closed form and far too many codes to sum over. The exact same integral blocks the other thing you would reach for, which is figuring out which code produced a given datapoint. Both the headline likelihood and the bookkeeping you would use to get it are out of reach.

Auto-Encoding Variational Bayes, by Diederik Kingma and Max Welling, is the paper that made this work at scale. It does two things. It optimizes a tractable lower bound instead of the impossible likelihood, and it trains a small second network to guess the hidden code, so you never have to solve for it. When both networks are neural nets fit jointly, the result is the variational autoencoder. To see why it works you only need a short stack of ideas: the intractable integral, the lower bound that dodges it, the autoencoder that the bound turns into, and the rewrite that lets a random sampling step carry a gradient.

The data, and the wall

Start with the generative model itself. There is a hidden code $\mathbf{z}$, drawn from a fixed, featureless prior, a standard Gaussian $p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ with no parameters to learn. Then a decoder network with weights $\boldsymbol{\theta}$ turns that code into a distribution over data, $p_{\boldsymbol{\theta}}(\mathbf{x}\mid\mathbf{z})$. Sampling is a two-step recipe: draw $\mathbf{z}$, then draw $\mathbf{x}$ from the decoder. A simple cloud goes in, structured data comes out.

Drag nothing, just watch: simple codes on the left stream through the decoder and land on the data.

Now the question that trains it. How probable is one real datapoint $\mathbf{x}$ under this model? You have to account for every code that could have produced it, weighting each by how likely that code was to begin with. That sum over codes is the marginal likelihood, also called the evidence:

Read it as an average of the decoder over the whole prior. For any single pair $(\mathbf{x},\mathbf{z})$ the thing inside the integral, the joint $p_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{z}) = p_{\boldsymbol{\theta}}(\mathbf{z})\,p_{\boldsymbol{\theta}}(\mathbf{x}\mid\mathbf{z})$, is cheap to evaluate: run the decoder, multiply by the prior density. The integral is what kills you. The code $\mathbf{z}$ lives in a continuous, high-dimensional space, the decoder is a nonlinear network, and there is no antiderivative and no small grid of codes that covers it. So the evidence is intractable. You cannot evaluate it, and you certainly cannot differentiate it to do gradient ascent.

It gets worse in a useful way. By Bayes' rule, the posterior, the distribution over codes that could have produced a given $\mathbf{x}$, is

and the denominator is that same impossible integral. The numerator is the cheap joint; the normalizer is the wall. So the posterior is intractable too. This is why the classical tool for latent-variable models, the EM algorithm, stalls here: its E-step needs the exact posterior to fill in the hidden codes, and we just said we do not have it. (The usual workaround, mean-field variational inference, assumes a simple factorized posterior it can solve for in closed form, but that closed form needs a conjugate likelihood, which a neural network is not.) One integral, three dead ends.

Everything the paper does is a way around that one integral. The move is to stop demanding the exact posterior and instead learn a cheap stand-in for it, a second network, and to optimize a bound on the evidence that this stand-in makes computable. That stand-in has a name in the paper, the recognition model $q_{\boldsymbol{\phi}}(\mathbf{z}\mid\mathbf{x})$, an encoder: feed it a datapoint and it returns a distribution over the codes that might have made it.

A bound you can actually optimize

We cannot touch $\log p_{\boldsymbol{\theta}}(\mathbf{x})$. But there is an exact identity that splits it into two pieces, and one of those pieces is something we can compute. Introduce the encoder $q_{\boldsymbol{\phi}}(\mathbf{z}\mid\mathbf{x})$ and the log-evidence decomposes as

The first term is a KL divergence, a measure of how far one distribution is from another. It measures the gap between our encoder $q_{\boldsymbol{\phi}}$ and the true posterior $p_{\boldsymbol{\theta}}(\mathbf{z}\mid\mathbf{x})$. A KL divergence is never negative, and it is zero only when the two distributions match exactly. That single fact is the whole lever. Because the left side $\log p_{\boldsymbol{\theta}}(\mathbf{x})$ is a fixed number once $\boldsymbol{\theta}$ is fixed, and the KL gap on the right cannot go below zero, the second term $\mathcal{L}$ can never exceed the evidence. It is a lower bound on $\log p_{\boldsymbol{\theta}}(\mathbf{x})$, and the paper calls it the variational lower bound. Today it is usually called the evidence lower bound, or ELBO.

So we get a bound for free, and pushing the bound up is productive in two ways at once. Hold the decoder fixed and improve the encoder: the bound rises, and since the total is pinned, the only thing that can give is the KL gap, which shrinks. Better bound and a better posterior approximation, from the same gradient step. The bound touches the evidence exactly when the encoder equals the true posterior, the moment the gap hits zero. We never reach it with a limited encoder, but we get as close as the encoder can manage.

Drag the encoder toward the true posterior and watch the gap close. The full bar is the (fixed) evidence; the teal part is the bound you are lifting, the amber part is the slack you are squeezing out.

What is this bound, concretely? Two equivalent ways to write it, and you will use both. The paper's first form keeps the joint together:

If you have met Jensen's inequality, this is also where the bound comes from without any mention of a posterior: write $\log p_{\boldsymbol{\theta}}(\mathbf{x})$ as $\log \mathbb{E}_{q}\!\left[\,p_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{z})/q_{\boldsymbol{\phi}}\,\right]$, and since $\log$ is concave you can pull it inside the average to get a smaller number, $\mathbb{E}_{q}\!\left[\log\big(p_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{z})/q_{\boldsymbol{\phi}}\big)\right]$. That is eq (2) again, and the size of Jensen's slack is exactly the KL gap from eq (1). Two doors into the same room.

The second form is the one that makes everything compute. Split the joint $p_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{z}) = p_{\boldsymbol{\theta}}(\mathbf{x}\mid\mathbf{z})\,p(\mathbf{z})$and collect terms:

A warning that has tripped up many readers: this KL is not the one from eq (1). Eq (1) measured distance to the true posterior $p_{\boldsymbol{\theta}}(\mathbf{z}\mid\mathbf{x})$, the intractable thing. This one measures distance to the prior $p(\mathbf{z}) = \mathcal{N}(\mathbf{0},\mathbf{I})$, which is fixed and known. Same symbol, different second argument, and only this one is computable. The first KL is the bound's invisible slack; this second KL is the term you actually evaluate.

One bookkeeping note you have to keep straight, because it is the single most common sign error in this whole subject. The paper maximizes $\mathcal{L}$ by gradient ascent, and writes the regularizer with a minus sign, as in eq (3). Modern code minimizes a loss, and the loss is just the negative ELBO, so the same KL flips to a positive penalty:

Maximize the bound or minimize its negative, it is the same training. But you have to know which one is on screen before you trust any sign. From here on, when a quantity has a sign, the text says whether we are maximizing the bound or minimizing the loss.

An autoencoder falls out

Read eq (3) as a machine rather than a formula. The encoder $q_{\boldsymbol{\phi}}$ takes a datapoint $\mathbf{x}$ and produces a distribution over codes. You draw a code $\mathbf{z}$ from it. The decoder $p_{\boldsymbol{\theta}}$ takes that code and tries to rebuild $\mathbf{x}$. The first term in eq (3) rewards a good rebuild, and the second keeps the codes tidy. That is an autoencoder, an encoder feeding a decoder through a bottleneck, with one twist: the bottleneck is a distribution you sample from, not a single point, and a regularizer keeps that distribution honest.

The two terms pull against each other, and that tension is the point. The reconstruction term wants codes that carry as much about $\mathbf{x}$ as possible, ideally a sharp, distinctive code per datapoint so the decoder can rebuild it perfectly. The KL term wants every code distribution to look like the same bland prior, which blurs the codes together and throws information away. A plain autoencoder, with only the reconstruction term, is free to scatter its codes anywhere it likes. It will, and the space between codes decodes to garbage, which is why you cannot sample from a plain autoencoder. The KL term is what makes the latent space filled in: every encoding is a little Gaussian parked near the origin, the encodings overlap and cover the prior, and so a fresh draw from $\mathcal{N}(\mathbf{0},\mathbf{I})$ lands somewhere the decoder actually knows. The regularizer is not a tax on the model. It is the thing that makes generation possible at all.

The encoder also solves the problem from the last section. We needed the posterior and could not compute it. Instead of solving for it per datapoint, we train one network that predicts an approximate posterior for any $\mathbf{x}$ in a single forward pass. That shortcut has a name, amortized inference: the cost of inference is paid once, in the encoder's weights $\boldsymbol{\phi}$, and then reused for every datapoint, instead of running a fresh optimization for each one the way classical variational inference does. The payoff lands at test time: encoding a new, unseen $\mathbf{x}$ is one forward pass through $\boldsymbol{\phi}$, not an optimization, which is what makes the encoder fast and reusable.

The KL, in closed form

One of the two terms is a gift. If the encoder is a diagonal Gaussian, the standard VAE choice, then for each datapoint it outputs a mean vector $\boldsymbol{\mu}$ and a vector of standard deviations $\boldsymbol{\sigma}$:

Both the prior and this $q$ are Gaussian, so the KL between them has a closed form. You do not sample it, you do not approximate it, you just evaluate it. With the prior $\mathcal{N}(\mathbf{0},\mathbf{I})$, summing over the $J$ latent dimensions, the negative KL (the form the paper adds to its maximized bound) is

Look at what it wants. This expression is at its largest, exactly zero, when every $\mu_j$ is $0$ and every $\sigma_j$ is $1$, which is to say when $q$ is the prior. Anywhere else it is negative, and the further $q$ drifts from the prior, the more negative it gets. A mean far from the origin is penalized through $-\mu_j^{2}$; a variance that is too large or too small is penalized through $\log(\sigma_j^2) - \sigma_j^2$. In a minimized loss every sign flips and you penalize the positive KL instead, the form you will meet in code:

It is a soft spring tethering each code distribution to the standard Gaussian, fought by the reconstruction term that wants the codes spread out and informative.

Move the encoder's output for one latent dimension and watch the spring. The teal bell is $q$, the amber bell the fixed prior; the landscape below is the penalty, brightest where it vanishes.

With the KL handled exactly, only one term in the bound still involves an expectation you cannot integrate: the reconstruction $\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z}\mid\mathbf{x})}\!\big[\log p_{\boldsymbol{\theta}}(\mathbf{x}\mid\mathbf{z})\big]$. Drawing samples estimates its value well enough. The obstacle is its gradient, and that is where the real difficulty hides.

Backprop through a random sample

Write the reconstruction's integrand as $f(\mathbf{z}) = \log p_{\boldsymbol{\theta}}(\mathbf{x}\mid\mathbf{z})$. We need the gradient of $\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z}\mid\mathbf{x})}[f(\mathbf{z})]$ with respect to the encoder's parameters $\boldsymbol{\phi}$. (The decoder $\boldsymbol{\theta}$ is the easy half: it sits inside $f$ at a fixed code, so its gradient passes straight through the sample, the way any gradient does. The hard half is $\boldsymbol{\phi}$, which controls which codes get drawn at all.) And that is the subtlety: $\boldsymbol{\phi}$ does not sit inside $f$, it sits inside the distribution we are averaging over. Change $\boldsymbol{\phi}$ and you change which codes $\mathbf{z}$ get sampled, not the value of $f$ at a fixed code. You cannot push the derivative through a sampling step the way you would through a fixed sum.

There is a textbook estimator that handles this, the score-function estimator (also called REINFORCE, or the likelihood-ratio estimator):

It is correct. It is unbiased, and it even works for discrete codes. But it is famously noisy, and the paper dismisses it in one line for being too high-variance to use. The reason is in its shape. The cost enters multiplicatively: each estimate is $f(\mathbf{z})$ times a score, so its magnitude and its noise scale with how big $f$ is in absolute terms. Add a constant to $f$ and the true gradient does not move, yet the estimator gets noisier, because it never uses how $f$ changes with the code, only its raw height. It learns by correlation: codes that happened to score well should be a little more likely. With one sample per datapoint that correlation is mostly noise.

The fix is the paper's headline idea, the reparameterization trick. Instead of drawing $\mathbf{z}$ directly from $q_{\boldsymbol{\phi}}$, write the same draw as a deterministic function of the encoder's outputs and a fixed, parameter-free noise source. For the Gaussian encoder that is just

where $\odot$ is the elementwise product. The randomness now lives entirely in $\boldsymbol{\varepsilon}$, which does not depend on $\boldsymbol{\phi}$ at all. The encoder's job is the deterministic part, the shift $\boldsymbol{\mu}$ and the scale $\boldsymbol{\sigma}$. And because $\mathbf{z}$ is now a plain differentiable function of $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$, the gradient of $f(\mathbf{z})$ flows straight back through it into the encoder, using $f$'s own slope at the drawn code. That is the difference that kills the variance: this estimate depends on how steeply $f$ leans near $\mathbf{z}$, never on its raw height, so a smooth cost already points the right way from a single sample. Same distribution of $\mathbf{z}$, completely different gradient.

Toggle the two wirings. The forward pass is identical; the backward pass is the whole story.

The slogans usually skip two caveats that matter. Both estimators are unbiased for the same gradient; they agree in expectation. What the trick changes is variance, not correctness. And lower variance is a strong, reliable tendency, not a theorem: for a cost that varies very sharply with the code, the pathwise estimator can in principle be the noisier one. For the smooth reconstruction losses VAEs use, it wins comfortably, which is the whole reason a single sample per datapoint is enough to train.

How much does it actually help?

You can see the variance gap instead of taking it on faith. Take a one-dimensional toy: a code $z \sim \mathcal{N}(\mu, \sigma^2)$ and a smooth cost $f(z) = (z-a)^2$, a stand-in for the reconstruction $f$ from before, and estimate the one gradient $\partial_{\mu}\, \mathbb{E}[f]$, whose true value is $2(\mu - a)$. Each estimator gets the same random draws. Drag the sample count up: both clouds of estimates converge on the true gradient (both are unbiased), but the score-function cloud stays splattered across a wide band while the reparameterized one sits in a tight knot around the answer.

That tightness is what makes the whole scheme practical. The paper's estimator uses a single code per datapoint, and it trains, because at a minibatch of a hundred the reparameterized gradients are already clean enough to point the right way.

The whole algorithm

Now assemble the pieces into the objective the paper actually optimizes. Take the bound in its reconstruction- minus-KL form (3), keep the KL analytic, and estimate only the reconstruction by sampling codes through the reparameterization. For a single datapoint with a Gaussian encoder this is the paper's eq (10):

with $\mathbf{z}^{(l)} = \boldsymbol{\mu} + \boldsymbol{\sigma}\odot\boldsymbol{\varepsilon}^{(l)}$ and $\boldsymbol{\varepsilon}^{(l)} \sim \mathcal{N}(\mathbf{0},\mathbf{I})$. This whole quantity is a single differentiable number. What you just built, the analytic KL plus a sampled reconstruction, is the paper's lower-variance estimator, SGVB ^B. (There is also a fully generic SGVB ^A that Monte-Carlo-estimates both terms, for a non-Gaussian prior or encoder where the KL has no closed form.) Differentiate it with respect to both $\boldsymbol{\theta}$ and $\boldsymbol{\phi}$ at once and take a gradient step. To scale to a large dataset you do this on a minibatch and rescale, eq (8):

The paper finds $L = 1$ sample per datapoint is enough as long as the minibatch is reasonably large, around $M = 100$, because the minibatch already averages the per-sample noise. So in practice you draw exactly one $\boldsymbol{\varepsilon}$ per datapoint. Written out, one training step is short:

# one AEVB step on a minibatch X of M datapoints; prior p(z) = N(0, I)
mu, logvar = encoder(X)                    # q(z|x) = N(mu, exp(logvar))
eps        = randn_like(mu)                # external noise, no gradient
z          = mu + exp(0.5 * logvar) * eps  # reparameterize  (eq 4)
x_hat      = decoder(z)                     # p(x|z)
recon = log_p_x_given_z(X, x_hat)          # +log-likelihood (eq 11 / 12)
kl    = -0.5 * sum(1 + logvar - mu**2 - exp(logvar))   # = +KL  (eq 10)
loss  = mean(kl - recon)                   # minimize -ELBO  (= -L)
loss.backward(); opt.step()                # grad -> encoder AND decoder

The listing hides two details. The encoder does not output $\boldsymbol{\sigma}$ directly; it outputs $\log \boldsymbol{\sigma}^2$, the log-variance, and you recover $\boldsymbol{\sigma} = \exp(\tfrac{1}{2}\log\boldsymbol{\sigma}^2)$. That keeps the variance positive for free and behaves under gradients. And the single loss.backward() sends gradients into the decoder $\boldsymbol{\theta}$ through the reconstruction term and into the encoder $\boldsymbol{\phi}$ through both terms, the KL directly and the reconstruction through the reparameterized code. One backward pass trains the entire model.

Make it concrete. Say $\mathbf{x}$ is a binarized MNIST digit, a 784-dimensional vector of zeros and ones. The encoder is an MLP with a 500-unit hidden layer that reads those 784 numbers and outputs two 20-dimensional vectors, $\boldsymbol{\mu}$ and $\log\boldsymbol{\sigma}^2$, so the latent is $J = 20$ dimensional. Draw one $\boldsymbol{\varepsilon} \in \mathbb{R}^{20}$ from $\mathcal{N}(\mathbf{0},\mathbf{I})$, form $\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma}\odot\boldsymbol{\varepsilon}$, and a decoder MLP (also one 500-unit hidden layer) maps that 20-vector back to the reconstruction the pipeline figure called $\hat{\mathbf{x}}$, here a vector of 784 per-pixel Bernoulli means $\hat{\mathbf{y}} \in (0,1)^{784}$. The reconstruction term is the Bernoulli log-likelihood of the real pixels under those probabilities, eq (11):

and the KL term is the 20-term positive sum from above. The loss is that KL minus this reconstruction, a single scalar, and backprop updates both MLPs. (For continuous data like the Frey Face images the decoder instead outputs a mean and a variance and the reconstruction is a Gaussian log-density, eq (12); everything else is identical.) Once it is trained you can throw the encoder away and just generate:

# generation: no encoder, no data, just ancestral sampling
z     = randn(n, J)        # z ~ prior N(0, I)
x_hat = decoder(z)         # p(x|z): a brand-new sample, never seen in training

Sample a code from the prior, run the decoder, and you have a new datapoint the model invented. The encoder existed only to make training tractable. At generation time the latent space is the prior, and the decoder is the artist.

A latent space you can walk

The regularizer promised something the reconstruction alone could not: a latent space with no holes, where every point decodes to something sensible and nearby points decode to similar things. That is not just a sampling convenience, it is a learned, organized coordinate system for the data. The paper makes the point with a picture. Train a VAE with a two-dimensional latent so you can see all of it, then sweep a grid of codes across the plane and decode each one.

Move the cursor through the latent plane. The decoded digit morphs smoothly, with no seams and no gaps.

The figure invites one careful detail. To make a grid that matches the prior, the paper (unlike this evenly-spaced demo grid) does not space the codes uniformly. It lays a regular grid on the unit square and passes it through the inverse Gaussian CDF, so the grid is dense where the prior puts its mass and sparse in the tails. That is the honest way to tile a Gaussian latent. The takeaway is the smoothness: because every datapoint was encoded to a Gaussian pinned near the origin, the codes pack together and overlap, and the decoder learns a continuous map off of them. Walking a straight line in latent space gives you a smooth interpolation in data space, which is the property that made VAEs a workhorse for representation learning.

What it actually showed

The paper is from 2013, and its results read best in their own terms, not through a modern lens. There is no FID, no Inception score, no CIFAR or ImageNet; those came years later. The datasets are MNIST and the Frey Face video frames. The metric is the variational lower bound itself, measured in nats, plus an estimate of the marginal likelihood for a small latent. The optimizer is Adagrad, not Adam, which Kingma would only publish the following year.

The comparison is against the two reasonable alternatives of the day, the wake-sleep algorithm and Monte Carlo EM with a Hybrid Monte Carlo sampler. On both datasets AEVB converges faster and reaches a better bound, and it keeps scaling where Monte Carlo EM, which needs an expensive sampling loop per datapoint, cannot follow onto the full dataset. The settings are modest by today's standards: minibatches of $M = 100$, one latent sample each, 500 hidden units for MNIST and 200 for the smaller Frey Face set, weights initialized from $\mathcal{N}(0, 0.01)$.

One finding surprised the authors. Adding more latent dimensions did not cause overfitting. A plain autoencoder with a fat bottleneck would happily memorize; this one did not, because the KL term regularizes the model for free, and superfluous latent dimensions just collapse back onto the prior and stop carrying information. The bound polices its own capacity.

What the trick opened up is larger than the experiments. The reparameterization estimator turned variational inference into something you could drop into any gradient-based pipeline, and the years after this paper are a long list of things built on it. There are honest limits, most of which became their own research programs. The encoder is only an approximation, so there is an amortization gap between it and the best per-datapoint posterior, which later work attacks with tighter bounds and richer encoders. The KL term can be too eager and shut a latent dimension off entirely, the failure called posterior collapse; turning its weight down with a single coefficient gives the β-VAE. The trick needs a continuous code, so discrete latents had to wait for the score-function estimator or a relaxation like Gumbel-Softmax. And the Gaussian decoder returns a blurry mean, which is part of why VAEs alone make soft samples.

None of that has dislodged the core idea. A latent-variable model has an intractable likelihood. Bound it, learn an encoder to approximate the posterior, and reparameterize the one sampling step so the whole thing is differentiable. That recipe is still how a great deal of modern generative modeling is trained, right down to the autoencoder that compresses images into the latent space that diffusion models like Stable Diffusion actually run in. The variational autoencoder was one of the first times you could take the clean probabilistic story, sample a code and decode it, and train it end to end with nothing fancier than backprop.