DATASCI 447 Lecture 19: Functional and Fast VAEs

Administrative Stuff

WHERE WE ARE

The VAE with perceptual loss produces diverse, structured outputs — different breeds, poses, colors, backgrounds.

But everything is blurry.

THREE SIGNALS TO THE DECODER

Keep the full VAE. Add a PatchGAN discriminator that judges the decoder’s reconstructions.

The decoder now receives three gradient signals simultaneously:

• L1: “match the input pixels”
• LPIPS: “match the input features”
• GAN: “look realistic”

THREE SIGNALS TO THE DECODER

THE PATCHGAN DISCRIMINATOR

GAN discriminators produced a single real/fake score per image. Need more information for reconstruction

THE PATCHGAN DISCRIMINATOR

THE PATCHGAN DISCRIMINATOR

Each position in the \(15 \times 15\) output sees a roughly \(34 \times 34\) patch of the \(128 \times 128\) input. Adjacent positions overlap heavily — offset by only 8 pixels.

• Before final 1x1 convolution on the channels (weighted average of the channel values), each position in the 15x15 grid is associated with a vector of features (perceptual)

• The weighted average links to the real/fake score by saying “Is this collection of features in this broad area realistic?”

PATCHGAN ARCHITECTURE

class PatchDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 64, 4, stride=2, padding=1),    # 128→64, RF: 4
            nn.LeakyReLU(0.2),

            nn.Conv2d(64, 128, 4, stride=2, padding=1),  # 64→32,  RF: 10
            nn.GroupNorm(32, 128),
            nn.LeakyReLU(0.2),

            nn.Conv2d(128, 256, 4, stride=2, padding=1),  # 32→16, RF: 22
            nn.GroupNorm(32, 256),
            nn.LeakyReLU(0.2),

            nn.Conv2d(256, 1, 4, stride=1, padding=1),    # 16→15, RF: 34
        )

• The first three layers downsample with stride=2, building up a \(34 \times 34\) receptive field. The final layer uses stride=1 — it doesn’t downsample further, just produces the real/fake score at each position.

• LeakyReLU, not ReLU — the discriminator needs gradients for “obviously fake” regions.

THE GENERATOR LOSS

\[\mathcal{L}_G = \underbrace{\text{L1}(\mathbf{x}, \hat{\mathbf{x}})}_{\text{pixel}} + \underbrace{\lambda \cdot \text{LPIPS}(\mathbf{x}, \hat{\mathbf{x}})}_{\text{perceptual}} + \underbrace{\beta \cdot D_{KL}(Q \| P)}_{\text{regularization}} + \underbrace{w_{\text{gan}} \cdot \mathcal{L}_{\text{GAN}}}_{\text{sharpness}}\]

• \(\lambda\) is set by gradient balancing (L1 vs LPIPS, same as Lecture 17).

DELAYED DISCRIMINATOR START

If we activate the discriminator from epoch 1, it overwhelms the VAE before the encoder and decoder have learned useful representations.

Train as a pure VAE (L1 + LPIPS + KL) for the first N epochs. Then activate the discriminator.

disc_factor = 1.0 if epoch >= 30 else 0.0
gan_term = disc_factor * w_gan * gan_loss

The VAE establishes good reconstructions first. The GAN sharpens on top of an already-working model.

• Best outcome for me was a linear ramp of the gan term weight

WHY THE VAE-GAN DOESN’T COLLAPSE

The standalone GAN collapsed because the generator had no anchor — its only objective was fooling D. It found one good image and repeated it.

• Training a GAN is quite difficult because of this lack of “correctness”

The VAE-GAN can’t collapse because of the reconstruction loss. Every training image demands its own reconstruction. Diversity is mandatory, not optional!

• The GAN loss is a secondary signal on top of a primary reconstruction objective. Sharpness without sacrificing diversity.

THE TRAINING LOOP

for data, *labels in train_loader:
    # ── Discriminator step ──
    recon, mu, logvar = model(data)
    real_logits = D(data)
    fake_logits = D(recon.detach())    # detach: don't update VAE
    d_loss = discriminator_loss(real_logits, fake_logits)
    opt_D.zero_grad(); d_loss.backward(); opt_D.step()

    # ── Generator (VAE) step ──
    recon, mu, logvar = model(data)    # fresh forward pass
    fake_logits = D(recon)             # no detach: grads to VAE

    recon_loss = l1 + adaptive_weight * lpips
    gan_loss = generator_loss(fake_logits)
    w_gan = compute_adaptive_weight(recon_loss, gan_loss, last_layer)

    total = recon_loss + beta * kl + disc_factor * w_gan * gan_loss
    opt_G.zero_grad(); total.backward(); opt_G.step()

WHAT’S STILL WRONG

The images are sharper, but look closely:

• Eyes not quite symmetric
• Heads slightly dented or lopsided
• Ears at wrong angles
• Occasional merged features

WHERE SPATIAL INFORMATION DIES

The current bottleneck:

\[\text{Encoder} \to (B, 512, 4, 4) \to \text{Flatten} \to \mathbb{R}^{8192} \to \text{fc\_mu} \to \mathbb{R}^{256}\]

The encoder produces a spatial feature map — position \((0,0)\) represents the top-left region, position \((3,3)\) represents the bottom-right.

Then we flatten and project — mixing every spatial position into one vector.

WHERE SPATIAL INFORMATION DIES

KEEP IT SPATIAL

What if we skip the flatten entirely?

\[\text{Encoder} \to (B, D, 4, 4)\]

Each position in the \(4 \times 4\) grid corresponds to a specific region of the image. Top-left stuff stays in the top-left. The decoder starts with a rough spatial map and refines it.

KEEP IT SPATIAL

BUT WE STILL HAVE THE KL PROBLEM

The spatial bottleneck helps with structure. But each of those \(D \times 4 \times 4\) continuous values still needs KL regularization toward \(\mathcal{N}(0, 1)\).

• That’s \(D \times 16\) KL terms pulling each latent dimension toward zero — fighting the reconstruction loss at every step.

• For D = 256, that’s 4096 KL terms! If each term contributes even a small amount to the loss, it adds up quickly and overwhelms the reconstruction loss

• We’re left at an impasse - even weak spatial structure will kill our VAE!

THE BILLION DOLLAR IDEA: CLUSTERING

What if instead of pushing continuous latents toward a Gaussian, we snap each spatial position to the nearest prototype from a learned dictionary?

THE BILLION DOLLAR IDEA: CLUSTERING

THE CODEBOOK

A codebook \(\mathbf{E} \in \mathbb{R}^{K \times D}\) — a matrix of \(K\) learned vectors, each of dimension \(D\).

Each row is a visual word — a prototype feature that captures a commonly occurring pattern in the encoder’s output.

• One visual word might represent “fur texture”, another “grass background”, and so on.

Think of it like a palette. An artist doesn’t use infinitely many colors — they use a finite palette and paint from those. The codebook is a palette of visual features.

self.codebook = nn.Embedding(K, D)  # K entries, each D-dimensional
# Typical: K = 512 entries, D = 64 dimensions per entry

THE CODEBOOK

QUANTIZATION: HOW IT WORKS

The encoder outputs a continuous spatial feature map: \(4 \times 4 \times D\)

For each of the 16 spatial positions \((h, w)\):

1. Take the continuous feature vector \(\mathbf{z}_e(h, w) \in \mathbb{R}^D\)
2. Compute distance to every codebook entry: \(\|\mathbf{z}_e(h, w) - \mathbf{e}_j\|_2\) for \(j = 1, \ldots, K\)
3. Replace with the nearest entry: \(\mathbf{z}_q(h, w) = \mathbf{e}_{k^*}\) where \(k^* = \arg\min_j \|\mathbf{z}_e(h,w) - \mathbf{e}_j\|_2\)

QUANTIZATION: HOW IT WORKS

QUANTIZATION: HOW IT WORKS

Once the visual vocabulary for an image has been created in the latent space, decode to an image!

• Locations in the 4x4 grid are spatially meaningful!

• Decoder is highly nonlinear and can figure out how grass background, gray fur, cat eyes, and cat ears should go together!

CODEBOOK SIZE IS EXPLICIT REGULARIZATION

With a continuous VAE, regularization (e.g. generalization) comes from the KL divergence — forcing the latent distribution toward \(\mathcal{N}(0, I)\).

THE REPRESENTATIONAL CAPACITY IS ENORMOUS

How much variation can we really have with only 512 vectors, though?

• Remember that the decoder figures out how to combine the actual vector with the positional context.

NO KL. AT ALL.

The prior over codebook entries is uniform — each entry is equally likely.

\[D_{KL}(Q(\mathbf{z} | \mathbf{x}) \| P(\mathbf{z})) = \log K\]

This is a constant. It drops out of the optimization entirely.

VQ-VAE

THE DIFFERENTIABILITY PROBLEM

\(\arg\min\) selects a discrete index. Discrete operations have zero gradient everywhere!

THE DIFFERENTIABILITY PROBLEM

THE STRAIGHT-THROUGH ESTIMATOR

The fix: during the forward pass, use \(\mathbf{z}_q\) (the actual codebook entry). During the backward pass, pretend the quantization didn’t happen — copy gradients directly from \(\mathbf{z}_q\) to \(\mathbf{z}_e\).

z_q = z_e + (z_q - z_e).detach()

WHY THE STRAIGHT-THROUGH ESTIMATOR WORKS

If \(\mathbf{z}_e\) is close to its nearest codebook entry \(\mathbf{e}_{k^*}\), then:

\[\mathbf{z}_q = \mathbf{e}_{k^*} \approx \mathbf{z}_e\]

• The gradient of the loss with respect to \(\mathbf{z}_q\) is approximately the same as the gradient with respect to \(\mathbf{z}_e\) — because they’re nearly the same point.

The approximation is good when encoder outputs stay close to codebook entries.

If the encoder output drifts far from all codebook entries, the straight-through gradient becomes poor. The encoder gets misleading gradient signals.

THE CODEBOOK LOSS

\[\mathcal{L}_{\text{codebook}} = \|\text{sg}[\mathbf{z}_e] - \mathbf{e}_{k^*}\|_2^2\]

\(\text{sg}[\cdot]\) is the stop-gradient operator. It says: “use this value in the forward pass, but treat it as a constant (zero gradient) in the backward pass.”

THE COMMITMENT LOSS

\[\mathcal{L}_{\text{commit}} = \beta_c \|\mathbf{z}_e - \text{sg}[\mathbf{e}_{k^*}]\|_2^2\]

Now the stop-gradient is on the codebook entry. The codebook is treated as a fixed target.

THE TWO LOSSES TOGETHER

Two losses pulling toward each other, but with stop-gradients so they don’t interfere.

The codebook loss only updates the codebook. The commitment loss only updates the encoder. Each treats the other’s output as a fixed target.

• Without stop-gradients: both losses would tug on both networks, creating unstable dynamics where the encoder and codebook chase each other in circles.

• With stop-gradients: coordinated, stable convergence.

EMA UPDATES: AN ALTERNATIVE FOR THE CODEBOOK

Instead of gradient-based codebook updates, track the running mean of encoder outputs assigned to each entry:

# For each codebook entry k:
N_k = decay * N_k + (1 - decay) * n_k       # count
m_k = decay * m_k + (1 - decay) * sum(z_e)   # sum of assigned vectors
e_k = m_k / N_k                               # updated entry

CODEBOOK COLLAPSE

The VQ equivalent of posterior collapse: some codebook entries never get selected.

• The encoder’s output distribution drifts away from unused entries. Those entries receive no updates. They become dead — wasting capacity.

• If \(K = 512\) but only 100 entries are active, you effectively have a much smaller codebook than intended.

# Every N steps:
usage_count = count_assignments_per_entry()
dead = usage_count == 0
if dead.any():
    random_z_e = z_e_flat[torch.randint(0, len(z_e_flat), (dead.sum(),))]
    codebook.weight.data[dead] = random_z_e

THE FULL VQ-VAE ARCHITECTURE

THE FULL LOSS

\[\mathcal{L} = \underbrace{\text{L1} + \lambda \cdot \text{LPIPS} + w_{\text{gan}} \cdot \mathcal{L}_{\text{GAN}}}_{\text{reconstruction + sharpness}} + \underbrace{\|\text{sg}[\mathbf{z}_e] - \mathbf{e}\|^2}_{\text{codebook}} + \underbrace{\beta_c \|\mathbf{z}_e - \text{sg}[\mathbf{e}]\|^2}_{\text{commitment}}\]

• The KL term is gone. Replaced by two simple L2 penalties that keep encoder and codebook aligned.

• This architecture — VQ-VAE with perceptual loss and a PatchGAN discriminator — is called a VQGAN (Esser, Rombach, & Ommer, “Taming Transformers for High-Resolution Image Synthesis,” CVPR 2021).

• With conditioning (next class), this is DALL-E 1.0!

WHAT EACH COMPONENT DOES

TRAINING COMPARISON

WHY THIS WORKS SO WELL FOR IMAGES

Images have strong spatial structure — nearby pixels are correlated, objects have consistent shapes, backgrounds are coherent. The spatial bottleneck preserves this.

THE SPATIAL STRUCTURE ADVANTAGE

The flat bottleneck forces the decoder to reconstruct spatial layout from a vector with no spatial organization.

The VQ-VAE’s spatial bottleneck preserves the layout. Position \((1,1)\) in the \(4 \times 4\) grid corresponds to the top-left of the image. The decoder starts with a rough spatial map and refines it.

• This is why the structural distortions — dented heads, asymmetric eyes — improve.

SUMMARY

We built two improvements to the VAE today:

VAE-GAN: Added a PatchGAN discriminator for sharpness. Three gradient signals keep the decoder honest and sharp. Delayed start prevents collapse.

VQ-VAE-GAN (VQGAN): Replaced the flat Gaussian bottleneck with a spatial discrete codebook. Eliminated KL entirely. Codebook size \(K\) controls compression. Spatial structure preserved. Faster training, no hyperparameter balancing.

Key references:

• van den Oord et al. (2017) — VQ-VAE (codebook, straight-through estimator)
• Esser, Rombach, & Ommer (2021) — VQGAN (perceptual + GAN loss)

NEXT TIME

We can reconstruct sharply with good spatial structure. But we can’t yet control what we generate.

Next lecture: interpreting VQ-VAE codebooks, conditioning, editing, the ethical implications, and the generation problem that bridges us to sequences and transformers.