Variational Autoencoders

Imagine training a convolutional autoencoder on human faces and then sampling from its latent space by scribbling down a random vector. What pops out of the decoder are grotesque hybrids—neither face nor noise but some mush where eyes float disconnected from a nose. The culprit is simple: the encoder has only seen a handful of latent vectors during training, so all the other points in the high-dimensional latent space are “dead zones” that never saw a decoder gradient. Variational Autoencoders (VAEs) confront that monster in the gaps by turning the deterministic autoencoder latent space into a smooth, probabilistic terrain where every coordinate has been visited during training and therefore decodes to something plausible. By the end of this page you will understand how VAEs use a latent prior, the ELBO, and the reparameterization trick to make that terrain traversable, how that changes training dynamics compared to plain autoencoders, and how to ship a working VAE build that visualizes smooth interpolations rather than random grotesqueries.

The territory

Generative modeling sits in a taxonomy of approaches that includes explicit likelihood models (autoregressive flows), implicit density models (GANs), and latent-variable models. VAEs land in the latent-variable corner: rather than modeling \(p(x)\) directly, they posit that each datum \(x\) arises from sampling a latent code \(z\) from a structured prior \(p(z)\) followed by a decoder \(p_\theta(x \mid z)\). The decoder is often parameterized with a neural network, so the key challenge is how to learn both the decoder and an encoder that “inverts” the generative process. What distinguishes VAEs from vanilla autoencoders is the insistence that the latent space not just compress data but conform to a smooth, tractable distribution. Without that structure, sampling fails and interpolation is meaningless, just as the opening scenario illustrated.

VAEs therefore borrow ideas from variational inference and information theory. The encoder is not a deterministic bottleneck; it is an approximate posterior \(q_\phi(z \mid x)\) that must stay close to a simple prior \(p(z)\), usually a standard normal. The learning signal then becomes the Evidence Lower Bound (ELBO), which trades off reconstruction fidelity with how well \(q_\phi\) fits \(p\). That trade-off is the mechanism that turns the fractured archipelago of latent points into a continuous continent: every nearby coordinate has a non-negligible probability mass and therefore, by design, decodes to something sensible. How does that work in practice? The mechanism is best understood by starting from the probabilistic model and following the ELBO into the reparameterization trick that makes gradients flow through sampling.

How it works

We begin with a latent-variable model \(p_\theta(x, z) = p_\theta(x \mid z) p(z)\), where \(p(z)\) is the chosen prior—typically a diagonal Gaussian \( \mathcal{N}(z; 0, I)\)—and \(p_\theta(x \mid z)\) is the decoder parameterized by \(\theta\). The marginal likelihood of a datum \(x\) is [p_\theta(x) = \int p_\theta(x \mid z) p(z) dz,] where the integral sums over all possible latent codes. The integral is intractable for neural decoders, so VAEs introduce an approximate posterior \(q_\phi(z \mid x)\) to sidestep it and define the Evidence Lower Bound (ELBO) as [\mathcal{L}(x; \theta, \phi) = \mathbb{E}{z\sim q\phi(z \mid x)}\left[\log p_\theta(x \mid z)\right] - \mathrm{KL}!\left(q_\phi(z \mid x)\;|\;p(z)\right).] Here the first term is the expected log-likelihood (reconstruction term) and the second is the Kullback-Leibler divergence between the encoder distribution \(q_\phi\) and the prior \(p(z)\), which penalizes deviations from the smooth prior.

The ELBO is a lower bound on \(\log p_\theta(x)\); optimizing it thus pushes the generative model to explain the data while keeping latents structured. The KL term is the bridge between reconstruction and structure: \(q_\phi(z \mid x)\) is encouraged to concentrate near the prior, which prevents the encoder from scattering points across disjoint regions. That geometric intuition is spelled out by Doersch’s tutorial on VAEs (Doersch 2016 arxiv:1606.05908), where he likens the KL term to a spring pulling the variational posterior into \(p(z)\). Without this spring, the points are isolated islands and sampling between them leads to the grotesque faces described earlier.

Training the ELBO requires differentiating through the expectation \(\mathbb{E}_{z\sim q_\phi}\). Kingma and Welling (2013) observed that if \(q_\phi\) is a Gaussian with mean \(\mu_\phi(x)\) and diagonal covariance \(\Sigma_\phi(x)\), then one can rewrite a draw from \(q_\phi\) as \(z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon\), where \(\epsilon \sim \mathcal{N}(0, I)\) is independent noise and \(\sigma_\phi(x)\) is the element-wise standard deviation. This reparameterization moves the randomness outside the neural network, allowing gradients to flow through \(\mu_\phi\) and \(\sigma_\phi\) during backpropagation because the sampling now only involves the fixed \(\epsilon\). The ELBO becomes differentiable with respect to \(\phi\), so both encoder and decoder can be trained end-to-end using SGD.

When \(q_\phi\) is Gaussian, the KL divergence has a closed form: [\mathrm{KL}\left(q_\phi(z \mid x) \parallel \mathcal{N}(0, I)\right) = -\frac{1}{2} \sum_{j=1}^{J} \left(1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2\right).] Here \(J\) is the latent dimensionality, \(\mu_j\) and \(\sigma_j^2\) are the \(j\)-th mean and variance, and the identity prior contributes no additional parameters. This term penalizes latents with large means or variances and rewards variances close to one; the KL therefore keeps the variational posterior anchored near the origin while still allowing per-point deviation for reconstruction. The resulting geometry is a continuous blob rather than dispersed islands.

Encoder-decoder interplay

Because the decoder is a neural network, it can in principle memorize the dataset and ignore the latents. This leads to the phenomenon of posterior collapse, where \(q_\phi(z \mid x)\) collapses to the prior and the decoder becomes a standalone generator. The KL term in the ELBO already resists this by punishing the encoder for too much deviation, but the decoder can still learn to ignore \(z\) if it is expressive enough. Practitioners combat this by weakening the decoder or by annealing the KL coefficient during training so that the reconstruction term dominates early, encouraging the encoder to produce informative latents before the KL pressure kicks in. When the decoder ignores \(z\), the latent space loses structure—posterior collapse is literally the archipelago reforming.

A complementary perspective is to view the VAE as learning a stochastic encoder \(q_\phi(z \mid x)\) and a deterministic decoder \(f_\theta(z) = \mathbb{E}[x \mid z]\). If \(f_\theta\) is too expressive relative to the encoder, it can model the data without relying on \(z\), so the encoder collapses. This tension is why later works pair VAEs with autoregressive decoders only after carefully regularizing the latent bottleneck.

Sampling and interpolation

Sampling from a VAE simply means drawing \(z \sim p(z)\) and passing it through the decoder. Because training keeps \(q_\phi(z \mid x)\) close to \(p(z)\), the decoder has seen points near every sample taken from the prior, producing realistic outputs instead of grotesques. Interpolation is similarly easy: linearly interpolate between \(\mu_\phi(x_1)\) and \(\mu_\phi(x_2)\) and decode each intermediate point. The continuity enforced by the KL term means the decoder transitions smoothly between the two data points without sudden artifacts, which is how VAEs become the go-to tool for latent-based explorations, morphing, or conditional generation.

Latent arithmetic and disentanglement

Because the latent prior is continuous and factorized, it also makes latent arithmetic possible. For example, if the variational posterior learns disentangled axes (say “smiling” vs “gender”), then moving along an axis corresponds to adding or subtracting a semantic feature. The ELBO does not guarantee disentanglement, but the smooth prior means such axes can be explored by adding small perturbations drawn from \(\mathcal{N}(0, \sigma^2 I)\), allowing practitioners to probe individual factors of variation using controlled sampling.

Conditioning and hierarchies

VAEs can be extended to conditional VAEs (CVAEs) by conditioning both the encoder and decoder on side information \(y\). The ELBO then becomes [\mathcal{L}(x; \theta, \phi \mid y) = \mathbb{E}{z\sim q\phi(z \mid x, y)}[\log p_\theta(x \mid z, y)] - \mathrm{KL}(q_\phi(z \mid x, y)\; |\; p(z \mid y)),] where \(p(z \mid y)\) is a conditional prior. The same reparameterization trick applies to the conditional posterior, and the KL term still enforces alignment with the conditional prior, ensuring that sampling from \(p(z \mid y)\) yields data consistent with the condition.

Hierarchical VAEs stack multiple latent variables to capture multi-scale structure. Each latent block has its own ELBO term and prior, and training alternates between encouraging each posterior to match its prior and reconstructing via the decoder conditioned on the hierarchy. The reparameterization trick generalizes by sampling each latent level conditioned on earlier ones, so gradient flow remains intact. The result is a more expressive latent space that can model fine-grained details and global structure simultaneously.

Failure modes to watch

Even with KL regularization, VAEs can suffer from blurry samples when the decoder assumes a Gaussian likelihood for pixel data. The Gaussian likelihood penalizes L2 error and therefore averages multiple plausible reconstructions, producing blur. Replacing the likelihood with more expressive distributions (autoregressive decoders, discretized logistic) or modeling logits directly can sharpen samples but increases complexity. Another failure mode is KL vanishing—the KL term collapses to zero even though the latents still carry information, causing optimization difficulties. Monitoring the KL contribution and adjusting the annealing schedule or decoder capacity usually resolves this.

Practical training recipe

Training a VAE from scratch means implementing the encoder and decoder networks, computing \(\mu_\phi(x)\) and \(\log \sigma_\phi^2(x)\), sampling \(z\) via the reparameterization trick, computing \(\log p_\theta(x \mid z)\) (e.g., cross-entropy for images), and adding the KL term. The optimizer minimizes \(-\mathcal{L}\), so gradient descent learns all parameters jointly. Keeping the learning rate modest (e.g., \(1 \times 10^{-3}\)), using batch normalization or layer normalization inside the encoder/decoder, and visualizing latent interpolations every few epochs are best practices to verify that the latent continent is forming as training progresses.

Where the field is now

The original VAE objective from Kingma and Welling (2013) forms the backbone of many current latent methods, but the frontier keeps pushing on the decoder side and on organizing the latent space for downstream use. From a research perspective, works such as VQ-VAE-2 (Razavi et al. 2019 arxiv:1906.00446) demonstrate that discrete latent bottlenecks combined with autoregressive priors scale to high-resolution image synthesis by quantizing the latent space and training the decoder to model discrete codes. More recently, approaches like the Gaussian-Quantized VAE family reinterpret the latent prior as a mixture of smooth Gaussians whose means align with learned discrete tokens, enabling continuous VAEs to plug into autoregressive token samplers without retraining the decoder. These papers show that the core ELBO-based training remains the anchor, but there is active work in reinterpreting the prior to serve modern discrete-generation pipelines.

On the systems side, engineers at Stability AI deploy VAE-derived latent compressors inside Stable Diffusion (Rombach et al. 2022) to map 1024×1024 pixel images into a 64×64 latent grid; the diffusion process then operates in that latent space, reducing both memory and compute. Stability AI’s engineering blog (stability.ai/research/blog) documents that the latent autoencoder at the front end runs in under 400 ms per 512×512 image on an A100, enabling the rest of the pipeline to scale to millions of requests daily. The latent-space compression provided by the VAE is also central to many LLM-vision pipelines, where a single VAE encoder feeds tokens into a transformer; without the VAE’s structured prior, those pipelines fall apart because the transformer receives inconsistent embedding distributions.

This combination of academic and production work suggests two frontiers: research still pushes on flexible priors and KL strategies that keep expressive decoders honest, while engineering focuses on speeding up the latent encoder/decoder pair so that the rest of the stack can scale. Both frontiers intersect on the same tension: how to keep the latent continent navigable without sacrificing the decoder’s power.

What's still open

1. Can a principled regularizer replace heuristic KL-annealing to prevent posterior collapse when pairing a VAE with an autoregressive or transformer decoder? Existing tricks tune the KL coefficient manually, but a theoretical understanding of decoder capacity versus encoder alignment is missing.

2. Does a structured prior—such as a mixture of Gaussians whose components are learned jointly with the encoder—allow VAEs to model multi-modal data without collapsing to the dominant mode, and can we train such priors without adding a prohibitive computational burden?

3. When a VAE encoder feeds a transformer, how can we guarantee that the transformer’s self-attention sees a stable latent distribution across domains, especially when the encoder is finetuned in a continual-learning scenario? The current practice is to freeze the encoder; an open path is to jointly regularize both.

Where to read next

If you want to see how ELBO and KL arise from first principles, → [[score-matching]] walks the probabilistic derivation that connects VAEs to diffusion and flow-based models. The engineering counterpart is → [[flash-attention]] because modern VAE decoders increasingly couple with transformer blocks whose efficiency depends on that kernel. For a different latent parametrization that sidesteps KL terms entirely, → [[flow-matching]] generalizes the noising process so you can compare how a continuous latent prior behaves versus a flow-trained latent.

Build it

This build proves that a VAE trained from scratch on Fashion-MNIST can produce smooth manifold interpolations when the KL term keeps the latent space regularized and the reparameterization trick allows gradients to flow through sampling.

What you're building: A PyTorch VAE trained on Fashion-MNIST whose decoder reconstructs full garments from sampled Gaussian latents and whose latent traversals interpolate between classes.

Why this is valuable: Implementing the encoder, decoder, KL term, and reparameterization from scratch forces you to engage with every term in the ELBO and to visualize how the KL prevents the latent islands from tearing the decoder output apart.

Stack: - Model: commit to no pretrained weights; architect a 3-layer convolutional encoder and decoder with latent size 64. - Dataset: huggingface: fashion_mnist — 60k training examples of 28×28 grayscale garments. - Framework: PyTorch 2.1 with torchmetrics 0.11 and torchvision 0.18. - Compute: Free Google Colab GPU (T4/RTX 4090 equivalent) — expect about 40 min for 60 epochs.

The recipe: 1. pip install torch torchvision torchmetrics matplotlib and clone a notebook environment; seed the RNG for reproducible samples. 2. Normalize Fashion-MNIST to \([-1, 1]\), stack training mini-batches of 128, and preprocess labels for interpolation (one-hot for visualization but not for loss). 3. Define encoder/decoder networks with convolutional blocks, compute \(\mu_\phi\) and \(\log \sigma_\phi^2\), sample \(z = \mu + \sigma \odot \epsilon\), reconstruct via decoder, and train with loss \( \text{reconstruction loss} + \beta \cdot \text{KL}\). Use Adam with learning rate \(1 \times 10^{-3}\), \(\beta=1\), and checkpoint the model every 10 epochs, watching both reconstruction and KL curves. 4. Evaluate by sampling 16 latents from \(\mathcal{N}(0, I)\) and decoding them, reporting the average binary cross-entropy per pixel and plotting a 4×4 grid of reconstructions; additionally, select two validation examples, interpolate their mean vectors in latent space over 10 steps, decode, and inspect the morph sequence. 5. What you now have is a trained VAE checkpoint plus visualization scripts that demonstrate the decoder producing realistic garments for random Gaussian latents and smooth interpolations between classes.

Expected outcome: A working VAE checkpoint that decodes from a structured Gaussian prior and a gallery of sampled/interpolated garments proving that the latent space is navigable.

• CS student: Run the same notebook on an RTX 3070 by reducing epochs to 30 and latent size to 32; the artifact is the interpolation GIF you save in under 3 hours.
• Applied engineer: Quantize the decoder with PyTorch 2.1 dynamic quantization, serve it via TorchServe, and demonstrate <80 ms p50 latency on an entry-level A10, showing the quantized decoder still decodes from Gaussian latents.
• Applied researcher: Hypothesize that KL scaling (tuning \(\beta\)) changes class overlap; sweep \(\beta = \{0.1, 0.5, 1.0, 2.0\}\), log the KL contribution, and compare the latent interpolation smoothness metric (e.g., mean squared difference between successive decoded images).
• Frontier researcher: Probe posterior collapse by pairing the VAE encoder with a lightweight transformer decoder and testing whether a principled regularizer (e.g., Fisher information-weighted KL) can keep the KL term non-zero without annealing; falsify the hypothesis if the regularizer fails to maintain KL > 0.1 across training.

---

If this build worked for you — a ⭐ on GitHub is the only signal we collect.