Latent Diffusion Models: High-Resolution Image Synthesis

TL;DR

Latent Diffusion Models reduce the cost of image diffusion by separating perceptual compression from generative modeling. A trained, frozen autoencoder supplies a spatial latent representation; diffusion learns to denoise that representation and a decoder returns the final latent to pixels. For one familiar later Stable Diffusion configuration, 512×512×3 pixels become 64×64×4 latents: 48× fewer tensor elements, not a measured 48× speedup. The original LDM paper studies both KL-regularized and VQ-regularized first stages, several downsampling factors, and task-specific conditioning. Cross-attention connects those latent features to conditioning embeddings.

The Problem: Diffusion Is Expensive

Diffusion models produce remarkable image quality by learning to reverse a gradual noising process. Starting from pure Gaussian noise, a neural network (typically a U-Net) iteratively predicts and removes noise over many steps, eventually producing a clean image. The training objective is elegant:

But there’s a catch. Pixel-space diffusion repeatedly evaluates a denoiser at image resolution. A 512×512×3 state contains 786,432 elements, and many sequential evaluations make training and inference expensive. Element counts alone do not determine FLOPs or time: network widths, layers, attention, sampler, and hardware matter. Doubling both spatial dimensions quadruples the number of state elements, while the computation depends on the architecture.

This computational burden meant that, before LDMs, high-resolution diffusion was practical only for well-resourced research labs. The question was: can we preserve diffusion’s quality while dramatically reducing its computational requirements?

The Solution: Move to Latent Space

The key insight of LDMs is a separation of concerns. Image generation involves two distinct phases: perceptual compression (learning a compact representation that captures visual structure) and semantic generation (learning the distribution of meaningful images). Pixel-space diffusion conflates these two phases — the model must simultaneously learn what images look like at a low level and how to generate semantically meaningful content.

LDMs separate these phases by introducing a two-stage approach:

1. Stage 1: Train a perceptual autoencoder, with either VAE-style KL regularization or VQ regularization. The encoder maps an image to a spatial tensor with task-dependent downsampling and channel count. Freeze the first stage for diffusion training; reuse it across new diffusion trainings.

2. Stage 2: Train a diffusion model in that first stage’s latent space. For a 512×512×3 → 64×64×4 configuration, the state has 48× fewer elements. Other first-stage configurations have different ratios.

The diffusion training objective becomes:

where z_t is the noised latent representation and \mathcal{E}(x) is the first-stage encoder output.

Why Latent Space Works

The critical question is how much compression preserves perceptually useful information. The learned first stage uses perceptual and adversarial objectives to balance reconstruction quality with a compact spatial representation. It can discard some details and still reconstruct useful images, but its reconstruction fidelity limits what downstream diffusion can generate.

For the 512×512×3 and 64×64×4 example, the states contain 786,432 and 16,384 elements. Multiplying by 50 gives 39,321,600 and 819,200 state-element visits. These are repeated tensor-element counts, not total operations, FLOPs, measured compute savings, or wall-clock predictions. The original paper measures task-specific training and inference efficiency separately.

The VAE: Perceptual Compression

For the KL-regularized variant, the following is a schematic grouping of reconstruction, perceptual, adversarial, and regularization terms, not the paper’s full first-stage training objective (Appendix G). The VQ-regularized variant uses a codebook regularizer instead:

The reconstruction loss promotes reconstruction fidelity. The KL divergence slightly regularizes the encoder’s diagonal Gaussian toward a standard normal to limit latent variance; it does not require the aggregate clean-latent distribution to be exactly Gaussian. Forward diffusion supplies the progressively Gaussian noisy distribution. The perceptual loss compares deep features, while the adversarial loss uses a patch discriminator to promote local realism. The schematic pixel MSE above is not a claim that all original first stages used that exact reconstruction term.

The paper studies several downsampling factors f. Its first-stage Table 8 and conditional-model Table 15 include different KL and VQ choices; f=4 and f=8 are useful in different experiments. There is no universal f=8 sweet spot. Stronger downsampling can reduce spatial computation but also degrade reconstructions; the suitable trade-off depends on the task and model.

The KL penalty is deliberately weak, so the first stage retains reconstruction capacity. Later Stable Diffusion releases use particular KL-style autoencoders, scales, and conditioning choices. Those release configurations should be distinguished from the original LDM family described in Sections 3.1–3.3 and Appendix G.

Conditioning with Cross-Attention

To enable conditional generation, LDMs introduce cross-attention layers into the U-Net. A task-specific conditioning encoder produces a sequence \tau_\theta(y) \in \mathbb{R}^{M \times d_\tau}, where M is its token count and d_\tau its feature width. In original LDM Eq. 3, the conditioner and denoiser are jointly optimized; Appendix E.2.1 describes trainable transformer and class-embedding conditioners. Frozen CLIP is a choice in particular later Stable Diffusion releases, not a universal original-LDM default.

Inside each U-Net block, cross-attention layers allow every spatial position in the latent to attend to all text token embeddings:

where the queries come from the latent features and the keys/values come from the text embeddings:

Here \varphi_i(z_t) is the intermediate representation of the U-Net at layer i, flattened to a sequence of spatial tokens. Each spatial position produces a query that attends to all text tokens, learning which parts of the prompt are relevant for generating content at that location.

Attention weights depend on the learned features and conditioning. Their spatial patterns can relate conditioning tokens to image locations, but an authored token heatmap alone does not demonstrate learned language grounding. The diagram below illustrates supplied weights; the executable section computes attention from numerical embeddings.

Cross-attention supports multiple conditioning modalities by designing and training a suitable encoder alongside the denoiser. Original LDMs also use concatenation for some image-conditioned tasks. Replacing a conditioning encoder after training can change the representation the denoiser expects and is not automatically interchangeable.

Implementation

The first stage and conditioning interface now meet in one executable path: encode a frozen latent, add known noise, compute conditioning attention, regress that noise, and sample entirely in latent space before decoding. This reduced numerical reference keeps autoencoder training separate. Guidance is a later addition, introduced after these mechanics.

class FrozenKLFirstStage(nn.Module):
    """Random Gaussian encoder/decoder interface, not perceptual compression.

    Attributes:
        downsampling: Spatial factor two, independent of latent channel count.
        latent_channels: Number of channels in a sampled diagonal Gaussian.
    """

    def __init__(self, image_channels=3, latent_channels=2):
        super().__init__()
        if min(image_channels, latent_channels) < 1:
            raise ValueError("Channel counts must be positive")
        self.image_channels = image_channels
        self.latent_channels = latent_channels
        self.downsampling = 2
        self.encoder = nn.Conv2d(
            image_channels, 2 * latent_channels, 2, stride=2
        )
        self.decoder = nn.ConvTranspose2d(
            latent_channels, image_channels, 2, stride=2
        )
        self.requires_grad_(False)
        self.eval()


def synthetic_fixture(generator, batch=2, height=8, width=8, tokens=3):
    """Return random images and supplied numerical conditions, without semantics."""
    if min(batch, height, width, tokens) < 1 or height % 2 or width % 2:
        raise ValueError(
            "Positive dimensions and even image sizes are required"
        )
    images = torch.randn(batch, 3, height, width, generator=generator)
    condition = torch.randn(batch, tokens, 3, generator=generator)
    return images, condition

@torch.no_grad()
def encode(first_stage, images, generator):
    """Sample z = mu + exp(logvar/2)*xi from the frozen KL-style interface."""
    require_tensor(images, 4, "images")
    if (
        images.shape[1] != first_stage.image_channels
        or images.shape[2] % 2
        or images.shape[3] % 2
    ):
        raise ValueError(
            "Image channels or spatial downsampling are incompatible"
        )
    moments = first_stage.encoder(images)
    mean, logvar = moments.chunk(2, dim=1)
    logvar = logvar.clamp(-30.0, 20.0)
    xi = torch.randn(mean.shape, generator=generator)
    latent = mean + (0.5 * logvar).exp() * xi
    return latent, mean, logvar

def noise_latent(latent, timesteps, noise, alpha_bar):
    """Apply the closed-form forward process; indices 0..T-1 mean paper t=1..T."""
    require_tensor(latent, 4, "latent")
    require_tensor(noise, 4, "noise")
    validate_schedule(alpha_bar)
    if (
        noise.shape != latent.shape
        or not isinstance(timesteps, torch.Tensor)
        or timesteps.dtype != torch.long
        or timesteps.device.type != "cpu"
        or timesteps.shape != (latent.shape[0],)
        or not ((timesteps >= 0) & (timesteps < len(alpha_bar))).all()
    ):
        raise ValueError(
            "Noise shape or per-example timestep indices are invalid"
        )
    coefficients = alpha_bar[timesteps].reshape(-1, 1, 1, 1)
    noisy = coefficients.sqrt() * latent + (1 - coefficients).sqrt() * noise
    return noisy

def spatial_queries(latent, projection):
    """Flatten N=h*w locations and retain shape for the inverse transform."""
    require_tensor(latent, 4, "latent")
    batch, channels, height, width = latent.shape
    if channels != projection.in_features:
        raise ValueError("Query input width must match latent channels")
    spatial = latent.flatten(2).transpose(1, 2)
    queries = projection(spatial)
    return queries, (batch, height, width)

def condition_keys_values(condition, key_projection, value_projection):
    """Project supplied numerical embeddings, without claiming text semantics."""
    require_tensor(condition, 3, "condition")
    if condition.shape[-1] != key_projection.in_features:
        raise ValueError("Condition width must match the key projection")
    if condition.shape[-1] != value_projection.in_features:
        raise ValueError("Condition width must match the value projection")
    keys = key_projection(condition)
    values = value_projection(condition)
    return keys, values

def cross_attention(queries, keys, values, spatial_shape, output_projection):
    """Normalize over M condition tokens, aggregate values, then restore space."""
    for name, tensor in (
        ("queries", queries),
        ("keys", keys),
        ("values", values),
    ):
        require_tensor(tensor, 3, name)
    batch, height, width = spatial_shape
    if (
        queries.shape[:2] != (batch, height * width)
        or keys.shape[0] != batch
        or values.shape[:2] != keys.shape[:2]
        or queries.shape[-1] != keys.shape[-1]
        or values.shape[-1] != output_projection.in_features
    ):
        raise ValueError(
            "Attention batch, token, feature, or spatial shapes mismatch"
        )
    scores = queries @ keys.transpose(-1, -2) / queries.shape[-1] ** 0.5
    weights = scores.softmax(dim=-1)
    attended = output_projection(weights @ values)
    residual = attended.transpose(1, 2).reshape(batch, -1, height, width)
    return residual, weights

class ConditionalDenoiser(nn.Module):
    """Reduced conditional epsilon network, not the paper's multiscale U-Net."""

    def __init__(self, latent_channels=2, condition_width=3, attention_width=4):
        super().__init__()
        if min(latent_channels, condition_width, attention_width) < 1:
            raise ValueError("Model widths must be positive")
        self.query = nn.Linear(latent_channels, attention_width, bias=False)
        self.key = nn.Linear(condition_width, attention_width, bias=False)
        self.value = nn.Linear(condition_width, attention_width, bias=False)
        self.output = nn.Linear(attention_width, latent_channels, bias=False)
        self.time = nn.Linear(1, latent_channels)
        self.conv = nn.Conv2d(latent_channels, latent_channels, 3, padding=1)

    def forward(self, noisy, timesteps, condition, total_steps):
        """Predict noise from current state, normalized time, and condition only."""
        require_tensor(noisy, 4, "noisy")
        if (
            not isinstance(total_steps, int)
            or total_steps < 1
            or not isinstance(timesteps, torch.Tensor)
            or timesteps.dtype != torch.long
            or timesteps.device.type != "cpu"
            or timesteps.shape != (noisy.shape[0],)
            or not ((timesteps >= 0) & (timesteps < total_steps)).all()
        ):
            raise ValueError("Denoiser time indices are invalid")
        queries, shape = spatial_queries(noisy, self.query)
        keys, values = condition_keys_values(condition, self.key, self.value)
        residual, weights = cross_attention(
            queries, keys, values, shape, self.output
        )
        time = ((timesteps.float() + 1) / total_steps).unsqueeze(1)
        time = self.time(time).unsqueeze(-1).unsqueeze(-1)
        predicted = self.conv(torch.tanh(noisy + residual + time))
        return predicted, weights


def epsilon_loss(predicted, target):
    """Use squared-vector norm per example, then batch mean (paper Eq. 3)."""
    require_tensor(predicted, 4, "predicted")
    require_tensor(target, 4, "target")
    if predicted.shape != target.shape:
        raise ValueError(
            "Noise target and prediction must have identical shapes"
        )
    return (predicted - target).square().flatten(1).sum(1).mean()


def train_step(
    model, optimizer, noisy, timesteps, condition, target, total_steps
):
    """Backpropagate noise regression and update only optimizer-owned parameters."""
    optimizer.zero_grad(set_to_none=True)
    predicted, weights = model(noisy, timesteps, condition, total_steps)
    loss = epsilon_loss(predicted, target)
    loss.backward()
    optimizer.step()
    return loss.detach(), predicted.detach(), weights.detach()

@torch.no_grad()
def reverse_and_decode(model, first_stage, initial, condition, alpha_bar):
    """Perform all T deterministic DDIM eta=0 updates, then decode once.

    The trace includes the initial z_T followed by each z_{t-1}. The final
    previous cumulative alpha is one, yielding the predicted clean latent.
    """
    require_tensor(initial, 4, "initial")
    validate_schedule(alpha_bar)
    if initial.shape[1] != first_stage.latent_channels:
        raise ValueError(
            "Latent channel count is incompatible with the decoder"
        )
    latent = initial.clone()
    trace = [latent.clone()]
    for index in reversed(range(len(alpha_bar))):
        timesteps = torch.full((latent.shape[0],), index, dtype=torch.long)
        epsilon, _ = model(latent, timesteps, condition, len(alpha_bar))
        current = alpha_bar[index]
        previous = (
            alpha_bar[index - 1] if index > 0 else current.new_tensor(1.0)
        )
        clean = (latent - (1 - current).sqrt() * epsilon) / current.sqrt()
        latent = previous.sqrt() * clean + (1 - previous).sqrt() * epsilon
        trace.append(latent.clone())
    decoded = first_stage.decoder(latent)
    return decoded, torch.stack(trace)

python -m pip install -r requirements-latentdiffusion.txt
python latentdiffusion.py --smoke --seed 7

Classifier-Free Guidance

At inference time, LDMs use classifier-free guidance (CFG) to control the trade-off between sample diversity and fidelity to the text prompt. During training, the text condition is randomly dropped (replaced with a null token) with probability p_{uncond} (a configuration-dependent choice). This trains the model to generate both conditionally and unconditionally.

At inference, each denoising step runs the model twice: once with the text prompt (\epsilon_\theta(z_t, t, y)) and once without (\epsilon_\theta(z_t, t, \varnothing)). The final noise prediction extrapolates beyond the conditional prediction:

The guidance scale w controls extrapolation. At w=1, this convention equals conditional sampling; larger values can improve condition fidelity at the expense of diversity and artifacts. The useful value depends on the trained model, task, and sampler. w=7.5 is familiar from particular Stable Diffusion examples, not a universal original-LDM optimum.

The cost of CFG is doubling the computation per denoising step (two forward passes). But because LDMs already operate in the compressed latent space, this overhead is far more manageable than it would be in pixel space.

How LDM Compares

Key Takeaways

1. Separate compression from generation. Train the first stage to retain perceptually useful structure, then train diffusion in its frozen spatial latent representation. A 48× tensor-element ratio in one configuration is not a universal compute speedup.

2. Cross-attention connects learned representations. It injects conditioning embeddings into latent features. Different modalities require compatible conditioning representations and training; some original image-conditioned LDMs use concatenation instead.

3. Choose the first stage for the task. Downsampling, latent channels, KL versus VQ regularization, and reconstruction fidelity all matter. Original LDM experiments use multiple choices, including f=4 and f=8.

4. Classifier-free guidance trades diversity for fidelity. By training with random condition dropout and extrapolating at inference, CFG gives users a single knob to control how closely outputs match the text prompt. This is more flexible than classifier guidance (which requires a separately trained classifier) and produces higher-quality results.

5. Reuse a first stage across diffusion trainings. Two-stage training supports modular development, but an arbitrary replacement autoencoder changes the latent distribution, scale, and potentially shape. A denoiser trained on the old representation generally needs compatibility validation or retraining; a better autoencoder is not automatically a drop-in improvement.

Impact and Legacy

The LDM paper, published at CVPR 2022, established a practical two-stage approach to high-resolution diffusion. Later Stable Diffusion releases built on latent-space denoising with specific autoencoders, conditioning encoders, and training choices. The original LDM code release and the later Stable Diffusion model releases are distinct artifacts.

The architecture has spawned an enormous ecosystem of improvements and applications:

• Stable Diffusion XL (SDXL) scaled the approach with a larger U-Net and a two-stage pipeline (base + refiner), producing higher-quality outputs at 1024×1024 resolution while maintaining the latent-space efficiency.

• Stable Diffusion 3 replaced the U-Net with a Diffusion Transformer (DiT) architecture, using the same latent-space principle but with transformer blocks instead of convolutional layers. This enabled better scaling and improved text rendering.

• ControlNet added trained control branches to diffusion networks for spatial control signals — edge maps, depth maps, pose skeletons — enabling precise compositional control over generated images.

• IP-Adapter, LoRA, and Textual Inversion showed that the modular architecture allowed efficient fine-tuning and personalization without retraining the full model, enabling users to adapt Stable Diffusion to specific styles, subjects, and domains.

• Latent Video Diffusion extended the approach to video generation, applying diffusion in a spatiotemporal latent space to produce coherent video sequences.

The fundamental insight is to learn a suitable image representation before modeling its distribution through diffusion. This has influenced many later image and video generators. It does not imply that every subsequent text-to-image system uses the original LDM architecture or even a latent diffusion process.

Related Reading

• VAE — the variational autoencoder whose encoder-decoder compresses images into the latent space that latent diffusion runs its denoising process in
• Attention Is All You Need — The transformer architecture underlying the text encoders and U-Net attention layers
• CLIP — The contrastive vision-language model used as the text encoder in Stable Diffusion
• Vision Transformer — The ViT architecture that inspired Diffusion Transformers (DiT)