Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

Latent Diffusion: Denoise Where the Information Is

6 minute read

Published:

TL;DR: Pixel-space diffusion spends most of its capacity and nearly all of its compute modelling high-frequency detail that contributes almost nothing perceptually. Latent diffusion splits the problem in two: train an autoencoder once to strip that detail, then run diffusion on its compact latent code. At a downsampling factor of 8, a \(512\times512\times3\) image becomes a \(64\times64\times4\) latent — exactly 48× fewer dimensions, and 4096× cheaper for a full self-attention layer. Conditioning enters through cross-attention. The price is a hard reconstruction ceiling you can never sample your way past.

Where the compute goes

A diffusion model on raw pixels has to learn everything at once: which objects exist, how they are arranged, and also the exact grain of a wall’s plaster. Rate–distortion analysis of natural images says the last of these dominates the bit budget while carrying very little perceptual weight. Rombach et al. call the two regimes semantic compression and perceptual compression, and the point of latent diffusion is that only the first needs a generative model with a thousand sequential forward passes. The second is a job for a single-pass autoencoder.

Squeezing this out matters because the cost is not linear in the pixel count. The U-Net’s attention blocks are quadratic in the number of spatial tokens, so halving each side cuts attention cost sixteenfold, and the cost is paid again at every one of the sampling steps.

Stage one: the autoencoder

Train an encoder \(\mathcal{E}\) and decoder \(\mathcal{D}\) on images so that \(\mathcal{D}(\mathcal{E}(\mathbf{x})) \approx \mathbf{x}\), with a downsampling factor \(f\): an \(H\times W\times 3\) image maps to \(\tfrac{H}{f}\times\tfrac{W}{f}\times c\). The reconstruction objective is not plain L2 — that produces blur — but a combination of a perceptual (LPIPS) loss and a patch discriminator, which is what keeps texture crisp at high compression.

The latent also needs a regulariser, or its scale drifts and the diffusion model’s noise schedule stops making sense. Two choices are used: a very weak KL penalty towards a standard Gaussian (the “KL-f8” autoencoder in Stable Diffusion) or a vector-quantisation layer. In practice the KL variant is also rescaled by a fixed constant so the latents have roughly unit variance before diffusion sees them.

 Pixel spaceLatent space, \(f = 8\)
Tensor shape\(512\times512\times3\)\(64\times64\times4\)
Dimensions786,43216,384
Ratio—48× fewer
Spatial tokens for attention262,1444,096
Cost of one full self-attention—4096× cheaper (\(64^2\))

Rombach et al. swept \(f\) and found the useful window to be roughly 4 to 8. Below that, too little is compressed and diffusion still wastes capacity on texture; above it, the autoencoder starts discarding structure the diffusion model would have needed, and sample quality falls no matter how long you train.

Stage two: diffusion on latents

Nothing about the diffusion maths changes. Freeze the autoencoder, encode the dataset once, and run the ordinary training objective with \(\mathbf{z}_0 = \mathcal{E}(\mathbf{x})\) in place of the image:

\[ \mathcal{L} = \mathbb{E}_{t,\ \mathbf{z}_0,\ \boldsymbol{\epsilon}} \left[\ \left\lVert \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{z}_t,\ t,\ \tau_\theta(y)) \right\rVert^2\ \right], \qquad \mathbf{z}_t = \sqrt{\bar{\alpha}_t}\,\mathbf{z}_0 + \sqrt{1-\bar{\alpha}_t}\,\boldsymbol{\epsilon} \]

with \(y\) the conditioning input (a caption, a layout, a class) and \(\tau_\theta\) a domain-specific encoder — for text, a frozen CLIP or T5 text encoder. Sampling runs the reverse chain in latent space and decodes once at the end: one call to \(\mathcal{D}\) per image, not per step.

The two-stage latent diffusion pipelineLeft to right: a large box labelled image, 512 by 512 by 3, equals 786,432 dimensions, feeds an encoder marked "runs once". The encoder produces a much smaller box labelled latent, 64 by 64 by 4, equals 16,384 dimensions, which is 48 times smaller. A loop arrow above the latent box is labelled "all diffusion steps happen here", with a note that this is repeated 20 to 50 times. The latent then feeds a decoder marked "runs once" producing the output image. A separate box on the right labelled text encoder feeds into the diffusion loop through cross-attention. 512×512×3 786,432 dims 𝓔 once 64×64×4 16,384 dims 48× smaller denoising U-Net 20–50 steps, in latent space every step stays here 𝓓 once output image text enc. τ(y) cross-attention Encode once, denoise many times, decode once
Notice where the repetition is. The autoencoder is evaluated exactly twice per image regardless of step count, so its cost is negligible; the 20 to 50 repeated evaluations all happen on the 48× smaller tensor.

Cross-attention as the conditioning interface

The mechanism that makes one architecture serve text, layouts and segmentation maps is cross-attention inserted into the U-Net’s intermediate blocks. Flatten the spatial feature map \(\varphi_i(\mathbf{z}_t)\) into tokens and attend from those to the conditioning tokens:

\[ \mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)V, \quad Q = W_Q\,\varphi_i(\mathbf{z}_t),\;\; K = W_K\,\tau_\theta(y),\;\; V = W_V\,\tau_\theta(y) \]

Queries come from the image, keys and values from the condition. Changing modality means changing only \(\tau_\theta\); the U-Net is untouched. This is also why per-token attention maps are inspectable, and why so much downstream editing work — attention re-weighting, prompt-to-prompt edits, region control — hooks into exactly this layer.

Key Insight — the split is a division of labour, not just a speed-up: the autoencoder is trained once, on a reconstruction objective, and reused across every model trained on top of it. The diffusion model then never has to allocate parameters to reproducing plausible texture, because the decoder supplies it unconditionally. That is why a latent model with far fewer effective degrees of freedom can match a pixel model: it is solving a strictly easier problem, and the hard-but-cheap half has been amortised.

What the autoencoder costs you

The reconstruction is a ceiling. Whatever \(\mathcal{D}(\mathcal{E}(\mathbf{x}))\) loses is unreachable by any amount of sampling, because the diffusion model has never seen it. At \(f = 8\) this shows up in the well-known failure modes: small faces, fine lettering, regular high-frequency patterns that pick up moiré or smearing. Encode a photograph of printed text and decode it, with no diffusion involved at all, and you will see the limit directly.

There are second-order costs too. Errors in latent space do not map uniformly to pixel error, so the L2 loss on \(\mathbf{z}\) implies a strange, decoder-dependent metric on images. Compounding is real: autoencoder artefacts and diffusion artefacts interact rather than adding. And the frozen decoder is a fixed component of every model built on it, so improving it means retraining downstream — which is why later releases ship refreshed VAEs with more latent channels rather than sticking with the original.

The trap: latent diffusion is often described as "diffusion on a compressed image". It is diffusion on the autoencoder's distribution, which is not the image distribution. Anything the encoder cannot represent has probability zero under the model, and this is a modelling limitation, not a sampling one.

The design is what makes the rest practical: with an 8× smaller field to integrate, fast samplers and step distillation push generation into the low-single-digit step range, and classifier-free guidance’s doubled forward pass stops being prohibitive.

References

  1. Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022.
  2. Esser, P., Rombach, R., & Ommer, B. Taming Transformers for High-Resolution Image Synthesis. CVPR 2021.
  3. Zhang, R., Isola, P., Efros, A. A., Shechtman, E., & Wang, O. The Unreasonable Effectiveness of Deep Features as a Perceptual Metric. CVPR 2018.
  4. Podell, D., English, Z., Lacey, K., et al. SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. ICLR 2024.

You May Also Enjoy

Noether’s Theorem: Symmetry, Conservation, and Equivariant Networks

7 minute read

Published:

Energy is conserved because the laws of physics do not care what time it is. That single sentence is Noether’s theorem, and its machine learning descendant is the reason an equivariant network needs less data than one that must learn the symmetry from examples.

Entropy and Free Energy: From the Second Law to the ELBO

7 minute read

Published:

Thermodynamic entropy and Shannon entropy differ by a constant with units. Once you accept that, the ELBO stops being an inference trick and becomes a free energy — and diffusion models stop being a clever architecture and become a driven non-equilibrium process.

Statistical Mechanics: The Boltzmann Distribution and the Cost of Z

6 minute read

Published:

Counting microstates gives you the Boltzmann distribution, and the Boltzmann distribution gives you softmax, simulated annealing and energy-based models. The partition function is not a bookkeeping constant — it is the object that contains every thermodynamic quantity, and it is intractable for exactly that reason.