Autoencoders & Variational Autoencoders (VAE)

The Problem

Compress a 784-pixel MNIST digit to a 16-number code, then reconstruct. A plain autoencoder will ace reconstruction MSE but the code space is a lumpy mess. Pick a random point in the code space, decode it, and you get noise. It has no sampler. It is a compression model dressed up.

What you actually want is: (a) the code space is a clean, smooth distribution you can sample from — say an isotropic Gaussian $\mathcal{N}(0, I)$, (b) decoding any sample produces a plausible digit, and (c) the encoder and decoder still compress well. Three goals, one architecture, one loss.

Kingma's 2013 VAE solves this by training the encoder to output a distribution $q(z \mid x) = \mathcal{N}(\mu(x), \sigma(x)^2)$, pulling that distribution toward the prior $\mathcal{N}(0, I)$ via a KL penalty, and then sampling $z$ from $q(z \mid x)$ before decoding. At inference time, drop the encoder, sample $z \sim \mathcal{N}(0, I)$, decode. The KL penalty is what forces the code space to be structured.

In 2026 VAEs rarely ship standalone — they have been outclassed by diffusion for raw image quality — but they are the encoder of choice for every latent-diffusion model (SD 1/2/XL/3, Flux, AudioCraft). Learn the VAE and you learn the invisible first layer of every image pipeline you use.

The Concept

Autoencoder. $z = \text{encoder}(x)$, $\hat{x} = \text{decoder}(z)$, loss = $\lVert x - \hat{x} \rVert^2$. Code space unstructured.

VAE encoder. Outputs two vectors: $\mu(x)$ and $\log \sigma^2(x)$. These define $q(z \mid x) = \mathcal{N}(\mu, \mathop{\mathrm{diag}}(\sigma^2))$.

Reparameterization trick. Sampling from $q(z \mid x)$ is not differentiable. Rewrite the sample as $z = \mu + \sigma \cdot \epsilon$ where $\epsilon \sim \mathcal{N}(0, I)$. Now $z$ is a deterministic function of $(\mu, \sigma)$ plus a non-parameter noise — gradients flow through μ and $\sigma$.

Loss. Evidence Lower BOund (ELBO), two terms:

\begin{aligned}
\text{loss} &= \text{reconstruction} + \beta \cdot \mathrm{KL}[q(z \mid x) \,\Vert\, \mathcal{N}(0, I)] \\
&= \lVert x - \hat{x} \rVert^2 + \beta \cdot \sum_i \frac{\sigma_i^2 + \mu_i^2 - \log \sigma_i^2 - 1}{2}
\end{aligned}

Reconstruction pushes $\hat{x}$ toward $x$. KL pushes $q(z \mid x)$ toward the prior. They trade off. Small $\beta$ ($<1$) = sharper samples, code space less Gaussian. Large $\beta$ ($>1$) = cleaner code space, blurrier samples. β-VAE (Higgins 2017) made this knob famous and kicked off disentanglement research.

Sampling. At inference: draw $z \sim \mathcal{N}(0, I)$, forward through decoder. One forward pass — no iterative sampling like diffusion.

vae-latent-grid

Build It

code/main.py implements a tiny VAE without numpy or torch. Input is 8-dimensional synthetic data drawn from a 2-component Gaussian mixture in 8-D. Encoder and decoder are single hidden-layer MLPs. We implement tanh activation, forward pass, loss, and a hand-written backward pass. Not production — pedagogy.

Step 1: encoder forward

def encode(x, enc):
    h = tanh(add(matmul(enc["W1"], x), enc["b1"]))
    mu = add(matmul(enc["W_mu"], h), enc["b_mu"])
    log_sigma2 = add(matmul(enc["W_sig"], h), enc["b_sig"])
    return mu, log_sigma2

$\log \sigma^2$ instead of $\sigma$ so the network output is unconstrained (softplus of σ is a trap — gradients die at $\sigma \approx 0$).

Step 2: reparameterize and decode

def reparameterize(mu, log_sigma2, rng):
    eps = [rng.gauss(0, 1) for _ in mu]
    sigma = [math.exp(0.5 * lv) for lv in log_sigma2]
    return [m + s * e for m, s, e in zip(mu, sigma, eps)]

def decode(z, dec):
    h = tanh(add(matmul(dec["W1"], z), dec["b1"]))
    return add(matmul(dec["W_out"], h), dec["b_out"])

Step 3: the ELBO

def elbo(x, x_hat, mu, log_sigma2, beta=1.0):
    recon = sum((a - b) ** 2 for a, b in zip(x, x_hat))
    kl = 0.5 * sum(math.exp(lv) + m * m - lv - 1 for m, lv in zip(mu, log_sigma2))
    return recon + beta * kl, recon, kl

Exact closed-form KL because both distributions are Gaussian. Do not integrate numerically. People still ship code with monte-carlo KL estimates in 2026 — it is 3x slower for no reason.

Step 4: generate

def sample(dec, z_dim, rng):
    z = [rng.gauss(0, 1) for _ in range(z_dim)]
    return decode(z, dec)

That is the generative model. Five lines.

Pitfalls

• Posterior collapse. KL term drives $q(z \mid x) \to \mathcal{N}(0, I)$ so aggressively that $z$ carries no info about $x$. Fix: β-annealing (start $\beta=0$, ramp to 1), free bits, or skip the KL on inactive dimensions.
• Blurry samples. The Gaussian decoder likelihood implies MSE reconstruction, which is Bayes-optimal for L2 (the mean) — the mean of a set of plausible digits is a fuzzy digit. Fix: discrete decoder (VQ-VAE, NVAE), or use the VAE only as an encoder and stack diffusion on the latents (this is what Stable Diffusion does).
• β too large, too early. See posterior collapse. Start at $\beta \approx 0.01$ and ramp.
• Latent dim too small. 16-D works for MNIST, 256-D for ImageNet 256², 2048-D for ImageNet 1024². Stable Diffusion's VAE compresses 512×512×3 → 64×64×4 (8x downsample per side, 64x in spatial area; channels go from 3 to 4, so about 48x fewer values overall).

Use It

The 2026 VAE stack:

A latent-diffusion model is a VAE with a diffusion model living between encoder and decoder. The VAE does coarse compression, the diffusion model does the heavy lifting. Same pattern for video (VAE + video-diffusion DiT) and audio (Encodec + MusicGen transformer).

Ship It

Save outputs/skill-vae-trainer.md.

Skill takes: dataset profile + latent-dim target + downstream use (reconstruction, sampling, or latent-diffusion input) and outputs: architecture choice (plain/β/VQ/RVQ), β schedule, latent dim, decoder likelihood (Gaussian vs categorical), and evaluation plan (recon MSE, KL per dim, Fréchet distance between $q(z \mid x)$ and $\mathcal{N}(0, I))$.

Exercises

1. Easy. Change $\beta$ in code/main.py to 0.01, 0.1, 1.0, 5.0. Record the final reconstruction MSE and KL. Which $\beta$ is Pareto-best for your synthetic data?
2. Medium. Replace the Gaussian decoder likelihood with a Bernoulli likelihood (cross-entropy loss). Compare sample quality on a binarized version of the same synthetic data.
3. Hard. Extend code/main.py into a mini VQ-VAE: replace the continuous $z$ with a nearest-neighbour lookup in a codebook of K=32 entries. Compare reconstruction MSE and report how many codebook entries get used (codebook collapse is real).

Key Terms

Production note: the VAE is the hottest path in a diffusion server

In a Stable Diffusion / Flux / SD3 pipeline the VAE is called twice per request — once to encode (if doing img2img / inpainting) and once to decode. At 1024² the decoder pass is often the single largest activation-memory peak in the whole pipeline because it upsamples 128×128×16 latents back to 1024×1024×3. Two practical consequences:

• Slice or tile the decode. diffusers exposes pipe.vae.enable_slicing() and pipe.vae.enable_tiling(). Tiling trades a small seam artifact for $O(\text{tile}^2)$ memory instead of $O(H \cdot W)$. Essential for 1024²+ on consumer GPUs.
• bf16 decoder, fp32 numerics for the final resize. The SD 1.x VAE was released in fp32 and silently produces NaNs when cast to fp16 at 1024²+. SDXL ships madebyollin/sdxl-vae-fp16-fix — always prefer the fp16-fix variant or use bf16.

Further Reading

• Kingma & Welling (2013). Auto-Encoding Variational Bayes — the VAE paper.
• Higgins et al. (2017). β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework — disentangled β-VAE.
• van den Oord et al. (2017). Neural Discrete Representation Learning — VQ-VAE.
• Vahdat & Kautz (2021). NVAE: A Deep Hierarchical Variational Autoencoder — state-of-the-art image VAE.
• Rombach et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models — Stable Diffusion; VAE as encoder.
• Défossez et al. (2022). High Fidelity Neural Audio Compression — Encodec, the audio VAE standard.