Lesson 57: Generative and Diffusion Models¶
Every network so far has mapped an image to something smaller: a label, a box, a mask, an embedding. Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) run the idea in reverse: learn to map pure noise to a realistic image. The core trick is deceptively simple — train a network to undo one small step of noise-corruption at a time, then chain many such steps together, starting from noise and ending at something that looks like the training data.
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
The forward process: gradually destroying an image¶
Define a fixed schedule of $T$ steps, each adding a small amount of Gaussian noise: $x_t = \sqrt{\alpha_t}\,x_{t-1} + \sqrt{1-\alpha_t}\,\varepsilon$. Applying this repeatedly is slow to simulate one step at a time, but because sums of independent Gaussians are themselves Gaussian, there's a closed form that jumps straight from the original image $x_0$ to any step $t$: $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon$, where $\bar\alpha_t = \prod_{s=1}^t \alpha_s$.
SIZE = 16
T = 100
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)
def make_image(shape_type, cx, cy, size=SIZE):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
return img * 2 - 1 # scale to [-1, 1], the usual diffusion-model convention
def make_dataset(rng, n):
imgs = []
for _ in range(n):
shape_type = rng.choice(['plus', 'circle'])
cx, cy = rng.integers(5, 11), rng.integers(5, 11)
imgs.append(make_image(shape_type, cx, cy))
return np.array(imgs, dtype=np.float32)
rng = np.random.default_rng(2)
X = make_dataset(rng, 500)
X_t = torch.tensor(X).unsqueeze(1)
# validate the closed form against literally simulating T small steps, statistically
x0_single = X_t[0:1]
t_idx = 50
n_trials = 3000
x0_rep = x0_single.expand(n_trials, -1, -1, -1)
torch.manual_seed(1)
x_t_iterative = x0_rep.clone()
for step in range(t_idx):
x_t_iterative = torch.sqrt(alphas[step]) * x_t_iterative + torch.sqrt(betas[step]) * torch.randn_like(x0_rep)
torch.manual_seed(2)
eps = torch.randn_like(x0_rep)
x_t_closed = torch.sqrt(alpha_bars[t_idx - 1]) * x0_rep + torch.sqrt(1 - alpha_bars[t_idx - 1]) * eps
print(f'{t_idx} iterative small steps: mean={x_t_iterative.mean().item():.4f}, std={x_t_iterative.std().item():.4f}')
print(f'closed-form single jump: mean={x_t_closed.mean().item():.4f}, std={x_t_closed.std().item():.4f}')
fig, axes = plt.subplots(1, 6, figsize=(11, 2))
x0_demo = X_t[0:1]
for ax, t_show in zip(axes, [0, 10, 25, 50, 75, 99]):
torch.manual_seed(0)
if t_show == 0:
img = x0_demo
else:
img = torch.sqrt(alpha_bars[t_show]) * x0_demo + torch.sqrt(1 - alpha_bars[t_show]) * torch.randn_like(x0_demo)
ax.imshow(img[0, 0], cmap='gray'); ax.set_title(f't={t_show}', fontsize=8); ax.axis('off')
plt.suptitle('Forward process: the same image at increasing noise levels')
plt.show()
Training: predict the noise, not the image¶
The network's job is simple to state: given a noisy image $x_t$ and the noise level $t$, predict the noise $\varepsilon$ that was added. The training signal comes for free — pick a random image and a random $t$, add known noise via the closed form above, and check whether the network can recover it. No labels needed at all, only unlabeled images (the same self-supervised spirit as Lesson 47's contrastive learning, applied to a completely different task). The network needs to know $t$ because the right amount of correction to apply is very different for a barely-noised image than for one that's almost pure noise — this is injected via a sinusoidal time embedding, the same construction as Lesson 45's positional encoding, indexing "how far along" instead of "where in the sequence."
class TimeEmbedding(nn.Module):
def __init__(self, dim=32):
super().__init__()
self.dim = dim
def forward(self, t):
half = self.dim // 2
freqs = torch.exp(-np.log(10000) * torch.arange(half).float() / half)
args = t[:, None].float() * freqs[None, :]
return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
class DenoiseNet(nn.Module):
def __init__(self, ch=32, time_dim=32):
super().__init__()
self.time_embed = TimeEmbedding(time_dim)
self.time_mlp = nn.Linear(time_dim, ch)
self.conv1 = nn.Conv2d(1, ch, 3, padding=1)
self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)
self.conv3 = nn.Conv2d(ch, 1, 3, padding=1)
def forward(self, x, t):
temb = self.time_mlp(self.time_embed(t))[:, :, None, None]
h = F.relu(self.conv1(x) + temb)
h = F.relu(self.conv2(h) + temb)
return self.conv3(h)
torch.manual_seed(0)
model = DenoiseNet()
opt = torch.optim.Adam(model.parameters(), lr=0.001)
n = len(X_t)
for epoch in range(800):
idx = np.random.default_rng(epoch).permutation(n)[:64]
x0 = X_t[idx]
t = torch.randint(0, T, (x0.shape[0],))
noise = torch.randn_like(x0)
ab = alpha_bars[t][:, None, None, None]
x_t_batch = torch.sqrt(ab) * x0 + torch.sqrt(1 - ab) * noise
pred_noise = model(x_t_batch, t)
loss = F.mse_loss(pred_noise, noise)
opt.zero_grad()
loss.backward()
opt.step()
print(f'final training loss (predicted-vs-true noise MSE): {loss.item():.3f}')
Sampling: reversing the process, one predicted-noise-removal at a time¶
Starting from pure Gaussian noise, repeatedly ask the trained network "what noise was added to get here?", subtract a scaled version of that prediction, and add back a small amount of fresh randomness (except on the very last step) — this last part matters: without it, the process is deterministic and tends to produce blurry averages rather than sharp, varied samples.
@torch.no_grad()
def sample(model, n_samples, seed):
torch.manual_seed(seed)
x = torch.randn(n_samples, 1, SIZE, SIZE)
for t in reversed(range(T)):
t_batch = torch.full((n_samples,), t, dtype=torch.long)
pred_noise = model(x, t_batch)
alpha_t, alpha_bar_t, beta_t = alphas[t], alpha_bars[t], betas[t]
mean = (1 / torch.sqrt(alpha_t)) * (x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * pred_noise)
x = mean + torch.sqrt(beta_t) * torch.randn_like(x) if t > 0 else mean
return x
gen_samples = sample(model, 8, seed=42)
fig, axes = plt.subplots(1, 8, figsize=(11, 2))
for ax, im in zip(axes, gen_samples):
ax.imshow(im[0], cmap='gray'); ax.axis('off')
plt.suptitle('Generated samples, starting from pure noise')
plt.show()
Do the samples actually resemble the training shapes?¶
Every real training image is, by construction, exactly one solid connected blob (a plus or a circle). Pure random noise, thresholded the same way, is a scatter of many small disconnected fragments. Check where generated samples land on that spectrum.
def connectivity_stats(imgs):
n_components_list = []
for im in imgs:
binary = (im[0].numpy() > 0).astype(np.uint8)
n_comp, _ = cv2.connectedComponents(binary)
n_components_list.append(n_comp - 1) # subtract the background label
return np.array(n_components_list)
n_eval = 50
gen_eval = sample(model, n_eval, seed=7)
noise_eval = torch.randn(n_eval, 1, SIZE, SIZE)
gen_ncomp = connectivity_stats(gen_eval)
noise_ncomp = connectivity_stats(noise_eval)
real_ncomp = connectivity_stats(X_t[:n_eval])
print(f'mean # connected foreground blobs, real training images: {real_ncomp.mean():.1f} (always exactly 1, by construction)')
print(f'mean # connected foreground blobs, generated samples: {gen_ncomp.mean():.1f}')
print(f'mean # connected foreground blobs, pure random noise: {noise_ncomp.mean():.1f}')
Generated samples land much closer to real images' single-blob structure than pure noise does — the reverse process has learned to reassemble coherent, connected shapes from nothing, not just to locally smooth random pixels.
From this toy model to Stable Diffusion¶
Three changes separate this lesson's model from a real text-to-image system like Stable Diffusion:
- Conditioning: feed a text embedding (Lesson 49's CLIP text encoder, or a dedicated language model) into the denoising network at every step — e.g. via cross-attention (Lesson 44), the same query/key/value mechanism used throughout Part 4, just with the queries coming from image features and the keys/values from the text embedding. This is what turns "generate a plausible image" into "generate an image matching this prompt."
- Classifier-free guidance: train the network both with and without the text conditioning (randomly dropping it during training), then at sampling time extrapolate away from the unconditional prediction toward the conditional one — this single trick is responsible for a large fraction of the prompt-adherence quality jump in modern text-to-image models.
- Latent diffusion: run the entire diffusion process in a compressed latent space (a pretrained autoencoder's bottleneck, not raw pixels) rather than on full-resolution images — the reason Stable Diffusion can generate a 512x512 image in a reasonable number of steps on a single GPU, instead of the pixel-space diffusion this lesson used.
The forward process, the noise-prediction training objective, and the iterative reverse sampling loop — the three pieces built from scratch above — are unchanged by any of these additions. They're still the mechanism underneath.
Exercise¶
- Reduce
Tfrom 100 to 20. Does the connected-blob metric for generated samples get better, worse, or stay about the same — and what does that suggest about the tradeoff between number of diffusion steps and sample quality? - The sampling loop adds fresh noise at every step except the last (
torch.sqrt(beta_t) * torch.randn_like(x) if t > 0 else mean). Remove that noise addition entirely (always usemean) and compare the connected-blob metric. Do fully deterministic samples look more or less like the training shapes than the stochastic version? - Class-conditional generation (without full text conditioning) is simpler to add: append a one-hot class vector to the time embedding before the
time_mlplayer, train on labeled plus/circle data, and sample separately for each class. Does conditioning on the class make the connected-blob metric better for either shape specifically?