Module 4 · Reading 4.1 · read before S23: Autoencoders & VAEs

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

Everything you have trained so far maps an input to a target: an image to a label, a prefix to the next token. Generative modeling asks a different question: can a network learn the data distribution itself, well enough to produce new samples from it? The three-session generative arc this reading opens (VAEs, then GANs, then diffusion — the models powering today’s image generators) starts somewhere much smaller: a network trained to copy its input through a bottleneck.

That sounds pointless (the identity function needs no neural network) until you make the middle of the network narrower than the input. Then the network cannot copy; it must compress, and to compress well it must discover the structure of the data. The autoencoder gives us the idea of a latent space. The variational autoencoder (VAE) then turns that latent space into something you can sample from, and in doing so introduces a trick, reparameterization, that you will meet again in every corner of modern generative modeling.

How to read this

The bottleneck forces structure

An autoencoder is two networks composed: an encoder \(z = f(x)\) that maps the input to a low-dimensional code, and a decoder \(\hat{x} = g(z)\) that maps the code back. Training minimizes reconstruction error, typically \(\lVert x - g(f(x))\rVert^2\), on the training data. No labels are involved; the data supervises itself.

The idea is old. Networks trained to reproduce their input under a bottleneck (then called autoassociative networks) were studied in the 1980s, and Bourlard and Kamp showed in 1988 that with linear layers the optimal solution spans the same subspace as principal component analysis (PCA). In other words, a linear autoencoder cannot learn anything PCA does not already give you. The interesting case is nonlinear: a curved data manifold that no linear projection can flatten.

You can see the difference on a dataset small enough to reason about: 2,000 points lying on a noisy arc in 2D. The arc is intrinsically one-dimensional (a single angle parameterizes it), but it is not a line, so a 1D linear projection must lose information. A nonlinear encoder can “unbend” it.

import torch
import torch.nn as nn

torch.manual_seed(0)

# Synthetic 2D data: a noisy arc (a 1D curve embedded in 2D)
n = 2000
t = torch.rand(n, 1) * 3.14159
x = torch.cat([torch.cos(t), torch.sin(t)], dim=1) + 0.05 * torch.randn(n, 2)

linear_ae = nn.Sequential(nn.Linear(2, 1), nn.Linear(1, 2))
nonlinear_ae = nn.Sequential(
    nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 1),   # encoder -> 1D code
    nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 2),   # decoder -> 2D
)

def train(model, steps=2000):
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
    for step in range(steps):
        loss = ((model(x) - x) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 500 == 0 or step == steps - 1:
            print(f"  step {step:4d}  reconstruction MSE {loss.item():.4f}")
    return loss.item()

print("Linear autoencoder (2 -> 1 -> 2):")
lin = train(linear_ae)
print("Nonlinear autoencoder (2 -> 32 -> 1 -> 32 -> 2):")
non = train(nonlinear_ae)
print(f"final MSE: linear {lin:.4f}  nonlinear {non:.4f}")
print(f"data noise floor (0.05^2 per coord): {0.05**2:.4f}")

Running this prints:

Linear autoencoder (2 -> 1 -> 2):
  step    0  reconstruction MSE 0.4376
  step  500  reconstruction MSE 0.0494
  step 1000  reconstruction MSE 0.0494
  step 1500  reconstruction MSE 0.0494
  step 1999  reconstruction MSE 0.0494
Nonlinear autoencoder (2 -> 32 -> 1 -> 32 -> 2):
  step    0  reconstruction MSE 0.6963
  step  500  reconstruction MSE 0.0015
  step 1000  reconstruction MSE 0.0014
  step 1500  reconstruction MSE 0.0014
  step 1999  reconstruction MSE 0.0013
final MSE: linear 0.0494  nonlinear 0.0013
data noise floor (0.05^2 per coord): 0.0025

The linear autoencoder plateaus at MSE 0.0494 (the best any 1D linear projection of an arc can do), while the nonlinear one reaches 0.0013, roughly 38 times better. The nonlinear network has learned a coordinate system along the curve: its 1D code is, in effect, the angle t that generated the data.

One detail in those numbers is worth a pause. The noise we added has variance \(0.05^2 = 0.0025\) per coordinate, yet the autoencoder’s MSE is 0.0013 — below the noise floor. That is not magic: noise has a component along the curve’s tangent and a component perpendicular to it. The 1D code can absorb the tangent component (a noisy point slid along the curve is still on the curve), so only the perpendicular half of the noise variance remains as irreducible error. Averaged over the two coordinates, that predicts roughly \(0.0025 / 2 = 0.00125\), almost exactly what we observe. Small experiments reward this kind of arithmetic.

Historical note. Autoencoders had a starring role in deep learning’s revival. Hinton and Salakhutdinov’s 2006 Science paper showed that a deep autoencoder, pretrained layer by layer, could learn far better low-dimensional codes than PCA — at a time when training deep networks end to end was widely considered hopeless. Layer-wise pretraining was later made obsolete by better initializations and normalization (see regularization and initialization), but it helped convince the field that deep representations were worth pursuing.

Denoising: structure without a bottleneck

The bottleneck is the oldest way to force an autoencoder to learn structure, but it is not the only one, and the alternatives turn out to matter for where this chapter is heading. The denoising autoencoder (DAE), introduced by Vincent, Larochelle, Bengio, and Manzagol at ICML 2008, keeps the architecture but changes the task: corrupt the input (add Gaussian noise, zero out random entries) and train the network to reconstruct the clean original. Now even a wide network cannot learn the identity, because the identity maps corrupted input to corrupted output. To succeed, the network must learn what typical data looks like, so it can pull a corrupted point back toward the nearest plausible clean one.

That phrase, pull a noisy point back toward the data, should sound familiar. For small Gaussian corruption, under regularity conditions, the denoising residual scaled by the noise variance approaches the score \(\nabla_x \log p(x)\); at finite noise it corresponds to a smoothed density. This asymptotic qualification matters when connecting denoising autoencoders to score-based diffusion models.

Two other regularization routes deserve a sentence each. Sparse autoencoders keep a wide code but penalize how many units activate at once, forcing each unit to specialize; a decade after their heyday as pretraining devices, they resurfaced as an interpretability tool, used by several research groups from around 2023 onward to decompose the internal activations of large language models into more interpretable features. Contractive autoencoders penalize the encoder’s sensitivity to input perturbations directly. The shared lesson across all variants: an autoencoder is only as interesting as the constraint that prevents it from cheating, and the choice of constraint decides what the representation is good for.

Why a plain autoencoder does not generate

A trained autoencoder gives you compression and a representation, but not a generator. To generate, you would pick a code \(z\) and decode it. Which \(z\)? The encoder places training points somewhere in code space, but you have no map of where. The codes might occupy a thin, curved region with gaps; a \(z\) drawn from the gaps decodes to garbage. Nothing in the training objective ever constrained what the decoder does off the data’s codes, or arranged the codes into a region you could sample from.

This is the precise gap the variational autoencoder fills. What we want is (1) a known distribution over codes, say a standard normal \(\mathcal{N}(0, I)\), that we can sample from, and (2) a decoder trained so that codes drawn from that distribution decode to realistic data. The VAE gets both by changing what the encoder outputs and adding one term to the loss.

The VAE: encode distributions, not points

The VAE was introduced by Kingma and Welling in the paper “Auto-Encoding Variational Bayes” (posted in late 2013, published at ICLR 2014); Rezende, Mohamed, and Wierstra arrived at essentially the same method independently (ICML 2014). The framing is probabilistic: assume data is generated by first drawing a latent \(z \sim \mathcal{N}(0, I)\) from a prior \(p(z)\), then drawing \(x \sim p(x \mid z)\) through a decoder network. Training should maximize the likelihood of the data under this model — but the exact likelihood is intractable, because it sums over every \(z\) that could have produced each \(x\).

The VAE sidesteps this with an encoder \(q(z \mid x)\) that outputs, for each input, a distribution over codes — in practice a mean vector \(\mu(x)\) and a variance vector \(\sigma(x)^2\) defining a Gaussian. Instead of the true likelihood, you maximize a tractable lower bound on it, the evidence lower bound (ELBO):

\[\log p(x) \;\ge\; \mathbb{E}_q [\log p(x \mid z)] - \mathrm{KL}\bigl(q(z \mid x) \,\Vert\, p(z)\bigr)\]

Where does this bound come from? The likelihood you want is \(p(x) = \mathbb{E}_{z \sim p(z)}[p(x \mid z)]\), an average over every possible code. Estimating that average by sampling \(z\) from the prior is hopeless in practice: for any particular \(x\), almost every \(z\) you draw decodes to something unrelated, contributing nearly nothing, so the estimate has enormous variance. The encoder \(q(z \mid x)\) is an importance sampler: it proposes codes in the region that plausibly generated \(x\), where the integrand actually lives. The price of sampling from \(q\) instead of the prior is a correction term, and applying Jensen’s inequality to the corrected estimate yields exactly the ELBO. The gap between the bound and the true \(\log p(x)\) is itself a KL divergence, \(\mathrm{KL}\bigl(q(z \mid x) \,\Vert\, p(z \mid x)\bigr)\) (the mismatch between the encoder and the true posterior over codes), so training the encoder tightens the bound at the same time as training the decoder raises it.

One more piece of vocabulary earns its keep here: the VAE’s encoder performs amortized inference. Classical variational methods fit a separate approximate posterior to each data point, an optimization per example. The VAE instead trains one network that maps any \(x\) directly to its posterior parameters — the cost of inference is amortized across the dataset into the encoder’s weights. Amortization is why VAE training scales, and its imperfection (the network cannot fit every point’s posterior exactly, a shortfall called the amortization gap) is one of the quiet reasons reconstructions are not sharper.

You do not need the derivation to use it; you need to internalize the two terms.

  • The first term is reconstruction: encode \(x\), sample \(z\), and score \(\log p(x \mid z)\). With an isotropic Gaussian decoder whose variance is fixed, negative log-likelihood is squared error up to scale and constants. Learned or non-isotropic variance changes the expression.
  • The second term is a regularizer on the code space: the KL divergence measures how far each encoding distribution strays from the prior \(\mathcal{N}(0, I)\). It pushes all encodings toward the origin, with unit scale, overlapping one another.

The two terms force a compromise: reconstruction preserves information while KL regularizes each approximate posterior toward the prior. This encourages a smoother, more sampleable latent space, but does not guarantee that the aggregate posterior exactly matches the prior or that all “holes” disappear. Prior–posterior mismatch and low-density gap samples remain practical VAE failure modes.

The reparameterization trick

There is one obstacle between the ELBO and gradient descent. The reconstruction term contains a sampling step: \(z\) is drawn from \(q(z \mid x)\). Drawing a sample is not a differentiable operation — the loss depends on \(\mu\) and \(\sigma\) only through a random draw, and autograd cannot push a gradient through “randomness happened here.” If you sample naively, the encoder gets no gradient at all.

The fix, from the same 2013 paper, is disarmingly small. Instead of

\(z \sim \mathcal{N}(\mu, \sigma^2)\), a draw that depends on the parameters,

write

\(z = \mu + \sigma \varepsilon\), with \(\varepsilon \sim \mathcal{N}(0, 1)\), a draw that does not.

The two produce identical distributions over \(z\). But in the second form all the randomness lives in \(\varepsilon\), which involves no learned parameters, while \(z\) is a plain differentiable function of \(\mu\) and \(\sigma\). The gradient path is restored: \(dz/d\mu = 1\) and \(dz/d\sigma = \varepsilon\). PyTorch exposes both forms: .sample() is the naive draw, .rsample() (r for reparameterized) is the trick. Watch what autograd sees in each case:

import torch
from torch.distributions import Normal

torch.manual_seed(0)

mu = torch.tensor([0.5, -1.0], requires_grad=True)
log_sigma = torch.tensor([0.0, 0.0], requires_grad=True)
dist = Normal(mu, log_sigma.exp())

# Naive sampling: draw z ~ N(mu, sigma^2) directly
z_naive = dist.sample()
print(f"sample():  z = {z_naive.tolist()}")
print(f"sample():  requires_grad = {z_naive.requires_grad}, grad_fn = {z_naive.grad_fn}")
try:
    (z_naive ** 2).sum().backward()
except RuntimeError as e:
    print(f"sample():  backward() fails: {e}")

# Reparameterized sampling: z = mu + sigma * eps, eps ~ N(0, 1)
z_rep = dist.rsample()   # same as mu + log_sigma.exp() * torch.randn(2)
print(f"rsample(): z = {[round(v, 4) for v in z_rep.tolist()]}")
print(f"rsample(): requires_grad = {z_rep.requires_grad}")
(z_rep ** 2).sum().backward()
print(f"rsample(): mu.grad        = {[round(v, 4) for v in mu.grad.tolist()]}")
print(f"rsample(): log_sigma.grad = {[round(v, 4) for v in log_sigma.grad.tolist()]}")

# Check against the closed form: for loss z^2 with z = mu + sigma*eps,
# dL/dmu = 2z and dL/dlog_sigma = 2z * sigma * eps
eps = (z_rep - mu) / log_sigma.exp()
print(f"expected mu.grad        = {[round(v, 4) for v in (2 * z_rep).tolist()]}")
print(f"expected log_sigma.grad = {[round(v, 4) for v in (2 * z_rep * log_sigma.exp() * eps).tolist()]}")

Running this prints:

sample():  z = [2.0409960746765137, -1.293428897857666]
sample():  requires_grad = False, grad_fn = None
sample():  backward() fails: element 0 of tensors does not require grad and does not have a grad_fn
rsample(): z = [-1.6788, -0.4316]
rsample(): requires_grad = True
rsample(): mu.grad        = [-3.3576, -0.8631]
rsample(): log_sigma.grad = [7.3155, -0.4906]
expected mu.grad        = [-3.3576, -0.8631]
expected log_sigma.grad = [7.3155, -0.4906]

The naive sample is a dead tensor: requires_grad = False, no grad_fn, and backward() raises an error. The reparameterized sample carries a full computation graph, and the gradients autograd computes match the hand-derived chain rule exactly. This is the same graph-based reasoning you built in the autograd section: a gradient flows only where a differentiable path exists, and reparameterization is the art of rearranging where the randomness sits so that a path exists.

People. Diederik (“Durk”) Kingma developed the VAE as a PhD student with Max Welling at the University of Amsterdam. Remarkably, his other famous paper from the same period is Adam (with Jimmy Ba, ICLR 2015) — the optimizer you have used in nearly every example in this course. Few PhD theses contain two ideas this load-bearing; Kingma went on to work at OpenAI (he was on its founding team) and later Google Brain, continuing work on generative models.

Reparameterization is not the only way to differentiate through sampling — it is the low-variance way. The older alternative, the score-function or REINFORCE estimator, needs only the ability to compute \(\log q(z \mid x)\) and works for any distribution, discrete ones included; the price is gradient estimates so noisy that training often crawls. Reparameterization trades generality for variance: it requires a continuous distribution you can write as a deterministic transform of parameter-free noise, and in exchange the gradient flows through the actual computation rather than through a correlation. This trade returns at the end of the section: when the latent space is made discrete, reparameterization is off the table, and the workaround (a straight-through estimator) is one of the defining engineering choices of the VQ-VAE.

A tiny VAE, end to end

Here is the whole machine on data we can inspect: a mixture of two Gaussian clusters at \((-2, 0)\) and \((+2, 0)\). A good generative model trained on this data should produce samples near both clusters and few in between.

import torch
import torch.nn as nn

torch.manual_seed(0)

# Data: mixture of two 2D Gaussians at (-2, 0) and (+2, 0), std 0.5
def real_batch(n):
    side = (torch.rand(n, 1) < 0.5).float() * 4 - 2      # -2 or +2
    return torch.cat([side, torch.zeros(n, 1)], dim=1) + 0.5 * torch.randn(n, 2)

enc = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 4))  # -> mu (2), logvar (2)
dec = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 2))
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=1e-3)

for step in range(4000):
    x = real_batch(256)
    h = enc(x)
    mu, logvar = h[:, :2], h[:, 2:]
    z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)   # reparameterization
    recon = ((dec(z) - x) ** 2).sum(dim=1).mean()          # reconstruction term
    kl = (-0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum(dim=1)).mean()
    loss = recon + 0.1 * kl                                # beta = 0.1
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 1000 == 0 or step == 3999:
        print(f"step {step:4d}  recon {recon.item():.4f}  KL {kl.item():.4f}")

# Generate: sample z from the prior N(0, I), decode
with torch.no_grad():
    gen = dec(torch.randn(4000, 2))
frac_left = ((gen[:, 0] + 2).abs() < 1).float().mean()
frac_right = ((gen[:, 0] - 2).abs() < 1).float().mean()
print(f"generated samples near left mode:  {frac_left.item():.4f}")
print(f"generated samples near right mode: {frac_right.item():.4f}")
print(f"generated samples in the gap (|x1| < 1): "
      f"{(gen[:, 0].abs() < 1).float().mean().item():.4f}")

Running this prints:

step    0  recon 5.1227  KL 0.1519
step 1000  recon 0.1019  KL 2.8391
step 2000  recon 0.0883  KL 2.6634
step 3000  recon 0.1046  KL 2.6021
step 3999  recon 0.1053  KL 2.5339
generated samples near left mode:  0.3865
generated samples near right mode: 0.4120
generated samples in the gap (|x1| < 1): 0.1787

Read the training trace first: at step 0 the KL is near zero (the untrained encoder outputs roughly the prior) and reconstruction is terrible; by step 1000 the model has bought reconstruction accuracy by spending KL — encodings moved away from the prior, to about 2.5 nats, exactly the tension described above. Then read the samples: drawing \(z \sim \mathcal{N}(0, I)\) and decoding (with no encoder involved) puts about 39% of samples at the left cluster and 41% at the right (real data, same windows: 48% / 48%). That is generation, and a plain autoencoder cannot do it.

But 18% of samples land in the gap between clusters: a prior/aggregate-posterior mismatch and connected-latent issue in this toy. Image blur has an additional, often dominant cause: pixelwise Gaussian likelihoods and uncertain reconstructions reward conditional averages of multiple plausible sharp outputs. Keep “gap samples” and “blurry reconstructions” related but distinct failure mechanisms.

Pitfall. The single most common VAE implementation bug is using .sample() where .rsample() is needed. In our small demo the mistake fails loudly, but inside a larger loss that touches the encoder through other paths, or after a stray .detach(), it can fail silently: training runs, the decoder learns something, but the encoder’s \(\mu\) and \(\sigma\) never receive gradients from the reconstruction term. If your VAE trains but the KL stays near zero and reconstructions look like an average of the dataset, check the gradient path through z first — print z.grad_fn and confirm it is not None.

Turning the beta dial

The tiny VAE used beta = 0.1 without justification. That number deserves a whole experiment, because the weight on the KL term is the single most consequential hyperparameter a VAE has: it sets the exchange rate between reconstruction fidelity and prior fidelity, and every VAE behavior you will read about (blurry samples, great samples, collapsed latents) lives somewhere on this one dial. The sweep below retrains the same model at five values of beta spanning three orders of magnitude, and evaluates each the only way that matters: by decoding prior samples. Two statistics summarize the result: the fraction of generated points that land near either real cluster, and the spread of the generated x1 coordinate (real data has standard deviation about 2.07; a collapsed model has nearly none).

import torch
import torch.nn as nn

torch.manual_seed(0)

# Same two-cluster data as the tiny VAE above
def real_batch(n):
    side = (torch.rand(n, 1) < 0.5).float() * 4 - 2
    return torch.cat([side, torch.zeros(n, 1)], dim=1) + 0.5 * torch.randn(n, 2)

def train_vae(beta, steps=4000):
    torch.manual_seed(0)
    enc = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 4))
    dec = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 2))
    opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=1e-3)
    for step in range(steps):
        x = real_batch(256)
        h = enc(x)
        mu, logvar = h[:, :2], h[:, 2:]
        z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)
        recon = ((dec(z) - x) ** 2).sum(dim=1).mean()
        kl = (-0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum(dim=1)).mean()
        (recon + beta * kl).backward()
        opt.step(); opt.zero_grad()
    with torch.no_grad():
        gen = dec(torch.randn(4000, 2))
        near = (((gen[:, 0] + 2).abs() < 1) | ((gen[:, 0] - 2).abs() < 1)).float().mean()
        spread = gen[:, 0].std()
    return recon.item(), kl.item(), near.item(), spread.item()

print(f"{'beta':>6} {'recon':>8} {'KL':>8} {'near a mode':>12} {'sample std(x1)':>15}")
for beta in [0.01, 0.1, 1.0, 5.0, 20.0]:
    r, k, near, spread = train_vae(beta)
    print(f"{beta:>6.2f} {r:>8.4f} {k:>8.4f} {near:>12.4f} {spread:>15.4f}")
ref = real_batch(4000)
near_ref = (((ref[:, 0] + 2).abs() < 1) | ((ref[:, 0] - 2).abs() < 1)).float().mean()
print(f"reference (real data): near-a-mode {near_ref.item():.4f}  std(x1) {ref[:, 0].std().item():.4f}")

Running this prints:

  beta    recon       KL  near a mode  sample std(x1)
  0.01   0.0111   5.1732       0.5890          1.7474
  0.10   0.1053   2.5339       0.7985          1.8853
  1.00   0.6592   0.9289       0.8673          1.8662
  5.00   2.1014   0.3707       0.6938          1.4817
 20.00   4.5083   0.0000       0.0000          0.0283
reference (real data): near-a-mode 0.9555  std(x1) 2.0708

Read the table as a journey along the dial. At beta = 0.01 the model is nearly a plain autoencoder: reconstruction is superb (0.011) and the KL has ballooned to 5.17 nats; the encoder has spread codes far from the prior to keep inputs separable. The punishment arrives at generation time: only 59% of prior samples land near a real cluster, because the prior is sampling regions the encoder never uses. At beta = 1.0, coverage peaks at 87% (the latent space is well matched to the prior), at the cost of visibly worse reconstructions. And at beta = 20, the KL term has won outright: it sits at exactly 0.0000, the encoder outputs the prior for every input, and every “sample” is the same point — the generated standard deviation is 0.028 against the data’s 2.07. The model has collapsed to producing the dataset mean.

Higgins and colleagues (ICLR 2017) turned this dial into a method: the beta-VAE deliberately sets beta > 1 (modestly, not at collapse levels) and reports that the extra pressure toward an isotropic prior encourages individual latent dimensions to align with individual generative factors of the data, a property called disentanglement. The mechanism, loosely: a diagonal-Gaussian posterior charges KL per dimension, so the cheapest way to encode information is along independent axes, one factor per axis. Their paper reports qualitative results on faces and 3D-shape datasets where single latents control single attributes (azimuth, scale, emotion) — with the trade-off you just measured, since higher beta always taxes reconstruction.

What the latent space learns

A generative latent space is not just a sampling device; it is a learned coordinate system for the data, and two classic experiments probe what those coordinates mean. The first is interpolation: encode two real points, walk the straight line between their codes, and decode along the way. In data space, the straight line between a left-cluster point and a right-cluster point passes through the empty gap. The question is what the latent-space line decodes to.

import torch
import torch.nn as nn

torch.manual_seed(0)

def real_batch(n):
    side = (torch.rand(n, 1) < 0.5).float() * 4 - 2
    return torch.cat([side, torch.zeros(n, 1)], dim=1) + 0.5 * torch.randn(n, 2)

# Train the same tiny VAE as above (beta = 0.1)
enc = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 4))
dec = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 2))
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=1e-3)
for step in range(4000):
    x = real_batch(256)
    h = enc(x)
    mu, logvar = h[:, :2], h[:, 2:]
    z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)
    recon = ((dec(z) - x) ** 2).sum(dim=1).mean()
    kl = (-0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum(dim=1)).mean()
    loss = recon + 0.1 * kl
    opt.zero_grad(); loss.backward(); opt.step()

# Encode one point from each cluster, then walk the straight line between codes
with torch.no_grad():
    x_left = torch.tensor([[-2.0, 0.0]])
    x_right = torch.tensor([[2.0, 0.0]])
    z_left = enc(x_left)[:, :2]
    z_right = enc(x_right)[:, :2]
    print(f"z_left  = ({z_left[0,0]:.4f}, {z_left[0,1]:.4f})")
    print(f"z_right = ({z_right[0,0]:.4f}, {z_right[0,1]:.4f})")
    print(f"{'alpha':>6} {'z':>20} {'decoded x':>20}")
    for a in [0.0, 0.2, 0.4, 0.5, 0.6, 0.8, 1.0]:
        z = (1 - a) * z_left + a * z_right
        xd = dec(z)
        print(f"{a:>6.1f} ({z[0,0]:>8.4f},{z[0,1]:>8.4f}) ({xd[0,0]:>8.4f},{xd[0,1]:>8.4f})")

Running this prints:

z_left  = (0.8202, 0.0062)
z_right = (-0.8541, 0.0178)
 alpha                    z            decoded x
   0.0 (  0.8202,  0.0062) ( -1.9741, -0.0044)
   0.2 (  0.4853,  0.0085) ( -1.6061, -0.0036)
   0.4 (  0.1505,  0.0109) ( -0.6543, -0.0210)
   0.5 ( -0.0169,  0.0120) (  0.0698, -0.0207)
   0.6 ( -0.1844,  0.0132) (  0.7951, -0.0206)
   0.8 ( -0.5192,  0.0155) (  1.6218, -0.0012)
   1.0 ( -0.8541,  0.0178) (  2.0531,  0.0165)

Three things are visible in this small table. The encoder has organized the latent space so that the first latent coordinate is the cluster identity: left-cluster points encode near z1 = +0.82, right-cluster points near z1 = -0.85, and the sign is arbitrary (nothing in the loss prefers left-positive). The decoded path is monotone and smooth: as alpha sweeps from 0 to 1, the output slides continuously from -1.97 to +2.05 with no detours. And the midpoint decodes to x1 = 0.07 — squarely inside the gap where real data is rare. That last observation is the same fact you saw when sampling: a connected latent space with a connected decoder must bridge disconnected data modes through low-density territory. On images, latent interpolations morph one face into another through plausible intermediate faces, and the same geometry explains both the charm and the occasional uncanny in-between frame.

The second probe is disentanglement. A diagonal posterior and unequal factor scales create an axis-aligned bias in this toy, so optimization may recover the generating axes. The isotropic Gaussian prior itself remains rotationally symmetric, and unsupervised VAEs do not generally identify the true factors without additional inductive bias or supervision. Read the result as a demonstration of a matched toy bias, not a recovery guarantee.

import torch
import torch.nn as nn

torch.manual_seed(0)

# Ground-truth factors: f1 ~ N(0,1), f2 ~ N(0,0.5^2), independent.
# Observations mix them through a 45-degree rotation plus small noise.
n = 4000
f = torch.randn(n, 2) * torch.tensor([1.0, 0.5])
c, s = 0.7071, 0.7071
R = torch.tensor([[c, -s], [s, c]])
x = f @ R.T + 0.01 * torch.randn(n, 2)

def factor_corr(z):
    """|correlation| between each latent (column of z) and each true factor."""
    m = torch.corrcoef(torch.cat([z, f], dim=1).T)[:2, 2:].abs()
    return m

def report(name, z):
    m = factor_corr(z)
    print(f"{name}: |corr(latent_i, factor_j)| =")
    for i in range(2):
        print(f"    z{i+1}:  f1 {m[i,0].item():.4f}   f2 {m[i,1].item():.4f}")

# Plain autoencoder, 2D code
torch.manual_seed(0)
ae = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2),
                   nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))
enc_ae = ae[:3]
opt = torch.optim.Adam(ae.parameters(), lr=1e-2)
for step in range(3000):
    loss = ((ae(x) - x) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
with torch.no_grad():
    report("plain AE           ", enc_ae(x))

# VAE with diagonal Gaussian posterior, beta = 0.05
torch.manual_seed(0)
enc = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 4))
dec = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=1e-2)
for step in range(3000):
    h = enc(x)
    mu, logvar = h[:, :2], h[:, 2:]
    z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)
    recon = ((dec(z) - x) ** 2).sum(dim=1).mean()
    kl = (-0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum(dim=1)).mean()
    loss = recon + 0.05 * kl
    opt.zero_grad(); loss.backward(); opt.step()
with torch.no_grad():
    report("VAE (beta = 0.05)  ", enc(x)[:, :2])

Running this prints:

plain AE           : |corr(latent_i, factor_j)| =
    z1:  f1 0.6214   f2 0.7868
    z2:  f1 0.9885   f2 0.1432
VAE (beta = 0.05)  : |corr(latent_i, factor_j)| =
    z1:  f1 0.9997   f2 0.0093
    z2:  f1 0.0062   f2 0.9995

The contrast is stark. The autoencoder’s first latent is a blend of both factors (correlations 0.62 and 0.79) — a perfectly good code for reconstruction, useless as a description of the underlying factors. The VAE’s latents align almost perfectly: z1 tracks factor 1 at correlation 0.9997 and ignores factor 2; z2 does the reverse. Nothing supervised this alignment; it fell out of the KL term’s preference for independent, axis-aligned codes.

Before extrapolating from this toy, know the sobering large-scale result. Locatello and colleagues (ICML 2019, best-paper award) trained more than ten thousand disentanglement models across methods, datasets, and seeds, and reported two things: unsupervised disentanglement is impossible in general without inductive biases (a theorem — for any disentangled model there is an entangled one with the identical data distribution), and in practice random seeds and hyperparameters mattered more than the choice of method. Our demo works because its inductive biases (diagonal posterior, factors with different scales, matched dimensionality) fit the data perfectly; real datasets are rarely so obliging. Treat disentanglement as a property to measure on your data, never one to assume from the method’s name.

When the latent is abandoned: posterior collapse

You saw one collapse mechanism in the beta sweep: crank the KL weight high enough and the encoder gives up. The subtler mechanism, the one that plagued real systems, needs no exaggerated beta. If the decoder is powerful enough to model the data without looking at \(z\), the cheapest optimum is to ignore the latent entirely: the reconstruction term does fine on its own, and the KL term goes to its minimum of zero, with \(q(z \mid x) = p(z)\) for every input. The latent code carries no information; the VAE has quietly become a plain (unconditional) density model.

The demonstration below gives the decoder a learned output variance — a modest superpower: it can explain data as “my mean plus Gaussian noise of whatever size fits.” Trained on the two-cluster data, that is not enough to explain the structure, so the model uses the latent (KL well above zero) and shrinks its output noise. Trained on pure Gaussian noise (data whose every property the decoder’s noise model can already express), the latent has nothing to add:

import torch
import torch.nn as nn

torch.manual_seed(0)

# Two datasets. "clusters": the two-mode data from above -- real structure to encode.
# "noise": pure N(0, I) -- nothing to encode beyond what a noise model explains.
def clusters(n):
    side = (torch.rand(n, 1) < 0.5).float() * 4 - 2
    return torch.cat([side, torch.zeros(n, 1)], dim=1) + 0.5 * torch.randn(n, 2)

def noise(n):
    return torch.randn(n, 2)

def train(data_fn, steps=4000):
    torch.manual_seed(0)
    enc = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 4))
    dec = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 2))
    log_s2 = torch.zeros(1, requires_grad=True)   # learned decoder output variance
    opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()) + [log_s2],
                           lr=1e-3)
    for step in range(steps):
        x = data_fn(256)
        h = enc(x)
        mu, logvar = h[:, :2], h[:, 2:]
        z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)
        # Gaussian NLL with learned variance s2 (up to a constant):
        # 0.5 * [ (x - x_hat)^2 / s2 + log s2 ] summed over coordinates
        nll = 0.5 * (((dec(z) - x) ** 2) / log_s2.exp() + log_s2).sum(dim=1).mean()
        kl = (-0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum(dim=1)).mean()
        loss = nll + kl
        opt.zero_grad(); loss.backward(); opt.step()
    # How much does the decoder's output actually depend on z?
    with torch.no_grad():
        zs = torch.randn(1000, 2)
        dec_spread = dec(zs).std(dim=0).mean()
    return kl.item(), log_s2.exp().item(), dec_spread.item()

for name, fn in [("two clusters", clusters), ("pure noise  ", noise)]:
    kl, s2, spread = train(fn)
    print(f"{name}:  final KL {kl:.4f}   decoder sigma^2 {s2:.4f}   "
          f"std of dec(z) over z {spread:.4f}")

Running this prints:

two clusters:  final KL 1.4967   decoder sigma^2 0.1526   std of dec(z) over z 1.1276
pure noise  :  final KL 0.0107   decoder sigma^2 0.9943   std of dec(z) over z 0.0993

On structured data, the model behaves like a healthy VAE: KL of 1.50 nats (information is flowing through \(z\)), output variance squeezed to 0.15, and the decoder’s output genuinely varies with \(z\) (spread 1.13). On noise data, all three symptoms of collapse appear together: KL is 0.01 (effectively zero), the learned output variance is 0.99, matching the data’s variance exactly, and the decoder is nearly a constant function of \(z\) (spread 0.10). The decoder alone is the model; the encoder has been abandoned.

This exact failure, at scale, nearly ended a research direction. Bowman and colleagues (CoNLL 2016) trained the first VAEs for text with LSTM decoders (decoders strong enough to model sentences autoregressively without any latent help) and reported that the KL term dropped to zero early in training unless countermeasures were taken. Their fixes became standard vocabulary: KL annealing (start training with beta = 0 and ramp it up, letting the encoder become useful before the KL tax kicks in) and word dropout (weaken the decoder’s autoregressive crutch so it must lean on \(z\)). A third fix, free bits (from Kingma and colleagues’ 2016 follow-up work on richer posteriors), exempts the first few nats of KL per dimension from the loss, so the optimizer has no incentive to squeeze the latent below a floor. All three are ways of re-tilting an optimization that has two ways to win toward the winner you actually want.

Going deeper

The VAE you built is the 2014 original. Three later developments matter enough to know by name, and each fixes a different weakness you have now seen in numbers.

Discrete latents: the VQ-VAE. VQ-VAE snaps encoder outputs to a learned codebook and uses a straight-through estimator through the non-differentiable choice. It avoids the Gaussian KL route to posterior collapse in the basic non-autoregressive setup, but it has its own collapse modes: only a few codes may be used and “dead” codebook entries may never recover. Codebook usage/perplexity is therefore a diagnostic, not an automatic guarantee.

Deeper hierarchies. A single latent layer asks one code to capture everything from global layout to texture. Hierarchical VAEs stack latent layers (coarse codes conditioning finer ones) and pushed VAE sample quality much closer to the competition: NVAE (Vahdat & Kautz, NeurIPS 2020) and very deep VAEs (Child, ICLR 2021) generate credible faces from dozens of stochastic layers. The reason to care goes beyond VAEs: a diffusion model is, in one standard reading, a hierarchical VAE with a fixed encoder (the noising process) and a thousand latent layers — the ELBO you learned here is literally the training objective DDPM starts its derivation from.

Tighter bounds and richer posteriors. The ELBO is a lower bound, and the gap is the price of a crude posterior. Importance-weighted autoencoders (Burda, Grosse & Salakhutdinov, ICLR 2016) tighten the bound by averaging over \(k\) samples from the encoder; normalizing-flow posteriors replace the diagonal Gaussian with a learned invertible transform. Both narrow the gap between what VAEs optimize and true likelihood — a reminder that “VAEs blur” was never a law of nature, just the signature of the simplest member of the family.

Where does that leave VAEs today? Mostly in a load-bearing supporting role. The autoencoder that defines Stable Diffusion’s latent space is trained with a small KL regularization (or as a VQ variant) — the section you just read is the design space that component came from. Meanwhile the representation-learning half of the autoencoder story continued into masked autoencoders for vision (He et al., CVPR 2022), which drop the probabilistic machinery and go back to reconstruction (of masked patches) as pure self-supervised pretraining, closing a loop that started with Hinton’s 2006 pretraining stack.

Historical note. The VQ-VAE’s discrete tokens turned out to be one of the quiet hinges of the generative decade: they are what let image generation borrow the transformer wholesale. DALL-E (Ramesh et al., 2021) modeled text tokens and VQ image tokens as one sequence; VQGAN (Esser, Rombach & Ommer, CVPR 2021) added an adversarial loss to sharpen the tokenizer, and its authors’ next system, latent diffusion, swapped the transformer for a diffusion model in a closely related latent space and became Stable Diffusion. Three model families, one shared idea: learn a good compressed space first, then generate inside it.

Common pitfalls

  • Reconstructions are excellent, but decoded random \(z\) vectors are garbage. Cause: you trained a plain autoencoder and are sampling from a latent space that was never shaped to match any prior. Fix: add the KL term (make it a VAE) or restrict sampling to encodings of real data.
  • KL term collapses to zero and samples all look like the dataset mean. Cause: the KL weight is too high relative to reconstruction (or the decoder is strong enough to ignore \(z\)), so the encoder gives up and outputs the prior for every input — “posterior collapse.” Fix: lower beta, or warm it up from 0 over training, and check that reconstruction error actually depends on \(z\).
  • Loss scale looks wildly different from a reference implementation. Cause: mixing sum and mean reductions — summing reconstruction over pixels while averaging KL over the batch (or vice versa) silently rescales the effective beta by the data dimensionality. Fix: pick one reduction convention for both terms and treat beta as meaningful only relative to it.
  • Gradients flow, but \(\sigma\) goes to zero and the VAE behaves like a plain autoencoder. Cause: beta is too small, so the model buys reconstruction with unbounded KL. Symptom: KL keeps climbing throughout training. Fix: raise beta until KL stabilizes and prior samples decode sensibly.
  • A “disentangled” model is selected because its latent traversals look clean on one run. Cause: disentanglement varies heavily across seeds and hyperparameters, and cherry-picking a lucky run overstates the method. Fix: quantify alignment (correlation or a disentanglement metric) across several seeds before claiming the property; treat single-run traversal grids as anecdotes.
  • Text or sequence VAEs train “successfully” with KL pinned at zero. Cause: posterior collapse — an autoregressive decoder models the data alone and the latent is ignored. Fix: KL annealing, weakening the decoder (word dropout), or a free-bits floor; verify the fix by checking that decoding with shuffled \(z\) actually changes outputs.

Try it yourself

  1. In the arc autoencoder, widen the bottleneck from 1 to 2 latent dimensions. Predict what happens to the linear model’s MSE before you run it (the arc lives in 2D — what is the best a 2D linear map can do?), then verify.
  2. Rerun the tiny VAE with beta set to 0.001, 0.1, and 1.0. For each, record the final reconstruction and KL values and the fraction of prior samples in the gap. Which beta gives the best trade-off between mode coverage and gap samples?
  3. Implement a latent traversal for the trained VAE: decode z = (a, 0) for a from -3 to 3 in steps of 0.5 and print the outputs. Where in latent space does the decoder switch from the left cluster to the right?
  4. Replace rsample()-style sampling in the tiny VAE with mu alone (no noise at all). Training still runs — but what happened to the KL term’s meaning, and what do prior samples look like now?
  5. In the disentanglement experiment, change the rotation angle from 45 to 10 degrees and equalize the two factor scales to 1.0. Rerun both models. Does the VAE still recover the factors? Connect what you observe to the Locatello et al. impossibility argument.
  6. In the posterior-collapse experiment, implement KL annealing: multiply the KL term by min(1, step / 2000). Does the two-cluster model reach a different final KL? Then try the noise dataset — can annealing rescue a latent that has nothing to encode?

Further reading

  • Kingma & Welling, “Auto-Encoding Variational Bayes,” ICLR 2014 — the original VAE paper; the ELBO derivation and reparameterization trick in eleven pages. arxiv.org/abs/1312.6114
  • Rezende, Mohamed & Wierstra, “Stochastic Backpropagation and Approximate Inference in Deep Generative Models,” ICML 2014 — the independent co-discovery, with a more general treatment of pathwise gradients. arxiv.org/abs/1401.4082
  • Kingma & Welling, “An Introduction to Variational Autoencoders,” Foundations and Trends in Machine Learning, 2019 — a readable book-length treatment by the original authors; the best single follow-up to this section. arxiv.org/abs/1906.02691
  • Vincent et al., “Stacked Denoising Autoencoders,” JMLR 2010 — the denoising-autoencoder program in full, the ancestor of diffusion’s denoising objective. jmlr.org/papers/v11/vincent10a.html
  • Higgins et al., “beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework,” ICLR 2017 — the KL-weight dial turned into a disentanglement method. openreview.net/forum?id=Sy2fzU9gl
  • Locatello et al., “Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations,” ICML 2019 — the large-scale study and impossibility result that recalibrated disentanglement claims. arxiv.org/abs/1811.12359
  • Bowman et al., “Generating Sentences from a Continuous Space,” CoNLL 2016 — where posterior collapse was diagnosed in the wild, with the KL-annealing and word-dropout fixes. arxiv.org/abs/1511.06349
  • van den Oord, Vinyals & Kavukcuoglu, “Neural Discrete Representation Learning” (VQ-VAE), NeurIPS 2017 — discrete latent codebooks and the straight-through estimator; the tokenizer behind DALL-E-style generation. arxiv.org/abs/1711.00937
  • Prince, Understanding Deep Learning, Chapter 17 (Variational autoencoders) — a free, careful textbook chapter with full derivations and figures. udlbook.github.io/udlbook/
  • PyTorch torch.distributions documentation — see rsample and the pathwise-derivative discussion for how the trick is exposed in the API. docs.pytorch.org/docs/stable/distributions.html

Before class you should be able to

  • Explain why a bottleneck (or noise, in a denoising autoencoder) is what forces an autoencoder to learn structure rather than the identity.
  • State why a trained plain autoencoder cannot generate, and which two things the VAE adds to fix that.
  • Name the two ELBO terms and describe the tension between them in one sentence each.
  • Given code using .sample() where .rsample() is needed, predict the failure and explain what the reparameterization \(z = \mu + \sigma \varepsilon\) restores.
  • Read a beta-sweep table and identify which rows show near-autoencoder behavior, a healthy VAE, and posterior collapse.
  • Describe posterior collapse: its symptoms (KL near zero, samples near the dataset mean) and one standard countermeasure such as KL annealing.
  • Say in one sentence how VAE-style autoencoders remain load-bearing inside Stable Diffusion.

Connects to

  • Quiz 3 takes the final 30 minutes of this class (6:45–7:15 PM, in Canvas), covering Module 3 (S13–S22). This session is itself covered by Quiz 4, given at S27.
  • Project proposal is due this week, so budget time across it and this new module.
  • The latent-variable and reparameterization ideas here are the foundation for S24 (GANs) and S25 (diffusion); VAE-style encoders reappear if your semester project touches generative modeling.

← S22 · ViT & Multimodal · All readings · S24 · GANs →