Autoencoders & VAEs

Background

Autoencoders as dimensionality reduction: Rumelhart et al. (1986) used backpropagation for non-linear PCA
Variational Autoencoders (VAEs) by Kingma & Welling (2013) added a probabilistic latent space enabling generation
Foundation for latent diffusion models (Stable Diffusion) and representation learning

The Concept

Autoencoders compress input to a latent code (encoder) then reconstruct it (decoder), learning a compact representation. VAEs add a constraint that the latent space follows a Gaussian distribution, enabling sampling and interpolation - you can generate new examples by sampling from the latent space.

How it Links

Latent diffusion models (Stable Diffusion) use VAE encoders/decoders around a diffusion process
The VAE ELBO loss connects to variational inference and KL divergence
Representation learning connects to SSL (self-supervised learning), contrastive learning, and CLIP

Concrete Example

A VAE trained on face images: the latent space organizes by age, gender, hairstyle, glasses. Sampling a point between two face latents produces a face that blends their features. Moving along the "age" direction in latent space makes the face look older or younger continuously.

Prerequisites

Multi-Layer Perceptron, Backpropagation & Optimizers, Math you need: linear algebra & calculus

Learning outcomes

Define the vanilla, denoising, and sparse autoencoders and explain when each helps.
Derive the ELBO from the marginal log-likelihood and the KL divergence between the variational posterior and the true posterior.
Justify the reparameterization trick: why we sample z = mu + sigma * eps instead of z ~ q(z|x), and how it enables backpropagation through stochastic nodes.
Derive the closed-form KL divergence between two diagonal Gaussians and apply it as the regularizer in a VAE.
Explain the beta-VAE disentanglement trade-off and the KL-vanishing failure mode (Bowman 2015).

TL;DR

An autoencoder learns a bottlenecked encoder-decoder pair $(g_\phi, f_\theta)$ that minimizes reconstruction error on its inputs [1]. Denoising autoencoders [2] corrupt the input and reconstruct the clean target, learning robust features. Sparse autoencoders add an L1 penalty on activations to learn sparse codes. The variational autoencoder (VAE) [3][4] places a probabilistic prior $p(z) = \mathcal{N}(0, I)$ on the latent and trains by maximizing the evidence lower bound (ELBO)

\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z)),

with the reparameterization trick $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$ for $\epsilon \sim \mathcal{N}(0, I)$ enabling gradients through the sample. beta-VAE [5] upweights the KL term to encourage disentangled factors at the cost of reconstruction sharpness. VAE samples are blurry because the per-pixel Gaussian likelihood is the wrong noise model for natural images, a fact diffusion models partially fix.

Intuition

An autoencoder is a network that learns to copy its input to its output through a bottleneck: the encoder $g_\phi : x \mapsto z$ compresses the input into a low-dimensional latent code $z$ , and the decoder $f_\theta : z \mapsto \hat x$ reconstructs $x$ from $z$ . The bottleneck forces the network to discover a compressed representation; once trained, the latent code $z$ is a learned feature vector for downstream tasks.

The plain autoencoder is deterministic: one input maps to one latent, and reconstructing $\hat x = f_\theta(g_\phi(x))$ is just a non-linear PCA-like operation. The variational autoencoder makes everything stochastic: the encoder outputs a distribution $q_\phi(z \mid x)$ over latents (typically a diagonal Gaussian), the decoder defines a distribution $p_\theta(x \mid z)$ (typically Gaussian or Bernoulli), and we train by maximizing a lower bound on $\log p_\theta(x)$ . The reward is a generative model: sample $z \sim p(z) = \mathcal{N}(0, I)$ , pass through the decoder, and you get a new “thought” from the same distribution as the training data.

What makes a VAE different from a plain autoencoder

A plain autoencoder produces an arbitrary code per example, with no constraint on what the code distribution looks like across the dataset. The codes can land anywhere in the latent space. To sample new data you would need to know which regions of latent space are plausible.

A VAE imposes a prior $p(z) = \mathcal{N}(0, I)$ on the latent distribution and pulls the encoder’s output toward it via a KL term. After training, the marginal posterior $\int q_\phi(z \mid x) p_{\text{data}}(x)\, dx$ is approximately $\mathcal{N}(0, I)$ , so to generate new examples you sample from $\mathcal{N}(0, I)$ and decode. The KL term is what turns “compressor” into “generative model” [3].

Vanilla and denoising autoencoders

A vanilla autoencoder with squared loss is

\mathcal{L}_{\text{vanilla}} = \frac{1}{N} \sum_i \lVert x_i - f_\theta(g_\phi(x_i)) \rVert_2^2.

If the encoder and decoder are linear and the loss is squared error, the optimum recovers PCA: $g_\phi$ projects onto the top- $k$ principal directions and $f_\theta$ is the transpose. With non-linear encoders/decoders the autoencoder generalizes PCA to arbitrary manifolds [1].

A denoising autoencoder [2] introduces a corruption distribution $\tilde x \sim q(\tilde x \mid x)$ (Gaussian noise, masking, salt-and-pepper) and trains the model to reconstruct the clean $x$ from the corrupted $\tilde x$ :

\mathcal{L}_{\text{denoising}} = \frac{1}{N} \sum_i \mathbb{E}_{\tilde x_i \sim q(\cdot \mid x_i)}\!\bigl[\, \lVert x_i - f_\theta(g_\phi(\tilde x_i)) \rVert_2^2 \bigr].

The model cannot just memorize the input; it has to learn structure that survives noise. Empirically denoising autoencoders learn features that transfer better to downstream tasks than vanilla autoencoders, and BERT’s masked-language-modelling objective is exactly denoising at the level of tokens.

A sparse autoencoder keeps the vanilla reconstruction loss but adds an L1 penalty on the latent activations:

\mathcal{L}_{\text{sparse}} = \frac{1}{N} \sum_i \lVert x_i - f_\theta(g_\phi(x_i)) \rVert_2^2 + \lambda \lVert g_\phi(x_i) \rVert_1.

The L1 pulls most latent units to zero, leaving only a few active per example. The result is a sparse, often more interpretable code; this is the textbook setup that motivated dictionary-learning and the mainstream sparse-coding literature.

The variational autoencoder

The full ELBO derivation (interview gold)

Setup

We have a generative model $p_\theta(x, z) = p(z)\, p_\theta(x \mid z)$ with prior $p(z) = \mathcal{N}(0, I)$ and decoder $p_\theta(x \mid z)$ (e.g. Gaussian or Bernoulli). We want maximum likelihood: maximize $\log p_\theta(x)$ over $\theta$ . The trouble is that

p_\theta(x) = \int p_\theta(x, z)\, dz = \int p(z)\, p_\theta(x \mid z)\, dz

is intractable, and the true posterior $p_\theta(z \mid x)$ is intractable for the same reason.

Step 1: introduce a variational posterior

Introduce a tractable approximation $q_\phi(z \mid x)$ to the true posterior. The encoder network produces the parameters of $q_\phi$ (e.g. mean and log-variance of a diagonal Gaussian).

Step 2: identity for $\log p_\theta(x)$

For any distribution $q$ that has the same support as $p$ ,

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x)],

because the expectation of a constant (in $z$ ) is the constant. Now use $p_\theta(x) = p_\theta(x, z) / p_\theta(z \mid x)$ :

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{p_\theta(z \mid x)} \right].

Step 3: insert and split with $q_\phi$

Multiply and divide by $q_\phi(z \mid x)$ inside the log:

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{q_\phi(z \mid x)} \cdot \frac{q_\phi(z \mid x)}{p_\theta(z \mid x)} \right].

Split the log into two terms:

\log p_\theta(x) = \underbrace{\mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{q_\phi(z \mid x)} \right]}_{\mathcal{L}(\theta, \phi; x)\, =\, \text{ELBO}} \,+\, \underbrace{\mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{q_\phi(z \mid x)}{p_\theta(z \mid x)} \right]}_{\mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p_\theta(z \mid x))}.

Both pieces have a name. The first is the evidence lower bound, $\mathcal{L}(\theta, \phi; x)$ ; the second is the KL divergence from the variational posterior to the true posterior.

Step 4: KL is non-negative, so ELBO $\le \log p_\theta(x)$

By Gibbs’ inequality $\mathrm{KL}(q \,\Vert\, p) \ge 0$ , with equality iff $q = p$ almost everywhere. Therefore

\boxed{\; \log p_\theta(x) \,\ge\, \mathcal{L}(\theta, \phi; x) \;=\; \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{q_\phi(z \mid x)} \right]. \;}

This is the evidence lower bound (ELBO). Maximizing the ELBO over $(\theta, \phi)$ does two things at once: it pushes $\log p_\theta(x)$ up (the gap closes from below) and it tightens the bound by minimizing $\mathrm{KL}(q_\phi \,\Vert\, p_\theta(z \mid x))$ (the variational posterior approaches the true posterior).

Step 5: split the ELBO into reconstruction and prior-matching

Use $p_\theta(x, z) = p_\theta(x \mid z)\, p(z)$ :

\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x \mid z) p(z)}{q_\phi(z \mid x)} \right] = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z)).

This is the standard form of the ELBO that the VAE optimizes:

$\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)]$ is the reconstruction term: how well the decoder explains $x$ given a code drawn from the encoder.
$\mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z))$ is the prior-matching term: how close the encoder’s output distribution is to the prior $\mathcal{N}(0, I)$ .

The VAE training objective for a minibatch is

\hat{\mathcal{L}} = \frac{1}{|B|}\sum_{x \in B} \Bigl( \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z)) \Bigr).

The reconstruction expectation is estimated with a single Monte-Carlo sample $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$ ; the KL has a closed form for diagonal Gaussians (next DeepDive).

Reparameterization trick: gradients through a stochastic node

The problem

The ELBO contains $\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)]$ . To optimize over $\phi$ we need

\nabla_\phi \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)].

If we sample $z \sim q_\phi(z \mid x)$ directly, the sampling step has no gradient (it is just an index into a random number stream), so the standard chain rule does not apply through $z$ .

REINFORCE-style estimator (high variance)

The score-function (REINFORCE) estimator handles this with the identity $\nabla_\phi q = q\, \nabla_\phi \log q$ :

\nabla_\phi \mathbb{E}_q[f(z)] = \mathbb{E}_q[\, f(z) \nabla_\phi \log q_\phi(z \mid x) \,].

This works for any $q$ , including discrete distributions, but its variance is large: the gradient signal is whatever the log-likelihood ratio happens to be at the sample, scaled by $f$ . Plain REINFORCE without baselines is too noisy to train a VAE in reasonable time.

Reparameterization

For continuous $q_\phi$ that can be written as a deterministic transformation of a parameter-free noise distribution, do the change of variables. For a diagonal Gaussian:

q_\phi(z \mid x) = \mathcal{N}(z; \mu_\phi(x), \mathrm{diag}(\sigma_\phi(x)^2)) \;\Longleftrightarrow\; z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I).

The noise $\epsilon$ does not depend on $\phi$ ; the transformation $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$ does. Now

\mathbb{E}_{q_\phi(z \mid x)}[f(z)] = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)}[f(\mu_\phi(x) + \sigma_\phi(x) \odot \epsilon)],

and the gradient passes through $f$ and the transformation by ordinary chain rule:

\nabla_\phi \mathbb{E}_q[f(z)] = \mathbb{E}_\epsilon\!\bigl[ \nabla_z f(z) \cdot \nabla_\phi (\mu_\phi(x) + \sigma_\phi(x) \odot \epsilon) \bigr].

This is the reparameterization trick [3][4]. The variance is much lower than REINFORCE because the gradient uses the shape of $f$ (its derivative), not just its value at the sample.

Why $\sigma$ is parameterized as $\log \sigma^2$ in practice

To keep $\sigma$ positive, the encoder outputs $\log \sigma_\phi^2(x)$ (a real-valued network output) and we recover $\sigma_\phi(x) = \exp(\tfrac{1}{2} \log \sigma^2)$ . This avoids softplus or absolute-value tricks and gives a well-behaved gradient near $\sigma \to 0$ .

Beyond Gaussian

The same trick works for any distribution that is the pushforward of a parameter-free base distribution under a smooth map: Gumbel-Softmax for categorical, planar/radial flows, normalizing flows in general. Discrete distributions need REINFORCE or relaxations.

Closed-form KL between two diagonal Gaussians

Statement

For $q = \mathcal{N}(\mu_1, \mathrm{diag}(\sigma_1^2))$ and $p = \mathcal{N}(\mu_2, \mathrm{diag}(\sigma_2^2))$ in $\mathbb{R}^d$ :

\mathrm{KL}(q \,\Vert\, p) = \frac{1}{2} \sum_{j=1}^{d} \left( \log \frac{\sigma_{2,j}^2}{\sigma_{1,j}^2} + \frac{\sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2}{\sigma_{2,j}^2} - 1 \right).

For the VAE special case $p = \mathcal{N}(0, I)$ , so $\mu_2 = 0$ and $\sigma_{2,j}^2 = 1$ :

\boxed{\; \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, \mathcal{N}(0, I)) = \frac{1}{2} \sum_{j=1}^{d} \bigl( \mu_{1,j}^2 + \sigma_{1,j}^2 - 1 - \log \sigma_{1,j}^2 \bigr). \;}

This is what the VAE actually computes for the KL term; no Monte Carlo needed.

Derivation

Start from the definition of KL and substitute the Gaussian density:

\mathrm{KL}(q \,\Vert\, p) = \int q(z) \log \frac{q(z)}{p(z)}\, dz = \mathbb{E}_q[\log q(z)] - \mathbb{E}_q[\log p(z)].

For a diagonal Gaussian, $\log q(z) = -\tfrac{1}{2} \sum_j [\log(2\pi \sigma_{1,j}^2) + (z_j - \mu_{1,j})^2 / \sigma_{1,j}^2]$ . Take the expectation under $q$ , using $\mathbb{E}_q[(z_j - \mu_{1,j})^2] = \sigma_{1,j}^2$ :

\mathbb{E}_q[\log q(z)] = -\frac{1}{2} \sum_j (\log(2\pi \sigma_{1,j}^2) + 1) = -\frac{d}{2}\log(2\pi) - \frac{1}{2} \sum_j (\log \sigma_{1,j}^2 + 1).

For $\log p(z)$ , expand the squared term: $\mathbb{E}_q[(z_j - \mu_{2,j})^2] = \sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2$ (variance plus mean shift). So

\mathbb{E}_q[\log p(z)] = -\frac{1}{2} \sum_j \log(2\pi \sigma_{2,j}^2) - \frac{1}{2} \sum_j \frac{\sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2}{\sigma_{2,j}^2}.

Subtract:

\mathrm{KL}(q \,\Vert\, p) = \frac{1}{2} \sum_j \log \frac{\sigma_{2,j}^2}{\sigma_{1,j}^2} + \frac{1}{2} \sum_j \frac{\sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2}{\sigma_{2,j}^2} - \frac{d}{2}.

Plug in $p = \mathcal{N}(0, I)$ to recover the boxed formula [9][6].

beta-VAE and the disentanglement trade-off

The standard VAE objective is the ELBO. Higgins et al. [5] proposed multiplying the KL term by a tunable coefficient $\beta$ :

\mathcal{L}_\beta = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \beta\, \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z)).

$\beta = 1$ is the standard VAE.
$\beta > 1$ strengthens the prior-matching pressure. Each latent dimension is pushed to be more independent from the others (because $p(z) = \mathcal{N}(0, I)$ is itself factorized), which encourages the latent code to align with disentangled factors of variation in the data: rotation, scale, lighting.
$\beta < 1$ relaxes the prior; the code carries more information about $x$ , reconstructions sharpen, but sampling quality suffers.

The trade-off is sharp: at $\beta = 4$ on dSprites, $\beta$ -VAE produces a code where each latent corresponds to one interpretable factor, but reconstructions are blurrier than at $\beta = 1$ . In practice, $\beta$ is tuned on the downstream task; for representation learning $\beta \in [2, 8]$ is typical, for image generation $\beta = 1$ .

Connection to the information bottleneck

The $\beta$ -VAE objective is mathematically equivalent to a Lagrangian on the information bottleneck principle [7]:

\max_\phi I(z; x) - \beta\, I(z; \text{nuisance}),

where $I$ denotes mutual information. The KL term in the VAE is an upper bound on $I(z; x) - \mathrm{const}$ measured against the prior, and pushing the KL down corresponds to compressing $z$ to remove variance unrelated to the reconstruction objective. The ” $\beta$ as compression dial” intuition is the practical takeaway.

Implementations

A canonical PyTorch VAE on MNIST and the equivalent in TensorFlow Keras.

python

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

class VAE(nn.Module):
  """Diagonal-Gaussian VAE on flattened MNIST."""

  def __init__(self, in_dim: int = 784, h_dim: int = 400, z_dim: int = 20) -> None:
      super().__init__()
      # Encoder.
      self.fc1 = nn.Linear(in_dim, h_dim)
      self.fc_mu = nn.Linear(h_dim, z_dim)
      self.fc_logvar = nn.Linear(h_dim, z_dim)
      # Decoder.
      self.fc2 = nn.Linear(z_dim, h_dim)
      self.fc3 = nn.Linear(h_dim, in_dim)

  def encode(self, x):
      h = F.relu(self.fc1(x))
      return self.fc_mu(h), self.fc_logvar(h)

  def reparameterize(self, mu, logvar):
      # z = mu + sigma * eps, where sigma = exp(0.5 * logvar).
      std = torch.exp(0.5 * logvar)
      eps = torch.randn_like(std)
      return mu + std * eps

  def decode(self, z):
      h = F.relu(self.fc2(z))
      return torch.sigmoid(self.fc3(h))   # Bernoulli mean over pixels.

  def forward(self, x):
      mu, logvar = self.encode(x)
      z = self.reparameterize(mu, logvar)
      x_recon = self.decode(z)
      return x_recon, mu, logvar


def vae_loss(x, x_recon, mu, logvar, beta: float = 1.0):
  # Bernoulli reconstruction; sum over pixels, mean over batch.
  bce = F.binary_cross_entropy(x_recon, x, reduction='sum')
  # Closed-form KL between q and N(0, I), summed over latent dims.
  kl = -0.5 * torch.sum(1.0 + logvar - mu.pow(2) - logvar.exp())
  return (bce + beta * kl) / x.shape[0]


# Training step (assuming x is a flattened MNIST minibatch in [0, 1]):
# vae = VAE().to(device); opt = torch.optim.Adam(vae.parameters(), lr=1e-3)
# for x, _ in loader:
#     x = x.to(device).view(x.size(0), -1)
#     x_recon, mu, logvar = vae(x)
#     loss = vae_loss(x, x_recon, mu, logvar, beta=1.0)
#     opt.zero_grad(); loss.backward(); opt.step()

Two heads on the encoder for mu and log-variance; reparameterize via z = mu + std * eps; closed-form KL against N(0, I).

Run a tiny denoising autoencoder in NumPy on a flat random vector. The model has to remove additive Gaussian noise from a fixed clean target:

Idle (packages: numpy)

import numpy as np

rng = np.random.default_rng(0)

# Clean signal: a random 16-dim vector that the autoencoder should reconstruct
# from noisy versions of itself. With a 6-dim bottleneck and a single example,
# this is a toy denoising task.
d_in = 16
d_h = 6
clean = rng.standard_normal(d_in)

# Tiny linear-ish autoencoder with a tanh non-linearity.
def init():
  W1 = rng.normal(0, 0.5, size=(d_h, d_in))
  b1 = np.zeros(d_h)
  W2 = rng.normal(0, 0.5, size=(d_in, d_h))
  b2 = np.zeros(d_in)
  return [W1, b1, W2, b2]

def forward(x, params):
  W1, b1, W2, b2 = params
  z = np.tanh(W1 @ x + b1)
  xr = W2 @ z + b2
  return xr, z

params = init()
lr = 0.05
losses = []
for it in range(2000):
  # Resample noise each step: this is what makes it a DENOISING autoencoder.
  noisy = clean + 0.3 * rng.standard_normal(d_in)
  xr, z = forward(noisy, params)
  err = xr - clean                    # MSE wrt the CLEAN target.
  # Manual backprop.
  dW2 = np.outer(err, z); db2 = err
  dz = params[2].T @ err
  dpre = dz * (1 - z * z)             # tanh' = 1 - tanh^2
  dW1 = np.outer(dpre, noisy); db1 = dpre
  grads = [dW1, db1, dW2, db2]
  params = [p - lr * g for p, g in zip(params, grads)]
  if it % 400 == 0:
      losses.append(np.mean(err * err))
      print(f"step {it:4d}  recon MSE on clean = {losses[-1]:.4f}")

# Evaluate on a freshly noisy version of the clean signal.
noisy = clean + 0.3 * rng.standard_normal(d_in)
xr, _ = forward(noisy, params)
print()
print("clean[:6] :", np.round(clean[:6], 3))
print("noisy[:6] :", np.round(noisy[:6], 3))
print("recon[:6] :", np.round(xr[:6], 3))

(no output yet)

Comparison

Model	Loss	Supports sampling?	Disentanglement	Output sharpness	Notes
Vanilla AE	reconstruction MSE	no (no prior on z)	no	sharp on training data	Reduces to PCA when linear; non-linear AEs generalize PCA.
Denoising AE	reconstruction MSE on (noisy -> clean) pairs	no (deterministic)	no	sharp	Robust features; basis of masked-LM (BERT-style) pretraining.
Sparse AE	reconstruction + L1 on z	no	partial (active units)	sharp	Learns sparse codes; classical motivation for dictionary learning.
VAE	ELBO = recon - KL	yes (sample from prior)	no (entangled latents)	blurry (Gaussian likelihood)	Standard generative AE; mode-covering due to KL(q \|\| p) form.
beta-VAE (beta > 1)	recon - beta * KL	yes	yes (compresses + factorizes)	blurrier than VAE	Higgins 2017; trade reconstruction for disentanglement.

Real-world uses

Domain	Use	Why an autoencoder works
Denoising	image, audio, signal	Bottleneck enforces a clean reconstruction; the network learns the data manifold and projects noise off it.
Anomaly detection	fraud, manufacturing defects, log analytics	Train on "normal" only; high reconstruction error at inference flags an outlier.
Representation learning	pretraining for downstream classifiers	Latent code captures a compressed, often more linearly separable, view of the input.
Image generation (pre-2020)	sampling from a learned latent	VAEs were one of the two main generative models pre-diffusion; now mostly displaced by diffusion + Latent Diffusion (which uses a VAE encoder under the hood).
Recommender / embedding compression	compress sparse user/item vectors	Dense bottleneck makes downstream similarity search and ranking cheap.

Pitfalls

VAE samples are blurry

The standard VAE decoder $p_\theta(x \mid z)$ is a diagonal Gaussian, which is equivalent to per-pixel MSE between the decoder mean and the target. MSE assumes Gaussian noise, which is wrong for natural images: the actual residuals are heavy-tailed and structured. The MSE-optimal decoder under uncertainty is the conditional mean, which averages over plausible images and produces a blurry result.

Two fixes beyond plain VAE:

Stronger likelihoods: pixelCNN/PixelSNAIL decoders (autoregressive over pixels), or a VAE-GAN hybrid where a discriminator provides perceptual loss.
Use the VAE as a feature compressor: train a VAE to land in a clean latent and put a separate generative model (diffusion, autoregressive) on the latent. This is exactly what Latent Diffusion / Stable Diffusion does.

KL-vanishing (posterior collapse)

If the decoder $p_\theta(x \mid z)$ is too powerful, the model can learn to ignore $z$ entirely: it sets $q_\phi(z \mid x) \approx p(z)$ (so the KL term goes to zero) and uses the decoder’s autoregressive capacity to fit $p_\theta(x)$ directly. The latent code carries no information; the VAE is degenerate. Bowman et al. [8] first documented this on text VAEs with LSTM decoders.

Standard mitigations: KL annealing (start training with the KL term scaled down and ramp it up), free-bits (clip the per-dimension KL to a minimum so the model cannot collapse all dimensions to the prior), or weaker decoders. KL-vanishing is a known recurring failure mode that requires explicit attention.

beta-VAE trade-off is sharp

Tuning $\beta > 1$ buys disentanglement but spends reconstruction sharpness; $\beta < 1$ buys reconstruction sharpness but spends sample quality and disentanglement. Always evaluate both objectives (reconstruction error AND a disentanglement metric like MIG) when choosing $\beta$ , not one alone [5].

Use log-variance, not variance

Parameterize the encoder to output $\log \sigma^2$ , not $\sigma^2$ directly: it keeps $\sigma$ positive without softplus tricks, the gradient is well-behaved near $\sigma \to 0$ , and the closed-form KL formula is in $\log \sigma^2$ already.

Pop quiz

Check yourself

Score: 0 / 5(0 answered)

1. Starting from log p(x), the ELBO is derived by:
2. Why do we need the reparameterization trick for VAEs?
3. Why are VAE samples blurry?
4. In beta-VAE, increasing beta:
5. The closed-form KL between q = N(mu, diag(sigma^2)) and p = N(0, I) is:

Interview Q&A

Derive the ELBO from the marginal log-likelihood and KL divergence.research

We want to maximize $\log p_\theta(x)$ but the marginal $p_\theta(x) = \int p(z)\, p_\theta(x \mid z)\, dz$ is intractable. Introduce a variational posterior $q_\phi(z \mid x)$ and use that the expectation of a constant-in- $z$ quantity under any distribution is the constant:

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x)].

Substitute $p_\theta(x) = p_\theta(x, z) / p_\theta(z \mid x)$ :

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{p_\theta(z \mid x)} \right].

Insert and divide by $q_\phi(z \mid x)$ :

\log p_\theta(x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{p_\theta(x, z)}{q_\phi(z \mid x)} \right] + \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log \frac{q_\phi(z \mid x)}{p_\theta(z \mid x)} \right].

The first term is the ELBO $\mathcal{L}(\theta, \phi; x)$ ; the second is $\mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p_\theta(z \mid x)) \ge 0$ . Therefore $\log p_\theta(x) \ge \mathcal{L}(\theta, \phi; x)$ , with equality iff $q_\phi$ matches the true posterior.

Splitting the joint $p_\theta(x, z) = p_\theta(x \mid z) p(z)$ rewrites the ELBO as

\mathcal{L}(\theta, \phi; x) = \underbrace{\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)]}_{\text{reconstruction}} - \underbrace{\mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z))}_{\text{prior matching}}.

The reconstruction term is estimated by a single Monte Carlo sample via the reparameterization trick; the KL term has a closed form when $q$ and $p$ are both diagonal Gaussian [3][6].

Why do we need the reparameterization trick instead of just sampling z?onsite

The ELBO contains $\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)]$ . To optimize over the encoder parameters $\phi$ we need

\nabla_\phi \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)].

If we sample $z \sim q_\phi$ directly, the sample is a black-box random number with no gradient path back through $\phi$ .

REINFORCE would fix this with $\nabla_\phi \mathbb{E}_q[f(z)] = \mathbb{E}_q[f(z) \nabla_\phi \log q]$ , but the variance is enormous because the gradient depends on the value of $f$ at the sample, not on its shape. Training a VAE with REINFORCE without baselines is impractical.

Reparameterization rewrites $z$ as a deterministic function of an external noise variable: $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$ , $\epsilon \sim \mathcal{N}(0, I)$ . The randomness lives in $\epsilon$ , which is independent of $\phi$ . The expectation becomes

\mathbb{E}_{q_\phi(z \mid x)}[f(z)] = \mathbb{E}_\epsilon[f(\mu_\phi(x) + \sigma_\phi(x) \odot \epsilon)],

and ordinary backprop flows through $f$ , then through the linear transformation, then into $\phi$ . The gradient uses the derivative of $f$ , not its raw value, which gives much lower variance.

The trick generalizes to any distribution that can be written as the pushforward of a parameter-free base distribution under a smooth map: planar/radial flows, Gumbel-Softmax for categorical, normalizing flows in general. Discrete or heavy-tailed distributions still need REINFORCE or relaxations [4].

Why are VAE samples blurry, and what fixes the problem?onsite

The standard VAE decoder is $p_\theta(x \mid z) = \mathcal{N}(\mu_\theta(z), \sigma^2 I)$ , equivalent to per-pixel MSE between the decoder mean $\mu_\theta(z)$ and the target $x$ . Per-pixel MSE assumes pixel residuals are independent Gaussian with constant variance. That is wrong for natural images: residuals are heavy-tailed, spatially correlated, and structured (edges, textures).

Under the wrong noise model, the loss-optimal decoder at any point in latent space is the conditional mean of the data conditional on $z$ . The conditional mean averages over plausible images consistent with the latent code (e.g. a digit at slightly different positions/styles), and the average of plausible images is blurry. There is no “blurry photograph” in the data; blur is an artifact of MSE optimality.

Three classes of fix:

Stronger likelihoods. Replace the diagonal Gaussian decoder with an autoregressive PixelCNN/PixelSNAIL or a discretized logistic mixture. The model now puts probability mass on actual sharp images, not their average.
Adversarial / perceptual loss (VAE-GAN). Train a discriminator that distinguishes real from reconstructed images, and replace pixel MSE with a feature-space loss from a fixed perceptual network. This pushes reconstructions toward the true image manifold rather than its mean.
Use the VAE as a compressor, put a generative model on the latent. This is the Latent Diffusion / Stable Diffusion architecture: a VAE compresses pixels to a small latent, a diffusion model is trained on the latent, and decoding back to pixels uses the (sharp) VAE decoder. The diffusion handles the mode-seeking; the VAE handles the compression.

The blurriness is a known property of the standard VAE objective, not a bug to debug.

What is the beta-VAE trade-off and how do you tune beta in practice?onsite

The beta-VAE objective [5] is

\mathcal{L}_\beta = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \beta\, \mathrm{KL}(q_\phi(z \mid x) \,\Vert\, p(z)).

The KL term is upper-bounded by mutual information $I(z; x)$ measured against the prior. Multiplying by $\beta$ changes the trade-off:

$\beta = 1$ : standard VAE, optimizes the ELBO exactly.
$\beta > 1$ : stronger compression of $z$ . Forces the encoder’s output distribution toward the factorized prior $\mathcal{N}(0, I)$ , encouraging each latent dimension to be independent and align with one disentangled factor of variation in the data. Reconstruction sharpness drops because the code carries less information.
$\beta < 1$ : weaker compression. Code carries more information, reconstructions sharpen, but sampling from the prior becomes less faithful and disentanglement disappears.

Tuning in practice:

Pick the metric that matters: reconstruction error for compression, sample FID/IS for generation, mutual information gap (MIG) or modularity for disentanglement.
Sweep $\beta$ across $\{0.5, 1, 2, 4, 8\}$ on the train+val split. Plot the trade-off curve.
Beware KL-vanishing at very small $\beta$ : even $\beta = 0$ does not guarantee informative latents because the decoder may simply ignore $z$ (Bowman 2015 [8]). Always check the per-dim KL after training; if many dims are at zero, the latent is collapsing.
Production recipe for representation learning: $\beta \in [2, 8]$ with KL annealing (start near 0, ramp to target).
Production recipe for generation: $\beta = 1$ with a stronger decoder (PixelCNN, GAN, diffusion on the latent).

The tradeoff is sharp: there is no $\beta$ that gives both perfect reconstruction and full disentanglement on natural data, only a Pareto curve.

Derive the closed-form KL divergence between two diagonal Gaussians.onsite

For $q = \mathcal{N}(\mu_1, \mathrm{diag}(\sigma_1^2))$ and $p = \mathcal{N}(\mu_2, \mathrm{diag}(\sigma_2^2))$ in $\mathbb{R}^d$ , the KL is

\mathrm{KL}(q \,\Vert\, p) = \mathbb{E}_q[\log q(z)] - \mathbb{E}_q[\log p(z)].

For a diagonal Gaussian,

\log q(z) = -\frac{d}{2} \log(2\pi) - \frac{1}{2} \sum_j \log \sigma_{1,j}^2 - \frac{1}{2} \sum_j \frac{(z_j - \mu_{1,j})^2}{\sigma_{1,j}^2}.

Take the expectation under $q$ , using $\mathbb{E}_q[(z_j - \mu_{1,j})^2] = \sigma_{1,j}^2$ :

\mathbb{E}_q[\log q(z)] = -\frac{d}{2} \log(2\pi) - \frac{1}{2} \sum_j \log \sigma_{1,j}^2 - \frac{d}{2}.

For $\log p(z)$ , use $\mathbb{E}_q[(z_j - \mu_{2,j})^2] = \sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2$ (variance plus squared mean shift):

\mathbb{E}_q[\log p(z)] = -\frac{d}{2} \log(2\pi) - \frac{1}{2} \sum_j \log \sigma_{2,j}^2 - \frac{1}{2} \sum_j \frac{\sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2}{\sigma_{2,j}^2}.

Subtract:

\mathrm{KL}(q \,\Vert\, p) = \frac{1}{2} \sum_{j=1}^{d} \left( \log \frac{\sigma_{2,j}^2}{\sigma_{1,j}^2} + \frac{\sigma_{1,j}^2 + (\mu_{1,j} - \mu_{2,j})^2}{\sigma_{2,j}^2} - 1 \right).

For the VAE special case $p = \mathcal{N}(0, I)$ ( $\mu_2 = 0$ , $\sigma_{2,j}^2 = 1$ ):

\mathrm{KL}(q \,\Vert\, \mathcal{N}(0, I)) = \frac{1}{2} \sum_{j=1}^{d} \bigl( \mu_{1,j}^2 + \sigma_{1,j}^2 - 1 - \log \sigma_{1,j}^2 \bigr).

This is what the VAE actually computes; no Monte Carlo needed. The corresponding code is -0.5 * (1 + logvar - mu**2 - logvar.exp()).sum() [9][3].

Build it yourself

A weekend project that takes a VAE from textbook to working:

Implement the PyTorch VAE above on MNIST. Train for 20 epochs at $\beta = 1$ . Check that the per-dim KL after training is non-zero on most dimensions (no collapse).
Plot a 2D latent slice: encode a few training images, then decode points along a line in latent space. The decoder should produce a smooth interpolation (e.g. 3 morphing into 8). This is the qualitative test that the latent is well-organized.
Sweep $\beta \in \{0.5, 1, 2, 4, 8\}$ . Record reconstruction error and sample quality at each. You should see the Pareto curve appear.
Trigger KL-vanishing on purpose: replace the decoder with a powerful autoregressive PixelCNN (or a deep ResNet) and train at $\beta = 1$ . Watch the per-dim KL collapse to ~0 and the latent become uninformative. Add KL annealing as the fix.
Stretch: train a denoising autoencoder on MNIST with masked input (set 50% of pixels to 0). Compare features extracted from the bottleneck against the VAE’s mean encoding for downstream classification. Denoising features often transfer better.

References

Hinton, G. E. and Salakhutdinov, R. R. (2006). Reducing the Dimensionality of Data with Neural Networks. Science 313(5786), 504-507. doi:10.1126/science.1127647
Vincent, P., Larochelle, H., Bengio, Y., and Manzagol, P.-A. (2008). Extracting and composing robust features with denoising autoencoders. ICML 2008. doi:10.1145/1390156.1390294
Kingma, D. P. and Welling, M. (2013). Auto-Encoding Variational Bayes. ICLR 2014 (arXiv preprint 2013). arXiv:1312.6114
Rezende, D. J., Mohamed, S., and Wierstra, D. (2014). Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML 2014. arXiv:1401.4082
Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. (2017). beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017. https://openreview.net/forum?id=Sy2fzU9gl
Doersch, C. (2016). Tutorial on Variational Autoencoders. arXiv preprint. arXiv:1606.05908
Tishby, N. and Zaslavsky, N. (2015). Deep Learning and the Information Bottleneck Principle. IEEE Information Theory Workshop. arXiv:1503.02406
Bowman, S. R., Vilnis, L., Vinyals, O., Dai, A. M., Jozefowicz, R., and Bengio, S. (2015). Generating Sentences from a Continuous Space. CoNLL 2016 (arXiv preprint 2015). arXiv:1511.06349
Bishop, C. M. (2006). Pattern Recognition and Machine Learning. Springer. ISBN 978-0-387-31073-2

Interactive Demo

Input

Latent Space (2D)

Latent noise: 0.00

Output (reconstructed)

Circle

Tip: Use Interpolate mode and slide between two shapes to see a smooth latent space transition

RNNs, LSTMs, GRUs Generative Adversarial Networks