Diffusion Models from Scratch: DDPM Training in PyTorch

The Forward Process: Adding Noise Until Nothing Remains

Diffusion models work by destroying data through gradual noise injection, then learning to reverse that destruction. That’s it. No adversarial training, no mode collapse nightmares, no discriminator to babysit. Just a deterministic noising schedule and a neural network that learns to denoise.

The math is surprisingly elegant once you stop trying to understand it from VAE analogies.

The forward diffusion process takes an image $x_0$ and progressively adds Gaussian noise over $T$ timesteps until you’re left with pure noise $x_T \sim \mathcal{N}(0, I)$. Each step follows:

$$x_t = \sqrt{1 – \beta_t} \cdot x_{t-1} + \sqrt{\beta_t} \cdot \epsilon_t$$

where $\beta_t$ is a small variance schedule (typically 0.0001 to 0.02) and $\epsilon_t \sim \mathcal{N}(0, I)$. The beauty is you can sample $x_t$ directly from $x_0$ without computing all intermediate steps:

$$x_t = \sqrt{\bar{\alpha}_t} \cdot x_0 + \sqrt{1 – \bar{\alpha}_t} \cdot \epsilon$$

where $\alpha_t = 1 – \beta_t$ and $\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s$. This closed-form sampling is what makes training efficient — you can jump to any timestep in constant time.

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class DiffusionSchedule:
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps
        # Linear schedule (DDPM paper uses this)
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), self.alphas_cumprod[:-1]])

        # Pre-compute values for closed-form sampling
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)

        # For reverse process
        self.sqrt_recip_alphas = torch.sqrt(1.0 / self.alphas)
        self.posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)

    def q_sample(self, x_start, t, noise=None):
        """Forward diffusion: sample x_t from x_0 directly."""
        if noise is None:
            noise = torch.randn_like(x_start)

        sqrt_alpha_prod = self.sqrt_alphas_cumprod[t].reshape(-1, 1, 1, 1)
        sqrt_one_minus_alpha_prod = self.sqrt_one_minus_alphas_cumprod[t].reshape(-1, 1, 1, 1)

        return sqrt_alpha_prod * x_start + sqrt_one_minus_alpha_prod * noise

The reshaping here matters — broadcasting failures are the #1 source of silent bugs in diffusion code. Always verify tensor shapes after each operation.

The U-Net: Why Everyone Copies the Same Architecture

The denoising model is typically a U-Net with time embeddings, attention layers, and residual connections. The original DDPM paper used a specific architecture that works well enough that most implementations just copy it.

The time embedding is crucial. The network needs to know which noise level it’s denoising, so you encode the timestep $t$ using sinusoidal positional embeddings (borrowed from Transformers):

$$PE(t, 2i) = \sin\left(\frac{t}{10000^{2i/d}}\right), \quad PE(t, 2i+1) = \cos\left(\frac{t}{10000^{2i/d}}\right)$$

These embeddings get projected and injected into each residual block.

class SinusoidalPositionEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, timesteps):
        device = timesteps.device
        half_dim = self.dim // 2
        embeddings = np.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = timesteps[:, None] * embeddings[None, :]
        embeddings = torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
        return embeddings

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, dropout=0.1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.time_mlp = nn.Linear(time_emb_dim, out_channels)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.norm2 = nn.GroupNorm(8, out_channels)
        self.dropout = nn.Dropout(dropout)

        # Skip connection needs dimension matching
        self.skip = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()

    def forward(self, x, time_emb):
        h = self.conv1(x)
        h = self.norm1(h)
        h = nn.functional.silu(h)  # Swish activation

        # Inject time embedding
        time_emb = self.time_mlp(time_emb)
        h = h + time_emb[:, :, None, None]  # Broadcasting trick

        h = self.conv2(h)
        h = self.norm2(h)
        h = self.dropout(h)
        h = nn.functional.silu(h)

        return h + self.skip(x)

I’m simplifying the full U-Net here — a production version needs downsampling blocks, upsampling blocks, skip connections between encoder/decoder, and self-attention layers at lower resolutions (typically 16×16). The full architecture is ~200 lines. What matters is the time embedding injection pattern.

Training: The Simplified Objective That Actually Works

The DDPM paper derives a complex variational lower bound, but the training objective simplifies to:

$$L_{\text{simple}} = \mathbb{E}_{t, x_0, \epsilon} \left[ \| \epsilon – \epsilon_{\theta}(x_t, t) \|^2 \right]$$

You’re just training the network to predict the noise $\epsilon$ that was added at timestep $t$. That’s it. No KL divergence terms, no weighting by timestep (though you can add that back for better sample quality).

Here’s the training loop:

def train_diffusion(model, dataloader, schedule, epochs=100, device='cuda'):
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
    model.to(device)

    for epoch in range(epochs):
        total_loss = 0
        for batch_idx, (images, _) in enumerate(dataloader):
            images = images.to(device)
            batch_size = images.shape[0]

            # Sample random timesteps
            t = torch.randint(0, schedule.timesteps, (batch_size,), device=device, dtype=torch.long)

            # Sample noise
            noise = torch.randn_like(images)

            # Forward diffusion
            x_t = schedule.q_sample(images, t, noise=noise)

            # Predict noise
            predicted_noise = model(x_t, t)

            # Simple MSE loss
            loss = nn.functional.mse_loss(predicted_noise, noise)

            optimizer.zero_grad()
            loss.backward()

            # Gradient clipping is CRITICAL — without it you'll hit NaN around epoch 5-10
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

        # The loss should drop from ~0.5 to ~0.02-0.05 on MNIST/CIFAR-10
        if avg_loss < 1e-4:
            print("Warning: Loss suspiciously low, check for bugs")

Gradient clipping saved me hours of debugging. Without it, you’ll see smooth training for a few epochs, then sudden NaN losses when large noise predictions cause gradient explosion. PyTorch 2.x is more stable than 1.x here, but clipping is still mandatory.

Sampling: The Reverse Process That Takes 1000 Steps

Generation starts from pure noise $x_T \sim \mathcal{N}(0, I)$ and iteratively denoises using the learned model:

$$x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t – \frac{1 – \alpha_t}{\sqrt{1 – \bar{\alpha}_t}} \epsilon_{\theta}(x_t, t) \right) + \sigma_t z$$

where $z \sim \mathcal{N}(0, I)$ and $\sigma_t^2 = \beta_t$ (or the posterior variance for better results). This is the DDPM sampling algorithm.

@torch.no_grad()
def sample(model, schedule, image_size=(1, 28, 28), batch_size=16, device='cuda'):
    model.eval()
    # Start from pure noise
    x_t = torch.randn(batch_size, *image_size, device=device)

    # Reverse process: t=T down to t=1
    for t in reversed(range(schedule.timesteps)):
        t_batch = torch.full((batch_size,), t, device=device, dtype=torch.long)

        # Predict noise
        predicted_noise = model(x_t, t_batch)

        # Compute alpha values
        alpha_t = schedule.alphas[t]
        alpha_t_cumprod = schedule.alphas_cumprod[t]

        # Mean of p(x_{t-1} | x_t)
        mean = (1.0 / torch.sqrt(alpha_t)) * (
            x_t - ((1.0 - alpha_t) / torch.sqrt(1.0 - alpha_t_cumprod)) * predicted_noise
        )

        # Add noise (except at t=0)
        if t > 0:
            noise = torch.randn_like(x_t)
            variance = schedule.posterior_variance[t]
            x_t = mean + torch.sqrt(variance) * noise
        else:
            x_t = mean

    return x_t

This takes 1000 forward passes through the network. On my setup (RTX 3090, PyTorch 2.1), generating a 32×32 image takes ~2 seconds. DDIM sampling reduces this to 50-100 steps with minimal quality loss, but that’s a separate algorithm.

The Training Instability No One Warns You About

Around epoch 15-20 on CIFAR-10, you might see the loss suddenly plateau or even increase slightly. This isn’t divergence — it’s the model learning to denoise harder timesteps (high noise levels) where the task is genuinely difficult. The early timesteps (low noise) are already nearly perfect.

I wasted a day thinking my learning rate was too high. The fix is weighted loss by timestep, but honestly, just training longer works too.

Another gotcha: GroupNorm with too few groups. If you use nn.GroupNorm(32, channels) but channels=16, PyTorch will silently fail or produce garbage. Always ensure num_groups divides num_channels.

Why This Beats GANs for Image Generation

No mode collapse. No discriminator hyperparameter search. No balancing act between generator and discriminator learning rates. You just… train the model. The loss goes down. The samples improve.

The tradeoff is inference speed. GANs generate images in one forward pass. Diffusion models need 50-1000 steps. For real-time applications, GANs still win (face filters, live video style transfer). For highest quality samples where you can wait 5 seconds, diffusion dominates.

Stable Diffusion, DALL-E 2, Imagen — all diffusion-based. The architecture has won the generative modeling wars for static image synthesis.

Classifier-Free Guidance: The Secret Sauce in Modern Models

Conditional generation (text-to-image, class-conditioned) uses classifier-free guidance. You train two models simultaneously: one conditioned on labels, one unconditional. During sampling, you extrapolate:

$$\tilde{\epsilon}_{\theta}(x_t, t, c) = \epsilon_{\theta}(x_t, t, \emptyset) + s \cdot \left( \epsilon_{\theta}(x_t, t, c) – \epsilon_{\theta}(x_t, t, \emptyset) \right)$$

where $s > 1$ is the guidance scale (typically 7-9). Higher $s$ means stronger conditioning but less diversity. I covered the original classifier-free guidance paper in detail — it’s the technique that unlocked high-fidelity text-to-image generation.

Memory Requirements and Training Time

On MNIST (28×28 grayscale), a minimal U-Net (4 downsample blocks, 64 base channels) trains in ~30 minutes on a single RTX 3090 (50 epochs, batch size 128). Peak GPU memory: 2GB.

CIFAR-10 (32×32 RGB) with a standard U-Net (128 base channels, attention at 16×16) takes ~6 hours for 200 epochs. Peak memory: 8GB. The loss converges to ~0.03-0.05 (lower is better, but don’t expect below 0.02 unless you add loss weighting).

ImageNet at 256×256 resolution? You’re looking at multiple A100s and days of training. Latent diffusion (Stable Diffusion’s approach) fixes this by diffusing in a compressed latent space instead of pixel space — cuts training cost by 10x.

Common Bugs That Waste Hours

NaN losses after epoch 5-10: Gradient explosion. Add torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) before optimizer step.

Blurry samples: Check your variance schedule. If $\beta_{\text{end}}$ is too small (< 0.01), the model doesn’t add enough noise at high timesteps. Increase it to 0.02.

Training loss stuck at 0.5: You’re probably not normalizing input images to [-1, 1]. Diffusion models expect zero-mean data. Use transforms.Normalize((0.5,), (0.5,)) for grayscale or transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) for RGB.

Samples are pure noise: The model didn’t train at all. Check that timestep embeddings are being injected correctly. Print time_emb.shape inside ResidualBlock — it should match (batch_size, time_emb_dim).

What I Still Don’t Fully Understand

Why does the linear noise schedule work so well? The original DDPM paper tried cosine schedules, learned schedules, all sorts of variants. Linear beat everything empirically, but there’s no strong theoretical justification.

And the weighting issue: the simplified objective treats all timesteps equally, but intuitively, denoising at $t=999$ (pure noise) should be easier than $t=10$ (nearly clean). Yet equal weighting works. My best guess is that the network’s capacity automatically balances the difficulty — it spends more parameters on hard timesteps.

But honestly, I’m not entirely sure. The math checks out, the samples look good, and that’s good enough for now.

Debugging Diffusion Training at 2am

When your loss refuses to drop below 0.3 and you’ve checked the schedule, the normalization, the gradients, and the architecture — try visualizing $x_t$ at different timesteps. Just save q_sample(x_0, t) for $t \in \{10, 100, 500, 999\}$ and plot them.

If $t=500$ still looks recognizable, your noise schedule is broken. If $t=10$ is already unrecognizable, you’re adding noise too aggressively. This saved me once when I accidentally swapped sqrt_alphas_cumprod and sqrt_one_minus_alphas_cumprod in the sampling function. The bug was invisible in the loss curve.

And if you’re debugging at 2am, Dark Chocolate Espresso Beans are genuinely the only thing that keeps the brain functional. The caffeine-to-sugar ratio is perfect for gradient descent debugging.

Latent Diffusion: Why Stable Diffusion Doesn’t Diffuse Pixels

The pixel-space approach I described works for 32×32 or 64×64 images, but collapses at 512×512. The memory and compute requirements scale quadratically with resolution.

Stable Diffusion solves this by training a separate autoencoder (VAE) to compress images into a low-dimensional latent space (e.g., 512×512 → 64x64x4), then running diffusion in that latent space. The VAE decoder upsamples back to pixels at the end.

This is a huge architectural win — you get high-resolution outputs without the compute cost. The tradeoff is training the VAE first, which adds complexity. For research and learning, pixel-space diffusion is simpler.

DDIM: Faster Sampling Without Retraining

DDPM sampling is slow because it’s a Markov chain — each step depends on the previous one, so you can’t parallelize. DDIM (Denoising Diffusion Implicit Models) reparameterizes the sampling process to skip timesteps:

$$x_{t-\Delta t} = \sqrt{\bar{\alpha}_{t-\Delta t}} \cdot \hat{x}_0 + \sqrt{1 – \bar{\alpha}_{t-\Delta t}} \cdot \epsilon_{\theta}(x_t, t)$$

where $\hat{x}_0 = \frac{x_t – \sqrt{1 – \bar{\alpha}_t} \cdot \epsilon_{\theta}(x_t, t)}{\sqrt{\bar{\alpha}_t}}$ is the predicted clean image. You can set $\Delta t = 20$ and sample in 50 steps instead of 1000.

The math is non-trivial (involves solving an ODE instead of an SDE), but the key insight is that DDPM’s Markov assumption is unnecessary. You can use the same trained model with either sampler.

When to Use Diffusion vs Autoregressive Models

For images: diffusion wins. PixelCNN and other autoregressive models are too slow and don’t scale to high resolution.

For discrete data (text, code): autoregressive models (GPT, LLaMA) still dominate. Diffusion on discrete spaces requires awkward workarounds (continuous relaxations, absorbing states). It’s an active research area, but as of 2026, transformers own the language modeling niche.

For video: still unclear. Diffusion models can generate coherent short clips (Runway, Pika), but temporal consistency over 10+ seconds is hard. Autoregressive models struggle with the same issue. Hybrid approaches might be the future.

FAQ

Q: How many timesteps should I use for training?

1000 is standard (from the DDPM paper). You can go lower (250-500) for faster training, but sample quality drops slightly. Going higher (2000+) doesn’t help much and wastes compute. For DDIM sampling, the training timestep count doesn’t matter — you can still sample in 50 steps.

Q: Why does my model generate the same image repeatedly?

You’re probably not re-sampling the initial noise $x_T$ for each generation. Make sure x_t = torch.randn(...) is inside your sampling loop, not outside. If $x_T$ is fixed, the output will always be identical (diffusion is deterministic given the initial noise and model weights).

Q: Can I fine-tune a pretrained diffusion model on custom data?

Yes, and it’s way easier than fine-tuning GANs. Just load the pretrained weights and continue training on your dataset. The noise schedule and architecture can stay unchanged. I’d recommend lowering the learning rate to ~1e-5 to avoid catastrophic forgetting. Stable Diffusion fine-tuning (DreamBooth, LoRA) follows this principle but adds parameter-efficient tricks.

The Real Advantage: You Can Actually Train These

GANs require babysitting. VAEs produce blurry samples. Autoregressive models are slow. Diffusion models just… work. The training is stable, the loss correlates with sample quality, and the architecture is well-understood.

The 1000-step sampling is annoying, but DDIM cuts it to 50 with minimal quality loss. And for static image generation, 2 seconds per sample is acceptable for most applications.

The field is moving fast — flow matching, consistency models, and rectified flows are newer variants that promise better speed-quality tradeoffs. But DDPM is still the foundation. Learn this first, then explore the optimizations.

I’m still curious whether score-based diffusion (the Langevin dynamics formulation) has practical advantages over DDPM beyond theoretical elegance. The math is prettier, but the code looks identical. Need to run side-by-side experiments to know for sure.