Diffusion Models

1. Forward Diffusion Process

1.1 Definition

The forward process adds Gaussian noise to data over timesteps:

where is the noise schedule.

1.2 Closed-Form Marginals

Define and . Then:

This allows sampling at any timestep directly:

1.3 Noise Schedules

Linear schedule:

Cosine schedule (Nichol & Dhariwal, 2021):

Sigmoid schedule:

2. Reverse Diffusion Process

2.1 Reverse Process Definition

The reverse process removes noise to generate data:

2.2 Optimal Reverse Process

When is Gaussian, the true reverse posterior is:

where

2.3 Reparameterisation with Noise Prediction

Using :

where is the noise prediction network.

3. DDPM Training Objective

3.1 ELBO Derivation

The variational lower bound (ELBO) is:

3.2 Simplified Objective

Ho et al. (2020) showed that predicting the noise is equivalent to minimising:

This is the denoising score matching objective.

3.3 Connection to Score Matching

The score function is . The noise prediction network relates to the score as:

So the diffusion model learns the score function at each noise level.

4. Score Matching

4.1 Score Function Definition

The score function is the gradient of the log-density:

4.2 Denoising Score Matching

The denoising score matching objective is:

For Gaussian noise :

4.3 NCSN (Noise Conditional Score Networks)

Train a single network for multiple noise levels:

5. Classifier-Free Guidance

5.1 Formulation

Classifier-free guidance (Ho & Salimans, 2022) trains a conditional model that can be interpolated with unconditional generation:

where is the guidance scale and is the null condition (unconditional).

5.2 Guidance Scale

• : Standard conditional generation
• : Amplifies conditional signal (higher quality, lower diversity)
• : Dampens conditional signal
• : Unconditional generation

5.3 Theoretical Interpretation

Classifier-free guidance can be viewed as implicit classifier guidance:

This is equivalent to:

6. Flow Matching

6.1 Continuous Normalising Flows

Flow matching (Lipman et al., 2023) defines a continuous transformation from noise to data:

where is the velocity field.

6.2 Probability Flow ODE

The probability flow ODE is:

This provides an exact likelihood computation.

6.3 Flow Matching Objective

where .

6.4 Advantages over DDPM

7. Consistency Models

7.1 Definition

Consistency models (Song et al., 2023) map any point on the probability flow ODE trajectory to the origin:

This enables one-step generation.

7.2 Consistency Property

The model must satisfy:

7.3 Training

Consistency distillation:

where is obtained from by one step of the diffusion ODE, and is the EMA of .

8. Code Examples

8.1 DDPM Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SinusoidalPositionEmbedding(nn.Module):
    """Sinusoidal embedding for timestep conditioning."""
    
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
    
    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None] * emb[None, :]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
        return emb


class ResBlock(nn.Module):
    """Residual block with timestep conditioning."""
    
    def __init__(self, in_ch, out_ch, time_emb_dim):
        super().__init__()
        
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        
        self.time_mlp = nn.Linear(time_emb_dim, out_ch)
        
        self.bn1 = nn.BatchNorm2d(out_ch)
        self.bn2 = nn.BatchNorm2d(out_ch)
        
        if in_ch != out_ch:
            self.shortcut = nn.Conv2d(in_ch, out_ch, 1)
        else:
            self.shortcut = nn.Identity()
    
    def forward(self, x, t_emb):
        h = F.gelu(self.bn1(self.conv1(x)))
        
        # Add time embedding
        time_emb = F.gelu(self.time_mlp(t_emb))[:, :, None, None]
        h = h + time_emb
        
        h = F.gelu(self.bn2(self.conv2(h)))
        
        return h + self.shortcut(x)


class SimpleUNet(nn.Module):
    """Simplified U-Net for diffusion models."""
    
    def __init__(self, in_channels=3, base_channels=64, time_emb_dim=256):
        super().__init__()
        
        # Time embedding
        self.time_embed = nn.Sequential(
            SinusoidalPositionEmbedding(base_channels),
            nn.Linear(base_channels, time_emb_dim),
            nn.GELU(),
            nn.Linear(time_emb_dim, time_emb_dim)
        )
        
        # Encoder
        self.enc1 = ResBlock(in_channels, base_channels, time_emb_dim)
        self.enc2 = ResBlock(base_channels, base_channels * 2, time_emb_dim)
        self.enc3 = ResBlock(base_channels * 2, base_channels * 4, time_emb_dim)
        
        self.pool = nn.MaxPool2d(2)
        
        # Bottleneck
        self.bottleneck = ResBlock(base_channels * 4, base_channels * 8, time_emb_dim)
        
        # Decoder
        self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, 2)
        self.dec3 = ResBlock(base_channels * 8, base_channels * 4, time_emb_dim)
        
        self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, 2)
        self.dec2 = ResBlock(base_channels * 4, base_channels * 2, time_emb_dim)
        
        self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, 2)
        self.dec1 = ResBlock(base_channels * 2, base_channels, time_emb_dim)
        
        # Output
        self.out = nn.Conv2d(base_channels, in_channels, 1)
    
    def forward(self, x, t):
        # Time embedding
        t_emb = self.time_embed(t)
        
        # Encoder
        e1 = self.enc1(x, t_emb)
        e2 = self.enc2(self.pool(e1), t_emb)
        e3 = self.enc3(self.pool(e2), t_emb)
        
        # Bottleneck
        b = self.bottleneck(self.pool(e3), t_emb)
        
        # Decoder with skip connections
        d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1), t_emb)
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1), t_emb)
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1), t_emb)
        
        return self.out(d1)


class GaussianDiffusion(nn.Module):
    """Gaussian Diffusion for DDPM."""
    
    def __init__(self, model, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        super().__init__()
        self.model = model
        self.timesteps = timesteps
        
        # Noise schedule
        betas = torch.linspace(beta_start, beta_end, timesteps)
        alphas = 1 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        
        self.register_buffer('betas', betas)
        self.register_buffer('alphas_cumprod', alphas_cumprod)
        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod))
    
    def q_sample(self, x0, t, noise=None):
        """Sample from q(x_t | x_0)."""
        if noise is None:
            noise = torch.randn_like(x0)
        
        sqrt_alpha = self.sqrt_alphas_cumprod[t][:, None, None, None]
        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t][:, None, None, None]
        
        return sqrt_alpha * x0 + sqrt_one_minus_alpha * noise
    
    def compute_loss(self, x0):
        """Compute simplified DDPM loss."""
        batch_size = x0.shape[0]
        
        # Sample random timesteps
        t = torch.randint(0, self.timesteps, (batch_size,), device=x0.device)
        
        # Add noise
        noise = torch.randn_like(x0)
        x_t = self.q_sample(x0, t, noise)
        
        # Predict noise
        predicted_noise = self.model(x_t, t)
        
        # MSE loss
        loss = F.mse_loss(predicted_noise, noise)
        
        return loss
    
    @torch.no_grad()
    def p_sample(self, x_t, t):
        """Sample from p(x_{t-1} | x_t)."""
        betas_t = self.betas[t][:, None, None, None]
        sqrt_one_minus_alpha_t = self.sqrt_one_minus_alphas_cumprod[t][:, None, None, None]
        sqrt_alpha_t = torch.sqrt(1 - betas_t)
        
        # Predict noise
        predicted_noise = self.model(x_t, t)
        
        # Compute mean
        mean = (x_t - betas_t / sqrt_one_minus_alpha_t * predicted_noise) / sqrt_alpha_t
        
        # Add noise (except at t=0)
        if t[0] > 0:
            noise = torch.randn_like(x_t)
            sigma = torch.sqrt(betas_t)
            return mean + sigma * noise
        else:
            return mean
    
    @torch.no_grad()
    def sample(self, shape):
        """Generate samples from noise."""
        x = torch.randn(shape, device=self.betas.device)
        
        for t in reversed(range(self.timesteps)):
            t_batch = torch.full((shape[0],), t, device=x.device, dtype=torch.long)
            x = self.p_sample(x, t_batch)
        
        return x


# Example: Create and train DDPM
model = SimpleUNet(in_channels=3, base_channels=64)
diffusion = GaussianDiffusion(model, timesteps=1000)

# Dummy training
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(5):
    x0 = torch.randn(4, 3, 32, 32)  # Dummy data
    loss = diffusion.compute_loss(x0)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print(f"Epoch {epoch+1}: Loss = {loss.item():.4f}")

# Generate samples
samples = diffusion.sample((4, 3, 32, 32))
print(f"Generated samples shape: {samples.shape}")

8.2 Flow Matching

class FlowMatchingModel(nn.Module):
    """
    Flow Matching for generative modelling.
    
    Learns velocity field v_θ(x_t, t) where x_t = (1-t)x_0 + t*x_1
    """
    
    def __init__(self, model, sigma_min=1e-4):
        super().__init__()
        self.model = model
        self.sigma_min = sigma_min
    
    def compute_loss(self, x0):
        """
        Compute flow matching loss.
        
        L = E[||v_θ(x_t, t) - (x_1 - x_0)||²]
        """
        batch_size = x0.shape[0]
        
        # Sample timestep t ~ U(0, 1)
        t = torch.rand(batch_size, device=x0.device)
        
        # Sample noise x_1 ~ N(0, I)
        x1 = torch.randn_like(x0)
        
        # Interpolate: x_t = (1-t)x_0 + t*x_1
        t_ = t[:, None, None, None]
        xt = (1 - t_) * x0 + t_ * x1
        
        # Target velocity: v = x_1 - x_0
        target = x1 - x0
        
        # Predict velocity
        predicted = self.model(xt, t * 999)  # Scale t to [0, 999] for model
        
        # MSE loss
        loss = F.mse_loss(predicted, target)
        
        return loss
    
    @torch.no_grad()
    def sample(self, shape, num_steps=50):
        """Generate samples using Euler integration."""
        device = self.model.time_embed[0].weight.device
        
        # Start from noise
        x = torch.randn(shape, device=device)
        
        # Euler integration
        dt = 1.0 / num_steps
        
        for i in range(num_steps):
            t = torch.full((shape[0],), i * dt, device=device)
            v = self.model(x, t * 999)
            x = x + v * dt
        
        return x
    
    @torch.no_grad()
    def sample_rk4(self, shape, num_steps=20):
        """Generate samples using 4th-order Runge-Kutta."""
        device = self.model.time_embed[0].weight.device
        
        x = torch.randn(shape, device=device)
        dt = 1.0 / num_steps
        
        for i in range(num_steps):
            t = torch.full((shape[0],), i * dt, device=device)
            
            k1 = self.model(x, t * 999)
            k2 = self.model(x + dt/2 * k1, (t + dt/2) * 999)
            k3 = self.model(x + dt/2 * k2, (t + dt/2) * 999)
            k4 = self.model(x + dt * k3, (t + dt) * 999)
            
            x = x + dt/6 * (k1 + 2*k2 + 2*k3 + k4)
        
        return x


# Example: Flow matching
model = SimpleUNet(in_channels=3, base_channels=64)
flow = FlowMatchingModel(model)

# Training
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(5):
    x0 = torch.randn(4, 3, 32, 32)
    loss = flow.compute_loss(x0)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print(f"Epoch {epoch+1}: Loss = {loss.item():.4f}")

# Generation
samples = flow.sample((4, 3, 32, 32))
print(f"Generated samples shape: {samples.shape}")

9. Summary

1. DDPM learns to denoise data through a Markov chain, with a simplified loss equivalent to denoising score matching.

2. Score matching trains a network to estimate the gradient of the log-density at multiple noise levels.

3. Classifier-free guidance enables controllable generation by interpolating conditional and unconditional predictions.

4. Flow matching defines a continuous transformation from noise to data, enabling faster sampling and exact likelihoods.

5. Consistency models enable one-step generation by mapping any noisy sample to the clean data.

References

• Ho, J., et al. (2020). Denoising diffusion probabilistic models. NeurIPS.
• Song, Y., & Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. NeurIPS.
• Ho, J., & Salimans, T. (2022). Classifier-free diffusion guidance. NeurIPS Workshop.
• Lipman, Y., et al. (2023). Flow matching for generative modeling. ICLR.
• Song, Y., et al. (2023). Consistency models. ICML.
• Nichol, A., & Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. ICML.
• Karras, T., et al. (2022). Elucidating the design space of diffusion-based generative models. NeurIPS.