Variational Autoencoders

An ordinary autoencoder learns a deterministic code $z=f(x)$ and a reconstruction $\hat{x}=g(z)$, and nothing constrains where in code-space the encodings land. The trained codes form isolated islands; the gaps between them decode to garbage, so the autoencoder compresses but cannot generate: there is no distribution to sample from. The variational autoencoder (VAE) repairs exactly this. It makes the encoder output a distribution over codes, trains that distribution to match a fixed prior $p(z)$, and thereby turns the latent space into a smooth density we can sample. The cost is a detour through variational inference; the result is a principled, gradient-trainable generative model.¹

The latent-variable model

A VAE is a latent-variable generative model: each datum $x$ is generated by first drawing a hidden code $z$ from a simple prior, then mapping it through a decoder $p_{\theta }(x\mid z)$.

We take the standard Gaussian prior and a decoder network $g_{\theta }$ that maps a code to the parameters of an output distribution:

so the data density is the marginal over all codes,

To train this by maximum likelihood we would maximize $\log p_{\theta }(x)$. The integral, however, is intractable: it ranges over the entire latent space, and the decoder $g_{\theta }$ is a nonlinear network, so there is no closed form. The posterior we would need to invert the generative process,

inherits the same intractable normalizer $p_{\theta }(x)$ in its denominator. We cannot compute it, and so we cannot directly do maximum likelihood. Variational inference works around this.²

Variational inference and the ELBO

The idea of variational inference is to give up on the exact posterior and instead introduce a tractable approximate posterior $q_{\varphi }(z\mid x)$ (a second network, the encoder) chosen from a family we can sample and differentiate. We pick a Gaussian with diagonal covariance whose mean and variance are emitted by the encoder network $f_{\varphi }$:

The encoder is also called the inference network because it performs approximate inference: given $x$, it guesses which codes $z$ could have produced it. The decoder is the generative network. We now derive a tractable objective that bounds the log-evidence from below.

The bound $\mathcal{L}(\theta ,\varphi ;x)$ is the evidence lower bound (ELBO). Maximizing it does double duty: because $\log p_{\theta }(x)=\mathcal{L}+D_{KL}(q_{\varphi }\Vert p_{\theta }(z\mid x))$ with the evidence fixed in $\varphi$, pushing $\mathcal{L}$ up simultaneously raises a lower bound on the data likelihood and squeezes the inference gap, tightening $q_{\varphi }$ toward the true posterior.³

Reading the two terms

The ELBO reads as a reconstruction term minus a regularizer, and this is the entire intuition behind the architecture:

With a unit-variance Gaussian decoder, $-\log p_{\theta }(x\mid z)$ is squared error up to a constant, so the reconstruction term is the familiar autoencoder loss. The KL term is the new ingredient: it is precisely what an ordinary autoencoder lacks, the term that packs the encodings against the prior so the gaps close.

The reparameterization trick

The ELBO contains an expectation over $z\sim q_{\varphi }(z\mid x)$, and we must differentiate it with respect to $\varphi$ — but $\varphi$ controls the very distribution we are sampling from. The naive estimator differentiates the sampling density directly (the score-function or REINFORCE estimator),

which is unbiased but notoriously high-variance: $h(z)$ swings widely and the ${\nabla }_{\varphi }\log q_{\varphi }$ weight does not cancel that swing, so estimates are noisy and training crawls.

The reparameterization trick removes the randomness from the gradient path. A Gaussian sample can be written as a deterministic, differentiable transform of a fixed, parameter-free noise source:

Now the only stochastic node is $ϵ$, which carries no parameters; the sample $z$ is a smooth function of ${\mu }_{\varphi }$ and ${\sigma }_{\varphi }$. The expectation becomes an expectation over the fixed distribution of $ϵ$, and the gradient slips inside it:

a low-variance estimate we can compute with a single sample of $ϵ$ and ordinary backpropagation.⁴

The trick is what makes the VAE trainable end-to-end: it lets a single gradient pass through the sampled code and into the encoder's weights.

The closed-form KL term

The KL regularizer between two Gaussians has a closed form, so the second ELBO term needs no sampling at all. For a diagonal posterior $q=\mathcal{N}(\mu ,\mathrm{diag}{\sigma }^2)$ against the prior $p=\mathcal{N}(0,I)$ in $k$ latent dimensions,

Each coordinate is penalized for drifting its mean off $0$ (${\mu }_j^2$) and for a variance that strays from $1$ (the ${\sigma }_j^2-\log {\sigma }_j^2-1$ term, which is $\ge 0$ and vanishes only at ${\sigma }_j=1$). Substituting the reparameterized reconstruction term, the per-example objective the network maximizes is fully explicit:

In practice the encoder emits $\log {\sigma }^2$ rather than $\sigma$ (the log keeps the variance positive without a constraint and stabilizes the gradient), and the expectation is approximated by the single noise draw $ϵ$ per example. The whole VAE is now a single differentiable graph.

A worked ELBO computation

For example, take a two-dimensional latent space ($k=2$) and an encoder that, for a particular input $x$, emits

The KL term first, coordinate by coordinate, using $\frac{1}{2}({\mu }_j^2+{\sigma }_j^2-\log {\sigma }_j^2-1)$:

So $D_{KL}(q_{\varphi }\Vert p)=0.417+0.054=0.471$ nats. Coordinate $1$ pays most of it: its mean sits well off the origin ($0.8$) and its variance is squeezed to $0.5$, both of which the prior penalizes. Coordinate $2$ is nearly free because its mean is small and its variance is close to $1$ — almost exactly the prior already.

Now the reconstruction term. Draw one noise sample $ϵ=(1.0,-0.5)$ and reparameterize:

Push $z$ through the decoder to get $g_{\theta }(z)=\hat{x}$. With a unit-variance Gaussian decoder over, say, a $D=4$-pixel output, the reconstruction log-density is $\log p_{\theta }(x\mid z)=-\frac{1}{2}\Vert x-\hat{x}{\Vert }^2-\frac{D}{2}\log (2\pi )$. If the squared reconstruction error came out to $\Vert x-\hat{x}{\Vert }^2=0.6$, then

The ELBO is their difference:

The optimizer minimizes $-\mathcal{L}=4.447$. The KL of $0.471$ is small next to the reconstruction cost, which is typical early in training — the network first learns to reconstruct, and the KL only bites once the codes are informative enough that packing them against the prior costs fidelity. The $\beta$ weight below controls exactly that ratio.

Training the VAE

Training maximizes the ELBO by minimizing its negation, averaged over minibatches, with a single reparameterized sample per example. Every line is an ordinary differentiable operation, so a standard optimizer trains both networks jointly.

Sampling, interpolation, and the smooth latent space

Once trained, the encoder can be discarded for generation: because $q_{\varphi }$ was pressed against $p(z)=\mathcal{N}(0,I)$, an unconditional sample from the prior lands where the decoder expects data.

This is what the KL term buys. In a plain autoencoder the codes occupy an unknown, holey region, and a prior sample almost surely misses it; the VAE's regularizer guarantees the prior is the code distribution, so sampling works. The same smoothness makes latent-space interpolation meaningful: a straight line between two codes $z_a,z_b$ decodes to a continuous morph because nearby codes decode to nearby data.⁵

The grid is the canonical VAE diagnostic: sweeping the two latent coordinates and decoding each point traces a continuous atlas of the data manifold, with no torn seams between encodings.

Two distributions, one pull

It helps to picture the KL term geometrically. The approximate posterior $q_{\varphi }(z\mid x)$ is a bell centered at ${\mu }_{\varphi }(x)$ with width ${\sigma }_{\varphi }(x)$; the prior $p(z)$ is the unit bell at the origin. The KL regularizer pulls every per-example bell toward the prior; the reconstruction term opposes it, placing each $x$ at a distinct, confident code.

The weight $\beta$ generalizes the objective to ${\mathcal{L}}_{\beta }={\mathbb{E}}_q[\log p_{\theta }(x\mid z)]-\beta D_{KL}(q_{\varphi }\Vert p)$: $\beta <1$ favors reconstruction (sharper samples, looser prior match), $\beta >1$ favors disentangled, prior-aligned codes (the $\beta$-VAE), and $\beta =1$ is the exact ELBO.⁶

What VAEs became

The VAE of Kingma and Welling was published in 2014 alongside an independent derivation by Rezende, Mohamed, and Wierstra, who called the same object a deep latent Gaussian model and its reparameterized gradient stochastic backpropagation.⁷ Both papers solved the same problem — how to backpropagate through a sampled latent — and the trick has since become standard machinery well outside generative modeling, anywhere an expectation over a learned distribution appears in a loss.

Discrete latents needed a new trick. Reparameterization as stated requires a continuous, differentiable transform, which rules out categorical codes. The Gumbel-softmax (or Concrete) relaxation fixed this by reparameterizing a categorical draw as a temperature-controlled softmax over Gumbel noise, recovering a differentiable path to discrete latents.⁸ The VQ-VAE took the opposite route — a hard nearest-neighbor lookup into a learned codebook, with a straight-through gradient — and its discrete codes are what let autoregressive and, later, diffusion models operate in a compact latent space rather than on raw pixels.⁹

The $\beta$-VAE turned the KL weight into a research program. Raising $\beta$ above $1$ was found to encourage disentangled representations, where individual latent coordinates track independent factors of variation such as pose or lighting; the same paper reframed the whole objective as a constrained optimization with the KL as an information bottleneck.¹⁰ That information-theoretic reading — the KL term is the number of nats the code is allowed to carry about the input — is the cleanest way to understand both disentanglement and posterior collapse.

The blur problem drove the field toward other models. VAEs reconstruct through a likelihood term that, with a Gaussian decoder, is squared error, and averaging over the posterior produces the characteristic soft, blurry samples. This is the practical reason GANs and diffusion models overtook the plain VAE for high-fidelity image synthesis — though the VAE's encoder-decoder latent space survives inside them, most visibly as the compression front-end of latent diffusion.

Takeaways

• A VAE is a latent-variable model $p_{\theta }(x)=\int p_{\theta }(x\mid z)p(z)dz$ with prior $p(z)=\mathcal{N}(0,I)$; the integral and the posterior $p_{\theta }(z\mid x)$ are intractable, so we cannot do direct maximum likelihood.
• Variational inference introduces an encoder $q_{\varphi }(z\mid x)$ and the ELBO $\log p_{\theta }(x)\ge {\mathbb{E}}_q[\log p_{\theta }(x\mid z)]-D_{KL}(q_{\varphi }\Vert p)$, whose gap is exactly $D_{KL}(q_{\varphi }\Vert p_{\theta }(z\mid x))$ — maximizing the bound both fits the data and tightens the approximation.
• The ELBO reads as reconstruction minus a KL regularizer; the KL term is the force a plain autoencoder lacks, packing codes against the prior so the latent space becomes smooth and samplable.
• The reparameterization trick $z={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)\odot ϵ$ moves randomness onto a parameter-free $ϵ$, replacing the high-variance score-function estimator with a low-variance pathwise gradient that backpropagates through the sample.
• The Gaussian-vs-Gaussian KL is closed-form, sampling from the prior generates data, latent interpolation morphs smoothly, and posterior collapse is the characteristic failure when the decoder ignores the code.