Skip to content

Autoencoders and VAEs

What This Is

An autoencoder is a pair of networks — an encoder f: x → z and a decoder g: z → x̂ — trained so that g(f(x)) ≈ x. The latent z is smaller or otherwise constrained, so the only way to reconstruct x well is for the encoder to compress the signal in x into z.

A variational autoencoder (VAE) is an autoencoder with a probabilistic bottleneck. The encoder outputs a distribution over z rather than a single point, and the decoder is trained so that sampling z from that distribution and running the decoder gives plausible x. The VAE's extra machinery buys a generative model: you can sample z ~ N(0, I) and decode it.

This topic covers plain autoencoders as a bridge from classical dimensionality reduction to generative models, and VAEs as the minimum generative model every student should understand before opening a GAN or diffusion paper.

When You Use It

  • you want a nonlinear dimensionality reduction that respects structure PCA cannot see
  • you need a dense, low-dimensional representation for downstream tasks (anomaly detection, retrieval, clustering)
  • you want a latent space you can sample from (VAE)
  • you want a lightweight encoder that feeds a latent diffusion model (Stable Diffusion's first stage is exactly this)

Do Not Use It When

  • you need sharp, high-fidelity samples as a final product — VAEs produce characteristic blur; diffusion / GANs are sharper
  • you only need a linear projection — PCA is cheaper and usually sufficient
  • labels are available — supervised representation learning will almost always beat unsupervised autoencoders on downstream tasks

The Plain Autoencoder

A minimal autoencoder in PyTorch:

import torch
from torch import nn

class Autoencoder(nn.Module):
    def __init__(self, in_dim, latent_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(in_dim, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent_dim),
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, in_dim),
        )

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z), z

Train with reconstruction loss (MSE for continuous data, BCE for 0–1 data):

loss = ((x_hat - x) ** 2).mean()     # MSE for continuous features

Two things the plain AE teaches:

  • what is in the latent — inspect z for a batch; cluster it; project it with t-SNE/UMAP. If the latent is degenerate (all points collapse), the bottleneck was not tight enough or the decoder is memorizing
  • what reconstruction error means — plot the per-sample MSE. The top decile of reconstruction error is often where anomalies live; this is the cheapest anomaly-detection model you can build

Variants Worth Knowing

variant twist use case
denoising AE train to reconstruct x from a corrupted x + noise robust features, pretext for SSL
sparse AE L1 penalty on the latent activations interpretable features
contractive AE penalize the Jacobian of the encoder stable features on nearby inputs
convolutional AE convolution + transposed convolution images
vector-quantized AE (VQ-VAE) discrete latents via codebook speech synthesis, tokenizers for diffusion

For continuous-data generative work today, the two variants that actually ship are VAEs (and their β variants) and VQ-VAEs.

Why Plain Autoencoders Are Not Generative

Train a plain AE. Sample a random z from the same shape as your latent. Decode it. The output is almost always garbage. Why?

  • the plain AE learns which z values correspond to real data, but nothing constrains the rest of the latent space
  • nearby z values may map to wildly different outputs — the latent is not smooth
  • the distribution of z across the training data is not matched to any prior, so you cannot sample from it

Fixing these problems is the exact job of the VAE.

VAE — The Two Ingredients

A VAE makes two changes:

  1. the encoder outputs a distribution q(z | x) = N(μ(x), σ²(x) I) instead of a point
  2. training penalizes the KL divergence between q(z | x) and a prior p(z) = N(0, I)

The loss is the evidence lower bound (ELBO):

L_VAE = reconstruction_loss(x, decoded)  +  β · KL( q(z|x) || N(0, I) )

β = 1 is the classic VAE. Higher β pushes the posterior closer to the prior (smoother latent, blurrier reconstruction); lower β preserves reconstruction at the cost of a less usable latent.

The reparameterization trick keeps this differentiable:

z = μ + σ · ε,    ε ~ N(0, I)

so gradients flow through μ and σ rather than through the non-differentiable sampling step.

Minimal VAE In PyTorch

import torch
from torch import nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, in_dim, latent_dim):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(in_dim, 256), nn.ReLU())
        self.mu = nn.Linear(256, latent_dim)
        self.log_var = nn.Linear(256, latent_dim)
        self.dec = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.ReLU(),
            nn.Linear(256, in_dim),
        )

    def forward(self, x):
        h = self.enc(x)
        mu, log_var = self.mu(h), self.log_var(h)
        z = mu + torch.exp(0.5 * log_var) * torch.randn_like(mu)
        return self.dec(z), mu, log_var

def vae_loss(x, x_hat, mu, log_var, beta=1.0):
    recon = F.mse_loss(x_hat, x, reduction="sum") / x.size(0)
    kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / x.size(0)
    return recon + beta * kl

Sampling from a trained VAE:

with torch.no_grad():
    z = torch.randn(16, latent_dim)
    samples = vae.dec(z)

Two inspections you should run the first time:

  • the KL term over training should rise at first and then stabilize; if it collapses to zero, you have posterior collapse (see below)
  • samples from N(0, I) should be recognizable, blurry versions of your data

β-VAE And The Disentanglement Story

β-VAE (β > 1) was proposed as a way to force each latent dimension to capture a single interpretable factor of variation. The story is partly oversold — strict disentanglement is hard to guarantee — but β is a useful knob:

  • β small (0.1–0.5): sharper reconstructions, less structured latent
  • β = 1: classical VAE
  • β large (4–10): more structure, noticeably blurrier reconstructions, more sample diversity

Sweep β when the downstream use of the VAE depends on how usable the latent is.

VQ-VAE And The Diffusion Connection

Vector-quantized VAEs replace the continuous latent with a codebook lookup: the encoder produces a vector, the vector is snapped to its nearest codebook entry, and the decoder works off the codebook entry. The codebook is learned jointly with the encoder and decoder.

Why this matters today: latent diffusion models (Stable Diffusion) use a VAE (not VQ, but the spiritual descendant) as their first stage. The diffusion runs in the latent space of the VAE, not in pixel space, and the quality gain is large. The lesson: VAEs, used correctly, are the compression stage of modern generative systems rather than the final generator.

Posterior Collapse

The characteristic VAE failure is posterior collapse: the encoder stops using z and the decoder ignores it entirely. The KL term goes to zero, the reconstruction uses no latent, and the model becomes a glorified unconditional decoder.

Diagnostics:

  • KL term near zero while reconstruction loss is stable
  • decoder outputs are insensitive to the latent z
  • μ(x) and σ(x) are near-constant across different x

Fixes:

  • KL warmup — start with β ≈ 0 and ramp to 1 over the first few epochs
  • free bits — clip each dimension's KL contribution at a floor (e.g., 0.1 nats per dim) so the model cannot zero out the latent entirely
  • stronger encoder or weaker decoder — decoders that are too expressive can reconstruct well without z

What To Inspect

  • reconstruction loss — should plateau well above zero for complex data (zero means memorization or a too-weak bottleneck)
  • KL term — should be non-trivial and stable; zero signals collapse
  • latent usage — per-dimension KL contributions; dead dimensions (KL ≈ 0) indicate unused capacity
  • sample quality — qualitative check on decoded N(0, I) draws; VAE samples are characteristically soft
  • interpolation smoothness — decode z_t = (1 - t) z_1 + t z_2 for t ∈ [0, 1]; smooth interpolations are a sign of a well-shaped latent
  • anomaly-detection usefulness — reconstruction error percentile is a strong cheap anomaly score

Failure Pattern

Beyond posterior collapse, the classical "failure" of VAEs is blurry samples. This is not a bug; it is a direct consequence of the Gaussian likelihood and KL regularization. If sharp samples are required, a VAE is the wrong choice. A VQ-VAE + autoregressive prior or a diffusion model is the right replacement.

Common Mistakes

  • training a plain AE and then sampling random latents — produces garbage, not generations
  • skipping KL warmup — leads to posterior collapse
  • using a very deep decoder with a weak encoder — the decoder ignores the latent
  • reporting only reconstruction loss — says nothing about the usefulness of the latent
  • expecting sharp VAE samples without moving to a VAE + discrete prior or VAE + diffusion pipeline
  • treating PCA and a plain AE as interchangeable — they are not on nonlinear data
  • comparing AE/VAE representations only by reconstruction, never by downstream linear probing

Decision: Autoencoder vs. VAE vs. PCA vs. Diffusion

option primary use strength limitation
PCA linear compression, quick exploration fast, interpretable linear only
plain AE nonlinear compression, anomaly detection simple not generative
VAE smooth latent, soft generative model latent is structured blurry samples
VQ-VAE + autoregressive prior tokenizer for generation sharp samples through discrete prior more moving parts
diffusion (in VAE latent) state-of-the-art continuous generation best quality complex pipeline

If the task is classification-adjacent, prefer SSL (see Self-Supervised and Representation Learning) over AEs for representation learning; SSL encoders usually dominate downstream.

Practice

  1. Train a plain AE on MNIST with a 2-D latent. Scatter-plot the encoded test set coloured by digit. You should see clear clusters.
  2. Sample 16 random 2-D latents from a uniform distribution and decode them. Note how unstructured the outputs are — the latent is not generative.
  3. Swap the AE for a VAE with the same architecture and latent size. Retrain. Sample 16 latents from N(0, 1). Confirm that the samples now look like recognizable (if blurry) digits.
  4. Sweep β ∈ {0.1, 1, 4, 8}. For each, report reconstruction loss, KL, and one qualitative sample panel. Pick the one that looks most like a usable latent.
  5. Deliberately induce posterior collapse: drop the KL term entirely (set β = 0) for a short training run, then re-enable it fully. Watch the KL stay at zero because the decoder has learned to ignore the latent.
  6. Use reconstruction error as an anomaly score on a binary "digit vs. letter" test — train the VAE on digits only, score a mixed test set. Plot the ROC curve.

Runnable Example

A small VAE on MNIST trains in minutes on a laptop CPU and produces enough signal to run the full set of inspections. Install with pip install torch torchvision matplotlib. The snippet above is self-contained; wire it to MNIST via torchvision.datasets.MNIST.

Longer Connection

Autoencoders and VAEs sit at the crossroads of several academy topics:

For the decision frame — "should I be using a VAE at all?" — the answer is usually "as a component, not as the final model." VAEs are excellent latent compressors and mediocre final generators. Knowing both uses is the point of this topic.