A complete diffusion model project — implementing a simplified DDPM (Denoising Diffusion Probabilistic Model) to generate images, directly experiencing the forward-noising/reverse-denoising process from first principles.
Problem Statement
Implement and train a simplified diffusion model that generates new images by learning to reverse a gradual noising process, evaluated qualitatively by visual inspection of generated samples.
Dataset
MNIST again works well as a starting dataset here — small images keep training time and the diffusion process itself (many denoising steps) tractable on modest hardware.
Architecture & Approach
A U-Net (the same architecture family from the segmentation project, here predicting noise instead of a segmentation mask) is trained to predict the noise added to an image at a given noise level — at generation time, this noise-predicting model is applied repeatedly, starting from pure random noise, gradually denoising toward a realistic image.
Step-by-Step Build
import torch
import torch.nn as nn
# 1. The forward noising process -- add noise according to a fixed schedule
T = 200 # number of diffusion steps
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1 - betas
alpha_bars = torch.cumprod(alphas, dim=0) # cumulative product -- how much signal remains at step t
def forward_noise(x0, t):
noise = torch.randn_like(x0)
sqrt_alpha_bar = alpha_bars[t].sqrt().view(-1, 1, 1, 1)
sqrt_one_minus_alpha_bar = (1 - alpha_bars[t]).sqrt().view(-1, 1, 1, 1)
noisy_x = sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise
return noisy_x, noise # noisy_x is the input; noise is what the model must learn to predict
# 2. A simplified noise-prediction U-Net (reusing the segmentation project's structure,
# with a time-step embedding added so the model knows which noise level it's denoising)
class SimpleDenoiser(nn.Module):
def __init__(self):
super().__init__()
self.time_embed = nn.Embedding(T, 32)
self.conv1 = nn.Conv2d(1, 64, 3, padding=1)
self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
self.conv3 = nn.Conv2d(64, 1, 3, padding=1)
self.time_proj = nn.Linear(32, 64)
def forward(self, x, t):
t_embed = self.time_proj(self.time_embed(t)).view(-1, 64, 1, 1)
h = torch.relu(self.conv1(x)) + t_embed # inject time information into the feature map
h = torch.relu(self.conv2(h))
return self.conv3(h) # predicts the noise that was added
model = SimpleDenoiser()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0002)
loss_fn = nn.MSELoss()
# 3. Training -- sample a random timestep and noise level for each training example
for epoch in range(10):
for x0, _ in train_loader:
t = torch.randint(0, T, (x0.size(0),))
noisy_x, true_noise = forward_noise(x0, t)
optimizer.zero_grad()
predicted_noise = model(noisy_x, t)
loss = loss_fn(predicted_noise, true_noise) # the model just learns to predict the noise
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}: loss={loss.item():.4f}")
# 4. Generation -- start from pure noise, iteratively denoise
@torch.no_grad()
def generate(model, num_images=4):
x = torch.randn(num_images, 1, 28, 28)
for t in reversed(range(T)):
t_batch = torch.full((num_images,), t, dtype=torch.long)
predicted_noise = model(x, t_batch)
alpha = alphas[t]; alpha_bar = alpha_bars[t]; beta = betas[t]
x = (1 / alpha.sqrt()) * (x - (beta / (1 - alpha_bar).sqrt()) * predicted_noise)
if t > 0:
x += beta.sqrt() * torch.randn_like(x) # add a small amount of noise back, except at the final step
return x
generated_images = generate(model)Expected Results
After sufficient training, the reverse process should transform pure random noise into recognizable digit-like images — with a simplified architecture and modest training time, expect results noticeably rougher than a production diffusion model, but the core mechanism (learning to predict and remove noise, applied iteratively) will be genuinely working and directly observable.
Key Learnings & Extensions
- This project makes the forward/reverse process from Diffusion Forward Process and Diffusion Reverse Process concrete — you're not just reading the math, you're watching random noise become a recognizable image, one denoising step at a time.
- Extension: Add class conditioning (like the GAN project's suggested extension) so you can control which digit gets generated.
- Extension: Compare training stability and sample diversity against the GAN project — diffusion models are generally noted for training more stably than GANs, and this project gives you a direct, hands-on basis for that comparison.
