Building a Latent Diffusion Model

Hello! Welcome back to your study of generative models.

In our last lesson, we built a solid understanding of Denoising Diffusion Probabilistic Models (DDPMs). We saw how they generate high-quality images by reversing a gradual noising process. However, we also identified their main weakness: the immense computational cost of running the U-Net denoiser on high-resolution images for hundreds of steps.

Today, we'll tackle that problem head-on. Your learning outcome is to implement a Latent Diffusion Model (LDM) like the one used in Stable Diffusion. This is the architecture that made high-resolution text-to-image generation accessible.

We will learn how LDMs achieve their remarkable efficiency by moving the entire diffusion process from the high-dimensional pixel space into a much smaller, compressed latent space. To do this, we'll explore three main components:

1. A Perceptual Autoencoder: A powerful encoder-decoder pair that learns to compress images into a meaningful latent space without losing important details.
2. Diffusion in Latent Space: How to adapt the DDPM process from our last lesson to operate on these compact latent representations.
3. The LDM U-Net: The specific U-Net architecture, including the crucial addition of cross-attention layers, which is the key to conditioning the model later on.

By the end of this lesson, you will understand the complete architecture of an LDM and be ready to see how it's guided by text prompts to become Stable Diffusion.

Let's start with a high-level view.

1. The Need for a Better "Compression"

The core idea of an LDM is simple: instead of running the expensive diffusion process on a 512x512 pixel image, we first compress the image into a smaller latent representation (e.g., 64x64) and run the diffusion there. This drastically reduces computational demand.

This compression is handled by an autoencoder. However, a standard autoencoder trained with a simple pixel-wise loss (like L1 or L2) tends to produce blurry, overly smooth reconstructions. It averages out high-frequency details, which are precisely what make an image look sharp and realistic. This would lead to a diffusion model that generates blurry images, defeating the purpose.

The Latent Diffusion Model paper introduced a clever solution: train the autoencoder with a combination of losses to ensure it learns a perceptually rich latent space.

In summary, to get a high-quality latent space, the LDM autoencoder is trained with three objectives:

1. Reconstruction Loss (L1/L2): Ensures the overall structure of the image is correct.
2. Perceptual Loss (LPIPS): Enforces that the reconstructed image is perceptually similar to the original, preserving textures and details.
3. Adversarial Loss: A GAN discriminator pushes the decoder to produce realistic, "unfakeable" high-frequency details.

This combination creates a VAE that can compress an image into a smaller latent space and reconstruct it with high fidelity.

2. The Autoencoder Architecture

With a robust training objective defined, let's examine the architecture of the VAE components. Given your background in Python and PyTorch, you'll find that the building blocks are familiar. The encoder and decoder in an LDM are typically symmetric, using a series of downsampling and upsampling blocks.

For a concrete code implementation, the following article provides a clear, from-scratch implementation of the VAE used in Stable Diffusion.

Once this powerful autoencoder is trained, we can freeze its weights. We now have two key functions:

• encode(image) -> latent
• decode(latent) -> image

These allow us to move freely between the pixel and latent spaces.

3. Diffusion in Latent Space: The U-Net

Now for the main event. We take the DDPM logic from our previous lesson and apply it entirely within the latent space.

1. Forward Process: Instead of adding noise to an image $x_0$, we first encode it to get its latent representation $z_0=\text{encode}(x_0)$. Then we apply the noising process to $z_0$ to get a noisy latent $z_t$.
   $z_t=\sqrt{{\bar{\alpha }}_t}{z}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$
2. Training: We train a U-Net, ${ϵ}_{\theta }$, to predict the noise $ϵ$ from the noisy latent $z_t$ and the timestep $t$. The loss function is the same MSE loss as before, just on latent variables:
   $L_{\text{LDM}}={\mathbb{E}}_{z_0,t,ϵ}\left[\Vert ϵ-{ϵ}_{\theta }(z_t,t){\Vert }^2\right]$

This is elegantly simple. The core diffusion algorithm doesn't change, only the data it operates on.

The Key Architectural Change: Cross-Attention

While the diffusion process is the same, the U-Net architecture used in LDMs like Stable Diffusion has one critical addition: cross-attention layers.

In the U-Net from our DDPM lesson, we likely saw self-attention, where the model relates different parts of the image to each other. Cross-attention allows the model to relate parts of the image to an external context, such as a text prompt embedding.

Even when training an unconditional LDM, the cross-attention layers are included in the architecture. They are simply fed a null or empty context. This makes the model ready for conditional training later.

4. The Full Pipeline

We have now defined all the pieces to build a Latent Diffusion Model. The full generation process looks like this:

1. (Optional) Text Encoding: A text prompt is converted into an embedding vector c using a text encoder like CLIP. For unconditional generation, this context is empty.
2. Latent Denoising: Start with a random noise tensor z_T in the latent space. Iteratively denoise it from t=T to 1 using the U-Net. At each step, the U-Net takes z_t, the timestep t, and the context c (via cross-attention) to predict the noise \epsilon_\theta.
3. Final Decoding: Once the loop finishes, you have a clean latent representation z_0. Pass this through the VAE's decoder to get the final, full-resolution image: image = decode(z_0).

This entire workflow is beautifully orchestrated in the generate function from the article we've been studying.

Conclusion

Excellent work! You have now dissected the architecture that powers one of the most influential generative models today. By moving diffusion to a compressed latent space, LDMs made high-resolution image synthesis practical.

Key Takeaways:

• Efficiency via Latent Space: LDMs dramatically reduce computational cost by performing the iterative diffusion process on small latent representations, not large pixel-based images.
• Perceptual VAE: A high-quality VAE is essential. It's trained with a combination of reconstruction, perceptual (LPIPS), and adversarial losses to ensure the latent space retains fine details.
• Diffusion in Latent Space: The core DDPM algorithm for training and sampling remains the same but is applied to latent vectors (z) instead of images (x).
• Cross-Attention is Key: The U-Net in an LDM is augmented with cross-attention layers. This allows the model to incorporate external conditioning information, such as text embeddings, making it a highly controllable generator.

Preview of the next lesson:
We have built and understood the complete LDM architecture, including the cross-attention layers that are waiting for a context. In the next lesson, "Apply classifier-free guidance for conditional image generation," we will finally plug something into that context input. You will learn the elegant technique of classifier-free guidance, which allows us to steer the generation process toward a text prompt by running the U-Net in two modes—conditional and unconditional—and combining the results. This is the final step in turning our LDM into a true text-to-image model like Stable Diffusion.