Diffusion Models for CV: DDPM, Noise Scheduling, U-Net for Denoising, Sampling

🎯 Learning Objectives

• Understand the forward and reverse diffusion processes
• Implement the U-Net architecture for denoising
• Configure noise schedules (linear, cosine, sigmoid)
• Train and sample from a DDPM
• Apply diffusion models for image generation tasks

📋 Prerequisites

• CNN Fundamentals (Week 1): Convolution, pooling
• Segmentation (Week 5): U-Net architecture
• Basic probability: Gaussian distributions

1. 📖 Core Content

1.1 Intuition: Destroy and Learn to Undo

1. Forward process: Gradually add Gaussian noise to an image until it's pure noise
2. Reverse process: Learn to denoise step by step, starting from pure noise to generate a new image This is like taking a photo, slowly crumpling it into a ball of paper, then learning to uncrumple it.

1.2 The DDPM Framework

Forward Process

Reverse Process (Learned)

Training Objective (simplified)

1.3 U-Net for Denoising

python
Run Inline
# runnable
import torch
import torch.nn as nn
class TimeEmbedding(nn.Module):
    """Sinusoidal time embedding as in Transformer."""
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
    def forward(self, t):
        half_dim = self.dim // 2
        embeddings = torch.log(torch.tensor(10000.)) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim) * -embeddings)
        embeddings = t[:, None].float() * embeddings[None, :]
        return torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
class DenoiseUNet(nn.Module):
    """Simplified U-Net with time conditioning."""
    def __init__(self, in_channels=3, time_dim=256):
        super().__init__()
        self.time_mlp = nn.Sequential(
            TimeEmbedding(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.ReLU()
        )
        # Encoder
        self.enc1 = nn.Conv2d(in_channels, 64, 3, padding=1)
        self.enc2 = nn.Conv2d(64, 128, 3, padding=1)
        self.enc3 = nn.Conv2d(128, 256, 3, padding=1)
        # Decoder (with skip connections)
        self.dec3 = nn.Conv2d(256 + 256, 128, 3, padding=1)
        self.dec2 = nn.Conv2d(128 + 128, 64, 3, padding=1)
        self.dec1 = nn.Conv2d(64 + 64, in_channels, 3, padding=1)
        # Time conditioning projections
        self.time_proj1 = nn.Linear(time_dim, 64)
        self.time_proj2 = nn.Linear(time_dim, 128)
        self.time_proj3 = nn.Linear(time_dim, 256)
    def forward(self, x, t):
        t_emb = self.time_mlp(t)
        # Encoder
        x1 = torch.relu(self.enc1(x) + self.time_proj1(t_emb)[:, :, None, None])
        x2 = torch.relu(self.enc2(nn.MaxPool2d(2)(x1)) + self.time_proj2(t_emb)[:, :, None, None])
        x3 = torch.relu(self.enc3(nn.MaxPool2d(2)(x2)) + self.time_proj3(t_emb)[:, :, None, None])
        # Decoder
        x = torch.relu(self.dec3(torch.cat([x3, nn.Upsample(scale_factor=2)(x3)], dim=1)))
        x = torch.relu(self.dec2(torch.cat([x, x2], dim=1)))
        x = self.dec1(torch.cat([x, x1], dim=1))
        return x

1.4 Sampling Algorithm

python
Run Inline
# runnable
@torch.no_grad()
def sample_ddpm(model, image_size, channels=3, n_steps=1000, device='cuda'):
    """Generate an image by iteratively denoising."""
    # Start from pure noise
    x = torch.randn(1, channels, image_size, image_size).to(device)
    # Pre-compute noise schedule (cosine schedule recommended)
    betas = cosine_beta_schedule(n_steps).to(device)
    alphas = 1 - betas
    alphas_bar = torch.cumprod(alphas, dim=0)
    # Reverse diffusion
    for t in reversed(range(1, n_steps)):
        t_batch = torch.full((1,), t, device=device, dtype=torch.long)
        predicted_noise = model(x, t_batch)
        # Denoising step
        alpha_t = alphas[t]
        alpha_bar_t = alphas_bar[t]
        alpha_bar_prev = alphas_bar[t-1] if t > 0 else torch.tensor(1.0)
        # Predict x_0 from current x_t and predicted noise
        x_0_pred = (x - torch.sqrt(1 - alpha_bar_t) * predicted_noise) / torch.sqrt(alpha_bar_t)
        # Compute x_{t-1}
        sigma_t = torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar_t) * betas[t])
        x = (1 / torch.sqrt(alpha_t)) * (
            x - betas[t] / torch.sqrt(1 - alpha_bar_t) * predicted_noise
        )
        # Add noise (except for last step)
        if t > 1:
            x += sigma_t * torch.randn_like(x)
    return x

1.5 Why This Matters

2. 📐 Key Formulas / Concepts

3. ⚠️ Common Pitfalls

Pitfall 1: Incorrect Beta Schedule

Pitfall 2: Not Scaling Time Embeddings

4. 📝 Practice Questions

Q1: Your DDPM generates very blurry images after 1000 training steps. The loss is still decreasing. What should you do?

Diagnosis: The model hasn't converged. DDPM training typically requires 100K-500K steps for reasonable results. Loss decreasing is normal.

Fixes:

1. Train longer: DDPM converges slowly. Use 100K+ steps.
2. Check noise schedule: Use cosine schedule (not linear).
3. Larger model: U-Net with more channels and deeper layers generates sharper images.
4. EMA (Exponential Moving Average): Track EMA of model parameters and use EMA for sampling.
5. Reduce learning rate: Use LR=1e-4 with cosine decay.

5. 🔗 Cross-References

• Previous: Image Processing (Week 9)
• Related: GenAI Diffusion Models
• Related: GenAI DDIM
• External: Ho et al., "Denoising Diffusion Probabilistic Models" (NeurIPS 2020) Join Discord PreviousImage ProcessingNextBSDA5006 — Deep Learning for Computer Vision (DL-CV)