{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Lesson 57: Generative and Diffusion Models\n\nEvery 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."
},
{
"cell_type": "code",
"id": "54670a9e",
"source": "import numpy as np\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "0bb4a387",
"source": "## The forward process: gradually destroying an image\n\nDefine 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$.",
"metadata": {}
},
{
"cell_type": "code",
"id": "a51edae5",
"source": "SIZE = 16\nT = 100\n\nbetas = torch.linspace(1e-4, 0.02, T)\nalphas = 1.0 - betas\nalpha_bars = torch.cumprod(alphas, dim=0)\n\ndef make_image(shape_type, cx, cy, size=SIZE):\n img = np.zeros((size, size), dtype=np.float32)\n if shape_type == 'plus':\n img[cy-1:cy+2, cx-3:cx+4] = 1.0\n img[cy-3:cy+4, cx-1:cx+2] = 1.0\n else:\n yy, xx = np.mgrid[0:size, 0:size]\n img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0\n return img * 2 - 1 # scale to [-1, 1], the usual diffusion-model convention\n\ndef make_dataset(rng, n):\n imgs = []\n for _ in range(n):\n shape_type = rng.choice(['plus', 'circle'])\n cx, cy = rng.integers(5, 11), rng.integers(5, 11)\n imgs.append(make_image(shape_type, cx, cy))\n return np.array(imgs, dtype=np.float32)\n\nrng = np.random.default_rng(2)\nX = make_dataset(rng, 500)\nX_t = torch.tensor(X).unsqueeze(1)\n\n# validate the closed form against literally simulating T small steps, statistically\nx0_single = X_t[0:1]\nt_idx = 50\nn_trials = 3000\nx0_rep = x0_single.expand(n_trials, -1, -1, -1)\n\ntorch.manual_seed(1)\nx_t_iterative = x0_rep.clone()\nfor step in range(t_idx):\n x_t_iterative = torch.sqrt(alphas[step]) * x_t_iterative + torch.sqrt(betas[step]) * torch.randn_like(x0_rep)\n\ntorch.manual_seed(2)\neps = torch.randn_like(x0_rep)\nx_t_closed = torch.sqrt(alpha_bars[t_idx - 1]) * x0_rep + torch.sqrt(1 - alpha_bars[t_idx - 1]) * eps\n\nprint(f'{t_idx} iterative small steps: mean={x_t_iterative.mean().item():.4f}, std={x_t_iterative.std().item():.4f}')\nprint(f'closed-form single jump: mean={x_t_closed.mean().item():.4f}, std={x_t_closed.std().item():.4f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"id": "54a8663d",
"source": "fig, axes = plt.subplots(1, 6, figsize=(11, 2))\nx0_demo = X_t[0:1]\nfor ax, t_show in zip(axes, [0, 10, 25, 50, 75, 99]):\n torch.manual_seed(0)\n if t_show == 0:\n img = x0_demo\n else:\n img = torch.sqrt(alpha_bars[t_show]) * x0_demo + torch.sqrt(1 - alpha_bars[t_show]) * torch.randn_like(x0_demo)\n ax.imshow(img[0, 0], cmap='gray'); ax.set_title(f't={t_show}', fontsize=8); ax.axis('off')\nplt.suptitle('Forward process: the same image at increasing noise levels')\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "655a33e5",
"source": "## Training: predict the noise, not the image\n\nThe 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.\"",
"metadata": {}
},
{
"cell_type": "code",
"id": "f1f264c5",
"source": "class TimeEmbedding(nn.Module):\n def __init__(self, dim=32):\n super().__init__()\n self.dim = dim\n\n def forward(self, t):\n half = self.dim // 2\n freqs = torch.exp(-np.log(10000) * torch.arange(half).float() / half)\n args = t[:, None].float() * freqs[None, :]\n return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)\n\nclass DenoiseNet(nn.Module):\n def __init__(self, ch=32, time_dim=32):\n super().__init__()\n self.time_embed = TimeEmbedding(time_dim)\n self.time_mlp = nn.Linear(time_dim, ch)\n self.conv1 = nn.Conv2d(1, ch, 3, padding=1)\n self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)\n self.conv3 = nn.Conv2d(ch, 1, 3, padding=1)\n\n def forward(self, x, t):\n temb = self.time_mlp(self.time_embed(t))[:, :, None, None]\n h = F.relu(self.conv1(x) + temb)\n h = F.relu(self.conv2(h) + temb)\n return self.conv3(h)\n\ntorch.manual_seed(0)\nmodel = DenoiseNet()\nopt = torch.optim.Adam(model.parameters(), lr=0.001)\nn = len(X_t)\nfor epoch in range(800):\n idx = np.random.default_rng(epoch).permutation(n)[:64]\n x0 = X_t[idx]\n t = torch.randint(0, T, (x0.shape[0],))\n noise = torch.randn_like(x0)\n ab = alpha_bars[t][:, None, None, None]\n x_t_batch = torch.sqrt(ab) * x0 + torch.sqrt(1 - ab) * noise\n pred_noise = model(x_t_batch, t)\n loss = F.mse_loss(pred_noise, noise)\n opt.zero_grad()\n loss.backward()\n opt.step()\n\nprint(f'final training loss (predicted-vs-true noise MSE): {loss.item():.3f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "49e3bbf1",
"source": "## Sampling: reversing the process, one predicted-noise-removal at a time\n\nStarting 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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "b9d8c558",
"source": "@torch.no_grad()\ndef sample(model, n_samples, seed):\n torch.manual_seed(seed)\n x = torch.randn(n_samples, 1, SIZE, SIZE)\n for t in reversed(range(T)):\n t_batch = torch.full((n_samples,), t, dtype=torch.long)\n pred_noise = model(x, t_batch)\n alpha_t, alpha_bar_t, beta_t = alphas[t], alpha_bars[t], betas[t]\n mean = (1 / torch.sqrt(alpha_t)) * (x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * pred_noise)\n x = mean + torch.sqrt(beta_t) * torch.randn_like(x) if t > 0 else mean\n return x\n\ngen_samples = sample(model, 8, seed=42)\n\nfig, axes = plt.subplots(1, 8, figsize=(11, 2))\nfor ax, im in zip(axes, gen_samples):\n ax.imshow(im[0], cmap='gray'); ax.axis('off')\nplt.suptitle('Generated samples, starting from pure noise')\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "4b761332",
"source": "## Do the samples actually resemble the training shapes?\n\nEvery 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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "f57b5a23",
"source": "def connectivity_stats(imgs):\n n_components_list = []\n for im in imgs:\n binary = (im[0].numpy() > 0).astype(np.uint8)\n n_comp, _ = cv2.connectedComponents(binary)\n n_components_list.append(n_comp - 1) # subtract the background label\n return np.array(n_components_list)\n\nn_eval = 50\ngen_eval = sample(model, n_eval, seed=7)\nnoise_eval = torch.randn(n_eval, 1, SIZE, SIZE)\n\ngen_ncomp = connectivity_stats(gen_eval)\nnoise_ncomp = connectivity_stats(noise_eval)\nreal_ncomp = connectivity_stats(X_t[:n_eval])\n\nprint(f'mean # connected foreground blobs, real training images: {real_ncomp.mean():.1f} (always exactly 1, by construction)')\nprint(f'mean # connected foreground blobs, generated samples: {gen_ncomp.mean():.1f}')\nprint(f'mean # connected foreground blobs, pure random noise: {noise_ncomp.mean():.1f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "15207c0d",
"source": "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.\n\n## From this toy model to Stable Diffusion\n\nThree changes separate this lesson's model from a real text-to-image system like Stable Diffusion:\n\n- **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.\"\n- **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.\n- **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.\n\nThe 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.\n\n### Exercise\n\n1. Reduce `T` from 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?\n2. 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 use `mean`) and compare the connected-blob metric. Do fully deterministic samples look more or less like the training shapes than the stochastic version?\n3. Class-conditional generation (without full text conditioning) is simpler to add: append a one-hot class vector to the time embedding before the `time_mlp` layer, 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?",
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}