Deep Learning · Generative Models

Latent Diffusion and the Cost of Compressing Before Generating

How moving the diffusion process into a compressed latent space changes the compute budget, what the autoencoder bottleneck removes, and which losses it silently introduces.

Diffusion in pixel space works, and its cost scales badly with resolution. Every sampling step evaluates the network on a tensor with three channels at full resolution, and the number of steps is in the hundreds. Doubling the linear resolution quadruples the per-step cost, and the cost of the whole sampling run multiplies by the same factor. At the resolutions where image synthesis is interesting, that product is the difference between a model that can be sampled interactively and one that cannot.

Latent diffusion addresses the cost by separating the problem into two stages. An autoencoder maps images to a smaller tensor and back; the diffusion process then runs entirely in the smaller space. The separation is not free, and the interesting engineering is in deciding what the autoencoder is allowed to throw away and what it must preserve.

The arithmetic of the compression ratio

An image of height H, width W, and three channels has 3HW values. An autoencoder with a spatial downsampling factor f and a latent channel count c maps it to c·(H/f)·(W/f) values, so the compression ratio in raw value count is 3f²/c. With f = 8 and c = 4 the ratio is 48; with f = 8 and c = 16 it is 12. Both configurations appear in practice, and the difference between them is the central design choice.

The ratio is not just a compute multiplier. It is the width of the channel through which all information about the image must pass. A ratio of 48 means the latent carries roughly one value for every 48 values in the image, and the autoencoder must decide which 47 values are predictable from the rest. Natural images are highly redundant at the pixel level, so a large fraction is predictable; the question is whether the fraction the bottleneck discards is the fraction the generator needs.

Compute scales with the latent’s spatial extent, so the downsampling factor is the dominant term. The channel count matters for the network’s width and therefore for its capacity, but it does not change the number of spatial positions the attention or convolution layers must process. This is why f is chosen first, on compute grounds, and c is then chosen as large as the compute budget allows.

What the bottleneck is allowed to discard

An autoencoder trained only on reconstruction error will use its capacity to reproduce whatever is easiest to reproduce, which is low-frequency structure. The residual it leaves behind is high-frequency detail, and the reconstructed images are visibly smooth. A generator trained in that latent space inherits the smoothness: it can produce the structure the autoencoder preserved, and the detail the autoencoder discarded is simply absent from its training targets.

The standard response is to change the autoencoder’s loss so that perceptual and adversarial terms supplement the reconstruction term. A perceptual loss compares features extracted from a pretrained network rather than raw values, so it penalises structural differences more than small shifts. An adversarial term pushes the decoder’s outputs toward the distribution of real images, which restores high-frequency texture that a pixel loss would leave out. Both are needed in practice: a pure perceptual loss produces images that are structurally right and texturally wrong, and a pure adversarial loss produces textures without structural fidelity.

The cost of these losses is that the autoencoder is no longer optimising a likelihood. Its reconstruction term, if it is a squared error, corresponds to a Gaussian observation model with a fixed variance; adding perceptual and adversarial terms changes the implied distribution to something with no closed form. Any likelihood reported for the full model is therefore a likelihood of the latent diffusion process, not of the image distribution, and the two are related by a Jacobian that is not computed.

The scaling factor and why the latent must be normalised

Diffusion’s forward process assumes the data have variance near one. A trained autoencoder’s latent does not satisfy this by construction: its per-channel variance depends on the encoder’s weights, the activation function, and the data. Feeding such a latent into a variance-preserving schedule means the signal-to-noise ratio at each timestep is not what the schedule says it is.

The usual fix is a single scalar, computed once from a batch of encoded examples, that rescales the latent so its overall standard deviation is one. The scalar is stored with the model and applied at both training and sampling. Getting it wrong by a factor of two shifts the effective ratio curve by a factor of four, which moves the region where the model spends its capacity and produces samples that are either too smooth or too noisy depending on the direction of the error.

The normalisation is global rather than per-channel in the common case, which means channels with different variances remain differently scaled relative to each other. A per-channel normalisation is more correct but requires the statistics to be stored per channel and complicates the architecture’s conditioning. The global scalar is a compromise that works because the encoder’s output channels tend to have similar scales after training with a normalised objective.

import numpy as np

def fit_scaling_factor(encoder, batches):
    # one pass over a sample of the data; the value is stored with the model
    total, count = 0.0, 0
    for x in batches:
        z = encoder(x)
        total += float((z ** 2).sum())
        count += z.size
    return 1.0 / np.sqrt(total / count)

def to_latent(encoder, x, scale):
    return encoder(x) * scale       # unit-variance latent for the diffusion process

def from_latent(decoder, z, scale):
    return decoder(z / scale)       # undo the scaling before decoding

The asymmetry in the two functions is the part that is easy to get wrong. The scale multiplies on the way in and divides on the way out, and an implementation that divides in both directions produces a latent that is either too large or too small by the square of the factor, with no error message.

Where the two-stage design changes the trade-offs

The compute saving is the headline, but three other properties change in ways that matter more for some applications.

The first is that the autoencoder can be trained once and reused. A latent diffusion model conditioned on text, on class labels, or on depth maps can share the same encoder and decoder, and only the diffusion network is retrained. This is a real engineering benefit: the autoencoder is the expensive part to train to convergence, and sharing it across conditions amortises that cost.

The second is that the latent space is a natural place to attach conditioning. Cross-attention from the diffusion network’s intermediate layers to an encoded condition operates on a tensor whose spatial extent is already reduced, so the attention cost is lower than it would be at pixel resolution. The same argument applies to any conditioning that is itself spatial, such as a segmentation map or a reference image.

The third is that the latent space is not a metric space in the way pixel space is. Distances between latents do not correspond to perceptual distances between images, because the decoder is not an isometry. Interpolating linearly between two latents and decoding produces a sequence that is smooth in latent coordinates and may not be smooth in image coordinates. Editing operations that add a direction to a latent rely on the decoder’s local linearity, which holds approximately for small perturbations and fails for large ones.

Failure modes specific to the two-stage construction

A latent diffusion model can fail in ways that a pixel-space model cannot, and the failures are informative.

The autoencoder can memorise. If it is trained on a small dataset and its latent is large enough, it can reconstruct training images by storing them rather than by learning a general encoding, and the diffusion model then samples from a latent space whose decoder has memorised the training set. The symptom is that generated images are near-duplicates of training images, and the diagnostic is to compare the autoencoder’s reconstruction error on held-out data against its error on training data.

The autoencoder’s decoder can hallucinate. Because the decoder was trained with an adversarial loss, it produces plausible detail in regions where the latent carries no information. When the diffusion model samples a latent that is outside the encoded distribution, the decoder fills the gaps with whatever its prior prefers, which is often a repeated texture. This is the mechanism behind the characteristic artefacts of latent diffusion at high guidance scales, where the sampled latents are pushed into low-density regions.

The two stages can disagree about what is important. The diffusion model’s loss weights all latent dimensions equally, because its objective is a unit-variance regression. The autoencoder may have allocated its dimensions unevenly, with a few carrying most of the information and the rest carrying small residuals. The diffusion model then spends equal capacity on dimensions that matter little, and the resulting samples have accurate coarse structure and unreliable detail.

The compression ratio sets a floor on the achievable quality that no amount of diffusion training can lift. If the autoencoder cannot represent a class of images, the generator cannot produce them, regardless of how well it models the latent distribution. This is the argument for evaluating the autoencoder’s reconstruction quality separately from the generator’s sample quality: a model with excellent samples and poor reconstructions is generating images that the decoder happens to render well, not images that the encoder could have represented.

The relationship between the latent’s distribution and the sampler is taken up in the reverse-time formulation, and the likelihood accounting that the two-stage design complicates is developed in bits per dimension.

References

  1. High-Resolution Image Synthesis with Latent Diffusion Models — Rombach, Blattmann, Lorenz, Esser & Ommer, CVPR 2022
  2. Taming Transformers for High-Resolution Image Synthesis — Esser, Rombach & Ommer, CVPR 2021
  3. Auto-Encoding Variational Bayes — Kingma & Welling, ICLR 2014