{
"cells": [
{
"cell_type": "markdown",
"id": "ffdf0c18",
"metadata": {},
"source": [
"# TP – DDPM on Two Moons \n",
"\n",
"This notebook walks through the implementation of a Denoising Diffusion Probabilistic Model (DDPM)\n",
"applied to the Two Moons dataset. The goal is to help students understand the core components\n",
"of diffusion models through hands-on coding and visualization. \n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "ea45dc7a",
"metadata": {},
"source": [
"\n",
"We will:\n",
"1. Build and visualize the Two Moons dataset (centered and rescaled).\n",
"2. Implement several beta schedules and a reusable `NoiseScheduler` class.\n",
"3. Visualize the noisy data for different timesteps.\n",
"4. Explore positional/time embeddings and inspect them graphically.\n",
"5. Assemble the MLP score network and verify intermediate tensors.\n",
"6. Train with a DDPM objective (MSE on predicted noise) using `tqdm`.\n",
"7. Sample new points, compare to the dataset, and export GIF animations of the reverse process.\n",
"8. Run a lightweight ablation on embeddings / schedules.\n",
"\n",
"The students are expected to fill in the missing code sections marked with `# TODO` comments."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6d9b5756",
"metadata": {},
"outputs": [],
"source": [
"import math\n",
"import random\n",
"from pathlib import Path\n",
"\n",
"import imageio\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import torch\n",
"from sklearn.datasets import make_moons\n",
"from torch import nn\n",
"from torch.nn import functional as F\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"import tqdm\n",
"\n",
"# ------------------------------------------------------------------\n",
"# Reproducibility and device\n",
"# ------------------------------------------------------------------\n",
"SEED = 42\n",
"random.seed(SEED)\n",
"np.random.seed(SEED)\n",
"torch.manual_seed(SEED)\n",
"\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"print(f\"Running on {device}\")\n",
"\n",
"cfg = {\n",
" \"num_samples\": 8000,\n",
" \"noise\": 0.03,\n",
" \"batch_size\": 512,\n",
" \"num_timesteps\": 1000,\n",
" \"beta_start\": 1e-4,\n",
" \"beta_end\": 2e-2,\n",
" \"beta_schedule\": \"linear\", # {linear, quadratic, cosine}\n",
" \"hidden_size\": 128,\n",
" \"hidden_layers\": 3,\n",
" \"emb_size\": 128,\n",
" \"lr\": 2e-4,\n",
" \"num_epochs\": 1000,\n",
" \"eval_every\": 250,\n",
" \"sample_points\": 1024,\n",
"}\n",
"print(cfg)"
]
},
{
"cell_type": "markdown",
"id": "869b1c3e",
"metadata": {},
"source": [
"## 1. Dataset preparation and visualization\n",
"We rely on `sklearn.make_moons`, then we **center and rescale** each axis so that\n",
"points cover approximately the square `[-1, 1]^2` (exactly as in `ddpm.ipynb`).\n",
"\n",
"\n",
"**TODO:**\n",
"- Load the Two Moons dataset with the specified noise level.\n",
"- Apply the exact shift/rescale used in `ddpm.ipynb` so both axes live in [-1, 1].\n",
"- Wrap the numpy array into a `TensorDataset` and prepare a shuffled `DataLoader`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "efec4f8e",
"metadata": {},
"outputs": [],
"source": [
"def moons_dataset(n_samples=cfg[\"num_samples\"], noise=cfg[\"noise\"], seed=SEED):\n",
" \"\"\"Return a TensorDataset with normalized Two Moons coordinates.\"\"\"\n",
" X, _ = make_moons(n_samples=n_samples, noise=noise, random_state=seed)\n",
" X = X.astype(np.float32)\n",
" X[:, 0] = X[:, 0] # TODO\n",
" X[:, 1] = X[:, 1] # TODO\n",
"\n",
" return TensorDataset(torch.from_numpy(X))\n",
"\n",
"\n",
"train_dataset = moons_dataset()\n",
"train_loader = DataLoader(\n",
" train_dataset,\n",
" batch_size=cfg[\"batch_size\"],\n",
" shuffle=True,\n",
" drop_last=True,\n",
")\n",
"\n",
"print(f\"Dataset size: {len(train_dataset)} samples\")\n",
"print(\n",
" f\"Mimimum and maximum of the samples:\\n {train_dataset.tensors[0].min(0).values}\\n {train_dataset.tensors[0].max(0).values}\"\n",
")"
]
},
{
"cell_type": "markdown",
"id": "813e2feb",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "33ea3bb4",
"metadata": {},
"source": [
"### 1.1. Quick scatter helper\n",
"We often want to compare the raw dataset with noisy samples or generated points.\n",
"`scatter_data_and_samples` overlays both sets and annotates the timestep.\n",
"Students will trigger it after prepping the dataset.\n",
"\n",
"**TODO:**\n",
"- Write a helper that scatters the dataset and optional generated samples on the same axes.\n",
"- Use the scatter function from `matplotlib.pyplot` to plot points.\n",
"- Fix the axis limits to `[-1.5, 1.5]` for better visibility.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b0ce028e",
"metadata": {},
"outputs": [],
"source": [
"def scatter_data_and_samples(\n",
" data,\n",
" samples=None,\n",
" title=\"Two Moons\",\n",
" label=None,\n",
" figsize=(5, 5),\n",
" alpha_real=0.5,\n",
" alpha_fake=0.6,\n",
"):\n",
" # TODO: plot dataset and optional samples on the same scatter chart\n",
" if isinstance(data, TensorDataset):\n",
" pts = data.tensors[0].cpu()\n",
" else:\n",
" pts = data.detach().cpu() if torch.is_tensor(data) else data\n",
"\n",
" pts # TODO\n",
"\n",
" plt.figure(figsize=figsize)\n",
" # TODO\n",
" if samples is not None:\n",
" samples = samples.detach().cpu()\n",
" # TODO\n",
" title_suffix = f\" {label}\" if label is not None else \"\"\n",
" plt.title(f\"{title}{title_suffix}\")\n",
" plt.legend()\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "d65d7aa5",
"metadata": {},
"source": [
"### 1.2. Visual inspection of the normalized data\n",
"\n",
"**TODO:**\n",
"- Call `scatter_data_and_samples` right after building the dataset to visually sanity check the scaling.\n"
]
},
{
"cell_type": "markdown",
"id": "ad9a602b",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "20d6a163",
"metadata": {},
"outputs": [],
"source": [
"# TODO: visualize the dataset to verify scaling\n",
"scatter_data_and_samples(train_dataset, title=\"Normalized Two Moons\", label=\" Dataset\")"
]
},
{
"cell_type": "markdown",
"id": "813e2feb",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "e4350c89",
"metadata": {},
"source": [
"## 2. Noise schedules and diffusion coefficients\n",
"The forward diffusion process is controlled by a sequence $\\{\\beta_t\\}_{t=1}^T$ called the noise schedule.\n",
"Given $\\alpha_t = 1 - \\beta_t$ and $\\bar{\\alpha}_t = \\prod_{s=1}^t \\alpha_s$, the\n",
"forward process satisfies:\n",
"$$\n",
"q(x_t \\mid x_{t-1}) = \\mathcal{N}(x_t; \\sqrt{\\alpha_t} x_{t-1}, \\beta_t I), \\qquad\n",
"q(x_t \\mid x_0) = \\mathcal{N}(x_t; \\sqrt{\\bar{\\alpha}_t} x_0, (1-\\bar{\\alpha}_t) I).\n",
"$$\n",
"Below we implement three schedules so students can compare them.\n",
"\n",
"Let use call `beta_start` and `beta_end` the initial and final beta values, respectively.\n",
"- **Linear schedule:** $\\beta_t$ increases linearly from `beta_start` to `beta_end`.\n",
"- **Quadratic schedule:** $\\beta_t$ increases quadratically from `beta_start` to `beta_end`.\n",
"- **Cosine schedule:** $\\beta_t$ follows a cosine schedule as in Nichol & Dhariwal (2021):\n",
" $$\n",
" \\bar{\\alpha}_t = \\frac{f(t/T)}{f(0)}, \\quad \\text{where } f(s) = \\cos\\left(\\frac{s + 0.008}{1.008} \\cdot \\frac{\\pi}{2}\\right)^2.\n",
" $$\n",
" Then $\\beta_t = 1 - \\frac{\\bar{\\alpha}_t}{\\bar{\\alpha}_{t-1}}$ with $\\bar{\\alpha}_0 = 1$. In practice, we clip $\\beta_t$ to a maximum of 0.999.\n",
" \n",
"\n",
"**TODO:**\n",
"- Implement linear, quadratic, and cosine beta schedules that return tensors of shape `(T,)`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "15ef4990",
"metadata": {},
"outputs": [],
"source": [
"# TODO: implement multiple beta schedules (linear, quadratic, cosine)\n",
"def linear_beta_schedule(num_timesteps, beta_start, beta_end):\n",
" return # TODO\n",
"\n",
"\n",
"def quadratic_beta_schedule(num_timesteps, beta_start, beta_end):\n",
" return # TODO\n",
"\n",
"\n",
"def cosine_beta_schedule(num_timesteps, s=0.008):\n",
" steps = num_timesteps + 1\n",
" x = torch.linspace(0, num_timesteps, steps, dtype=torch.float32)\n",
" # TODO"
]
},
{
"cell_type": "markdown",
"id": "1b247003",
"metadata": {},
"source": [
"**TODO:** Run the next cell to try your code."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7cb973b0",
"metadata": {},
"outputs": [],
"source": [
"plt.figure(figsize=(10, 4))\n",
"plt.subplot(1, 3, 1)\n",
"plt.plot(linear_beta_schedule(cfg[\"num_timesteps\"], cfg[\"beta_start\"], cfg[\"beta_end\"]))\n",
"plt.title(\"Linear Beta Schedule\")\n",
"plt.subplot(1, 3, 2)\n",
"plt.plot(\n",
" quadratic_beta_schedule(cfg[\"num_timesteps\"], cfg[\"beta_start\"], cfg[\"beta_end\"])\n",
")\n",
"plt.title(\"Quadratic Beta Schedule\")\n",
"plt.subplot(1, 3, 3)\n",
"plt.plot(cosine_beta_schedule(cfg[\"num_timesteps\"]))\n",
"plt.title(\"Cosine Beta Schedule\")\n",
"plt.tight_layout()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "c8bc194b",
"metadata": {},
"source": [
"Using the $\\beta_t$ schedules, we can compute and plot the corresponding $\\bar{\\alpha}_t$ sequences:\n",
"$$ \\bar{\\alpha}_t = \\prod_{s=1}^t (1 - \\beta_s). $$\n"
]
},
{
"cell_type": "markdown",
"id": "813e2feb",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "86bcb8b8",
"metadata": {},
"source": [
"### 2.1. NoiseScheduler helper class\n",
"\n",
"We will build a `NoiseScheduler` class to encapsulate the noise schedule and precompute all necessary diffusion coefficients. We will progressively add methods to this class as we implement the DDPM components.\n",
"\n",
"You will need to implement the following methods using the DDPM equations:\n",
"- `add_noise(x0, eps, t)` implements: $$x_t = \\sqrt{\\bar{\\alpha}_t} x_0 + \\sqrt{1-\\bar{\\alpha}_t} \\epsilon.$$\n",
"- `reconstruct_x0(x_t, t, eps_theta)` implements the denoising prediction:\n",
"$$\\hat{x}_0 = \\frac{1}{\\sqrt{\\bar{\\alpha}_t}} \\left( x_t - \\sqrt{1 - \\bar{\\alpha}_t} \\epsilon_\\theta(x_t, t) \\right).$$\n",
"- `q_posterior(x0, x_t, t)` computes the mean of $q(x_{t-1} \\mid x_t, x_0)$:\n",
"$$\\mu_q = \\frac{\\sqrt{\\bar{\\alpha}_{t-1}} \\beta_t}{1 - \\bar{\\alpha}_t} x_0 + \\frac{\\sqrt{\\alpha_t} (1 - \\bar{\\alpha}_{t-1})}{1 - \\bar{\\alpha}_t} x_t.$$\n",
"- `step(model_output, t, x_t)` performs a single reverse update:\n",
"$$x_{t-1} = \\mu_q + \\sigma_t \\epsilon, \\quad \\text{where } \\sigma_t^2 = \\frac{1 - \\bar{\\alpha}_{t-1}}{1 - \\bar{\\alpha}_t} \\beta_t.$$\n",
"\n",
"\n",
"We precompute all necessary coefficients in the `__init__` method for efficiency:\n",
"- `betas`: $\\beta_t$ for $t=1,\\ldots,T$\n",
"- `alphas`: $\\alpha_t = 1 - \\beta_t$ for $t=1,\\ldots,T$\n",
"- `alphas_cumprod`: $\\bar{\\alpha}_t = \\prod_{s=1}^t \\alpha_s$ for $t=1,\\ldots,T$\n",
"- `alphas_cumprod_prev`: $\\bar{\\alpha}_{t-1}$ for $t=1,\\ldots,T$ with $\\bar{\\alpha}_0 = 1$\n",
"- `sqrt_alphas_cumprod`: $\\sqrt{\\bar{\\alpha}_t}$ for $t=1,\\ldots,T$\n",
"- `sqrt_one_minus_alphas_cumprod`: $\\sqrt{1 - \\bar{\\alpha}_t}$ for $t=1,\\ldots,T$\n",
"- `sqrt_recip_alphas_cumprod`: $\\frac{1}{\\sqrt{\\bar{\\alpha}_t}}$ for $t=1,\\ldots,T$\n",
"- `sqrt_recipm1_alphas_cumprod`: $\\sqrt{\\frac{1}{\\bar{\\alpha}_t} - 1}$ for $t=1,\\ldots,T$\n",
"- `posterior_variance`: $\\sigma_t^2 = \\frac{1 - \\bar{\\alpha}_{t-1}}{1 - \\bar{\\alpha}_t} \\beta_t$ for $t=1,\\ldots,T$\n",
"- `posterior_mean_coef1`: $\\frac{\\sqrt{\\bar{\\alpha}_{t-1}} \\beta_t}{1 - \\bar{\\alpha}_t}$ for $t=1,\\ldots,T$\n",
"- `posterior_mean_coef2`: $\\frac{\\sqrt{\\alpha_t} (1 - \\bar{\\alpha}_{t-1})}{1 - \\bar{\\alpha}_t}$ for $t=1,\\ldots,T$\n",
"\n",
"\n",
"In this class, you will also need a helper method `_extract` to broadcast timestep-dependent scalars over arbitrary batch shapes:\n",
"```python\n",
"example_coeff = self._extract(self.alphas_cumprod, t, x_t.shape)\n",
"```\n",
"\n",
"**TODO:**\n",
"- Implement `add_noise`, `reconstruct_x0`, `q_posterior`, and `step` exactly as in the DDPM equations.\n",
"- Use a helper like `_extract` to broadcast timestep-dependent scalars over arbitrary batch shapes.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "66b9065d",
"metadata": {},
"outputs": [],
"source": [
"class NoiseScheduler:\n",
" def __init__(\n",
" self,\n",
" num_timesteps=cfg[\"num_timesteps\"],\n",
" beta_start=cfg[\"beta_start\"],\n",
" beta_end=cfg[\"beta_end\"],\n",
" schedule_type=cfg[\"beta_schedule\"],\n",
" ):\n",
" # TODO: pick the requested schedule and pre-compute diffusion coefficients\n",
" if schedule_type == \"linear\":\n",
" betas = linear_beta_schedule(num_timesteps, beta_start, beta_end)\n",
" elif schedule_type == \"quadratic\":\n",
" betas = quadratic_beta_schedule(num_timesteps, beta_start, beta_end)\n",
" elif schedule_type == \"cosine\":\n",
" betas = cosine_beta_schedule(num_timesteps)\n",
" else:\n",
" raise ValueError(f\"Unknown schedule: {schedule_type}\")\n",
"\n",
" self.num_timesteps = num_timesteps\n",
" self.betas = betas\n",
" self.alphas = 1.0 - betas\n",
" self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)\n",
" self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)\n",
"\n",
" self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)\n",
" self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)\n",
" self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)\n",
" self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)\n",
"\n",
" self.posterior_variance = (\n",
" betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)\n",
" ).clamp(min=1e-20)\n",
" self.posterior_mean_coef1 = (\n",
" betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)\n",
" )\n",
" self.posterior_mean_coef2 = (\n",
" (1.0 - self.alphas_cumprod_prev)\n",
" * torch.sqrt(self.alphas)\n",
" / (1.0 - self.alphas_cumprod)\n",
" )\n",
"\n",
" def _extract(self, arr, timesteps, x_shape):\n",
" batch = timesteps.shape[0]\n",
" out = arr.to(timesteps.device)[timesteps]\n",
" return out.view(batch, *([1] * (len(x_shape) - 1)))\n",
"\n",
" def add_noise(self, x0, noise, timesteps):\n",
" # TODO: apply forward diffusion formula\n",
" \n",
" def reconstruct_x0(self, x_t, timesteps, noise):\n",
" # TODO: reconstruct x0 from x_t and predicted noise\n",
" \n",
" def q_posterior(self, x0, x_t, timesteps):\n",
" # TODO: compute posterior mean for q(x_{t-1} | x_t, x0)\n",
" \n",
" def step(self, model_output, timesteps, sample):\n",
" # TODO: implement one reverse step (Eq. 11 in DDPM paper)\n",
" nonzero_mask = (timesteps != 0).float().view(-1, *([1] * (sample.dim() - 1)))\n",
"\n",
" def __len__(self):\n",
" return self.num_timesteps"
]
},
{
"cell_type": "markdown",
"id": "2acc06e6",
"metadata": {},
"source": [
"### 2.2. Visual comparison of schedules\n",
"\n",
"**TODO:**\n",
"- Plot every schedule on the same matplotlib figure to highlight how their alphas ramps differ.\n"
]
},
{
"cell_type": "markdown",
"id": "94e5b6d8",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3de81195",
"metadata": {},
"outputs": [],
"source": [
"# TODO: plot the beta curves for each schedule to compare shapes\n",
"lin = # TODO\n",
"quad = # TODO\n",
"cos = # TODO\n",
"plt.figure(figsize=(6, 4))\n",
"plt.plot(lin.alphas_cumprod, label=\"linear\")\n",
"plt.plot(quad.alphas_cumprod, label=\"quadratic\")\n",
"plt.plot(cos.alphas_cumprod, label=\"cosine\")\n",
"plt.title(\"Alpha cumulative product schedules\")\n",
"plt.xlabel(\"t\")\n",
"plt.ylabel(\"Cumulative product of alpha_t\")\n",
"plt.legend()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "0d47146a",
"metadata": {},
"source": [
"### 2.3. Inspect noising for different timesteps\n",
"For a batch of real samples we can directly compute $x_t$ at arbitrary timesteps:\n",
"$$\n",
"x_t = \\sqrt{\\bar{\\alpha}_t} x_0 + \\sqrt{1-\\bar{\\alpha}_t} \\epsilon.\n",
"$$\n",
"This diagnostic plot helps spot mistakes in the forward process.\n",
"\n",
"**TODO:**\n",
"- Write a function that takes a few real points, corrupts them at several timesteps, and shows the scatter side-by-side.\n",
"- Verify visually that larger `t` corresponds to heavier noise (nearly Gaussian at `t=T-1`).\n"
]
},
{
"cell_type": "markdown",
"id": "fcc543f4",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cd52e764",
"metadata": {},
"outputs": [],
"source": [
"# TODO: visualize forward diffusion at multiple timesteps\n",
"def visualize_forward_process(\n",
" dataset, scheduler, steps=(0, 50, 200, 400, 999), n_points=2048\n",
"):\n",
" base = dataset.tensors[0]\n",
" idx = torch.randperm(len(base))[:n_points]\n",
" base = base[idx].to(device)\n",
" for t in steps:\n",
" noise = torch.randn_like(base)\n",
" timesteps = torch.full((base.shape[0],), t, device=device, dtype=torch.long)\n",
" x_t = # TODO\n",
" scatter_data_and_samples(\n",
" base.cpu(), samples=x_t, title=\"Forward diffusion\", label=f\"t={t}\"\n",
" )\n",
"\n",
"visualize_forward_process(train_dataset, lin, steps=(0, 10, 100, 500, 999))"
]
},
{
"cell_type": "markdown",
"id": "c76629f4",
"metadata": {},
"source": [
"## 3. Positional and time embeddings\n",
"The MLP consumes $(x, y, t)$ but benefits from richer features. Following the standard\n",
"DDPM/Transformer recipe we combine several embeddings. For the sinusoidal one we use:\n",
"$$\n",
"\t\\text{Sinusoidal}(t, k) = \\sin\\left( \n",
"\\frac{t}{10000^{2k / d}} \n",
"\\right), \\quad\n",
"\t\\text{Sinusoidal}(t, k + d/2) = \\cos\\left( \n",
"\\frac{t}{10000^{2k / d}} \n",
"\\right).\n",
"$$\n",
"\n",
"Other embeddings include:\n",
"- **Linear embedding:** scales `t` linearly to `[-1, 1]`.\n",
"- **Learnable embedding:** a learnable `nn.Embedding` layer. \n",
"- **Identity embedding:** simply returns `t` as a float tensor.\n",
"- **Zero embedding:** returns a tensor of zeros.\n",
"\n",
"**TODO:**\n",
"- Ensure the sinusoidal embedding matches the Transformer-style equation provided above.\n",
"- Looking at the sinusoidal embedding code, implement the other positional embeddings (linear, learnable, identity, zero).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "584b0da7",
"metadata": {},
"outputs": [],
"source": [
"class SinusoidalEmbedding(nn.Module):\n",
" def __init__(self, size: int, scale: float = 1.0):\n",
" super().__init__()\n",
" self.size = size\n",
" self.scale = scale\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" x = x * self.scale\n",
" half = self.size // 2\n",
" freq = torch.exp(\n",
" -math.log(10000.0) * torch.arange(half, device=x.device) / (half - 1)\n",
" )\n",
" angles = x.unsqueeze(-1) * freq.unsqueeze(0)\n",
" emb = torch.cat((torch.sin(angles), torch.cos(angles)), dim=-1)\n",
" return emb\n",
"\n",
" def __len__(self):\n",
" return self.size\n",
"\n",
"\n",
"class LinearEmbedding(nn.Module):\n",
" def __init__(self, size: int, scale: float = 1.0):\n",
" super().__init__()\n",
" self.size = size\n",
" self.scale = scale\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" # TODO\n",
"\n",
" def __len__(self):\n",
" return 1\n",
"\n",
"\n",
"class LearnableEmbedding(nn.Module):\n",
" def __init__(self, size: int):\n",
" super().__init__()\n",
" self.size = size\n",
" self.linear = nn.Linear(1, size)\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" # TODO\n",
"\n",
" def __len__(self):\n",
" return self.size\n",
"\n",
"\n",
"class IdentityEmbedding(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.size = 1\n",
"\n",
" def __len__(self):\n",
" return 1\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
"\n",
"\n",
"\n",
"class ZeroEmbedding(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.size = 1\n",
"\n",
" def __len__(self):\n",
" return 1\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" # TODO\n",
"\n",
"\n",
"class PositionalEmbedding(nn.Module):\n",
" def __init__(self, size: int, emb_type: str, **kwargs):\n",
" super().__init__()\n",
" if emb_type == \"sinusoidal\":\n",
" self.layer = SinusoidalEmbedding(size, **kwargs)\n",
" elif emb_type == \"linear\":\n",
" self.layer = LinearEmbedding(size, **kwargs)\n",
" elif emb_type == \"learnable\":\n",
" self.layer = LearnableEmbedding(size)\n",
" elif emb_type == \"identity\":\n",
" self.layer = IdentityEmbedding()\n",
" elif emb_type == \"zero\":\n",
" self.layer = ZeroEmbedding()\n",
" else:\n",
" raise ValueError(f\"Unknown embedding type: {emb_type}\")\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" x = self.layer(x)\n",
" if x.shape[-1] != self.layer.size:\n",
" x = x.unsqueeze(-1)\n",
" return x\n",
"\n",
" def __len__(self):\n",
" return len(self.layer)"
]
},
{
"cell_type": "markdown",
"id": "813e2feb",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "99cc01d3",
"metadata": {},
"source": [
"### 3.1. Visualization utility for embeddings\n",
"\n",
"**TODO:**\n",
"- Plot a handful of embedding dimensions versus normalized time to understand their dynamics before training.\n"
]
},
{
"cell_type": "markdown",
"id": "12c40ff4",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6d802881",
"metadata": {},
"outputs": [],
"source": [
"# TODO: visualize a few embedding channels across time to understand their patterns\n",
"\n",
"\"\"\"Plot several channel of a time embedding over [0, 1].\"\"\"\n",
"for embedding in [\"sinusoidal\", \"linear\", \"learnable\", \"identity\", \"zero\"]:\n",
" emb_layer = PositionalEmbedding(size=16, emb_type=embedding, scale=10.0)\n",
" t = torch.linspace(0, 1, steps=100)\n",
" emb = emb_layer(t).detach().cpu().numpy()\n",
"\n",
" plt.figure(figsize=(6, 4))\n",
" for i in range(min(emb.shape[1], 8)):\n",
" plt.plot(t.numpy(), emb[:, i], label=f\"Channel {i}\")\n",
" plt.title(f\"{embedding.capitalize()} Embedding\")\n",
" plt.xlabel(\"t\")\n",
" plt.ylabel(\"Embedding value\")\n",
" plt.legend()\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "cbe914c5",
"metadata": {},
"source": [
"## 4. Score network (MLP)\n",
"The model predicts the added noise $\\epsilon$ conditioned on $(x, y, t)$ by doing the following:\n",
"- We concatenate embeddings for $x$, $y$, and $t$.\n",
"- Feed the concatenated embeddings through a linear layer and apply GELU activation.\n",
"- Pass the result through several residual `Block`s with skip connections and GELU activations. (`hidden_layers` times)\n",
"- Finally, map to a 2D noise prediction with a linear layer.\n",
"\n",
"**TODO:**\n",
"- Create a residual `Block` with a skip connection plus GELU activation. `y = x + GELU(Linear(x))`\n",
"- Instantiate embeddings for `x`, `y`, and `t`, concatenate them, and map to a 2D noise prediction with an MLP.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1b91b67e",
"metadata": {},
"outputs": [],
"source": [
"# TODO: implement residual blocks and concatenate (x, y, t) embeddings\n",
"class Block(nn.Module):\n",
" def __init__(self, size: int):\n",
" super().__init__()\n",
"\n",
" def forward(self, x):\n",
"\n",
"\n",
"class MLP(nn.Module):\n",
" def __init__(\n",
" self,\n",
" hidden_size=cfg[\"hidden_size\"],\n",
" hidden_layers=cfg[\"hidden_layers\"],\n",
" emb_size=cfg[\"emb_size\"],\n",
" time_emb_type=\"sinusoidal\",\n",
" input_emb_type=\"sinusoidal\",\n",
" ):\n",
" super().__init__()\n",
" self.time_emb = PositionalEmbedding(emb_size, time_emb_type)\n",
" self.input_emb_x = PositionalEmbedding(emb_size, input_emb_type, scale=25.0)\n",
" self.input_emb_y = PositionalEmbedding(emb_size, input_emb_type, scale=25.0)\n",
"\n",
" concat_dim = (\n",
" len(self.time_emb.layer)\n",
" + len(self.input_emb_x.layer)\n",
" + len(self.input_emb_y.layer)\n",
" )\n",
"\n",
" layers = [] # TODO\n",
" \n",
"\n",
" self.net = nn.Sequential(*layers)\n",
"\n",
" def forward(self, x, t):\n",
" # TODO"
]
},
{
"cell_type": "markdown",
"id": "29a7dab0",
"metadata": {},
"source": [
"### 4.1. Sanity check on the MLP\n",
"\n",
"**TODO:**\n",
"- Run a quick forward pass on a mini batch with random timesteps to confirm tensor shapes and device placement.\n"
]
},
{
"cell_type": "markdown",
"id": "07c9eee2",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5c0dab29",
"metadata": {},
"outputs": [],
"source": [
"# TODO: run a quick forward pass to check tensor shapes\n",
"model = MLP().to(device)\n",
"sample_pts = train_dataset.tensors[0][:5].to(device)\n",
"rand_t = torch.randint(0, cfg[\"num_timesteps\"], (5,), device=device).float()\n",
"model(sample_pts, rand_t / cfg[\"num_timesteps\"]).shape"
]
},
{
"cell_type": "markdown",
"id": "76b870c6",
"metadata": {},
"source": [
"## 5. DDPM loss and helpers\n",
"Given random timesteps $t \\sim \\mathcal{U}([0, T))$, noise $\\epsilon \\sim \\mathcal{N}(0, I)$,\n",
"and $x_t = \\sqrt{\\bar{\\alpha}_t} x_0 + \\sqrt{1-\\bar{\\alpha}_t} \\epsilon$, the training objective is:\n",
"$$\n",
"\\mathcal{L} = \\mathbb{E}_{x_0, \\epsilon, t} \\big[ \\lVert \\epsilon - \\epsilon_\\theta(x_t, t) \n",
"\\Vert_2^2 \\big].\n",
"$$\n",
"\n",
"**TODO:**\n",
"- Sample random timesteps per batch element. The timestep `t` should be an integer in $[0, T)$ given by `len(scheduler)`.\n",
"- Don't forget to put the random vector on the right device.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acc13320",
"metadata": {},
"outputs": [],
"source": [
"def sample_timesteps(batch_size, scheduler):\n",
" # TODO\n",
"\n",
"\n",
"timesteps = sample_timesteps(16, lin)\n",
"assert timesteps.shape == (16,), \"Timesteps shape is incorrect\"\n",
"assert timesteps.min() > 0, \"Timesteps minimum value is incorrect\"\n",
"assert timesteps.max() < 1000, \"Timesteps maximum value is incorrect\""
]
},
{
"cell_type": "markdown",
"id": "75017062",
"metadata": {},
"source": [
"Now we can create a function that takes a batch of real samples `x0`, samples random timesteps, adds noise accordingly, and computes the MSE loss between the true noise and the model prediction.\n",
"\n",
"**TODO:**\n",
"- Implement the `forward_diffusion_sample` function to add noise to `x0` at given timesteps.\n",
"- Make sure the function returns both the noisy samples `x_t` and the true noise `ε`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d995af7",
"metadata": {},
"outputs": [],
"source": [
"def forward_diffusion_sample(scheduler, x0, timesteps, noise=None):\n",
" if noise is None:\n",
" noise = torch.randn_like(x0)\n",
" # TODO\n",
" return xt, noise\n",
"\n",
"\n",
"x0 = train_dataset.tensors[0][:16].to(device)\n",
"output = forward_diffusion_sample(lin, x0, timesteps)\n",
"assert len(output) == 2, \"Output should be a tuple of (noisy samples, noise)\"\n",
"assert output[0].shape == x0.shape, \"Noisy samples shape is incorrect\"\n",
"assert output[1].shape == x0.shape, \"Noise shape is incorrect\""
]
},
{
"cell_type": "markdown",
"id": "3308f103",
"metadata": {},
"source": [
"And finally, we can implement the DDPM loss function. Given a batch of real samples `x0`, we:\n",
"1. Sample random timesteps `t` for each batch element.\n",
"2. Generate noisy samples `x_t` and the true noise `ε` using `forward_diffusion_sample`.\n",
"3. Normalize timesteps to `[0, 1]` for model input.\n",
"4. Predict noise `ε_θ(x_t, t)` using the model.\n",
"5. Compute and return the MSE loss between the true noise `ε` and the predicted noise `ε_θ(x_t, t)`.\n",
"\n",
"**TODO:**\n",
"- Implement the `ddpm_loss` function that computes the MSE loss between the true noise and the model's predicted noise."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2dc06c52",
"metadata": {},
"outputs": [],
"source": [
"def ddpm_loss(model, scheduler, batch):\n",
" # TODO"
]
},
{
"cell_type": "markdown",
"id": "813e2feb",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "2e375e97",
"metadata": {},
"source": [
"## 6. Training loop with tqdm\n",
"We iterate over the dataloader, backpropagate the MSE, and every `eval_every`\n",
"epochs we sample a batch to monitor progress.\n",
"\n",
"**TODO:**\n",
"- Build a training loop that iterates with `tqdm`, calls `ddpm_loss`, and steps an Adam optimizer.\n",
"- Keep track of epoch losses and every `eval_every` epochs draw fresh samples for monitoring.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "57ae1a07",
"metadata": {},
"outputs": [],
"source": [
"# TODO: implement training loop with optimizer, tqdm, and periodic sampling\n",
"def train(model, scheduler, dataloader, cfg):\n",
" optimizer = torch.optim.Adam(model.parameters(), lr=cfg[\"lr\"])\n",
" losses = []\n",
" # Use tqdm for progress bar to monitor training progress\n",
" pg = tqdm.trange(cfg[\"num_epochs\"], desc=\"Training\")\n",
" for epoch in pg:\n",
" epoch_losses = []\n",
"\n",
" mean_loss = float(np.mean(epoch_losses))\n",
" losses.append(mean_loss)\n",
" pg.set_postfix({\"loss\": mean_loss})\n",
" return losses"
]
},
{
"cell_type": "markdown",
"id": "1b2d2ddc",
"metadata": {},
"source": [
"### 6.1. First short training run\n",
"(For a quick check you can reduce `num_epochs` or the dataset size.)\n",
"\n",
"**TODO:**\n",
"- Instantiate the scheduler + model with the current hyper-parameters.\n",
"- Optionally launch a reduced training run (fewer epochs) to ensure everything converges before the long job.\n"
]
},
{
"cell_type": "markdown",
"id": "36ec1c71",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "503d612d",
"metadata": {},
"outputs": [],
"source": [
"# TODO: instantiate scheduler + model and launch training when ready\n",
"main_scheduler = NoiseScheduler(\n",
" num_timesteps=cfg[\"num_timesteps\"],\n",
" beta_start=cfg[\"beta_start\"],\n",
" beta_end=cfg[\"beta_end\"],\n",
" schedule_type=cfg[\"beta_schedule\"],\n",
")\n",
"model = MLP().to(device)\n",
"losses = train(model, main_scheduler, train_loader, cfg)"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "3a5f1b13",
"metadata": {},
"source": [
"## 7. Reverse diffusion and sampling\n",
"The reverse update uses:\n",
"$$\n",
"\\hat{x}_0 = \\frac{1}{\\sqrt{\\bar{\\alpha}_t}} \\Big( x_t - \\sqrt{1 - \\bar{\\alpha}_t} \\; \\epsilon_\\theta(x_t, t) \\Big),\n",
"$$\n",
"then\n",
"$$\n",
"q(x_{t-1} \\mid x_t, \\hat{x}_0) = \\mathcal{N}(x_{t-1}; \\mu_t(\\hat{x}_0, x_t), \\sigma_t^2 I).\n",
"$$\n",
"and \n",
"$$\n",
"\\mu_t(\\hat{x}_0, x_t) = \\frac{\\sqrt{\\bar{\\alpha}_{t-1}} \\beta_t}{1 - \\bar{\\alpha}_t} \\hat{x}_0 + \\frac{\\sqrt{\\alpha_t} (1 - \\bar{\\alpha}_{t-1})}{1 - \\bar{\\alpha}_t} x_t,\n",
"$$\n",
"with $\\sigma_t^2 = \\frac{1 - \\bar{\\alpha}_{t-1}}{1 - \\bar{\\alpha}_t} \\beta_t$.\n",
"\n",
"This step has been implemented in the `NoiseScheduler.step` method. We implement `p_sample_step` and `sample_ddpm`, and a helper to plot the final samples.\n",
"\n",
"**TODO:**\n",
"- Implement `p_sample_step` (single reverse update) using the `NoiseScheduler.step` method`.\n",
"- Implement `sample_ddpm` (full reverse chain) by starting from a standard Gaussian and iteratively applying `p_sample_step`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "82f1e9a1",
"metadata": {},
"outputs": [],
"source": [
"# TODO: code the reverse diffusion sampler\n",
"@torch.no_grad()\n",
"def p_sample_step(model, scheduler, x_t, timesteps):\n",
" # TODO\n",
"\n",
"@torch.no_grad()\n",
"def sample_ddpm(model, scheduler, n_samples=1024, device=device, plot=False):\n",
" model.eval()\n",
" x = torch.randn(n_samples, 2, device=device)\n",
" for t in reversed(range(len(scheduler))):\n",
" # TODO\n",
" if t % (len(scheduler) // 10) == 0 and plot:\n",
" scatter_data_and_samples(\n",
" train_dataset,\n",
" samples=x,\n",
" title=\"DDPM Sampling\",\n",
" label=f\"t={t}\",\n",
" )\n",
" return x.cpu()"
]
},
{
"cell_type": "markdown",
"id": "49c45120",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "84cb3d5a",
"metadata": {},
"outputs": [],
"source": [
"samples = sample_ddpm(model, main_scheduler, n_samples=1024, device=device, plot=True)"
]
},
{
"cell_type": "markdown",
"id": "cff3d3bf",
"metadata": {},
"source": [
"To visualize the final generated samples against the dataset, we can use the following helper function.\n",
"**TODO:**\n",
"- Let us visualize what `sample_ddpm` produces after training using the helper function `compare_samples`.\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ff86071b",
"metadata": {},
"outputs": [],
"source": [
"def compare_samples(model, scheduler, n_samples=1024, name=\"\"):\n",
" samples = sample_ddpm(model, scheduler, n_samples=n_samples, device=device)\n",
" scatter_data_and_samples(train_dataset, samples, title=\"DDPM samples\" + \"\")\n",
"\n",
"\n",
"compare_samples(model, main_scheduler, n_samples=cfg[\"sample_points\"])"
]
},
{
"cell_type": "markdown",
"id": "71216efc",
"metadata": {},
"source": [
"As seen in the lectures, the higher the number of denoising steps, the better the quality of the generated samples. You can experiment with different values of `num_sampling_steps` in `sample_ddpm` to observe this effect.\n",
"\n",
"**TODO:** Run next cell to try your code."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "daea339a",
"metadata": {},
"outputs": [],
"source": [
"for num_steps in [10, 25, 50, 100, 250, 500, 1000]:\n",
" print(f\"Sampling with {num_steps} steps\")\n",
" sub_scheduler = NoiseScheduler(\n",
" num_timesteps=num_steps,\n",
" beta_start=cfg[\"beta_start\"],\n",
" beta_end=cfg[\"beta_end\"],\n",
" schedule_type=cfg[\"beta_schedule\"],\n",
" )\n",
" compare_samples(\n",
" model,\n",
" sub_scheduler,\n",
" n_samples=cfg[\"sample_points\"],\n",
" name=f\" ({num_steps} steps)\",\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": true
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"id": "c71d7cb7",
"metadata": {},
"source": [
"## 8. Ablation study scaffolding\n",
"To encourage experimentation we expose a helper that trains short runs with different\n",
"embeddings and beta schedules.\n",
"\n",
"**TODO:**\n",
"- Write `run_ablation` to rerun short trainings with alternative embeddings or schedules and report the losses.\n",
"- Populate an `experiments` list (baseline + >=2 variants) and log/visualize the resulting samples.\n"
]
},
{
"cell_type": "markdown",
"id": "a6445394",
"metadata": {},
"source": [
"**TODO:** Run next cell to try your code.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e496b833",
"metadata": {},
"outputs": [],
"source": [
"# TODO: relaunch short experiments with variant embeddings / schedules\n",
"def run_ablation(\n",
" time_emb=\"sinusoidal\",\n",
" input_emb=\"sinusoidal\",\n",
" schedule=\"linear\",\n",
" epochs=400,\n",
" name=\"\",\n",
"):\n",
" local_cfg = cfg.copy()\n",
" local_cfg[\"num_epochs\"] = epochs\n",
" scheduler = NoiseScheduler(\n",
" num_timesteps=local_cfg[\"num_timesteps\"],\n",
" beta_start=local_cfg[\"beta_start\"],\n",
" beta_end=local_cfg[\"beta_end\"],\n",
" schedule_type=schedule,\n",
" )\n",
" model = MLP(time_emb_type=time_emb, input_emb_type=input_emb).to(device)\n",
" ablation_loader = DataLoader(\n",
" train_dataset,\n",
" batch_size=local_cfg[\"batch_size\"],\n",
" shuffle=True,\n",
" drop_last=True,\n",
" )\n",
" losses = train(model, scheduler, ablation_loader, local_cfg)\n",
" samples = sample_ddpm(\n",
" model, scheduler, n_samples=local_cfg[\"sample_points\"], device=device\n",
" )\n",
" scatter_data_and_samples(\n",
" train_dataset, samples, title=f\"{name} : {time_emb}/{schedule}\"\n",
" )\n",
" return {\"model\": model, \"losses\": losses, \"scheduler\": scheduler}\n",
"\n",
"\n",
"experiments = [\n",
" {\n",
" \"name\": \"baseline\",\n",
" \"time_emb\": \"sinusoidal\",\n",
" \"input_emb\": \"sinusoidal\",\n",
" \"schedule\": \"linear\",\n",
" },\n",
" {\n",
" \"name\": \"learnable-time\",\n",
" \"time_emb\": \"learnable\",\n",
" \"input_emb\": \"sinusoidal\",\n",
" \"schedule\": \"linear\",\n",
" },\n",
" {\n",
" \"name\": \"learnable-input\",\n",
" \"time_emb\": \"sinusoidal\",\n",
" \"input_emb\": \"learnable\",\n",
" \"schedule\": \"linear\",\n",
" },\n",
" {\n",
" \"name\": \"learnable-input-time\",\n",
" \"time_emb\": \"learnable\",\n",
" \"input_emb\": \"learnable\",\n",
" \"schedule\": \"linear\",\n",
" },\n",
"]\n",
"results = {exp[\"name\"]: run_ablation(**exp, epochs=200) for exp in experiments}"
]
},
{
"cell_type": "markdown",
"id": "7a53f9d6",
"metadata": {},
"source": [
"We can compare the final losses across experiments.\n",
"\n",
"**TODO:**\n",
"- Display the final losses for each experiment to compare performance across different configurations."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "39b0e2d9",
"metadata": {},
"outputs": [],
"source": [
"def plot_losses(results):\n",
" plt.figure(figsize=(8, 5))\n",
" for name, output in results.items():\n",
" losses = output[\"losses\"]\n",
" plt.plot(losses, label=name)\n",
" plt.xlabel(\"Epoch\")\n",
" plt.ylabel(\"MSE Loss\")\n",
" plt.title(\"Training Losses for Different Configurations\")\n",
" plt.legend()\n",
" plt.show()\n",
"\n",
"\n",
"plot_losses(results)"
]
},
{
"cell_type": "markdown",
"id": "1f807f4d",
"metadata": {},
"source": [
"To continue the evaluation, we can build a function to assess the model per noise level. This function should sample points from the dataset, add noise at a specific timestem and evaluate the model's denoising performance.\n",
"\n",
"**TODO:**\n",
"- Implement a function to evaluate the model's denoising performance at different noise levels by adding noise to sampled data points and measuring reconstruction quality.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1c572509",
"metadata": {},
"outputs": [],
"source": [
"def evaluate_denoising_performance(\n",
" model, scheduler, dataset, noise_levels, n_samples=256\n",
"):\n",
" model.eval()\n",
" x0 = dataset.tensors[0]\n",
" idx = torch.randperm(len(x0))[:n_samples]\n",
" x0 = x0[idx].to(device)\n",
"\n",
" normalized_timesteps = []\n",
" performance = []\n",
" for noise_level in noise_levels:\n",
" t = int(noise_level * (len(scheduler) - 1))\n",
" # TODO\n",
"\n",
" mse = F.mse_loss(noise_pred, noise).item()\n",
" performance.append(mse)\n",
" normalized_timesteps.append(noise_level)\n",
" return normalized_timesteps, performance"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5628aeca",
"metadata": {},
"outputs": [],
"source": [
"def plot_losses_vs_noise_levels(results):\n",
" noise_levels = np.linspace(0, 1, 100)\n",
" plt.figure(figsize=(8, 5))\n",
" for name, output in results.items():\n",
" model = output[\"model\"]\n",
" scheduler = output[\"scheduler\"]\n",
" timesteps, performance = evaluate_denoising_performance(\n",
" model, scheduler, train_dataset, noise_levels\n",
" )\n",
" plt.plot(timesteps, performance, label=name)\n",
" plt.xlabel(\"Normalized Noise Level\")\n",
" plt.ylabel(\"Denoising MSE Loss\")\n",
" plt.title(\"Denoising Performance vs. Noise Level\")\n",
" plt.legend()\n",
" plt.show()\n",
"\n",
"\n",
"plot_losses_vs_noise_levels(results)"
]
},
{
"cell_type": "markdown",
"id": "5c9abd0f",
"metadata": {},
"source": [
"Now compare the different scheduling stategies with a fixed time/ input embedding. \n",
"\n",
"**TODO:**\n",
"- Run the ablation experiments and plot the losses and denoising performance across noise levels for each configuration."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2245614f",
"metadata": {},
"outputs": [],
"source": [
"experiments = [] # TODO\n",
"\n",
"results = {exp[\"name\"]: run_ablation(**exp, epochs=200) for exp in experiments}\n",
"plot_losses(results)\n",
"plot_losses_vs_noise_levels(results)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2e36bba5",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}