Image Generation › Implementation: Building a Mini Diffusion Pipeline
Language AI HandbookAll chapters Image Generation › Implementation: Building a Mini Diffusion PipelineOne section · Sections in this chapter10 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a simplified diffusion training and sampling loop to make these concepts concrete. We will not train a full-scale text-to-image model (that requires significant compute and large datasets), but we will build a complete DDPM pipeline that learns to generate a structured 2D distribution. This allows us to visualize the full forward-reverse process and verify that the math works as described.
Setup and Data
We will work with a toy 2D dataset: a two-moons arrangement. This lets us visualize the full forward-reverse process clearly. Despite its simplicity, this dataset is non-trivial for generative models because it has a bimodal, non-Gaussian shape with a curved manifold.
import numpy as np
import torch
def make_moons_dataset(n_samples=2000, noise=0.05, seed=42):
"""Create a simple two-moons 2D dataset."""
rng = np.random.RandomState(seed)
n_samples_each = n_samples // 2
theta1 = np.linspace(0, np.pi, n_samples_each)
theta2 = np.linspace(np.pi, 2 * np.pi, n_samples_each)
x1 = np.column_stack([np.cos(theta1), np.sin(theta1)])
x2 = np.column_stack([np.cos(theta2) + 1, np.sin(theta2)])
x = np.vstack([x1, x2]) + rng.randn(n_samples, 2) * noise
return x.astype(np.float32)
data = make_moons_dataset(n_samples=3000)
X = torch.tensor(data)
n_samples, d = X.shapeDataset shape: torch.Size([3000, 2]) Data range: [-1.136, 2.134] Data mean: [4.9965525e-01 3.2643635e-05]
Defining the Noise Schedule
The DiffusionSchedule class precomputes all the quantities derived from the noise schedule that we need for both training (forward process) and sampling (reverse process).
class DiffusionSchedule:
"""Linear noise schedule for DDPM."""
def __init__(self, T=200, beta_start=1e-4, beta_end=0.02, device="cpu"):
self.T = T
self.device = device
# Linear noise schedule
self.betas = torch.linspace(beta_start, beta_end, T, device=device)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# Precompute useful quantities
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(
1.0 - self.alphas_cumprod
)
def q_sample(self, x0, t, noise=None):
"""Add noise to x0 at timestep t. Implements the forward process."""
if noise is None:
noise = torch.randn_like(x0)
sqrt_alpha = self.sqrt_alphas_cumprod[t].view(-1, 1)
sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1)
return sqrt_alpha * x0 + sqrt_one_minus_alpha * noise, noise
schedule = DiffusionSchedule(T=200)Noise schedule: T=200 Alpha(t=0) = 0.9999 (almost no noise) Alpha(t=100) = 0.5964 (partial noise) Alpha(t=199) = 0.132183 (almost pure noise)
At , so the noisy image is nearly identical to the clean data. At , so the signal has been destroyed and we have nearly pure Gaussian noise. The middle timesteps correspond to partial noise: the data has some structure but is significantly corrupted.
Denoising Network
For our 2D toy problem, a small MLP handles noise prediction. In full image diffusion, this would be a U-Net with millions of parameters. The architecture here captures the essential design: a network that takes both the noisy data and the timestep, and outputs a noise prediction of the same shape as the input.
import torch.nn as nn
class SinusoidalPositionEmbedding(nn.Module):
"""Embed timestep t as a sinusoidal vector."""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, t):
half_dim = self.dim // 2
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
emb = t[:, None].float() * emb[None, :]
return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
class DenoisingMLP(nn.Module):
"""MLP that predicts noise from noisy data and timestep."""
def __init__(self, data_dim=2, hidden_dim=256, time_dim=64):
super().__init__()
self.time_embed = SinusoidalPositionEmbedding(time_dim)
self.net = nn.Sequential(
nn.Linear(data_dim + time_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, data_dim),
)
def forward(self, x, t):
t_emb = self.time_embed(t)
return self.net(torch.cat([x, t_emb], dim=-1))
model = DenoisingMLP(data_dim=2, hidden_dim=256, time_dim=64)The sinusoidal timestep embedding encodes the scalar as a vector of alternating sines and cosines at different frequencies. This is the same approach used in transformer positional embeddings, and it ensures the model can distinguish all timesteps and interpolate smoothly between them. Without this embedding, the network would not know whether it is being asked to denoise heavily corrupted data (where only global structure matters) or lightly corrupted data (where fine details matter).
Model parameters: 149,250 Input shape: torch.Size([8, 2]), output shape: torch.Size([8, 2])
Training Loop
import torch.nn.functional as F
from torch.optim import Adam
def train_diffusion(model, schedule, X, epochs=3000, batch_size=256, lr=1e-3):
"""Train the denoising network using DDPM objective."""
optimizer = Adam(model.parameters(), lr=lr)
losses = []
for epoch in range(epochs):
# Sample random batch
idx = torch.randint(0, len(X), (batch_size,))
x0 = X[idx]
# Sample random timesteps
t = torch.randint(0, schedule.T, (batch_size,))
# Add noise according to schedule
xt, noise = schedule.q_sample(x0, t)
# Predict noise
noise_pred = model(xt, t)
# Loss: MSE between predicted and actual noise
loss = F.mse_loss(noise_pred, noise)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 200 == 0:
losses.append(loss.item())
return losses
losses = train_diffusion(model, schedule, X, epochs=3000, batch_size=256)Training epochs: 3000 Final loss: 0.5023 Initial loss: 1.1503 Loss reduction: 56.3%
The loss measures the mean squared error between the predicted noise and the actual noise added. As training progresses, the model learns to predict noise more accurately across all timesteps, and the loss decreases. A perfect model would achieve loss of 0, but in practice the irreducible loss reflects the stochasticity of the forward process.
DDPM Sampling
With a trained model, we run the reverse process to generate new samples. The sampling loop iterates from down to , denoising at each step.
@torch.no_grad()
def ddpm_sample(model, schedule, n_samples=1000, d=2):
"""Sample from the learned distribution using DDPM reverse process."""
model.eval()
# Start from pure noise
x = torch.randn(n_samples, d)
trajectory = [x.clone()]
for t in reversed(range(schedule.T)):
t_batch = torch.full((n_samples,), t, dtype=torch.long)
# Predict noise
eps_pred = model(x, t_batch)
# Compute denoised estimate
beta_t = schedule.betas[t]
alpha_t = schedule.alphas[t]
alpha_bar_t = schedule.alphas_cumprod[t]
# DDPM reverse step mean
coef = beta_t / torch.sqrt(1 - alpha_bar_t)
mu = (x - coef * eps_pred) / torch.sqrt(alpha_t)
if t > 0:
# Add stochastic noise for all but the last step
noise = torch.randn_like(x)
x = mu + torch.sqrt(beta_t) * noise
else:
x = mu
if t % 50 == 0:
trajectory.append(x.clone())
return x, trajectory
generated, trajectory = ddpm_sample(model, schedule, n_samples=2000)Generated 2000 samples in 2D Generated range: [-1.283, 2.232] Generated mean: [ 0.5127893 -0.04144546]
Notice that we add stochastic noise at every reverse step except the last (). This stochasticity is essential to DDPM sampling: without it, the reverse process would be deterministic from the initial noise, producing less diverse samples and sometimes getting stuck in local optima. At , we take the deterministic mean to obtain the final clean sample.
Visualizing the Noise Schedule
Visualizing Generated vs. Real Data
The generated distribution captures both crescent arcs, confirming that the model has learned the broad structure through denoising alone and avoided collapsing to one mode. The bridge and off-manifold points also show that this small denoising network is not a perfect density model. A GAN on this dataset would often collapse to one mode; diffusion models resist this failure mode more effectively.
Visualizing Training Loss
Visualizing Classifier-Free Guidance Scale
The guidance scale is one of the most impactful hyperparameters in text-to-image generation. The CFG formula decomposes the guided noise prediction into a linear combination of conditional and unconditional predictions. We can visualize exactly how the weights on each component change with different guidance scales.
At , the conditional noise prediction receives a weight of 7.5 and the unconditional prediction receives a weight of . The negative unconditional weight means the model actively moves away from unconditional predictions, amplifying whatever features the text condition adds. This is why high guidance scales can cause oversaturation: the model is maximally exploiting the direction of the text signal, sometimes beyond what produces photorealistic results.