Latent diffusion

Pixels are expensive

A 512 by 512 RGB image has 786,432 values. A diffusion model that works on all of them spends most of its compute on detail the eye will never resolve. Attention over 262,144 pixel tokens grows with the square of that number. Most of those values are redundant. Neighbouring pixels are highly correlated, and the eye is far more sensitive to some frequencies than others. Latent diffusion moves the whole process into a compressed representation. An autoencoder (a network that squeezes an input down and then rebuilds it) maps the image to a small tensor. Diffusion runs there, and a decoder turns the result back into pixels.

The compression is typically a factor of eight in each spatial dimension, and a change from three channels to four. A 512 image becomes a 64 by 64 by 4 latent (the small tensor the model works in), which is 16,384 values. That is forty-eight times smaller. The diffusion model never sees a pixel. It sees a tensor whose statistics are close to a standard normal, and the same UNet or transformer architecture works on it unchanged.

The autoencoder is not a detail

The autoencoder is trained first, on reconstruction, and then frozen. Its job is to keep what matters and discard what does not. A good one preserves edges, text and faces. A bad one turns them into mush that no amount of diffusion can recover. The reconstruction loss is usually a mix of pixel error, perceptual loss (a loss measured on features rather than raw pixels) and a patch discriminator. Pure pixel error produces blurry latents, which is why the mix is used.

The encoder’s output is scaled by a constant before diffusion, usually around 0.18, so that the latent distribution has roughly unit variance. The decoder divides by the same constant. That constant is part of the checkpoint, not a universal law. Using the wrong one gives washed-out or blown-out images. It is a small number with a large blast radius, and it belongs in the model card. The scaling factor is fitted on the training set and then frozen with the rest of the autoencoder.

Encode, diffuse, decode

import torch
from diffusers import AutoencoderKL

vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
image = torch.randn(1, 3, 512, 512)

latent = vae.encode(image).latent_dist.sample()
latent = latent * vae.config.scaling_factor
print(latent.shape)

decoded = vae.decode(latent / vae.config.scaling_factor).sample
print(decoded.shape)

What changes in the sampler

The diffusion maths is identical. The schedule, the reverse update, guidance and the solver all work on tensors. None of them care that the tensor is a latent rather than an image. What changes is the cost: attention over 4,096 latent tokens instead of 65,536 pixel patches, and a memory footprint that fits on a single accelerator. The decoder runs once, at the end, and it is the only part that knows about pixels.

Two consequences follow. First, the schedule is tuned for latents, not pixels. A pixel-space schedule moved into a latent model will be wrong at both ends. Second, the VAE’s blind spots become the model’s blind spots. The autoencoder here is a variational autoencoder, or VAE. Small text, fine patterns and regular grids are where reconstruction fails, and no sampler can fix a latent that never encoded them. Judge the autoencoder separately before you blame the diffusion model. The decoder is also where most of the perceptual quality lives. So latent models are judged on the whole pipeline, not only on the denoiser.

Transformers and patches

Latent diffusion does not need a UNet. A transformer can treat the latent as a sequence of patches, exactly as a language model treats tokens, and predict noise for each patch. This is the diffusion transformer, and it scales with the same tools as a language model: more layers, more width, more data. The patch size is the new resolution knob, and it trades sequence length against detail.

At a patch size of two, a 64 by 64 latent becomes 1,024 tokens. At a patch size of one it becomes 4,096. Attention cost grows with the square, so the choice is not cosmetic. Most modern systems use a mix: local attention within patches, global attention at low resolution, or a convolutional stem before the transformer. The principle is unchanged: do the expensive work where the tensor is small.