Variational Autoencoders
Imagine you want a machine that can dream up new, plausible handwritten digits, faces, or molecules. It should not just memorize the ones it has seen. This requires modeling a full probability distribution over complex, high-dimensional data. It also requires being able to sample fresh points from it. A latent variable model is a natural framework for this task. But the classical machinery for fitting such models, exact posterior inference, does not scale. It fails for deep neural networks with millions of parameters. The Variational Autoencoder (VAE), introduced by Kingma and Welling in 2013, solves this problem. It replaces exact inference with a learned approximation. Everything is trained jointly with backpropagation. The key engineering trick involves the two hardest parts of latent variable modeling: inference and learning. These become a pair of neural networks that look exactly like an autoencoder. An encoder network compresses data into a latent code. A decoder network reconstructs data from that code. Unlike an ordinary autoencoder, the VAE's encoder outputs a whole distribution over codes. It does not output a single point. The decoder is trained differently too. Sampling from this distribution and pushing the sample through the decoder should generate data that looks real. This is called amortized inference. Instead of solving a separate optimization problem for every data point's posterior, one neural network learns to output approximate posteriors. It does this for all data points at once. VAEs matter for an important reason. They gave deep learning its first scalable recipe for likelihood-based generative modeling with a continuous latent space. This recipe trains with a single stable, differentiable objective. It also produces a latent space with genuinely useful geometric structure. This structure includes interpolation, arithmetic, and disentanglement. This combination made VAEs foundational for several fields. These include image and audio generation, representation learning, anomaly detection, and semi-supervised learning. VAEs also serve as a conceptual stepping stone to diffusion models. Diffusion models are themselves best understood as deep hierarchical VAEs. This article gives the precise generative story and objective. It derives the Evidence Lower Bound (ELBO) that VAEs optimize. It walks through the encoder-decoder training loop with pseudocode. It explains why the reparameterization trick is necessary for backpropagation. Finally, it traces a concrete numerical example from a data point to a gradient update.
- A VAE models data as generated by a latent code `z` drawn from a prior, then pushed through a decoder network `p(x|z)`.
- Exact posterior inference `p(z|x)` is intractable, so an encoder network learns an approximate posterior `q(z|x)`.
- Training maximizes the Evidence Lower Bound (ELBO), a tractable surrogate for the true data log-likelihood.
- The ELBO decomposes into a reconstruction term and a KL-divergence regularizer pulling `q(z|x)` toward the prior.
- The reparameterization trick rewrites sampling as a deterministic function of noise, making the ELBO differentiable end-to-end.
- The result is a smooth, structured latent space that supports sampling, interpolation, and downstream representation learning.
A VAE trains an encoder and decoder jointly by maximizing a tractable lower bound on the data likelihood, using a reparameterized random sample as the bridge between the two networks.
The idea
A VAE posits a generative story for data `x`: first draw a latent code `z` from a simple prior `p(z)`, typically a standard Gaussian `N(0, I)`, then draw `x` from a conditional distribution `p_θ(x|z)` parameterized by a neural network called the decoder. The joint model is `p_θ(x, z) = p(z) p_θ(x|z)`, and the quantity we ultimately want to maximize is the marginal likelihood of the observed data, `p_θ(x) = ∫ p_θ(x|z) p(z) dz`. This integral is intractable for any nontrivial decoder network, since it requires integrating over all possible latent codes. The standard remedy from variational inference is to introduce an approximate posterior `q_φ(z|x)`, and use it to derive a tractable lower bound on `log p_θ(x)`. In classical variational inference `q` would be optimized separately, from scratch, for every single data point `x`, which is prohibitively expensive at deep learning scale. The VAE's central innovation is amortized inference: a single neural network, the encoder (or recognition network), takes `x` as input and directly outputs the parameters of `q_φ(z|x)` — typically a mean vector `μ(x)` and a diagonal covariance `σ²(x)` for a Gaussian. One forward pass replaces an entire per-datapoint optimization loop, and the encoder's parameters `φ` are shared and learned across the whole dataset, just like the decoder's parameters `θ`. Training proceeds by jointly optimizing `θ` and `φ` to maximize the Evidence Lower Bound (ELBO), a single scalar objective that is tractable to estimate and differentiate. Because the encoder produces a distribution rather than a single deterministic code, and because we need to sample from that distribution to evaluate the decoder, naive backpropagation through the sampling step is impossible — you cannot differentiate through a random number generator. The reparameterization trick, covered in depth elsewhere, resolves this by expressing the sampled `z` as a deterministic, differentiable function of `x`'s encoder outputs and an independent noise variable `ε`. This lets gradients flow from the reconstruction loss, through the decoder, through `z`, and into the encoder, all in one backward pass.
How it works
Deriving the ELBO. Start from the log marginal likelihood and introduce `q_φ(z|x)` via a multiply-and-divide trick, then apply Jensen's inequality: ``` log p_θ(x) = log ∫ q_φ(z|x) · [p_θ(x, z) / q_φ(z|x)] dz = log E_{q_φ(z|x)} [ p_θ(x, z) / q_φ(z|x) ] ≥ E_{q_φ(z|x)} [ log p_θ(x, z) - log q_φ(z|x) ] (Jensen's inequality) = ELBO(θ, φ; x) ``` Expanding `p_θ(x, z) = p_θ(x|z) p(z)` gives the standard two-term form: ``` ELBO(θ, φ; x) = E_{q_φ(z|x)}[log p_θ(x|z)] - KL( q_φ(z|x) || p(z) ) └── reconstruction term ──┘ └────── regularizer ──────┘ ``` The reconstruction term rewards decoders that assign high probability to `x` given codes sampled from the encoder. The KL term penalizes the encoder for straying from the prior, keeping the latent space well-organized and making the prior a valid sampling distribution at generation time. One can show `log p_θ(x) = ELBO + KL(q_φ(z|x) || p_θ(z|x))`, and since KL divergence is always non-negative, the ELBO is indeed a lower bound on the true log-likelihood; the gap is exactly the KL divergence between the approximate and true posteriors, which shrinks as the encoder becomes a better approximator.
Training loop. Both networks are trained jointly with stochastic gradient ascent on the ELBO (equivalently, gradient descent on the negative ELBO loss), using minibatches and the reparameterization trick to get low-variance gradient estimates. ``` PSEUDOCODE: VAE training step Input: minibatch X = {x_1, ..., x_B} Params: encoder φ (outputs μ(x), log σ²(x)), decoder θ (outputs distribution params for x given z) for each x in X:
# --- Encode --- μ, log_σ2 = Encoder_φ(x) # amortized inference σ = exp(0.5 * log_σ2)
# --- Reparameterized sample --- ε ~ N(0, I) # sample noise, NOT z directly z = μ + σ ⊙ ε # deterministic, differentiable in μ, σ
# --- Decode --- x_hat_params = Decoder_θ(z) # e.g. Bernoulli logits or Gaussian mean
# --- Loss terms --- recon_loss = -log p_θ(x | z) # e.g. binary cross-entropy for pixels kl_loss = 0.5 sum(exp(log_σ2) + μ^2 - 1 - log_σ2) # closed form vs N(0,I) loss(x) = recon_loss + β kl_loss # β=1 for standard ELBO; β-VAE varies it L = (1/B) * sum(loss(x) for x in X)
# --- Update --- g_θ, g_φ = backprop(L) # gradients flow through z because it's reparameterized θ, φ = θ - η g_θ, φ - η g_φ # SGD / Adam step ``` The closed-form KL term above is derived from the KL divergence between two Gaussians, `q_φ(z|x) = N(μ, diag(σ²))` and the prior `p(z) = N(0, I)`: ``` KL(N(μ,σ²) || N(0,1)) = 0.5 * (σ² + μ² - 1 - log σ²) ``` summed over latent dimensions. At generation time, the encoder is discarded entirely: new samples are produced by drawing `z ~ p(z) = N(0, I)` directly and passing it through the decoder.
Complexity & correctness
Time complexity. A single training step performs one forward and one backward pass through the encoder and the decoder. If the encoder and decoder are feedforward or convolutional networks with `P` total parameters, the cost per data point is `O(P)` for the forward pass and `O(P)` for backpropagation, by the standard cost analysis of Backpropagation where the backward pass costs a constant factor times the forward pass. For a minibatch of size `B`, one training step costs `O(B·P)`, and training for `E` epochs over a dataset of size `N` costs `O(E·N·P)` — identical, asymptotically, to training an ordinary autoencoder or feedforward classifier of the same size; the VAE's variational machinery (sampling `ε`, computing the closed-form KL) adds only `O(B·d)` overhead, where `d` is the latent dimension, which is negligible next to `O(B·P)`.
Space complexity. Memory is dominated by storing activations for backpropagation, `O(B·P)` in the worst case for a network with `P` neuron-activations per example (more precisely, proportional to the total width-times-depth of the network), plus `O(B·d)` for the latent samples and their reparameterization intermediates. This is exactly the memory profile of a standard autoencoder of the same depth, since the encoder and decoder here are ordinary networks; the only Wchange is that the encoder's output layer has `2d` units (mean and log-variance) instead of `d`.
Correctness argument. There are two separate things to verify: (1) that the ELBO is truly a valid lower bound on `log p_θ(x)`, and (2) that maximizing it with reparameterized, sampled gradients converges to a stationary point of the expected ELBO. For (1), the derivation above used Jensen's inequality on the concave `log` function applied to an expectation, which is valid unconditionally, so `ELBO ≤ log p_θ(x)` for any choice of `q_φ`, even a badly-fit one — the bound never breaks, it just becomes looser. The exact gap is: ``` log p_θ(x) - ELBO(θ,φ;x) = KL( q_φ(z|x) || p_θ(z|x) ) ≥ 0 ``` so maximizing the ELBO with respect to `φ` alone, holding `θ` fixed, is exactly equivalent to minimizing the KL divergence between the encoder's approximate posterior and the true (intractable) posterior — this is why the encoder is called an inference network: it is literally being trained to approximate Bayesian posterior inference. For (2), the Monte Carlo estimate of the ELBO using one reparameterized sample `z = μ + σ⊙ε` is an unbiased estimator of the true expectation, because the change-of-variables from `z ~ q_φ(z|x)` to `ε ~ N(0,I)` preserves the expectation exactly (a linear reparameterization has a Jacobian determinant that cancels correctly in the change-of-variables formula, and more directly, one can just verify `E_ε[f(μ+σε)] = E_{z~N(μ,σ²)}[f(z)]` by substitution). Because `ε` is drawn independently of `θ, φ`, the gradient of this Monte Carlo estimate is also an unbiased estimator of the gradient of the true expectation (gradient and expectation commute here since the sampling distribution no longer depends on the parameters being differentiated). Stochastic gradient ascent with unbiased gradient estimates and a suitable learning rate schedule converges to a stationary point of the ELBO objective, by the standard theory of stochastic approximation — the same guarantee underlying SGD training of any deep network.
Tradeoffs. The ELBO is only a lower bound, not the true likelihood, so a VAE trained to convergence still underestimates `log p_θ(x)` by exactly the encoder's residual approximation error; tighter bounds (e.g. importance-weighted autoencoders) close this gap at extra compute cost. VAEs are also well known to produce blurrier samples than GANs, because the pixelwise reconstruction loss (often a Gaussian or Bernoulli likelihood) averages over the encoder's posterior uncertainty rather than committing to sharp modes — this is a property of the objective, not a bug in the derivation.
Worked example
Take a toy VAE with 1-dimensional latent space, trained on scalar data `x ∈ ℝ` (imagine it's one pixel intensity, rescaled to `[0,1]`, for concreteness). Suppose at the current parameters, for input `x = 0.8`, the encoder computes: ``` Encoder_φ(x=0.8) → μ = 0.50, log σ² = -1.386 (so σ² = 0.25, σ = 0.5) ``` We draw one noise sample `ε = 1.2` from `N(0,1)` and reparameterize: ``` z = μ + σ·ε = 0.50 + 0.5 × 1.2 = 1.10 ``` Feed `z = 1.10` into the decoder, which (say) computes a Bernoulli reconstruction probability via a sigmoid output: ``` Decoder_θ(z=1.10) → logit = 1.5 → x̂ = σ(1.5) ≈ 0.8176 ``` Compute the two loss terms: ``` recon_loss = -[x·log(x̂) + (1-x)·log(1-x̂)] = -[0.8·log(0.8176) + 0.2·log(0.1824)] = -[0.8×(-0.2014) + 0.2×(-1.7016)] = -[-0.1611 - 0.3403] = 0.5014 kl_loss = 0.5·(σ² + μ² - 1 - log σ²) = 0.5·(0.25 + 0.25 - 1 - (-1.386)) = 0.5·(0.886) = 0.443 ``` With `β = 1`, the total per-example loss is: ``` loss = recon_loss + kl_loss = 0.5014 + 0.443 = 0.9444 ``` This scalar loss is backpropagated: the gradient with respect to `θ` flows through the decoder as usual; the gradient with respect to `φ` flows through the chain `loss → x̂ → z → (μ, σ) → φ`, which is only possible because `z` was written as `μ + σ·ε` with `ε` treated as a constant during this backward pass. A step of Adam or SGD nudges `μ` and `log σ²` in the direction that would have made `z` decode closer to `x = 0.8` while also nudging `q_φ(z|x)` closer to `N(0,1)`. Repeating this over the dataset, the encoder learns to place codes for "similar" inputs near each other and near the origin, and the decoder learns to map the whole `N(0,1)` ball to plausible data — which is exactly what lets one later sample `z ~ N(0,1)` with no `x` at all and decode a brand-new, never-seen data point.
| Quantity | Value |
|---|---|
| Input `x` | 0.8 |
| Encoder output `μ, σ²` | 0.50, 0.25 |
| Sampled noise `ε` | 1.2 |
| Latent code `z` | 1.10 |
| Decoder output `x̂` | 0.8176 |
| Reconstruction loss | 0.5014 |
| KL loss | 0.443 |
| Total loss | 0.9444 |
Common pitfalls & traps
| Pitfall | Why it happens | Fix | ||||
|---|---|---|---|---|---|---|
| Posterior collapse | Decoder is powerful enough to ignore `z` entirely; KL term pushes `q_φ(z\ | x) → p(z)` and the model degenerates to unconditional generation | Anneal the KL weight `β` from 0 up (KL warm-up), or use a weaker/limited decoder | |||
| Forgetting the reparameterization trick | Sampling `z` directly with `z ~ N(μ,σ²)` inside the graph blocks gradient flow into the encoder | Always express `z = μ + σ⊙ε` with `ε` sampled outside the differentiable path | ||||
| Using `σ` directly instead of `log σ²` | Networks can output negative values for what must be a positive variance | Predict `log σ²` (unconstrained) and exponentiate, guaranteeing positivity | ||||
| Treating ELBO as the true likelihood | The gap `KL(q_φ(z\ | x)\ | \ | p_θ(z\ | x))` is silently ignored when reporting "log-likelihood" | Report the ELBO explicitly, or use tighter bounds (importance weighting) for likelihood estimates |
| Wrong reconstruction likelihood for the data type | Using mean-squared error (implicitly Gaussian) on binary pixels, or vice versa, biases the learned decoder | Match the likelihood family to the data: Bernoulli/categorical for discrete pixels, Gaussian for continuous | ||||
| Blaming blurriness on a "broken" model | Blurring is an expected consequence of averaging over posterior uncertainty under a pixelwise likelihood | Accept the tradeoff, or move to hierarchical/autoregressive decoders, or use adversarial/perceptual losses |
Pioneers & history
Diederik Kingma, then a PhD student, and Max Welling published "Auto-Encoding Variational Bayes" in 2013, introducing the reparameterization trick as the key device that let stochastic variational inference be trained with ordinary backpropagation at the scale of deep networks. Danilo Rezende, Shakir Mohamed, and Daan Wierstra independently published a closely related "stochastic backpropagation" framework the same year, applying essentially the same reparameterization idea to deep generative models. Together these papers turned decades of variational Bayesian inference theory into a practical recipe that any deep learning practitioner could implement in a few dozen lines of code.
Origins · 1 minDiederik P. Kingma and how this was discovered
Where this matters
VAEs are used directly for anomaly detection (data points with unusually low reconstruction likelihood or ELBO are flagged as outliers), for molecule and drug candidate generation in computational chemistry, for collaborative filtering and recommendation via latent user/item embeddings, and for speech and music synthesis where a smooth continuous latent space enables controllable generation. The encoder-decoder-with-a-bottleneck-distribution pattern also underlies modern representation learning: pretraining an encoder to produce useful, compressed features is a direct ancestor of Self-Supervised Learning pipelines used in vision and language models. Perhaps most importantly, the mathematics of the VAE, a hierarchy of latent variables, an ELBO-style objective, and a learned "denoising" or generative process, is the direct conceptual precursor to diffusion models, which can be understood as VAEs with a fixed, many-step, non-learned encoder (the forward noising process) and a learned decoder (the reverse denoising process). Understanding VAEs thoroughly is therefore also the fastest route to understanding why diffusion models are trained the way they are.
Ask the tutor: How does the reparameterization trick generalize when the approximate posterior isn't Gaussian?
Why can't we just maximize log p_θ(x) directly instead of the ELBO?
Because `log p_θ(x) = log ∫ p_θ(x|z) p(z) dz` requires integrating the decoder's output over every possible latent code, an integral with no closed form once the decoder is a deep neural network. Even Monte Carlo estimation of this integral by sampling `z` from the prior is hopelessly high-variance in practice, since for most randomly sampled `z`, `p_θ(x|z)` is essentially zero — the vast majority of the latent space produces outputs nothing like `x`. The ELBO sidesteps this by sampling `z` from an encoder distribution `q_φ(z|x)` that is trained to concentrate its mass on codes actually likely to have produced `x`, making the reconstruction term's Monte Carlo estimate low-variance and tractable while still bounding the true objective we care about.
What is amortized inference, and why is it central to VAEs?
In classical variational inference, for each data point `x` you would run an iterative optimization (e.g. coordinate ascent) to find the best variational distribution `q(z)` approximating that point's specific posterior, which is expensive when repeated for every one of thousands or millions of data points. Amortized inference instead trains a single shared neural network, the encoder, whose weights `φ` are optimized once over the whole dataset, and which then produces a good approximate posterior for any `x` in a single forward pass. The cost of inference is "amortized" across the dataset: you pay the cost of training the encoder once, and afterward inference is essentially free, an `O(1)`-network-pass operation rather than an iterative per-point optimization.
Why does the KL term use the standard normal prior N(0, I), and what happens if we change it?
`N(0, I)` is chosen mainly for mathematical convenience: it makes the KL divergence term between two Gaussians have a simple closed form, and it makes sampling for generation trivial, just draw independent standard normals. It also encourages the encoder to spread codes for different inputs across a shared, roughly isotropic region, giving the latent space smooth, interpolatable geometry. Using a different or learned prior (e.g. a mixture of Gaussians, or a autoregressive prior as in VQ-VAE-2) can improve expressiveness and match complex data structure more closely, at the cost of a more complex or less tractable KL term and sometimes a harder-to-navigate latent space.
How does a VAE differ from a plain (deterministic) autoencoder?
A plain autoencoder maps `x` to a single deterministic code `z = Encoder(x)` and back, minimizing only reconstruction error, with no probabilistic story and no mechanism forcing the latent space to be well-organized or sample-able. A VAE's encoder instead outputs a distribution `q_φ(z|x)`, adds a KL-divergence penalty pulling that distribution toward a fixed prior, and this combination is precisely what makes the latent space usable for generation: sampling from the prior alone, with no input `x`, now decodes into realistic outputs, something a plain autoencoder's arbitrary, unconstrained latent space typically cannot support.
Why do VAE-generated images often look blurry compared to GAN-generated images?
The reconstruction term of the ELBO is typically a Gaussian or independent-pixel Bernoulli log-likelihood, which is maximized, for a given input, by predicting something close to the average of all plausible reconstructions consistent with the sampled `z`, rather than committing to one crisp, high-frequency mode. Combined with the KL regularizer compressing information through a bottleneck, fine detail that isn't cheaply encodable tends to get smoothed out. GANs instead use an adversarial loss that directly penalizes any output a discriminator can distinguish from real data, which pushes toward sharp, realistic-looking samples but sacrifices the stable likelihood-based training and mode-covering behavior that VAEs provide.
Is the ELBO objective the same as, or different from, the objective used to train ordinary variational inference in Bayesian statistics?
It is the same objective, evidence lower bound maximization, applied with a specific architectural choice for how `q_φ(z|x)` is parameterized. Classical variational Bayes typically restricts `q` to a simple, hand-derived family (e.g. mean-field factorized Gaussians updated via coordinate ascent equations) and refits it per data point. The VAE keeps the same mathematical objective but replaces the per-point optimization with a single amortized neural network trained by stochastic gradient ascent, which is what makes it scale to large, high-dimensional deep learning datasets rather than small, hand-modeled Bayesian problems.
What role does β play in the β-VAE variant, and why would you tune it?
β scales the KL term in the loss, `recon_loss + β·kl_loss`, so `β=1` recovers the exact ELBO while `β>1` more aggressively pushes `q_φ(z|x)` toward the prior. Larger β tends to produce a more disentangled latent space, where individual latent dimensions correspond to more interpretable, independent factors of variation, at the cost of worse reconstruction fidelity since the bottleneck becomes tighter. This is a genuine tradeoff rather than a free improvement: setting β too high can trigger posterior collapse, where the KL penalty dominates and the encoder stops encoding useful information about `x` at all.