DATASCI 447 Lecture 16: Variational Autoencoders

Administrative Stuff

THE GENERATIVE GOAL

We want to create a coherent space that we can sample from to produce images of dogs with high probability.

The problem: pixel space is enormous. A 64×64×3 image lives in \(\mathbb R^{64 \times 64 \times 3}\).

Almost all of that space is garbage — random static, melted faces, impossible textures. The set of “images that look like dogs” is a tiny, thin manifold winding through this vast space.

THE GENERATIVE GOAL

THE GENERATIVE SOLUTION

Project high-dimensional images down to a low-dimensional code vector in a simple space that uniquely maps back to the image.

Do it so that:

• Nearby codes → similar images

• We can sample new codes to generate new images

THE GENERATIVE SOLUTION

THE GENERATIVE SOLUTION

THE GENERATIVE SOLUTION

THE BASE METHOD: AUTOENCODERS

Minimize reconstruction error after passing through a low-dimensional bottleneck.

We project to a latent vector (not a tensor) with orthogonal dimensions:

• We know how to define distributions over vectors (like \(\mathcal N(0, \mathbf I)\))

• We want each dimension to capture an independent factor of variation

• Sampling from a structured tensor is much harder — what does a “random 512×4×4 feature map” even mean?

PROBLEM 1: HOLES AND MASS COLLAPSE

With infinite capacity and no constraints, the encoder maps each image to an isolated point — a Dirac delta.

The latent space becomes a scattered collection of dots with vast voids.

The decoder has never seen the voids. Sample there → garbage.

HOLES AND MASS COLLAPSE

PROBLEM 2: POSTERIOR COLLAPSE

An even worse failure mode: with enough capacity, the model learns to skip the bottleneck entirely.

The decoder learns to reconstruct without z — essentially an identity function that routes information around the bottleneck.

The latent space becomes meaningless. The encoder outputs the same z for every input.

PROBLEM 2: POSTERIOR COLLAPSE

THE FIX: POSTERIORS INSTEAD OF POINTS

Instead of learning a point mapping x → z, assume each image was generated as a draw from a posterior distribution over z.

Learn \(Q(\mathbf{z} | \mathbf{x})\) instead of a single z. Give it probabilistic structure.

This ensures the latent space is:

• Smooth — a continuum of dogs with no holes

• Compact — everything concentrated in the same region

• A priori orthogonal — independent dimensions via the prior N(0, I)

POSTERIORS INSTEAD OF POINTS

POSTERIORS REQUIRE PRIORS

To compute the posterior \(P(\mathbf{z} | \mathbf{x}, \theta)\), we need a prior \(P(\mathbf{z})\).

A natural choice is \(P(\mathbf{z}) = \mathcal{N}(\mathbf{z} | \mathbf{0}, \mathbf{I})\), a standard Gaussian.

• Induces compactness properties

• Easy to sample from

This is going to be the standard for any generation methods we show over this semester!

THE NEW PROBLEM: WHAT IS THE POSTERIOR?

The encoder \(f(\mathbf{x}) \to \mathbf{z}\) and decoder \(f(\mathbf{z}) \to \mathbf{x}\) are both nonlinear neural networks.

We don’t know the functional form of \(P(\mathbf{z} | \mathbf{x}, \theta)\).

THE MARGINALIZATION ATTEMPT

Since z is a proper probability distribution, marginalize it out:

\[P(\mathbf{x} | \theta) = \int P(\mathbf{x} | \mathbf{z}, \theta) \, P(\mathbf{z}) \, d\mathbf{z}\]

Choose \(\theta\) that maximizes this value

• Integrating over our prior belief about \(z\), maximize the likelihood that any candidate \(\theta\) would have generated the images that we see!

NOTATION

• \(\theta\): parameters of the model (e.g., weights of the neural networks - both encoder and decoder)

THE INTEGRAL IS IMPOSSIBLE

\[P(\mathbf{x} | \theta) = \int P(\mathbf{x} | \mathbf{z}, \theta) \, P(\mathbf{z}) \, d\mathbf{z}\]

• Neural network decoder → highly nonlinear

• z is high-dimensional → no grid method

• Naive Monte Carlo → astronomical variance

• \(\log \int \neq \int \log\)

• Bayes’ rule is circular — the denominator is the same integral

CAN WE JUST USE BAYES’ RULE DIRECTLY?

Tempting thought. We wrote the posterior \(P(\mathbf{z} | \mathbf{x}, \theta) \propto P(\mathbf{x} | \mathbf{z}, \theta) P(\mathbf{z})\). Conjugacy gave us a nice Gaussian formula.

Can we just compute it?

CAN WE JUST USE BAYES’ RULE DIRECTLY?

Our goal is \(P(\mathbf x | \theta)\)

Rearranging Bayes’ theorem:

\[P(\mathbf{z} | \mathbf{x}, \theta) = \frac{P(\mathbf{x} | \mathbf{z}, \theta) \; P(\mathbf{z})}{P(\mathbf{x} | \theta)}\]

\[P(\mathbf{x} | \theta) = \frac{P(\mathbf{x} | \mathbf{z}, \theta) \; P(\mathbf{z})}{P(\mathbf{z} | \mathbf x , \theta)}\]

If I knew \(P(\mathbf{z} | \mathbf{x}, \theta)\), then I could compute \(P(\mathbf{x} | \theta)\).

• But we don’t. That’s the whole issue!

THE KEY IDEA: APPROXIMATE THE POSTERIOR

We can’t compute the true posterior \(P(\mathbf{z} | \mathbf{x}, \theta)\).

So let’s approximate it.

THE KEY IDEA: APPROXIMATE THE POSTERIOR

We choose Q to be Gaussian:

\[Q_\phi(\mathbf{z} | \mathbf{x}) = \mathcal{N}(\mathbf{z} \;|\; \boldsymbol{\mu}_\phi(\mathbf{x}),\; \text{diag}(\boldsymbol{\sigma}^2_\phi(\mathbf{x})))\]

• Full support. Convenient.

• This is the encoder. The neural network that maps images to posterior parameters. Two heads that map \(\mathbf X\) to \(\boldsymbol{\mu}(\mathbf X)\) and \(\boldsymbol{\sigma}^2(\mathbf X)\).

• Defensible as \(N \to \infty\) due to Bayesian CLT - Bernstein-von Mises theorem.

Goal: Find the best normal approximation to the true posterior!

• But, we don’t know the true posterior…

THE KEY IDEA: APPROXIMATE THE POSTERIOR

DERIVING THE ELBO: STEP 1

Start from the log marginal likelihood — the thing we want to maximize but can’t compute:

\[\log P(\mathbf{x} | \theta) = \log P(\mathbf{x} | \theta) \cdot \underbrace{\int Q_\phi(\mathbf{z} | \mathbf{x}, \theta) \; d\mathbf{z}}_{= 1}\]

DERIVING THE ELBO: STEP 2

Now use Bayes’ rule: \(P(\mathbf{x} | \theta) = \frac{P(\mathbf{x} | \mathbf{z}, \theta) \; P(\mathbf{z})}{P(\mathbf{z} | \mathbf{x}, \theta)}\)

\[= \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{x} | \mathbf{z}, \theta) \; P(\mathbf{z})}{P(\mathbf{z} | \mathbf{x}, \theta)}\right]\]

DERIVING THE ELBO: STEP 3

Multiply and divide by \(Q(\mathbf{z} | \mathbf{x}, \theta)\) inside the log:

\[\log P(\mathbf{x} | \theta) = \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{x} | \mathbf{z}, \theta) \; P(\mathbf{z}) \; Q(\mathbf z | \mathbf x, \theta)}{P(\mathbf{z} | \mathbf{x}, \theta) \; Q(\mathbf z | \mathbf x, \theta)}\right]\]

DERIVING THE ELBO

\[ \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{z})}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] \]

• Recognize this thing? Maybe it would be more familiar in the integral form?

• An expected log density ratio…

\[ E_{Q}\left[ \log \frac{P(\mathbf z)}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] = \int Q(\mathbf{z} | \mathbf{x}, \theta) \; \log \frac{P(\mathbf z)}{Q(\mathbf{z} | \mathbf{x}, \theta)} \; d\mathbf{z} \]

DERIVING THE ELBO

\[ E_{Q}\left[ \log \frac{P(\mathbf z)}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] \]

This is the expected log ratio of the prior to the approximate posterior

DERIVING THE ELBO

THE KL HAS A CLOSED FORM

For \(Q(\mathbf{z}|\mathbf{x}) = \mathcal{N}_K(\boldsymbol{\mu}, \text{diag}(\boldsymbol{\sigma}^2))\) and \(P(\mathbf{z}) = \mathcal{N}_K(\mathbf{0}, I)\):

\[D_{KL} = -\frac{1}{2} \sum_{j=1}^{K} \left(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right)\]

No approximation required. Just \(\boldsymbol \mu \in \mathbb R^K\) and \(\boldsymbol \sigma^2 \in \mathbb R^K\) from the encoder, plugged into a formula.

THE THREE TERMS

\[\log P(\mathbf{x} | \theta) = \mathbb{E}_{Q}\left[\log P(\mathbf{x} | \mathbf{z}, \theta)\right] + \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{z})}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] + \mathbb{E}_{Q}\left[\log \frac{Q(\mathbf{z} | \mathbf{x}, \theta)}{P(\mathbf{z} | \mathbf{x}, \theta)}\right]\]

\[\log P(\mathbf{x}) = \underbrace{\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z}, \theta)] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}))}_{\text{ELBO: } \mathcal{L}(\theta, \theta; \mathbf{x})} + \underbrace{D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}|\mathbf{x}))}_{\geq 0, \text{ unknowable}}\]

Term 1 — Reconstruction: \(\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z}, \theta)]\)

“How well does the decoder reconstruct x from z’s sampled according to Q?”

• How close is our prediction to the truth?

• The expected log-likelihood of an image in our dataset given a latent space, encoder, and decoder.

• The autoencoder loss! Maximize to create best reconstruction of training data.

THE THREE TERMS

\[\log P(\mathbf{x} | \theta) = \mathbb{E}_{Q}\left[\log P(\mathbf{x} | \mathbf{z}, \theta)\right] + \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{z})}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] + \mathbb{E}_{Q}\left[\log \frac{Q(\mathbf{z} | \mathbf{x}, \theta)}{P(\mathbf{z} | \mathbf{x}, \theta)}\right]\]

\[\log P(\mathbf{x}) = \underbrace{\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z}, \theta)] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}))}_{\text{ELBO: } \mathcal{L}(\theta, \theta; \mathbf{x})} + \underbrace{D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}|\mathbf{x}))}_{\geq 0, \text{ unknowable}}\]

Term 2 — KL Regularizer: \(-D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}))\)

“How close is the approximate posterior to the prior?” This is the term that enforces the shared geometry.

• How far is the approximation from our prior?

• Regularization in the latent space

• Always greater than zero.

• Minimize by making latent posteriors for each image exactly N(0,I)

THE THREE TERMS

\[\log P(\mathbf{x} | \theta) = \mathbb{E}_{Q}\left[\log P(\mathbf{x} | \mathbf{z}, \theta)\right] + \mathbb{E}_{Q}\left[\log \frac{P(\mathbf{z})}{Q(\mathbf{z} | \mathbf{x}, \theta)}\right] + \mathbb{E}_{Q}\left[\log \frac{Q(\mathbf{z} | \mathbf{x}, \theta)}{P(\mathbf{z} | \mathbf{x}, \theta)}\right]\]

\[\log P(\mathbf{x}) = \underbrace{\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z}, \theta)] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}))}_{\text{ELBO: } \mathcal{L}(\theta, \theta; \mathbf{x})} + \underbrace{D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}|\mathbf{x}))}_{\geq 0, \text{ unknowable}}\]

Term 3 — Approximation Gap: \(D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}|\mathbf{x}))\)

“How good is our approximation?”

• Always \(\ge\) 0 (Gibbs inequality).

• Can’t compute it — it requires the true posterior. Drop it.

THE ELBO

\[\log P(\mathbf{x}) = \underbrace{\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z})] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}))}_{\text{ELBO: } \mathcal{L}(\theta, \theta; \mathbf{x})} + \underbrace{D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z}|\mathbf{x}))}_{\geq 0, \text{ unknowable}}\]

Since the gap term \(\geq\) 0:

\[\text{ELBO} = \mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z})] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z})) \;\leq\; \log P(\mathbf{x})\]

THE ELBO

\[\text{ELBO} = \mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z})] - D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| P(\mathbf{z})) \;\leq\; \log P(\mathbf{x})\]

L4 callback: the marginal likelihood was our “golden goose” — the quantity that automatically trades off fit and complexity, the thing that replaces cross-validation.

The ELBO is our computable approximation to it.

• And it’s a pretty great approximation in most circumstances!

THE ELBO IS RECONSTRUCTION VS. REGULARIZATION

\[\text{ELBO} = \underbrace{\mathbb{E}_{Q}[\log P(\mathbf{x} | \mathbf{z}, \theta)]}_{\text{Reconstruction}} - \underbrace{D_{KL}(Q(\mathbf{z}|\mathbf{x}, \theta) \| P(\mathbf{z}))}_{\text{Regularization}}\]

Look at this through the lens of every model we’ve built:

Reconstruction term → the encoder-decoder must still compress and reconstruct well. This is the autoencoder job. Push \(Q(\mathbf z | \mathbf x, \theta)\) to be informative about \(\mathbf x\).

KL term → the posteriors must stay close to \(\mathcal N(0, \mathbf I)\) - minimize the distance between the image posterior and the prior. Prevent mass collapse. Maintain shared geometry. Push \(Q(\mathbf z | \mathbf x , \theta)\) toward the prior.

THE KL TERM INDUCES THE SPACE WE WANT

The penalty \(D_{KL}(Q(\mathbf{z}|\mathbf{x}) \| \mathcal{N}(0, I))\) does three things simultaneously:

• Pulls posterior means toward the origin → compactness (no codes flying off to infinity)

• Prevents posterior variances from collapsing to zero → smoothness (no Dirac deltas)

• Encourages independence between dimensions → orthogonality (decorrelated latent factors)

A Quick History Lesson

Classical variational inference (pre-deep learning):

For each observation xᵢ, learn separate posterior parameters \((\mu_i, \sigma^2_i)\). N observations \(\to\) N separate optimization problems.

This was the standard approach for decades. No one did this…

A Quick History Lesson

Amortized inference:

Learn a single function \(Q_\theta(\mathbf{z} | \mathbf{x})\) — a neural network that takes any x and outputs \((\mu_\theta(\mathbf x), \sigma^2_\theta(\mathbf x))\).

One network replaces N optimizations.

WHY AMORTIZATION IS THE NATURAL APPROACH

Amortization seems obviously better. And it is.

• One neural network handles all observations — no per-observation optimization

• Generalizes to new x at test time — just run the encoder

• Scales to millions of observations without growing the parameter count

WHY AMORTIZATION IS THE NATURAL APPROACH

So why wasn’t this done from the start?

It required autodifferentiation frameworks like PyTorch to be practical. Classical VI was developed when computing gradients through a neural encoder was not feasible.

THE VAE IS AMORTIZED VI

Put it all together:

• Encoder network \(Q_\theta\): takes \(\mathbf x\), outputs \((\mu_\theta(\mathbf x), \sigma^2_\theta(\mathbf x))\). This is \(Q(\mathbf{z} | \mathbf{x}, \theta)\).

• Decoder network \(f_\theta\): takes \(\mathbf z \in \mathbb{R}^K\), outputs reconstructed image. This is \(P(\mathbf{x} | \mathbf{z}, \theta)\).

• Train both jointly by maximizing the ELBO (really minimizing the negative ELBO) with gradient descent.

TRAINING VAEs

The training strategy:

• Use CNN to project images from pixel space to latent posterior - \(\mathcal N(\boldsymbol \mu_\theta(\mathbf x), \text{diag}(\boldsymbol{\sigma}^2_\theta(\mathbf x)))\)

• Latent posterior represents the plausible locations of \(\mathbf x\) in \(\mathbf z \in \mathbb R^K\)

• Decode the latent posterior back to pixel space

• Minimize the negative ELBO. Under MSE reconstruction loss (the average sum of squared pixel misses in an image):

\[ \begin{aligned} -\text{ELBO} = \; & \frac{1}{N} \sum_{i=1}^N \|\mathbf{X}_i - \hat{\mathbf{X}}_i \|^2_F \\ + \; & \frac{1}{N} \sum_{i=1}^N D_{KL}\!\left(\mathcal{N}\!\left(\boldsymbol{\mu}_\theta(\mathbf{x}_i),\, \text{diag}(\boldsymbol{\sigma}^2_\theta(\mathbf{x}_i))\right) \,\Big\|\, \mathcal{N}(\mathbf{0}, \mathbf{I})\right) \end{aligned} \]

TRAINING VAEs

New idea: How do we deal with a distribution in the middle of an encoder/decoder?

• A distribution is a pair of values and densities.

• Not just a single number

How do we differentiate with respect to a distribution?

TRAINING VAEs

We don’t!

Instead, each iteration of the autodifferentiator we sample a single value:

\[ \mathbf{z}_i \sim \mathcal{N}\left(\boldsymbol{\mu}_\theta(\mathbf{x}_i),\, \text{diag}(\boldsymbol{\sigma}^2_\theta(\mathbf{x}_i))\right) \]

• Sampling forces the encoder and decoder to hedge their bets.

• Every forward pass, z lands at a slightly different location in the posterior.

• The decoder must learn to reconstruct from all of them — not just the posterior mean.

• This is what fills in the latent space.

THE PROBLEM: SAMPLING ISN’T DIFFERENTIABLE

We need to backpropagate through the encoder. But the sampling operation \(z \sim \mathcal N(\mu, \sigma^2 I)\) is stochastic.

Gradients can’t flow through a random draw.

THE REPARAMETERIZATION TRICK

Instead of sampling \(z \sim \mathcal{N}(\mu, \sigma^2 I)\) directly, rewrite it:

\[\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, I)\]

The randomness is moved into \(\epsilon\), which doesn’t depend on any parameters.

THE REPARAMETERIZATION TRICK

\(z\) is now a deterministic, differentiable function of \(\mu\) and \(\sigma\) (given \(\epsilon\)). Gradients flow cleanly.

def reparameterize(mu, log_var):
    std = torch.exp(0.5 * log_var)
    eps = torch.randn_like(std)        # randomness lives here
    return mu + std * eps              # differentiable in mu and std

ARCHITECTURE CHOICES

Key design decisions for our dog VAE:

• Replace MaxPool with Conv stride 2 — learnable downsampling gives the encoder more control over what information is preserved vs. discarded

• Symmetry between encoder and decoder — ConvTranspose stride 2 mirrors Conv stride 2. Balanced reconstruction problem.

• No skip connections — U-Net skip connections let information bypass the bottleneck. That’s cheating for generation. If the decoder gets pixel details from skip connections, it ignores z → posterior collapse.

BATCHNORM IS FIGHTING AGAINST YOU

We’ve used BatchNorm everywhere this semester. It stabilized training, enabled deeper networks. Why stop?

GROUPNORM: PER-IMAGE NORMALIZATION

The fix: GroupNorm. Instead of normalizing across the batch, GroupNorm normalizes across groups of channels within a single image.

# BatchNorm: statistics computed across batch dimension
nn.BatchNorm2d(num_channels)

# GroupNorm: statistics computed within each image
nn.GroupNorm(num_groups=8, num_channels=num_channels)

GROUPNORM: PER-IMAGE NORMALIZATION

GROUPNORM: PER-IMAGE NORMALIZATION

Each image’s normalization depends only on itself. No cross-contamination between samples. Identical behavior at training and test time — no running averages, no train/eval mode mismatch.

• This works because feature maps channels can be arranged in any order.

• Organize them so that each group contains channels that are highly correlated

Rule of thumb: use GroupNorm in generative models.

THE ENCODER OUTPUTS

The encoder produces two vectors, not one:

self.fc_mu     = nn.Linear(flat_dim, latent_dim)
self.fc_logvar = nn.Linear(flat_dim, latent_dim)

• \(\boldsymbol \mu\) = the “best guess” location in latent space for this image

• \(\log \boldsymbol{\sigma}^2\) = how uncertain the encoder is about that location

THE FULL VAE LOSS

def vae_loss(recon, target, mu, log_var, beta=1.0):
    recon_loss = nn.functional.mse_loss(recon, target, reduction='sum') / target.shape[0]
    kl_loss    = torch.sum(-0.5 * torch.mean(1 + log_var - mu.pow(2) - log_var.exp())) / target.shape[0]
    return recon_loss + beta * kl_loss, recon_loss, kl_loss

The entire ELBO in four lines. Reconstruction (MSE) + \(\beta\) × KL penalty.

• We’ll see \(\beta\) in a second.

THE VAE FORWARD PASS

def forward(self, x):
    # Encode
    mu, log_var = self.encode(x)
    # Sample (reparameterized)
    z = self.reparameterize(mu, log_var)
    # Decode
    recon = self.decode(z)
    return recon, mu, log_var

Encode → sample → decode. One forward pass. Backprop updates both encoder and decoder jointly.

Training Tricks

VAEs are notoriously picky at the beginning of training.

Now, I’ll outline some training tricks that help VAEs settle into decent generation machines.

KL ANNEALING

If we turn on the full KL penalty from the start, the model never learns to reconstruct — the penalty dominates and pushes everything to N(0, I) immediately. The decoder gets no useful signal.

FREE BITS: THE COLLAPSE PROBLEM

Even with annealing, some latent dimensions can collapse — their KL goes to zero.

The encoder has learned to ignore them entirely!

• \(\mu_j \approx 0\) and \(\sigma^2_j \approx 1\) for those dimensions — Zero penalty. Zero information.

FREE BITS: THE FIX

Free bits (Kingma et al., 2016): set a minimum KL threshold λ per dimension. If a dimension’s KL is below λ, don’t penalize it.

\[\mathcal{L}_{KL} = \sum_{j=1}^{d} \max\left(\lambda, \; D_{KL}(q_\phi(z_j | \mathbf{x}) \| p(z_j))\right)\]

Each dimension gets a budget of λ nats of information for free. The encoder isn’t penalized for using a dimension up to that budget.

LEARNING RATE: DON’T LET IT DIE

VAEs don’t converge to a fixed point — they keep improving the longer you train them. Reconstruction gets sharper, latent space gets more organized.

This means: don’t let the learning rate decay to zero. Standard schedulers (ReduceLROnPlateau, cosine annealing to 0) may kill training prematurely.

WARM RESTARTS

Cosine annealing with warm restarts: the LR drops smoothly, then jumps back up. Each restart lets the model escape the current basin and explore new regions.

scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer, T_0=50, T_mult=2, eta_min=1e-5
)

WHAT WE’LL SEE: RECONSTRUCTIONS

Training reconstructions will be a little blurry. This is expected.

• The Gaussian decoder averages over posterior samples, and MSE rewards the mean prediction. The mean of several sharp images is a soft, blurry compromise.

• But validation set reconstructions should still be recognizable dogs. The model generalizes — unlike the deterministic autoencoder, which memorized training images and produced garbage on everything else.

WHAT WE’LL SEE: GENERATION

The payoff moment.

Sample \(\mathbf z \sim \mathcal N(\mathbf 0, \mathbf I)\). Decode.

EVALUATION: USE THE MEAN

During training: sample z ~ N(μ, σ²I) to fill in the latent space. The stochastic noise is essential for learning.

During evaluation: skip the sampling, use μ directly.

# Training
z = self.reparameterize(mu, log_var)   # sample

# Evaluation
z = mu                                  # use the mean

This removes stochastic noise and gives the decoder the posterior’s best estimate. Helps a little with blurriness. But the fundamental issue remains — MSE treats pixels independently.

WHY MSE MAKES THINGS BLURRY

MSE treats each pixel independently:

\[\text{MSE} = \frac{1}{N} \sum_i (x_i - \hat{x}_i)^2\]

Pixel (32, 17) doesn’t know about pixel (32, 18). But we know neighboring pixels are highly correlated — edges, textures, and shapes span many pixels.

• When the decoder is uncertain, the MSE-optimal output is the pixel-wise mean of all plausible images.

• The mean of sharp images is blurry — like averaging 10 photos of a face and getting a smooth blob.

WHY MSE MAKES THINGS BLURRY

PERCEPTUAL LOSS: FEATURES, NOT PIXELS

Instead of comparing pixels, compare features.

Take a pretrained VGG network (same transfer learning backbone from L13). Pass both the real image and reconstruction through VGG. Compare the intermediate feature maps:

\[\mathcal{L}_{\text{perceptual}} = \sum_{\ell} \left\| \phi_\ell(\mathbf{x}) - \phi_\ell(\hat{\mathbf{x}}) \right\|^2\]

where \(\phi_\ell\) is the feature map at layer ℓ of VGG.

PERCEPTUAL LOSS: FEATURES, NOT PIXELS

WHY PERCEPTUAL LOSS WORKS

VGG features capture spatial relationships between pixels. Early layers detect edges and textures. Later layers detect object parts and compositions.

Two images can have low perceptual loss even if they differ pixel-by-pixel — as long as they have the same visual structure.

class VGGPerceptualLoss(nn.Module):
    def __init__(self):
        super().__init__()
        vgg = torchvision.models.vgg16(pretrained=True).features
        self.blocks = nn.ModuleList([
            vgg[:4],   # relu1_2
            vgg[4:9],  # relu2_2
            vgg[9:16], # relu3_3
        ])
        for p in self.parameters():
            p.requires_grad = False

THE MISMATCH

We sample z ~ N(0, I) because that’s our prior.

But the actual distribution of encoded z-values — the aggregate posterior — isn’t exactly N(0, I). The KL penalty pulls toward it, but doesn’t force an exact match.

THE MISMATCH

THE TWO-STAGE SOLUTION: A SECOND VAE

After training the first VAE, train a second, smaller VAE on the encoded latent vectors {μ₁, …, μₙ}.

The second VAE learns the actual distribution of latent codes.

THE TWO-STAGE SOLUTION: A SECOND VAE

WHY THIS WORKS

The first VAE learned a good latent space and a good decoder. The only problem was the mismatch between where we sample and where the data lives.

The second VAE corrects exactly that mismatch without touching the original VAE.

RESULTS: TWO-STAGE SAMPLING

Generated images should now cover a wider variety of dogs — different breeds, poses, colorings.

The samey-ness is reduced because we’re sampling from a distribution that matches the actual latent structure.

NEXT LECTURE

Two directions: