Hello there! It is wonderful to assist you today. We are about to embark on an exciting journey into the fascinating world of diffusion networks. These powerful generative models are truly revolutionizing how we think about creating new data, especially images. By the end of this tutorial, you will have a solid conceptual understanding of how they work and even a practical foundation to start building your own.
Let us dive right in!
Imagine you have a beautiful, clear photograph. Now, imagine someone starts adding a tiny bit of static, then a bit more, and then even more, until the photograph is completely obscured by random noise, like a fuzzy old TV screen. This process, where information is gradually destroyed by adding noise, is the core idea behind the "forward diffusion process" in diffusion models.
The truly magical part, and what diffusion models excel at, is the "reverse diffusion process." Here, we train a neural network to do the opposite: starting from pure noise, it learns to gradually remove the static, step by step, until the original clear photograph (or a brand new, similar one) emerges. It is like teaching an artist to reconstruct a masterpiece from a canvas that was initially just a random splatter of paint.
The goal of a diffusion model is to learn this reverse process. If we can accurately reverse the noise addition, we can then start with random noise and generate completely new, realistic data that resembles the data it was trained on.
Let us visualize this intuitive process:
Original Image -> Slightly Noisy -> More Noisy -> Even More Noisy -> Pure Noise (The Forward Diffusion Process)
Pure Noise -> Less Noisy -> Even Less Noisy -> Almost Clear -> Generated Image (The Reverse Denoising Process)
Diving Deeper: The Forward Diffusion Process (Noising)
The forward diffusion process, also known as the noising process, is not learned by the model; it is a fixed, predefined process. We start with an original data sample, let us call it (x_0), which could be an image. Over a series of discrete time steps, typically denoted as (t = 1, 2, \ldots, T), we progressively add Gaussian noise to the sample. Each step (t) generates a slightly noisier version, (x_t), from the previous step's sample, (x_{t-1}).
The mathematical formulation for adding noise at each step is governed by a variance schedule. Let (\beta_t) be a small, positive value that determines the amount of noise added at time step (t). This (\beta_t) typically increases over time, meaning more noise is added in later steps.
The transition from (x_{t-1}) to (x_t) is defined as:
$$ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t \mathbf{I}) $$
Here, (\mathcal{N}) denotes a Gaussian distribution. This equation tells us that (x_t) is sampled from a Gaussian distribution with a mean of (\sqrt{1 - \beta_t} x_{t-1}) and a variance of (\beta_t \mathbf{I}) (where (\mathbf{I}) is the identity matrix, meaning noise is added independently to each dimension). The term (\sqrt{1 - \beta_t}) scales down the previous sample, while (\beta_t \mathbf{I}) adds the new noise.
While this step-by-step process is intuitive, for training the denoising model, it is more efficient to be able to sample (x_t) directly from (x_0) at any arbitrary time step (t), rather than iteratively applying noise (t) times. This is where a clever reparameterization comes into play.
Let us define (\alpha_t = 1 - \beta_t) and (\bar{\alpha}t = \prod{s=1}^{t} \alpha_s). Using these, we can derive a direct way to sample (x_t) from (x_0):
$$ q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) \mathbf{I}) $$
This equation is incredibly powerful. It means that to get a noisy version of (x_0) at any time step (t), we simply scale (x_0) by (\sqrt{\bar{\alpha}_t}) and add noise scaled by (\sqrt{1 - \bar{\alpha}_t}). The noise itself is sampled from a standard Gaussian distribution, (\epsilon \sim \mathcal{N}(0, \mathbf{I})).
So, we can write:
$$ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon $$
This is often called the "reparameterization trick." It allows us to compute (x_t) for any (t) and any (x_0) by just sampling one noise vector (\epsilon). This is crucial because during training, we will randomly sample a time step (t) and then directly generate (x_t) from (x_0) to train our denoising network.
Here is a conceptual ASCII diagram of the forward process:
[x_0] | (add noise beta_1) V [x_1] | (add noise beta_2) V [x_2] | (add noise ... ) V [x_t] | (add noise beta_T) V [x_T] (Pure Gaussian Noise)
The Heart of the Matter: The Reverse Denoising Process
The real challenge, and where the neural network comes into play, is learning to reverse this forward process. Our goal is to train a model that can predict the noise that was added at any given step, or equivalently, predict the original image (x_0) from a noisy image (x_t).
More specifically, the reverse process involves estimating the mean and variance of the reverse diffusion step (q(x_{t-1} | x_t)). It turns out that if (\beta_t) is small, this reverse distribution is also Gaussian. However, its parameters depend on (x_0), which we do not know.
This is where our neural network, let us call it (\epsilon_\theta), comes in. We train (\epsilon_\theta) to predict the noise (\epsilon) that was added to (x_0) to get (x_t). If our network can accurately predict this noise (\epsilon), then we can use the reparameterization trick in reverse to estimate (x_0) or (x_{t-1}).
The denoising network (\epsilon_\theta) typically takes two inputs:
- The noisy image (x_t).
- The current time step (t).
And it outputs:
- The predicted noise (\epsilon_\theta(x_t, t)).
The architecture of choice for (\epsilon_\theta) is often a U-Net.
Why a U-Net? A U-Net is a type of convolutional neural network particularly well-suited for image-to-image translation tasks, such as image segmentation or, in our case, denoising. It has an encoder-decoder structure with "skip connections."
- Encoder: This part downsamples the image, extracting hierarchical features and capturing contextual information.
- Decoder: This part upsamples the features, reconstructing the image while incorporating the learned context.
- Skip Connections: These connections directly pass information from the encoder to the corresponding decoder layers. This is crucial because it allows the network to retain fine-grained spatial details that might otherwise be lost during downsampling, which is essential for generating high-quality images.
The time step (t) is also an important input. Since the amount of noise varies with (t), the network needs to know which specific noise level it is trying to denoise. This is usually incorporated by embedding (t) into a high-dimensional vector (similar to positional embeddings in Transformers) and then adding or concatenating this embedding to the feature maps at various points within the U-Net.
The training process for the denoising network is surprisingly straightforward. For each training step:
- Sample an original image (x_0) from your dataset.
- Randomly sample a time step (t) between 1 and (T).
- Generate a noisy image (x_t) by adding noise to (x_0) using the forward diffusion equation: (x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon), where (\epsilon) is randomly sampled Gaussian noise.
- Feed (x_t) and (t) into the denoising network (\epsilon_\theta) to get its prediction of the noise, (\epsilon_\theta(x_t, t)).
- Calculate the loss between the predicted noise (\epsilon_\theta(x_t, t)) and the actual noise (\epsilon) that was added. The most common loss function is the Mean Squared Error (MSE):
$$ L_t = ||\epsilon - \epsilon_\theta(x_t, t)||^2 $$
- Perform backpropagation and update the network's parameters (\theta) using an optimizer (e.g., Adam).
By minimizing this loss, the network learns to accurately predict the noise component at any given time step and noise level.
Here is a conceptual ASCII diagram of the denoising U-Net:
Input: Noisy Image (x_t)
Input: Time Step (t)
|
V
+-----------------+
| Encoder Path |
| (Downsampling) |
+-------+---------+
|
V
(Latent Representation)
|
+-------+---------+
| Decoder Path |
| (Upsampling) |
+-----------------+
|
V
Output: Predicted Noise (epsilon_theta)
Building Blocks of a Diffusion Model: Key Components
Let us break down the essential components we need to implement a diffusion model.
Variance Schedule The variance schedule, (\beta_t), dictates how much noise is added at each step in the forward process. A common choice is a linear schedule, where (\beta_t) increases linearly from a small value ((\beta_{start})) to a larger value ((\beta_{end})). Other schedules, like cosine schedules, can also be used for better performance.
We need to calculate (\beta_t), (\alpha_t = 1 - \beta_t), and (\bar{\alpha}t = \prod{s=1}^{t} \alpha_s) for all time steps (t). These values are typically pre-computed and stored in tensors for efficient access during training and sampling.
The Denoising Model (U-Net) As discussed, a U-Net is the backbone of the denoising network. Its key features include:
- Encoder-Decoder Structure: Convolutional layers with downsampling (e.g., stride-2 convolutions or max-pooling) in the encoder, and transposed convolutions or nearest-neighbor upsampling followed by convolutions in the decoder.
- Residual Connections: These are often integrated within the convolutional blocks (e.g., ResNet blocks). They help with training deep networks by allowing gradients to flow more easily.
- Skip Connections: Direct connections from encoder blocks to symmetrically positioned decoder blocks. These concatenate feature maps, providing the decoder with high-resolution information.
- Time Embeddings: The time step (t) needs to be encoded in a way that the network can understand its significance. A common approach is to use sinusoidal positional embeddings, similar to those in Transformer models. This embedding vector is then typically added or concatenated to the feature maps within the U-Net blocks.
- Attention Mechanisms (Optional but Beneficial): For higher-resolution images, self-attention layers can be incorporated into the U-Net, especially at lower resolutions (in the bottleneck or middle layers). These allow the network to capture global dependencies across the image.
Optimizer and Training Loop The training loop for a diffusion model largely follows standard neural network training practices:
- Optimizer: Adam or AdamW are popular choices.
- Loss Function: Mean Squared Error (MSE) between the predicted noise and the true noise.
- Data Loader: To feed batches of images to the model.
- Training Steps: Iterate over epochs, sampling images and time steps, performing forward pass, calculating loss, backpropagation, and optimizer step.
Step-by-Step Implementation Walkthrough
Let us walk through the implementation of a simple diffusion model using PyTorch. We will use a running example of generating grayscale images, similar to MNIST digits.
Step 1: Imports and Configuration
First, we need to import the necessary libraries and define some basic configuration parameters.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import os
from PIL import Image
import math
from tqdm import tqdm # For progress bars
# --- Configuration Parameters ---
IMG_SIZE = 28 # Size of the input images (e.g., 28 for MNIST)
BATCH_SIZE = 128 # Number of images per training batch
NUM_EPOCHS = 100 # Number of training epochs
LEARNING_RATE = 1e-4 # Learning rate for the optimizer
TIMESTEPS = 1000 # Total number of diffusion steps (T)
BETA_START = 1e-4 # Start value for the linear beta schedule
BETA_END = 0.02 # End value for the linear beta schedule
SAVE_DIR = "diffusion_results" # Directory to save generated images
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Use GPU if available
# Ensure save directory exists
os.makedirs(SAVE_DIR, exist_ok=True)
print(f"Using device: {DEVICE}")
Step 2: Defining the Variance Schedule
We will implement a linear variance schedule and pre-compute the (\alpha_t) and (\bar{\alpha}_t) values. These will be stored as PyTorch tensors on the chosen device.
def linear_beta_schedule(timesteps, beta_start, beta_end):
"""
Generates a linear schedule for beta values.
Args:
timesteps (int): The total number of diffusion steps (T).
beta_start (float): The starting value for beta.
beta_end (float): The ending value for beta.
Returns:
torch.Tensor: A tensor of beta values for each timestep.
"""
return torch.linspace(beta_start, beta_end, timesteps)
# Pre-compute the schedule values
betas = linear_beta_schedule(TIMESTEPS, BETA_START, BETA_END).to(DEVICE)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # alpha_bar_t = product(alpha_s from s=1 to t)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # alpha_bar_{t-1}
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # Used in reverse process
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # sqrt(alpha_bar_t)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # sqrt(1 - alpha_bar_t)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) # Variance for reverse step
Step 3: The Forward Diffusion Helper Functions
These functions help us apply noise to an image at a given time step (t), using the reparameterization trick.
def extract(a, t, x_shape):
"""
Extracts the values from a tensor 'a' at given indices 't' and reshapes them
to match the shape of 'x_shape'. This is used to select the correct
alpha_bar_t or sqrt(1 - alpha_bar_t) for each sample in a batch.
Args:
a (torch.Tensor): The tensor to extract values from (e.g., alphas_cumprod).
t (torch.Tensor): A tensor of time indices for each sample in the batch.
x_shape (torch.Size): The desired shape for the extracted values.
Returns:
torch.Tensor: The extracted and reshaped tensor.
"""
batch_size = t.shape[0]
out = a.gather(-1, t.cpu()) # Extract values using cpu indices
# Reshape to (batch_size, 1, 1, 1) for broadcasting with image tensors
return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
def q_sample(x_start, t, noise=None):
"""
Applies noise to the original image x_start at time step t.
This implements the forward diffusion process using the reparameterization trick.
Args:
x_start (torch.Tensor): The original, clean image (x_0).
t (torch.Tensor): The current time step for each sample in the batch.
noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
Returns:
torch.Tensor: The noisy image x_t.
"""
if noise is None:
noise = torch.randn_like(x_start) # Sample Gaussian noise
# Extract sqrt(alpha_bar_t) and sqrt(1 - alpha_bar_t) for the current batch and time steps
sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)
# Apply the reparameterization trick: x_t = sqrt(alpha_bar_t)*x_0 + sqrt(1 - alpha_bar_t)*epsilon
x_t = sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
return x_t
Step 4: The Denoising U-Net Model
This is the core neural network that learns to predict the noise. We will define helper blocks and then assemble them into a U-Net.
class SinusoidalPositionalEmbedding(nn.Module):
"""
Generates sinusoidal positional embeddings for time steps.
These embeddings allow the model to understand the current time step (noise level).
"""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, time):
device = time.device
half_dim = self.dim // 2
# Compute the sinusoidal arguments
embeddings = math.log(10000) / (half_dim - 1)
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
embeddings = time[:, None] * embeddings[None, :]
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
return embeddings
class Block(nn.Module):
"""
A basic convolutional block with Group Normalization and ReLU activation.
"""
def __init__(self, dim_in, dim_out, groups=8):
super().__init__()
self.proj = nn.Conv2d(dim_in, dim_out, 3, padding=1)
self.norm = nn.GroupNorm(groups, dim_out)
self.act = nn.SiLU() # Swish activation function
def forward(self, x, scale_shift=None):
x = self.proj(x)
x = self.norm(x)
# Apply optional scale and shift from time embedding
if scale_shift is not None:
scale, shift = scale_shift
x = x * (scale + 1) + shift
x = self.act(x)
return x
class ResnetBlock(nn.Module):
"""
A Residual Block, commonly used in U-Nets.
It includes two convolutional blocks and a skip connection.
Time embeddings are incorporated here.
"""
def __init__(self, dim_in, dim_out, time_emb_dim=None, groups=8):
super().__init__()
self.mlp = nn.Sequential(
nn.SiLU(),
nn.Linear(time_emb_dim, dim_out * 2)
) if time_emb_dim is not None else None
self.block1 = Block(dim_in, dim_out, groups=groups)
self.block2 = Block(dim_out, dim_out, groups=groups)
self.res_conv = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity()
def forward(self, x, time_emb=None):
scale_shift = None
if self.mlp is not None and time_emb is not None:
time_emb = self.mlp(time_emb)
# Split the output into scale and shift for normalization
time_emb = time_emb.reshape(time_emb.shape[0], time_emb.shape[1], 1, 1)
scale_shift = time_emb.chunk(2, dim=1) # Split into two along dimension 1
h = self.block1(x, scale_shift=scale_shift)
h = self.block2(h)
return h + self.res_conv(x) # Add residual connection
class Downsample(nn.Module):
"""
Downsampling layer using a stride-2 convolution.
"""
def __init__(self, dim):
super().__init__()
self.conv = nn.Conv2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1
def forward(self, x):
return self.conv(x)
class Upsample(nn.Module):
"""
Upsampling layer using a transposed convolution.
"""
def __init__(self, dim):
super().__init__()
self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1
def forward(self, x):
return self.conv(x)
class Unet(nn.Module):
"""
The Denoising U-Net model.
It takes a noisy image and a time step, and outputs the predicted noise.
"""
def __init__(self,
dim, # Base dimension for the model
image_channels=1, # Number of channels in the input image (e.g., 1 for grayscale)
dim_mults=(1, 2, 4, 8), # Multipliers for dimension increase at each downsampling level
time_emb_dim=128, # Dimension of the time embedding
groups=8): # Number of groups for Group Normalization
super().__init__()
self.time_mlp = nn.Sequential(
SinusoidalPositionalEmbedding(dim),
nn.Linear(dim, time_emb_dim),
nn.SiLU(),
nn.Linear(time_emb_dim, time_emb_dim)
)
# Initial convolution to project image channels to base dimension
self.init_conv = nn.Conv2d(image_channels, dim, 7, padding=3)
dims = [dim, *map(lambda m: dim * m, dim_mults)] # [dim, dim*1, dim*2, dim*4, ...]
in_out = list(zip(dims[:-1], dims[1:])) # [(dim, dim*1), (dim*1, dim*2), ...]
# Encoder (downsampling path)
self.downs = nn.ModuleList([])
for ind, (dim_in, dim_out) in enumerate(in_out):
self.downs.append(nn.ModuleList([
ResnetBlock(dim_in, dim_out, time_emb_dim, groups=groups),
Downsample(dim_out) if ind != (len(in_out) - 1) else nn.Identity() # No downsample at bottleneck
]))
# Bottleneck (middle layer)
mid_dim = dims[-1]
self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
# Decoder (upsampling path)
self.ups = nn.ModuleList([])
for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
self.ups.append(nn.ModuleList([
ResnetBlock(dim_out * 2, dim_in, time_emb_dim, groups=groups), # *2 for skip connection concatenation
Upsample(dim_in) if ind != 0 else nn.Identity() # No upsample at final layer
]))
# Final output convolution
self.final_conv = nn.Sequential(
Block(dim, dim, groups=groups),
nn.Conv2d(dim, image_channels, 1) # Output noise with original image channels
)
def forward(self, x, time):
# Time embedding
t = self.time_mlp(time)
# Initial convolution
x = self.init_conv(x)
h = [] # To store skip connections outputs
# Downsampling path
for resnet_block, downsample in self.downs:
x = resnet_block(x, t)
h.append(x) # Store for skip connection
x = downsample(x)
# Bottleneck
x = self.mid_block1(x, t)
x = self.mid_block2(x, t)
# Upsampling path
for resnet_block, upsample in self.ups:
# Concatenate with skip connection from encoder
x = torch.cat((x, h.pop()), dim=1)
x = resnet_block(x, t)
x = upsample(x)
# Final output
return self.final_conv(x)
Step 5: The Training Process
We will define the loss function and the training loop. The p_losses function encapsulates the core training step for a single batch.
def p_losses(denoise_model, x_start, t, noise=None):
"""
Calculates the loss for a given batch of original images and time steps.
Args:
denoise_model (nn.Module): The U-Net model to train.
x_start (torch.Tensor): The original, clean images (x_0).
t (torch.Tensor): The time steps for each image in the batch.
noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
Returns:
torch.Tensor: The mean squared error loss.
"""
if noise is None:
noise = torch.randn_like(x_start) # Sample noise for the forward process
# Apply noise to get x_t
x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
# Predict the noise using the denoising model
predicted_noise = denoise_model(x_noisy, t)
# Calculate MSE loss between actual noise and predicted noise
loss = F.mse_loss(noise, predicted_noise)
return loss
Step 6: Sampling (Generating New Images)
This is the reverse process, where we start from pure noise and iteratively denoise it to generate new images.
@torch.no_grad() # Disable gradient calculations for sampling
def p_sample(model, x, t, t_index):
"""
Performs one step of the reverse diffusion process (denoising).
Estimates x_{t-1} from x_t.
Args:
model (nn.Module): The trained U-Net model.
x (torch.Tensor): The noisy image at time t (x_t).
t (torch.Tensor): The current time step.
t_index (int): The integer index of the current time step.
Returns:
torch.Tensor: The denoised image at time t-1 (x_{t-1}).
"""
betas_t = extract(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = extract(
sqrt_one_minus_alphas_cumprod, t, x.shape
)
sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
# Predict the noise using the model
model_mean = sqrt_recip_alphas_t * (
x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
)
# If t is the first step (t=0), there's no more noise to add
if t_index == 0:
return model_mean
else:
posterior_variance_t = extract(posterior_variance, t, x.shape)
# Sample new noise for the reverse step (unless t=0)
noise = torch.randn_like(x)
return model_mean + torch.sqrt(posterior_variance_t) * noise
@torch.no_grad()
def p_sample_loop(model, shape):
"""
Generates a batch of images by iteratively denoising from pure noise.
Args:
model (nn.Module): The trained U-Net model.
shape (tuple): The shape of the images to generate (batch_size, channels, height, width).
Returns:
torch.Tensor: A batch of generated images.
"""
batch_size = shape[0]
# Start with pure noise
img = torch.randn(shape, device=DEVICE)
# Iterate backwards through time steps
for i in tqdm(reversed(range(0, TIMESTEPS)), desc='sampling loop time step', total=TIMESTEPS):
t = torch.full((batch_size,), i, device=DEVICE, dtype=torch.long)
img = p_sample(model, img, t, i) # Denoise one step
return img
@torch.no_grad()
def sample_and_save_images(model, epoch, num_samples=16):
"""
Generates and saves a grid of sample images.
"""
# Define the shape of the images to generate
sample_shape = (num_samples, 1, IMG_SIZE, IMG_SIZE)
# Generate images
samples = p_sample_loop(model, sample_shape)
# Normalize images to [0, 1] range and convert to PIL Image
samples = (samples + 1) * 0.5 # [-1, 1] to [0, 1]
samples = samples.clamp(0, 1)
# Create a grid of images
grid_size = int(math.sqrt(num_samples))
# Create a blank image to paste samples onto
combined_image_width = grid_size * IMG_SIZE
combined_image_height = grid_size * IMG_SIZE
combined_image = Image.new('L', (combined_image_width, combined_image_height)) # 'L' for grayscale
for i in range(num_samples):
row = i // grid_size
col = i % grid_size
# Convert tensor to numpy array, scale to 0-255, convert to uint8
img_array = (samples[i].squeeze().cpu().numpy() * 255).astype(np.uint8)
img = Image.fromarray(img_array)
# Paste onto the combined image
combined_image.paste(img, (col * IMG_SIZE, row * IMG_SIZE))
# Save the combined image
filepath = os.path.join(SAVE_DIR, f"epoch_{epoch:04d}_samples.png")
combined_image.save(filepath)
print(f"Saved {num_samples} samples to {filepath}")
Putting It All Together: Training and Generation
Now we have all the pieces. We will set up the data loading, initialize the model and optimizer, and run the training loop. After training, we can use the sample_and_save_images function to generate new images.
The overall workflow during training is:
- Load a batch of real images (x_0).
- Randomly choose a time step (t) for each image in the batch.
- Add noise to (x_0) to get (x_t) and simultaneously record the noise (\epsilon) that was added.
- Feed (x_t) and (t) into the U-Net to predict the noise, (\epsilon_\theta).
- Calculate the MSE loss between (\epsilon) and (\epsilon_\theta).
- Perform backpropagation and update the U-Net's weights.
After training, to generate new images:
- Start with a tensor of pure random noise.
- Iteratively apply the
p_samplefunction, moving backward from (T) down to 0. Each step removes a bit of noise based on the U-Net's prediction. - The final output is a newly generated image.
Conclusion
Diffusion models represent a significant leap forward in generative AI. By understanding the simple yet powerful concept of gradually adding and then learning to reverse noise, developers can unlock incredible capabilities for generating highly realistic and diverse data. We have explored the forward noising process, the crucial reverse denoising process powered by U-Nets, and the key mathematical and architectural components involved.
With the conceptual understanding and the step-by-step implementation guide provided, you are now equipped to delve deeper, experiment with different schedules, model architectures, and apply these fascinating models to your own creative and technical challenges. The field is rapidly evolving, and your journey into diffusion models has just begun!
Addendum: Full Running Example Code
This section provides a complete, runnable Python script that implements the diffusion model for generating MNIST-like grayscale images. This code integrates all the snippets and additional necessary components like data loading and the main training loop.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import os
from PIL import Image
import math
from tqdm import tqdm # For progress bars
# --- Configuration Parameters ---
IMG_SIZE = 28 # Size of the input images (e.g., 28 for MNIST)
BATCH_SIZE = 128 # Number of images per training batch
NUM_EPOCHS = 100 # Number of training epochs
LEARNING_RATE = 1e-4 # Learning rate for the optimizer
TIMESTEPS = 1000 # Total number of diffusion steps (T)
BETA_START = 1e-4 # Start value for the linear beta schedule
BETA_END = 0.02 # End value for the linear beta schedule
SAVE_DIR = "diffusion_results" # Directory to save generated images
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Use GPU if available
# Ensure save directory exists
os.makedirs(SAVE_DIR, exist_ok=True)
print(f"Using device: {DEVICE}")
# --- Data Loading and Preprocessing ---
def load_mnist_data(batch_size, img_size):
"""
Loads and preprocesses the MNIST dataset.
Images are normalized to the range [-1, 1] for better model stability.
"""
transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(), # Converts to [0, 1] range
transforms.Normalize((0.5,), (0.5,)) # Normalizes to [-1, 1] range
])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)
return dataloader
# --- Variance Schedule Pre-computation ---
def linear_beta_schedule(timesteps, beta_start, beta_end):
"""
Generates a linear schedule for beta values.
Args:
timesteps (int): The total number of diffusion steps (T).
beta_start (float): The starting value for beta.
beta_end (float): The ending value for beta.
Returns:
torch.Tensor: A tensor of beta values for each timestep.
"""
return torch.linspace(beta_start, beta_end, timesteps)
# Pre-compute the schedule values and move to device
betas = linear_beta_schedule(TIMESTEPS, BETA_START, BETA_END).to(DEVICE)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # alpha_bar_t = product(alpha_s from s=1 to t)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # alpha_bar_{t-1}
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # Used in reverse process
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # sqrt(alpha_bar_t)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # sqrt(1 - alpha_bar_t)
# Variance for the reverse step, used in p_sample
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
# --- Forward Diffusion Helper Functions ---
def extract(a, t, x_shape):
"""
Extracts the values from a tensor 'a' at given indices 't' and reshapes them
to match the shape of 'x_shape'. This is used to select the correct
alpha_bar_t or sqrt(1 - alpha_bar_t) for each sample in a batch.
Args:
a (torch.Tensor): The tensor to extract values from (e.g., alphas_cumprod).
t (torch.Tensor): A tensor of time indices for each sample in the batch.
x_shape (torch.Size): The desired shape for the extracted values.
Returns:
torch.Tensor: The extracted and reshaped tensor.
"""
batch_size = t.shape[0]
out = a.gather(-1, t.cpu()) # Extract values using cpu indices
# Reshape to (batch_size, 1, 1, 1) for broadcasting with image tensors
return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
def q_sample(x_start, t, noise=None):
"""
Applies noise to the original image x_start at time step t.
This implements the forward diffusion process using the reparameterization trick.
Args:
x_start (torch.Tensor): The original, clean image (x_0).
t (torch.Tensor): The current time step for each sample in the batch.
noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
Returns:
torch.Tensor: The noisy image x_t.
"""
if noise is None:
noise = torch.randn_like(x_start) # Sample Gaussian noise
# Extract sqrt(alpha_bar_t) and sqrt(1 - alpha_bar_t) for the current batch and time steps
sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)
# Apply the reparameterization trick: x_t = sqrt(alpha_bar_t)*x_0 + sqrt(1 - alpha_bar_t)*epsilon
x_t = sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
return x_t
# --- Denoising U-Net Model Definition ---
class SinusoidalPositionalEmbedding(nn.Module):
"""
Generates sinusoidal positional embeddings for time steps.
These embeddings allow the model to understand the current time step (noise level).
"""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, time):
device = time.device
half_dim = self.dim // 2
# Compute the sinusoidal arguments
embeddings = math.log(10000) / (half_dim - 1)
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
embeddings = time[:, None] * embeddings[None, :]
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
return embeddings
class Block(nn.Module):
"""
A basic convolutional block with Group Normalization and SiLU (Swish) activation.
Includes an optional scale and shift for time embedding integration.
"""
def __init__(self, dim_in, dim_out, groups=8):
super().__init__()
self.proj = nn.Conv2d(dim_in, dim_out, 3, padding=1)
self.norm = nn.GroupNorm(groups, dim_out)
self.act = nn.SiLU()
def forward(self, x, scale_shift=None):
x = self.proj(x)
x = self.norm(x)
# Apply optional scale and shift from time embedding
if scale_shift is not None:
scale, shift = scale_shift
x = x * (scale + 1) + shift
x = self.act(x)
return x
class ResnetBlock(nn.Module):
"""
A Residual Block, commonly used in U-Nets.
It includes two convolutional blocks and a skip connection.
Time embeddings are incorporated here by modulating the normalization layer.
"""
def __init__(self, dim_in, dim_out, time_emb_dim=None, groups=8):
super().__init__()
self.mlp = nn.Sequential(
nn.SiLU(),
nn.Linear(time_emb_dim, dim_out * 2)
) if time_emb_dim is not None else None
self.block1 = Block(dim_in, dim_out, groups=groups)
self.block2 = Block(dim_out, dim_out, groups=groups)
self.res_conv = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity()
def forward(self, x, time_emb=None):
scale_shift = None
if self.mlp is not None and time_emb is not None:
time_emb = self.mlp(time_emb)
# Split the output into scale and shift for normalization
time_emb = time_emb.reshape(time_emb.shape[0], time_emb.shape[1], 1, 1)
scale_shift = time_emb.chunk(2, dim=1) # Split into two along dimension 1
h = self.block1(x, scale_shift=scale_shift)
h = self.block2(h)
return h + self.res_conv(x) # Add residual connection
class Downsample(nn.Module):
"""
Downsampling layer using a stride-2 convolution.
"""
def __init__(self, dim):
super().__init__()
self.conv = nn.Conv2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1
def forward(self, x):
return self.conv(x)
class Upsample(nn.Module):
"""
Upsampling layer using a transposed convolution.
"""
def __init__(self, dim):
super().__init__()
self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1
def forward(self, x):
return self.conv(x)
class Unet(nn.Module):
"""
The Denoising U-Net model.
It takes a noisy image and a time step, and outputs the predicted noise.
"""
def __init__(self,
dim, # Base dimension for the model
image_channels=1, # Number of channels in the input image (e.g., 1 for grayscale)
dim_mults=(1, 2, 4, 8), # Multipliers for dimension increase at each downsampling level
time_emb_dim=128, # Dimension of the time embedding
groups=8): # Number of groups for Group Normalization
super().__init__()
# Time embedding MLP
self.time_mlp = nn.Sequential(
SinusoidalPositionalEmbedding(dim),
nn.Linear(dim, time_emb_dim),
nn.SiLU(),
nn.Linear(time_emb_dim, time_emb_dim)
)
# Initial convolution to project image channels to base dimension
self.init_conv = nn.Conv2d(image_channels, dim, 7, padding=3)
# Calculate dimensions for encoder/decoder blocks
dims = [dim, *map(lambda m: dim * m, dim_mults)] # [dim, dim*1, dim*2, dim*4, ...]
in_out = list(zip(dims[:-1], dims[1:])) # [(dim, dim*1), (dim*1, dim*2), ...]
# Encoder (downsampling path)
self.downs = nn.ModuleList([])
for ind, (dim_in, dim_out) in enumerate(in_out):
self.downs.append(nn.ModuleList([
ResnetBlock(dim_in, dim_out, time_emb_dim, groups=groups),
Downsample(dim_out) if ind != (len(in_out) - 1) else nn.Identity() # No downsample at bottleneck
]))
# Bottleneck (middle layer)
mid_dim = dims[-1]
self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
# Decoder (upsampling path)
self.ups = nn.ModuleList([])
for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
self.ups.append(nn.ModuleList([
ResnetBlock(dim_out * 2, dim_in, time_emb_dim, groups=groups), # *2 for skip connection concatenation
Upsample(dim_in) if ind != 0 else nn.Identity() # No upsample at final layer
]))
# Final output convolution
self.final_conv = nn.Sequential(
Block(dim, dim, groups=groups),
nn.Conv2d(dim, image_channels, 1) # Output noise with original image channels
)
def forward(self, x, time):
# Time embedding
t = self.time_mlp(time)
# Initial convolution
x = self.init_conv(x)
h = [] # To store skip connections outputs for concatenation
# Downsampling path
for resnet_block, downsample in self.downs:
x = resnet_block(x, t)
h.append(x) # Store for skip connection
x = downsample(x)
# Bottleneck
x = self.mid_block1(x, t)
x = self.mid_block2(x, t)
# Upsampling path
for resnet_block, upsample in self.ups:
# Concatenate with skip connection from encoder
x = torch.cat((x, h.pop()), dim=1)
x = resnet_block(x, t)
x = upsample(x)
# Final output
return self.final_conv(x)
# --- Loss Function for Training ---
def p_losses(denoise_model, x_start, t, noise=None):
"""
Calculates the loss for a given batch of original images and time steps.
Args:
denoise_model (nn.Module): The U-Net model to train.
x_start (torch.Tensor): The original, clean images (x_0).
t (torch.Tensor): The time steps for each image in the batch.
noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
Returns:
torch.Tensor: The mean squared error loss.
"""
if noise is None:
noise = torch.randn_like(x_start) # Sample noise for the forward process
# Apply noise to get x_t
x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
# Predict the noise using the denoising model
predicted_noise = denoise_model(x_noisy, t)
# Calculate MSE loss between actual noise and predicted noise
loss = F.mse_loss(noise, predicted_noise)
return loss
# --- Sampling (Image Generation) Functions ---
@torch.no_grad() # Disable gradient calculations for sampling
def p_sample(model, x, t, t_index):
"""
Performs one step of the reverse diffusion process (denoising).
Estimates x_{t-1} from x_t.
Args:
model (nn.Module): The trained U-Net model.
x (torch.Tensor): The noisy image at time t (x_t).
t (torch.Tensor): The current time step.
t_index (int): The integer index of the current time step.
Returns:
torch.Tensor: The denoised image at time t-1 (x_{t-1}).
"""
betas_t = extract(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = extract(
sqrt_one_minus_alphas_cumprod, t, x.shape
)
sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
# Predict the noise using the model
# This formula is derived from the reverse process mean estimation
model_mean = sqrt_recip_alphas_t * (
x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
)
# If t is the first step (t=0), there's no more noise to add
if t_index == 0:
return model_mean
else:
# Add noise sampled from the posterior distribution's variance
posterior_variance_t = extract(posterior_variance, t, x.shape)
noise = torch.randn_like(x)
return model_mean + torch.sqrt(posterior_variance_t) * noise
@torch.no_grad()
def p_sample_loop(model, shape):
"""
Generates a batch of images by iteratively denoising from pure noise.
Args:
model (nn.Module): The trained U-Net model.
shape (tuple): The shape of the images to generate (batch_size, channels, height, width).
Returns:
torch.Tensor: A batch of generated images.
"""
batch_size = shape[0]
# Start with pure Gaussian noise
img = torch.randn(shape, device=DEVICE)
# Iterate backwards through time steps, denoising at each step
for i in tqdm(reversed(range(0, TIMESTEPS)), desc='sampling loop time step', total=TIMESTEPS):
t = torch.full((batch_size,), i, device=DEVICE, dtype=torch.long)
img = p_sample(model, img, t, i) # Denoise one step
return img
@torch.no_grad()
def sample_and_save_images(model, epoch, num_samples=16):
"""
Generates and saves a grid of sample images.
Args:
model (nn.Module): The trained U-Net model.
epoch (int): The current epoch number, used for naming the saved file.
num_samples (int): The number of images to generate and save.
"""
# Define the shape of the images to generate (e.g., 16 samples, 1 channel, 28x28 pixels)
sample_shape = (num_samples, 1, IMG_SIZE, IMG_SIZE)
# Generate images using the reverse diffusion process
samples = p_sample_loop(model, sample_shape)
# Normalize images from [-1, 1] to [0, 1] range for saving as image files
samples = (samples + 1) * 0.5
samples = samples.clamp(0, 1) # Ensure values are within valid range
# Create a grid of images for visualization
grid_size = int(math.sqrt(num_samples))
# Create a blank image to paste samples onto
combined_image_width = grid_size * IMG_SIZE
combined_image_height = grid_size * IMG_SIZE
# 'L' mode for grayscale images
combined_image = Image.new('L', (combined_image_width, combined_image_height))
for i in range(num_samples):
row = i // grid_size
col = i % grid_size
# Convert tensor to numpy array, scale to 0-255, convert to uint8
img_array = (samples[i].squeeze().cpu().numpy() * 255).astype(np.uint8)
img = Image.fromarray(img_array)
# Paste the individual generated image onto the combined grid image
combined_image.paste(img, (col * IMG_SIZE, row * IMG_SIZE))
# Save the combined image to the specified directory
filepath = os.path.join(SAVE_DIR, f"epoch_{epoch:04d}_samples.png")
combined_image.save(filepath)
print(f"Saved {num_samples} samples to {filepath}")
# --- Main Training Loop ---
def train_diffusion_model():
"""
Main function to train the diffusion model.
Initializes the model, optimizer, loads data, and runs the training epochs.
"""
# Load data
dataloader = load_mnist_data(BATCH_SIZE, IMG_SIZE)
# Initialize model
model = Unet(
dim=64, # Base dimension for the U-Net
image_channels=1, # MNIST is grayscale
dim_mults=(1, 2, 4), # Dimension multipliers for downsampling path
time_emb_dim=256 # Dimension of time embeddings
).to(DEVICE)
# Initialize optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
print("Starting training...")
for epoch in range(NUM_EPOCHS):
model.train() # Set model to training mode
total_loss = 0
for step, (images, _) in enumerate(tqdm(dataloader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}")):
optimizer.zero_grad() # Clear gradients
images = images.to(DEVICE)
# Sample random time steps for the current batch
t = torch.randint(0, TIMESTEPS, (images.shape[0],), device=DEVICE).long()
# Calculate loss
loss = p_losses(model, images, t)
total_loss += loss.item()
# Backpropagation and optimization step
loss.backward()
optimizer.step()
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch+1} completed. Average Loss: {avg_loss:.4f}")
# Generate and save sample images periodically
if (epoch + 1) % 10 == 0 or epoch == 0: # Save samples at epoch 0 and every 10 epochs
model.eval() # Set model to evaluation mode for sampling
sample_and_save_images(model, epoch + 1, num_samples=16)
print("Training finished.")
# Optionally save the final model
torch.save(model.state_dict(), os.path.join(SAVE_DIR, "diffusion_model_final.pth"))
print(f"Final model saved to {os.path.join(SAVE_DIR, 'diffusion_model_final.pth')}")
if __name__ == "__main__":
train_diffusion_model()
No comments:
Post a Comment