Perceptual VAE for Image Compression

1. From Standard VAE to a High-Fidelity Autoencoder

You're likely familiar with the standard Variational Autoencoder. It consists of an encoder that maps an image to a distribution in the latent space (parameterized by a mean and variance), and a decoder that reconstructs the image from a sample of that distribution.

The traditional VAE loss function is a sum of a reconstruction term (like Mean Squared Error) and a KL divergence term that regularizes the latent space. As we discussed, relying on MSE alone leads to blurry reconstructions, which is unacceptable for our purposes.

To build a high-fidelity autoencoder for latent diffusion, we need to upgrade both the architecture and the loss function.

Architectural Foundation

The autoencoder used in latent diffusion models is often built with similar blocks to the U-Net you're familiar with from DDPMs (e.g., ResNet blocks, self-attention), but with two key differences:

1. No skip connections between the encoder and decoder. The latent space must be a true information bottleneck.
2. No time embeddings, as the autoencoder's job is purely spatial compression, not denoising over time.

The video below details this architecture and contrasts the VAE approach (which we're using) with the VQ-VAE variant.

2. The Loss Function Quartet

The real power of this autoencoder comes from a carefully constructed loss function that balances four distinct objectives. Let's break down each component.

The total loss can be expressed as a weighted sum:
${\mathcal{L}}_{AE}={\mathcal{L}}_{rec}+{\lambda }_{KL}{\mathcal{L}}_{KL}+{\lambda }_{perceptual}{\mathcal{L}}_{perceptual}+{\lambda }_{adv}{\mathcal{L}}_{adv}$

Where $\lambda$ are hyperparameters that balance the contribution of each term.

A. Reconstruction Loss (${\mathcal{L}}_{rec}$)

This is a straightforward pixel-wise L1 or Mean Squared Error (MSE) loss between the input image $x$ and its reconstruction $\hat{x}$.
${\mathcal{L}}_{rec}=||x-\hat{x}|{|}_2^2$
Its role is to provide a strong, stable gradient signal that ensures the overall structure and color of the image are preserved. As we've established, on its own it is insufficient and leads to blurriness, but it acts as an essential "guardrail" during training.

B. KL Regularization (${\mathcal{L}}_{KL}$)

This term comes from the VAE formulation and measures the KL divergence between the learned latent distribution $q(z|x)$ and a standard normal prior $p(z)=\mathcal{N}(0,I)$. While in a classic VAE this is a core part of the Evidence Lower Bound (ELBO), its role here is more pragmatic.

In latent diffusion models, the weight of this term, ${\lambda }_{KL}$, is typically very small. Its primary function is not to strictly enforce a Gaussian latent space, but rather to act as a mild regularizer. It keeps the latent codes from growing to extreme values and helps prevent the variance from collapsing to zero, which makes the subsequent training of the diffusion model on these latents more stable.

C. Perceptual Loss (${\mathcal{L}}_{perceptual}$)

This is our first major tool for fighting blurriness. Instead of comparing pixels, perceptual loss compares high-level features extracted from a pre-trained network (like VGG). The idea is that if two images have similar feature activations inside a network trained to understand images, they are perceptually similar to a human. The most common implementation is called LPIPS (Learned Perceptual Image Patch Similarity).

The following videos explain the concept and implementation of perceptual loss and LPIPS.

D. Adversarial Loss (${\mathcal{L}}_{adv}$)

To achieve maximum realism and sharpness, we introduce a GAN-style training setup.

• Generator: Our autoencoder's decoder, $G$, which generates the reconstructed image $\hat{x}$.
• Discriminator: A separate network, $D$, trained to distinguish between "real" images (from the dataset, $x$) and "fake" ones (the reconstructions, $\hat{x}$).

The autoencoder is then trained not just to reconstruct the input, but to produce reconstructions that are so realistic they can "fool" the discriminator. This adversarial pressure forces the decoder to learn how to generate the fine-grained textures and high-frequency details that a simple pixel loss would ignore. Typically, a PatchGAN discriminator is used, which classifies overlapping image patches as real or fake, providing a more localized and efficient gradient signal.

3. The Training Loop: An Adversarial Dance

Because we have a discriminator, the training loop is more complex than a standard model.fit(). It involves an alternating optimization process with two separate optimizers.

1. Instantiate Models:

   • The Autoencoder (Generator)
   • The Discriminator (e.g., PatchGAN)
   • The LPIPS model for perceptual loss (with frozen weights).

2. Instantiate Optimizers:

   • opt_ae: An optimizer for the autoencoder's parameters.
   • opt_disc: An optimizer for the discriminator's parameters.

3. In each training step:

   • Discriminator Update:
     • Get reconstructions $\hat{x}$ from the autoencoder.
     • Calculate the discriminator's loss: its ability to correctly label real images $x$ as real and fake reconstructions $\hat{x}$ as fake.
     • Perform a backward pass and update the discriminator's weights using opt_disc.
   • Autoencoder (Generator) Update:
     • Calculate the adversarial loss for the generator: its ability to fool the discriminator into labeling reconstructions $\hat{x}$ as real.
     • Calculate the other losses: reconstruction ${\mathcal{L}}_{rec}$, perceptual ${\mathcal{L}}_{perceptual}$, and KL divergence ${\mathcal{L}}_{KL}$.
     • Combine all losses into the total autoencoder loss ${\mathcal{L}}_{AE}$.
     • Perform a backward pass and update the autoencoder's weights using opt_ae.

A common and important practical trick is to delay the introduction of the adversarial loss. This allows the autoencoder to first learn a decent reconstruction with the more stable losses before the discriminator's influence begins.

This entire process is detailed exceptionally well in the following video.

Conclusion

In this lesson, we've detailed the process of training the high-fidelity VAE that forms the foundation of a latent diffusion model. You've seen how it moves beyond a simple autoencoder by incorporating a more powerful architecture and a sophisticated, multi-part loss function.

Key Takeaways:

• Architecture: The autoencoder uses proven building blocks (ResNet, Attention) in an encoder-decoder setup without inter-level skip connections to create a true information bottleneck.
• Loss Function: High-fidelity reconstruction is achieved by balancing four loss terms:
  1. Reconstruction Loss: For global structural integrity.
  2. Perceptual Loss (LPIPS): To ensure outputs are perceptually similar to inputs, preserving details that humans notice.
  3. Adversarial Loss: To force the generation of realistic, sharp textures by trying to fool a discriminator.
  4. KL Divergence: As a mild regularizer to keep the latent space well-behaved for the subsequent diffusion training.
• Training: The process requires a GAN-style loop with two separate optimizers, one for the autoencoder (generator) and one for the discriminator, updated in an alternating fashion.

Next Lesson Preview:

With our high-quality perceptual compression model trained, we are now ready for the second stage. In the next lesson, we will implement a latent diffusion model by adapting the training loop to operate on latent vectors. We'll take our pre-trained autoencoder, freeze its weights, and use it to convert our entire dataset into latent representations. Then, we will train a U-Net—very similar to the one you've built before—to perform diffusion in this compact and efficient new space.