Latent Diffusion Models (Stable Diffusion)

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Pixel-Space vs Latent-Space
5. The VAE Component
6. Latent Diffusion Training
7. Full Stable Diffusion Pipeline
8. Efficiency Analysis
9. SDXL and SD3 Improvements
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Explain why diffusion in latent space is more efficient than pixel space.
2. Derive the compression ratio and its effect on computational cost.
3. Describe the two-stage training pipeline (VAE then diffusion).
4. Trace the full generation pipeline: text → CLIP → U-Net → VAE decoder → image.
5. Analyze the computational savings quantitatively.

Notation

• $\mathcal{E}$ — VAE encoder
• $\mathcal{D}$ — VAE decoder
• $\mathbf{z} = \mathcal{E}(\mathbf{x}) \in \mathbb{R}^{h \times w \times c}$ — latent representation
• $f$ — downsampling factor ($f=8$ means $256 \times 256 \to 32 \times 32$)

Core Intuition

Running diffusion directly on high-resolution images ($512 \times 512 \times 3$) is computationally prohibitive — the U-Net must process 786K pixels. Latent diffusion first compresses images into a low-dimensional latent space using a pre-trained VAE ($512 \times 512 \to 64 \times 64 \times 4$), then runs diffusion in this compressed space. The result: 50–100x less compute with comparable quality.

Pixel-Space vs Latent-Space

Pixel-space diffusion (DDPM, ADM):

• Input dimension: $512 \times 512 \times 3 = 786{,}432$.
• U-Net processes full resolution.
• Self-attention at $64 \times 64$ costs $O(64^4) \approx 16M$ entries.

Latent-space diffusion (LDM/Stable Diffusion):

• Encode: $\mathbf{z} = \mathcal{E}(\mathbf{x}) \in \mathbb{R}^{64 \times 64 \times 4} = 16{,}384$ values.
• U-Net processes latent.
• Self-attention at $8 \times 8$ costs $O(8^4) = 4{,}096$ entries.

Compression ratio: $786{,}432 / 16{,}384 = 48\times$.

The VAE Component

Architecture: Encoder $\mathcal{E}$ maps images to latents; decoder $\mathcal{D}$ reconstructs images from latents.

Training objective:

KL regularization: Forces latents to be roughly standard normal — important for diffusion to start from $\mathcal{N}(0, \mathbf{I})$.

Perceptual loss (LPIPS): Ensures reconstructions are perceptually faithful, not just pixel-accurate.

Typical architecture: Downsampling factor $f=8$. Input $512 \times 512 \times 3$ → latent $64 \times 64 \times 4$. Latent channels $c=4$ (vs 3 for RGB).

Latent Diffusion Training

Two-stage pipeline:

Stage 1: Train VAE on images (autoencoder reconstruction).

• Train until reconstruction quality is high (PSNR > 30 dB).
• Freeze VAE after training.

Stage 2: Train diffusion U-Net in latent space.

• Encode training images: $\mathbf{z}_0 = \mathcal{E}(\mathbf{x})$.
• Run standard DDPM training on $\mathbf{z}_0$ instead of $\mathbf{x}$.
• Loss: $\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{z}_t, t, \mathbf{c})\|^2$.

Generation:

1. Sample $\mathbf{z}_T \sim \mathcal{N}(0, \mathbf{I})$.
2. Run reverse diffusion in latent space: $\mathbf{z}_T \to \mathbf{z}_0$.
3. Decode: $\hat{\mathbf{x}} = \mathcal{D}(\mathbf{z}_0)$.

Full Stable Diffusion Pipeline

Components:

1. Text encoder (CLIP ViT-L/14): Text → 77 × 768 embeddings.
2. U-Net ($\sim$860M params): Denoises latents conditioned on text.
3. VAE decoder: Latents → pixels.
4. Scheduler (DDIM/DPM-Solver): Controls sampling steps.

Generation flow:

Classifier-free guidance (CFG): At each denoising step:

where $s = 7.5$ is the guidance scale and $\emptyset$ is the null (unconditional) embedding.

Efficiency Analysis

Compute savings: For a 50-step DDIM with $512 \times 512$ images:

• Pixel-space U-Net: ~2000 GFLOPS per step × 50 = 100 TFLOPS.
• Latent U-Net: ~40 GFLOPS per step × 50 = 2 TFLOPS.
• VAE decode: ~50 GFLOPS (one-time).
• Total savings: ~50x.

Memory savings: Latent attention matrices are $(64/f)^2 \times (64/f)^2$ vs $(512/f)^2 \times (512/f)^2$ in pixel space. At the same relative resolution: 64x less attention memory.

SDXL and SD3 Improvements

SDXL:

• Larger U-Net (2.6B params).
• Two text encoders (CLIP + OpenCLIP).
• Higher latent resolution ($128 \times 128 \times 4$).
• Refiner model for high-frequency details.

Stable Diffusion 3 (SD3):

• Replaces U-Net with MMDiT (Multimodal Diffusion Transformer).
• Separate transformer streams for image and text, merged via cross-attention.
• Flow matching objective (rectified flow) instead of DDPM.
• Three text encoders (CLIP, OpenCLIP, T5-XXL).

Common Pitfalls

Pitfall 1. Assuming VAE is lossless. The VAE introduces reconstruction error — fine details may be lost. This limits the maximum quality achievable by latent diffusion.

Pitfall 2. Using a poorly-regularized VAE. If the latent space is not well-structured (e.g., holes, mismatched scale), the diffusion model struggles to learn proper distributions.

Pitfall 3. Ignoring the VAE decode cost for real-time applications. While the U-Net runs on small latents, the final VAE decode still operates at full resolution.

Summary

• Latent diffusion runs diffusion in compressed VAE space: ~50x more efficient.
• Two-stage training: (1) train VAE, (2) train U-Net in latent space.
• Stable Diffusion: CLIP encoder + latent U-Net + VAE decoder + CFG.
• Compression factor $f=8$: images $512^2$ → latents $64^2$.
• Modern evolution: DiT replacing U-Net (SD3, DALL-E 3).

Exercises

Exercise 1. Compute the compression ratio for $f=16$ with latent channels $c=16$ vs the original $512 \times 512 \times 3$ image.

Exercise 2. Estimate the minimum VAE PSNR needed for latent diffusion to achieve FID comparable to pixel-space diffusion.

Exercise 3. Derive the total FLOPs for generating a $1024 \times 1024$ image with Stable Diffusion (latent $128 \times 128 \times 4$, 50 steps).

Exercise 4. Explain why KL regularization of the VAE is necessary for the diffusion model to generate from $\mathcal{N}(0, \mathbf{I})$.

Exercise 5. Compare the inductive biases of U-Net vs DiT for latent diffusion at resolution $64 \times 64$.