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