Diffusion Models — Why Training Diverges at 10K Steps

Diffusion models have quietly staged a coup in generative AI. Stable Diffusion, DALL·E 2, Imagen, Sora — every one of these headline-grabbing systems is powered by the same elegant probabilistic idea first formalized in 2020. They've dethroned GANs as the dominant generative architecture not by being simpler, but by being more stable to train, more theoretically grounded, and dramatically better at capturing the full diversity of a data distribution without mode collapse.

The core problem every generative model must solve is: how do you learn to produce samples from a complex, high-dimensional distribution (e.g., all possible realistic photographs) when you only have a finite training set? GANs solved it with adversarial games that are notoriously hard to balance. VAEs solved it with a learned latent bottleneck that trades fidelity for tractability. Diffusion models solve it differently — by decomposing generation into thousands of tiny, individually tractable denoising steps, each one learned by a neural network. The math is cleaner, the training signal is more stable, and the results speak for themselves.

By the end of this article you'll understand the forward noising process and why it's designed the way it is, the reverse denoising process and the neural network that drives it, the mathematical connection to score matching and why that matters, the practical difference between DDPM and DDIM sampling, and how to implement a minimal but fully functional diffusion model in PyTorch. You'll also know the production gotchas that cost teams weeks to debug.

What is a Diffusion Model? — The Core Idea

A diffusion model is a generative model that learns to produce data from pure random noise through a sequential denoising process. The key insight is to decompose the complex task of generating a full image into thousands of small, tractable steps. Each step transforms a slightly noisy image into a slightly cleaner one. The model learns the reverse of a fixed forward process that gradually adds Gaussian noise.

The forward process (noising) is a Markov chain: given data x₀ ∼ q(x), we define q(x₁|x₀), q(x₂|x₁), ..., q(x_T|x_{T-1}) where each step adds small Gaussian noise. For T large enough, x_T is approximately isotropic Gaussian. The reverse process (denoising) is then learned: p_θ(x_{t-1}|x_t). The model is trained to maximize a variational lower bound on the data likelihood.

Why does this work? Because denoising a slightly noisy image is a much easier problem than generating a realistic image from scratch. The model can focus on local structure recovery, and the cumulative effect of many small corrections yields globally coherent outputs.

import torch
import torch.nn as nn
import math

def cosine_schedule(t, T, s=0.008):
    """Cosine variance schedule from Nichol & Dhariwal 2021."""
    return torch.cos(((t / T) + s) / (1 + s) * math.pi / 2) ** 2

class ForwardProcess:
    def __init__(self, betas):
        self.betas = betas
        self.alphas = 1. - betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)
    
    def q_sample(self, x_0, t, noise=None):
        """Sample x_t ~ q(x_t | x_0) in closed form."""
        if noise is None:
            noise = torch.randn_like(x_0)
        sqrt_alpha_bar = self.alpha_bars[t].sqrt()
        sqrt_one_minus_alpha_bar = (1. - self.alpha_bars[t]).sqrt()
        return sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise

The Forward (Noising) Process — Adding Chaos Methodically

The forward process is a fixed Markov chain that transforms data into noise over T steps. It's designed so that the distribution at any timestep can be computed directly from the original data without simulating all intermediate steps. This is crucial for efficient training.

Given a data point x₀, we define:

q(x_t | x_{t-1}) = N(x_t; sqrt(1 - β_t) x_{t-1}, β_t I)

where β_t is a predetermined variance schedule (e.g., linearly increasing from 1e-4 to 0.02). Using reparameterization, we can write:

x_t = sqrt(α_t) x_{t-1} + sqrt(1 - α_t) ϵ, where α_t = 1 - β_t and ϵ ∼ N(0,I)

By induction, we get the closed form:

x_t = sqrt(α̅_t) x₀ + sqrt(1 - α̅_t) ϵ, where α̅_t = ∏_{s=1}^{t} α_s

This means during training we can randomly sample a timestep t, compute the corresponding noisy image x_t from x₀ and ϵ, and train the model to predict ϵ from x_t. No iterative simulation needed.

def get_index_from_list(vals, t, x_shape):
    """Utility: get values at timestep t and broadcast to batch shape."""
    batch_size = t.shape[0]
    out = vals.gather(-1, t.cpu())
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def forward_diffusion_sample(x_0, t, device="cpu"):
    """
    Sample from q(x_t | x_0) given timestep t.
    Returns (x_t, noise) where noise = epsilon ~ N(0,1).
    """
    noise = torch.randn_like(x_0)
    sqrt_alpha_bar = get_index_from_list(sqrt_alpha_bar_t, t, x_0.shape)
    sqrt_one_minus_alpha_bar = get_index_from_list(sqrt_one_minus_alpha_bar_t, t, x_0.shape)
    x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise
    return x_t, noise

The Reverse (Denoising) Process — Learning to Unadd Noise

Now we need to learn the reverse: given a noisy image x_t, predict a slightly cleaner image x_{t-1}. The reverse process is also Gaussian with a learned mean:

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_t² I)

The variance σ_t² is fixed: σ_t² = β_t (or a learned diagonal). The mean μ_θ is parameterized as:

μ_θ(x_t, t) = 1/√α_t ( x_t - β_t / √(1-α̅_t) ϵ_θ(x_t, t) )

where ϵ_θ is the denoising U-Net that predicts the noise added between x₀ and x_t. This formulation reparameterizes the reverse step to predict noise instead of the clean image directly. Why noise? Because the noise has unit variance across all timesteps, making the loss well scaled.

Training uses a simple mean-squared error between the true noise ϵ and the predicted noise ϵ_θ:

L = ||ϵ - ϵ_θ(√{α̅_t} x₀ + √{1-α̅_t} ϵ, t)||²

class DenoisingUNet(nn.Module):
    """A simple U-Net for noise prediction. In practice, use a larger model."""
    def __init__(self, in_channels=3, base_ch=64):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, base_ch, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(base_ch, base_ch*2, 3, stride=2, padding=1),
            nn.SiLU(),
            nn.Conv2d(base_ch*2, base_ch*4, 3, stride=2, padding=1),
            nn.SiLU(),
        )
        self.decoder = nn.Sequential(\n            nn.Upsample(scale_factor=2),
            nn.Conv2d(base_ch*4, base_ch*2, 3, padding=1),
            nn.SiLU(),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(base_ch*2, base_ch, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(base_ch, in_channels, 3, padding=1),
        )
        # Time embedding: simple sinusoidal
        self.time_embed = nn.Sequential(\n            nn.Linear(1, base_ch*4),
            nn.SiLU(),
            nn.Linear(base_ch*4, base_ch*4),
        )

    def forward(self, x, t):
        # Add time embedding to each spatial position
        t_emb = self.time_embed(t.float().unsqueeze(-1))
        t_emb = t_emb.view(t_emb.shape[0], -1, 1, 1).expand(-1, -1, x.shape[2], x.shape[3])
        x = torch.cat([x, t_emb], dim=1)
        x = self.encoder(x)
        x = self.decoder(x)
        return x

Training: The Simplified Variational Loss

The full diffusion model is trained by minimizing the variational bound on the negative log-likelihood. This bound reduces to a sum of KL divergences between the true reverse conditional and the learned reverse conditional at each step. Remarkably, Ho et al. (2020) showed that a simplified loss — just the mean-squared error between true noise and predicted noise — works at least as well as the full bound:

L_simple = E_{t, x₀, ϵ} [ || ϵ - ϵ_θ(√{α̅_t} x₀ + √{1-α̅_t} ϵ, t) ||² ]

where t is uniformly sampled from {1, ..., T}. The uniform weighting over t works because the model sees all noise levels equally during training, which forces it to learn a consistent denoising function across the entire noise range.

In practice, we train with mini-batches, sampling a random t for each image in the batch. The U-Net takes both the noisy image x_t and the timestep t (as a sinusoidal embedding). This joint conditioning allows the model to behave differently at different noise levels.

import torch

def train_step(model, optimizer, x_0, forward_process):
    """Single training step for DDPM."""
    batch_size = x_0.shape[0]
    T = len(forward_process.betas)
    # Sample random timestep for each image
    t = torch.randint(0, T, (batch_size,), device=x_0.device).long()
    # Get noisy image and noise
    x_t, noise = forward_process.q_sample(x_0, t)
    # Predict noise
    predicted_noise = model(x_t, t)
    # Simple MSE loss
    loss = torch.nn.functional.mse_loss(predicted_noise, noise)
    optimizer.zero_grad()
    loss.backward()
    # Gradient clipping to prevent divergence
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    return loss.item()

Sampling: DDPM vs DDIM — The Speed-Quality Trade-off

Once trained, we generate new images by starting from pure noise x_T ∼ N(0,I) and iteratively applying the reverse step for t = T, T-1, ..., 1. This is the DDPM (Denoising Diffusion Probabilistic Models) sampler. It's stochastic: at each reverse step we sample from the predicted Gaussian:

x_{t-1} = μ_θ(x_t, t) + σ_t · z, where z ∼ N(0,I)

This stochasticity is what gives DDPM its high quality — it can correct errors from previous steps. The cost: we must run all T steps (typically 1000), making sampling slow.

DDIM (Denoising Diffusion Implicit Models) makes the process deterministic by setting σ_t = 0. This allows us to skip many steps during sampling. For example, we can sample only every 20th timestep (50 total steps). The quality degrades gracefully. DDIM also enables latent space interpolation: because the process is deterministic, you can travel between two generated images in noise space and get a smooth interpolation.

Which should you use? If quality is paramount and you have GPU time, use DDPM with T=1000. If you need fast sampling for deployment or experimentation, use DDIM with 50-200 steps.

@torch.no_grad()
def sample_ddpm(model, img_shape, forward_process, device='cpu'):
    """Generate an image using DDPM (stochastic, all T steps)."""
    x = torch.randn(img_shape, device=device)
    T = len(forward_process.betas)
    for t in reversed(range(T)):
        t_tensor = torch.full((img_shape[0],), t, device=device, dtype=torch.long)
        predicted_noise = model(x, t_tensor)
        alpha = forward_process.alphas[t]
        alpha_bar = forward_process.alpha_bars[t]
        # Compute predicted x_0
        pred_x0 = (x - (1 - alpha_bar).sqrt() * predicted_noise) / alpha_bar.sqrt()
        # Compute mean for x_{t-1}
        coef1 = alpha.sqrt() * (1 - alpha_bar) / (1 - alpha_bar)
        coef2 = (1 - alpha) / (1 - alpha_bar).sqrt() * alpha_bar.sqrt()  # simplified
        # Actually the standard formula:
        pred_mean = (1 / alpha.sqrt()) * (x - (1 - alpha) / (1 - alpha_bar).sqrt() * predicted_noise)
        if t == 0:
            x = pred_mean
        else:
            noise = torch.randn_like(x)
            sigma = forward_process.betas[t].sqrt()
            x = pred_mean + sigma * noise
    return x

@torch.no_grad()
def sample_ddim(model, img_shape, forward_process, steps=50, device='cpu'):
    """DDIM deterministic sampling with skipping."""
    T = len(forward_process.betas)
    skip = T // steps
    times = list(reversed(range(0, T, skip)))[:steps]
    x = torch.randn(img_shape, device=device)
    for i, t in enumerate(times):
        t_tensor = torch.full((img_shape[0],), t, device=device, dtype=torch.long)
        pred_noise = model(x, t_tensor)
        alpha_bar = forward_process.alpha_bars[t]
        pred_x0 = (x - (1 - alpha_bar).sqrt() * pred_noise) / alpha_bar.sqrt()
        # For DDIM, use next timestep's alpha_bar
        next_t = times[i+1] if i+1 < len(times) else 0
        alpha_bar_next = forward_process.alpha_bars[next_t]
        x = alpha_bar_next.sqrt() * pred_x0 + (1 - alpha_bar_next).sqrt() * pred_noise
    return x

Score Matching Connection — The Theoretical Foundation

Wait, there's a deeper connection. The noise prediction network ϵ_θ(x_t, t) is closely related to the score function of the data distribution — the gradient of the log-density at noise level t. Specifically:

ϵ_θ(x_t, t) ≈ -√{1 - α̅_t} · ∇_{x_t} log p(x_t)

This means that diffusion models are implicitly learning the score function at multiple noise levels. This perspective unifies them with score-based generative models (Song & Ermon, 2019). The denoising score matching objective (Vincent, 2011) is exactly what we're optimizing.

Why does this connection matter? Because it explains why diffusion models don't suffer from mode collapse: score-based models estimate the gradient of the data distribution, which is unique and identifies the full distribution. They can generate diverse samples without adversarial training.

Furthermore, the score matching view enables extensions like classifier-free guidance (where you combine conditional and unconditional score estimates) and accelerated sampling (e.g., via the Probability Flow ODE).

def score_from_noise(x_t, predicted_noise, alpha_bar):
    """Convert noise prediction to score estimate."""
    # score = - predicted_noise / sqrt(1 - alpha_bar)
    return -predicted_noise / torch.sqrt(1 - alpha_bar + 1e-8)

Latent Diffusion (LDM) — The Secret to High-Resolution Generation

Pixel-space diffusion is expensive: applying a U-Net to a 1024×1024 image is computationally prohibitive. Latent Diffusion Models (LDM), introduced by Rombach et al. (2022) and used in Stable Diffusion, solve this by compressing the image into a lower-dimensional latent space via a pretrained autoencoder. The diffusion process then runs in this latent space, which is 4× to 64× smaller in spatial dimensions.

The architecture consists of three components: 1. A VAE (vector quantized or continuous) that maps images to latents and back. The encoder compresses 256×256×3 to 64×64×4 (or 32×32×4). 2. A U-Net denoiser that operates on the latent representation. It is conditioned on the timestep and optionally on text embeddings via cross-attention. 3. A decoder that reconstructs the image from the denoised latent.

Because the latent space is much smaller, the U-Net can be shallower and the number of forward passes is drastically reduced. This makes training feasible on a single consumer GPU and enables high-resolution synthesis. For example, Stable Diffusion's U-Net has about 860M parameters but runs in seconds on an A100.

The key insight: the VAE's latent space is perceptually equivalent to the pixel space but with reduced spatial redundancy. The diffusion model learns the distribution of these perceptually compressed latents. Conditioning mechanisms (text, segmentation maps, etc.) are injected via cross-attention layers in the U-Net.

import torch
import torch.nn as nn

class LatentDiffusion(nn.Module):
    """Minimal LDM: VAE + Denoising U-Net on latents."""
    def __init__(self, vae_encoder, vae_decoder, denoiser_unet):
        super().__init__()
        self.encoder = vae_encoder  # pretrained; frozen during diffusion training
        self.decoder = vae_decoder  # pretrained; frozen
        self.denoiser = denoiser_unet  # trained on latent noise prediction
    
    def encode(self, x):
        # Returns latent z with shape [B, C_latent, H_latent, W_latent]
        return self.encoder(x)
    
    def decode(self, z):
        # Returns reconstructed image
        return self.decoder(z)
    
    def forward(self, x_0, t, noise=None):
        # Diffusion training in latent space
        z_0 = self.encode(x_0)
        noise = torch.randn_like(z_0) if noise is None else noise
        z_t, noise = self.q_sample(z_0, t, noise)
        noise_pred = self.denoiser(z_t, t)
        loss = nn.functional.mse_loss(noise_pred, noise)
        return loss

Generative Model Comparison — Stability, Speed, and Quality

Choosing the right generative architecture for a production application requires understanding the trade-offs between training stability, sampling speed, and output quality. The table below compares GANs, VAEs, Flows, and Diffusion models across these axes.

Key takeaways for production: If you need real-time generation (e.g., interactive avatars), GANs are still viable. For tasks requiring high fidelity and diversity (e.g., stock image generation), diffusion models are now the default. VAEs are unmatched for anomaly detection due to their reconstruction likelihood. Flows are rarely used in production due to large model sizes.

Visual ControlNet Guide — Structure-Conditioned Generation

ControlNet (Zhang & Agrawala, 2023) is a neural network architecture that adds spatial conditioning to pretrained diffusion models without requiring full fine-tuning. It works by copying the encoder blocks of the U-Net and connecting them via zero-initialised convolutional layers (zero convolutions). The copied weights are trainable side branches that learn to control the generation based on an input condition (e.g., edge maps, depth maps, pose skeletons).

The beauty of ControlNet is that it preserves the knowledge of the base model: the side branches start from zeros, so the model initially generates unconditionally. Only the side branch weights are updated during training, leaving the original U-Net untouched. This makes ControlNet extremely parameter-efficient: you can train a new condition with less than 5% of the data and time required for a full fine-tune.

Training ControlNet on canny edges: The user provides an edge map as input. The side branch encodes it into features at multiple resolutions, which are added to the U-Net skip connections via zero convs. The base U-Net remains frozen. After training on 50K–200K image–condition pairs, the model learns to generate images that respect the edge structure.

For production, ControlNet is typically used with Stable Diffusion. The pipeline is: input image → condition extractor (e.g., Canny edge detector, depth estimator) → ControlNet side branch → standard denoising steps. The result is a generated image that faithfully follows the input structure.

import torch
import torch.nn as nn

class ZeroConv2d(nn.Module):
    """Zero-initialised convolution: starts as identity, learns slowly."""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 1, bias=False)
        nn.init.zeros_(self.conv.weight)
    
    def forward(self, x):
        return self.conv(x)

class ControlNetSideBranch(nn.Module):
    """Copies encoder layers from the main U-Net, connected via zero convs."""
    def __init__(self, unet_encoder_blocks):
        super().__init__()
        # Deep-copy each encoder block (example for first block)
        self.blocks = nn.ModuleList()
        self.zero_convs = nn.ModuleList()
        for block in unet_encoder_blocks:
            # Shallow copy is enough for architecture; weights are separate
            self.blocks.append(block.__class__(*block.arg_params))  # simplified
            self.zero_convs.append(ZeroConv2d(block.out_channels, block.out_channels))
    
    def forward(self, x, feature_maps_to_skip):
        outputs = []
        for block, zc, skip in zip(self.blocks, self.zero_convs, feature_maps_to_skip):
            x = block(x)
            x = x + zc(skip)  # add conditioned features via zero conv
            outputs.append(x)
        return outputs

Keras/TensorFlow Implementation — Forward, Training, and Sampling

While PyTorch dominates research, TensorFlow and Keras are still widely used in production pipelines, especially for serving on Google Cloud, TFX, or mobile (TFLite). Below is a minimal Keras implementation of the key diffusion components: forward noising, a simple U-Net, and the training step.

Forward process in TensorFlow – the closed-form sampling works identically. We'll use TensorFlow's vectorised operations.

U-Net – a Keras model with time embedding. Note that Keras does not have a built-in SiLU (swish) in older versions, so we use tf.keras.activations.swish.

Training step – written as a custom training loop or compiled model. The code below shows a train_step for a custom fit override or standalone.

Keras' fit expects model inputs and outputs. We can create a functional model that takes [x_0, t] and outputs ϵ_pred. Then we compile with MSE loss and train with a custom data generator that samples t randomly.

The key difference from PyTorch is the need to handle device placement manually (or rely on tf.distribute for multi-GPU). Also, batch normalisation (BatchNormalization) is the default in Keras — remember to replace with GroupNormalization (available in Keras 3 or via tensorflow_addons).

import tensorflow as tf

def get_alpha_bar(betas):
    alphas = 1.0 - betas
    return tf.math.cumprod(alphas, axis=0)

def q_sample(x_0, t, alpha_bar, noise=None):
    noise = tf.random.normal(tf.shape(x_0)) if noise is None else noise
    sqrt_alpha_bar = tf.gather(alpha_bar, t)[:, tf.newaxis, tf.newaxis, tf.newaxis]
    sqrt_one_minus = tf.sqrt(1.0 - tf.gather(alpha_bar, t))[:, tf.newaxis, tf.newaxis, tf.newaxis]
    return sqrt_alpha_bar * x_0 + sqrt_one_minus * noise, noise

def sinusoidal_embedding(timesteps, embedding_dim):
    # Standard sinusoidal time embedding
    half_dim = embedding_dim // 2
    emb = tf.math.log(10000.0) / (half_dim - 1)
    emb = tf.exp(tf.range(half_dim, dtype=tf.float32) * -emb)
    emb = tf.cast(timesteps[:, tf.newaxis], tf.float32) * emb[tf.newaxis, :]
    return tf.concat([tf.sin(emb), tf.cos(emb)], axis=-1)

def build_unet(input_shape=(64,64,3), base_channels=64):
    # Input: noisy image + time embedding
    image_input = tf.keras.Input(shape=input_shape, name='noisy_image')
    t_input = tf.keras.Input(shape=(), name='timestep', dtype=tf.int32)
    
    # Time embedding -> Dense -> reshape to spatial
    t_emb = sinusoidal_embedding(t_input, base_channels * 4)
    t_dense = tf.keras.layers.Dense(base_channels * 4, activation='swish')(t_emb)
    t_dense = tf.keras.layers.Dense(base_channels * 4, activation='swish')(t_dense)
    # reshape to (batch, 1, 1, channels) for addition
    t_dense = tf.reshape(t_dense, (-1, 1, 1, base_channels * 4))
    
    # Encoder (using GroupNorm from tensorflow_addons or custom)
    x = tf.keras.layers.Conv2D(base_channels, 3, padding='same')(image_input)
    x = tfa.layers.GroupNormalization(groups=32)(x)  # requires tensorflow-addons
    x = tf.keras.layers.Activation('swish')(x)
    # ... more blocks
    
    # Decoder with skip connections (simplified)
    # For brevity, we return a placeholder output
    output = tf.keras.layers.Conv2D(3, 3, padding='same')(x)
    model = tf.keras.Model(inputs=[image_input, t_input], outputs=output)
    return model

# Training loop
@tf.function
def train_step(model, optimizer, x_0, betas):
    alpha_bar = get_alpha_bar(betas)
    batch_size = tf.shape(x_0)[0]
    t = tf.random.uniform((batch_size,), maxval=len(betas), dtype=tf.int32)
    x_t, noise = q_sample(x_0, t, alpha_bar)
    with tf.GradientTape() as tape:
        pred = model([x_t, t])
        loss = tf.reduce_mean(tf.square(noise - pred))
    grads = tape.gradient(loss, model.trainable_variables)
    # Gradient clipping
    grads, _ = tf.clip_by_global_norm(grads, 1.0)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss