Image Generation › Latent Diffusion Models

A key limitation of pixel-space diffusion is computational cost. Running a U-Net over full-resolution images for hundreds of timesteps is expensive. A $512 \times 512$ image has 786,432 pixels; processing this at each of 1000 timesteps during training requires enormous GPU memory and compute. Pixel-space diffusion models like the one used in DALL-E 2's decoder were trained on thousands of GPU-days.

Latent Diffusion Models (LDMs), introduced by Rombach et al. (2022) and forming the basis of Stable Diffusion, solve this by moving the diffusion process into a compressed latent space.

The approach uses a separate variational autoencoder (VAE) that:

1. Encodes images to a compact latent representation: $\mathbf{z} = \mathcal{E}(\mathbf{x})$, where $\mathbf{z}$ has spatial dimensions roughly $8\times$ smaller than $\mathbf{x}$
2. Decodes latent representations back to images: $\hat{\mathbf{x}} = \mathcal{D}(\mathbf{z})$

The VAE encoder-decoder is trained separately with a combination of losses. A reconstruction loss (typically L1 or L2 in pixel space) ensures the decoded image resembles the original. A perceptual loss (computing similarity in the feature space of a pre-trained VGG network rather than raw pixel space) ensures the reconstruction captures high-level semantic content, not just low-level pixel patterns. A GAN loss discriminates between real and decoded images, encouraging sharp textures and preventing blurry averaging. Additionally, a KL regularization term keeps the latent distribution close to a standard Gaussian, enabling sampling.

Once trained, the encoder-decoder is frozen. The diffusion model then operates exclusively on latent codes $\mathbf{z}$ rather than on pixels. Because latent codes are much smaller (typically $64 \times 64$ with 4 channels for a $512 \times 512$ image), the U-Net processes a $64^2 = 4096$ element feature map rather than a $512^2 = 262144$ pixel image. This is a 64x reduction in spatial resolution, reducing the computational cost of each forward pass by roughly two orders of magnitude.

The practical consequence is significant. Models like Stable Diffusion can be trained and run on consumer GPUs with 8-24 GB of VRAM, whereas pixel-space models of comparable quality required data center hardware. This democratization of image generation through latent compression was arguably as important as the diffusion algorithm itself.

The latent space also has a useful structural property: it is smoother and more semantically organized than pixel space. Nearby points in latent space correspond to semantically similar images, which means the diffusion model can learn a smoother, more coherent distribution compared to the raw pixel distribution that contains abrupt transitions between adjacent pixel values.