{
"cells": [
{
"cell_type": "markdown",
"id": "59b03327",
"metadata": {},
"source": [
"# Lab — Score-Based Diffusion: SDE, ODE and *Classifier Guidance*\n",
"\n",
"In this lab, we implement and compare **diffusion-based generative models** in 2D.\n",
"The goal is to understand:\n",
"\n",
"1. how a **forward diffusion SDE** progressively destroys information,\n",
"2. how to learn the **score**\n",
" $\n",
" \\nabla_x \\log p_t(x),\n",
" $\n",
"3. how to *invert* the process for sampling via:\n",
" - the **reverse-time SDE** (stochastic sampling),\n",
" - the **probability-flow ODE** (deterministic trajectories),\n",
" - **classifier guidance** (conditional sampling toward a target class).\n",
"\n",
"The dataset is intentionally simple: a mixture of **8 well-separated Gaussians** arranged on a circle.\n",
"This makes trajectories easy to visualize and allows us to quantify both **sample quality** and **mode coverage**.\n",
"\n",
"Prerequisites: basic probability (Gaussians), gradients, and PyTorch (MLPs, optimization).\n",
"\n",
"**TODO:**\n",
"- Run both samplers (reverse SDE and probability-flow ODE) and compare their behavior.\n",
"- Inspect the effect of guidance strength on class fidelity vs diversity.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7daf9d0",
"metadata": {},
"outputs": [],
"source": [
"import math\n",
"import random\n",
"\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import torch\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",
"if torch.cuda.is_available():\n",
" torch.cuda.manual_seed_all(SEED)\n",
" device = torch.device(\"cuda\")\n",
"elif torch.backends.mps.is_available():\n",
" device = torch.device(\"mps\")\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
"\n",
"\n",
"cfg = {\n",
" \"num_per_class\": 5000,\n",
" \"num_classes\": 8,\n",
" \"data_std\": 0.08,\n",
" \"radius\": 2.0,\n",
" \"batch_size\": 256,\n",
" \"lr\": 1e-3,\n",
" \"epochs\": 200,\n",
" \"beta_min\": 0.1,\n",
" \"beta_max\": 20.0,\n",
" \"num_steps\": 500,\n",
" \"eps\": 1e-3,\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "4db94dc0",
"metadata": {},
"source": [
"## 1. Dataset: Mixture of 8 Gaussians\n",
"\n",
"We build a 2D dataset with $K=8$ Gaussian modes.\n",
"Each sample $x\\in\\mathbb{R}^2$ is associated with a discrete label $y\\in\\{0,\\dots,K-1\\}$ indicating the mode index.\n",
"We will reuse these labels later to train a **noisy classifier** and perform **classifier guidance**.\n",
"\n",
"Means placed on a circle:\n",
"$$\n",
"\\mu_k = r\\,\\big(\\cos(2\\pi k/K),\\; \\sin(2\\pi k/K)\\big),\\qquad k=0,\\dots,K-1.\n",
"$$\n",
"Each class generates points around $\\mu_k$ with a small isotropic covariance $\\sigma_0^2 I$.\n",
"Finally, we apply a global normalization so that the data roughly lives in $[-3,3]^2$.\n",
"This stabilizes training by keeping consistent scales for both the neural network and the diffusion dynamics.\n",
"\n",
"**TODO:**\n",
"- Visualize the point cloud and verify that there are 8 distinct modes.\n",
"- Overlay the means $\\mu_k$ on the plot and verify they match the code (radius, orientation).\n",
"- Check the effect of normalization to make sure the data fits in $[-3,3]^2$.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "be13598f",
"metadata": {},
"outputs": [],
"source": [
"def make_gaussian_mixture(\n",
" n_per_class=cfg[\"num_per_class\"],\n",
" num_classes=cfg[\"num_classes\"],\n",
" std=cfg[\"data_std\"],\n",
" radius=cfg[\"radius\"],\n",
" seed=SEED,\n",
"):\n",
" rng = np.random.default_rng(seed)\n",
" angles = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)\n",
" means = np.stack([radius * np.cos(angles), radius * np.sin(angles)], axis=1)\n",
"\n",
" X_list, y_list = [], []\n",
" for k, m in enumerate(means):\n",
" pts = rng.normal(loc=m, scale=std, size=(n_per_class, 2))\n",
" X_list.append(pts)\n",
" y_list.append(np.full(n_per_class, k))\n",
"\n",
" X = np.concatenate(X_list, axis=0).astype(np.float32)\n",
" y = np.concatenate(y_list, axis=0).astype(np.int64)\n",
"\n",
" # TODO\n",
" scale = 1\n",
" X = X\n",
" # TODO\n",
"\n",
" return (\n",
" TensorDataset(\n",
" torch.from_numpy(X),\n",
" torch.from_numpy(y),\n",
" ),\n",
" means.astype(np.float32),\n",
" std / scale,\n",
" )\n",
"\n",
"\n",
"train_dataset, means, data_std = make_gaussian_mixture()\n",
"train_loader = DataLoader(train_dataset, batch_size=cfg[\"batch_size\"], shuffle=True)\n",
"\n",
"COLORS = [\n",
" \"#1f77b4\",\n",
" \"#ff7f0e\",\n",
" \"#2ca02c\",\n",
" \"#d62728\",\n",
" \"#9467bd\",\n",
" \"#8c564b\",\n",
" \"#e377c2\",\n",
" \"#7f7f7f\",\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "679eb118",
"metadata": {},
"outputs": [],
"source": [
"def plot_mixture(dataset, means=None, title=\"Mixture of 8 Gaussians\", figsize=(5, 5)):\n",
" if isinstance(dataset, TensorDataset):\n",
" X, y = dataset.tensors\n",
" X = X.numpy()\n",
" y = y.numpy()\n",
" else:\n",
" X, y = dataset\n",
" plt.figure(figsize=figsize)\n",
" for k in range(cfg[\"num_classes\"]):\n",
" pts = X[y == k]\n",
" plt.scatter(pts[:, 0], pts[:, 1], s=6, alpha=0.6, color=COLORS[k])\n",
" if means is not None:\n",
" plt.scatter(means[:, 0], means[:, 1], s=80, c=\"black\", marker=\"x\")\n",
" plt.title(title)\n",
" plt.axis(\"equal\")\n",
" plt.show()\n",
"\n",
"\n",
"def plot_gaussian(figsize=(5, 5)):\n",
" pts = np.random.normal(loc=0.0, scale=1, size=(1000, 2))\n",
" plt.figure(figsize=figsize)\n",
" plt.scatter(pts[:, 0], pts[:, 1], s=6, alpha=0.6, color=COLORS[0])\n",
" if means is not None:\n",
" plt.scatter(means[:, 0], means[:, 1], s=80, c=\"black\", marker=\"x\")\n",
" plt.title(\"gaussian distribution\")\n",
" plt.axis(\"equal\")\n",
" plt.show()\n",
"\n",
"\n",
"plot_mixture(train_dataset, means=means)\n",
"plot_gaussian()"
]
},
{
"cell_type": "markdown",
"id": "d5235f5c",
"metadata": {},
"source": [
"## 2. Variance Preserving SDE (VP-SDE)\n",
"\n",
"We define a **forward diffusion process** that gradually turns data $x_0\\sim p_0$ into Gaussian noise.\n",
"The continuous-time VP-SDE is:\n",
"$$\n",
"\\mathrm d x_t = -\\tfrac{1}{2}\\,\\beta(t)\\,x_t\\,\\mathrm dt + \\sqrt{\\beta(t)}\\,\\mathrm dW_t,\n",
"$$\n",
"where $W_t$ is a standard Brownian motion and $\\beta(t)>0$ controls the noise intensity.\n",
"\n",
"In this lab we use a linear schedule on $t\\in[0,1]$:\n",
"$$\n",
"\\beta(t) = \\beta_{\\min} + (\\beta_{\\max}-\\beta_{\\min})\\,t.\n",
"$$\n",
"\n",
"A key fact: the conditional marginal $p(x_t\\mid x_0)$ is Gaussian with a closed form:\n",
"$$\n",
"x_t = \\sqrt{\\bar\\alpha(t)}\\,x_0 + \\sqrt{1-\\bar\\alpha(t)}\\,\\varepsilon,\n",
"\\qquad \\varepsilon\\sim\\mathcal N(0,I),\n",
"$$\n",
"where\n",
"$$\n",
"\\bar\\alpha(t)=\\exp\\Big(-\\int_0^t \\beta(s)\\,\\mathrm ds\\Big).\n",
"$$\n",
"For the linear schedule,\n",
"$$\n",
"\\int_0^t\\beta(s)\\,\\mathrm ds\n",
"= \\beta_{\\min}t + \\tfrac12(\\beta_{\\max}-\\beta_{\\min})t^2,\n",
"$$\n",
"hence\n",
"$$\n",
"\\bar\\alpha(t)=\\exp\\Big(-\\beta_{\\min}t - \\tfrac12(\\beta_{\\max}-\\beta_{\\min})t^2\\Big).\n",
"$$\n",
"\n",
"Interpretation:\n",
"- At $t=0$, $\\bar\\alpha(0)=1$ and $1-\\bar\\alpha(0)=0$: we recover $x_0$.\n",
"- At $t=1$, $\\bar\\alpha(1)$ is small and $1-\\bar\\alpha(1)\\approx 1$: we are close to standard noise.\n",
"\n",
"**TODO:**\n",
"- Derive $\\int_0^t\\beta(s)\\,\\mathrm ds$ for the linear schedule and confirm the expression of $\\bar\\alpha(t)$ and complete the function in the code `beta_t`, `alpha_bar_t`.\n",
"- Complete the function `marginal_prob` that returns the mean $\\sqrt{\\bar\\alpha(t)}\\,x_0$ and standard deviation $\\sqrt{1-\\bar\\alpha(t)}$ of $p(x_t\\mid x_0)$.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8de66c96",
"metadata": {},
"outputs": [],
"source": [
"def beta_t(t):\n",
" # TODO\n",
" raise NotImplementedError\n",
"\n",
"\n",
"def alpha_bar(t):\n",
" # TODO\n",
"\n",
" raise NotImplementedError\n",
"\n",
"\n",
"def marginal_prob(x0, t):\n",
" # TODO\n",
" # Return mean and std of x_t given x0.\n",
" while t.dim() < x0.dim():\n",
" t = t.unsqueeze(-1)\n",
" mean = \n",
" std = \n",
" return mean, std"
]
},
{
"cell_type": "markdown",
"id": "3d9c9e85",
"metadata": {},
"source": [
"## 3. Score network (MLP + time embedding)\n",
"\n",
"The core object in score-based methods is the **score**\n",
"$$\n",
"s_\\theta(x,t) \\approx \\nabla_x \\log p_t(x),\n",
"$$\n",
"i.e., the gradient of the log-density at time $t$.\n",
"Intuitively, this vector field tells us **which direction to move** a point $x$ to increase its probability under $p_t$.\n",
"\n",
"Why does the network need the time $t$?\n",
"- For small $t$, $p_t$ looks like the data distribution (highly multi-modal).\n",
"- For large $t$, $p_t$ is close to a standard Gaussian.\n",
"Therefore the score changes dramatically with $t$, and the model must be time-conditioned.\n",
"\n",
"We use an **MLP** that takes $(x_t,t)$ as input and outputs a 2D vector.\n",
"Time is embedded with a sinusoidal positional encoding (similar to Transformers):\n",
"$$\n",
"\\gamma(t) = (\\sin(\\omega_1 t),\\cos(\\omega_1 t),\\dots,\\sin(\\omega_m t),\\cos(\\omega_m t)).\n",
"$$\n",
"This provides the MLP with a rich basis to represent non-linear dependencies in $t$.\n",
"\n",
"**TODO:**\n",
"- Complete the forward pass of the `MLP` class to concatenate the time and input embeddings and feed them to the joint MLP.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "39798b59",
"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_size = self.size // 2\n",
" emb = torch.log(torch.tensor([10000.0], device=x.device)) / (half_size - 1)\n",
" emb = torch.exp(-emb * torch.arange(half_size, device=x.device))\n",
" emb = x.unsqueeze(-1) * emb.unsqueeze(0)\n",
" emb = torch.cat((torch.sin(emb), torch.cos(emb)), dim=-1)\n",
" return emb\n",
"\n",
" def __len__(self):\n",
" return self.size\n",
"\n",
"\n",
"class Block(nn.Module):\n",
" def __init__(self, size: int):\n",
" super().__init__()\n",
" self.ff = nn.Linear(size, size)\n",
" self.act = nn.GELU()\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" return x + self.act(self.ff(x))\n",
"\n",
"\n",
"class MLP(nn.Module):\n",
" def __init__(\n",
" self,\n",
" hidden_size: int = 128,\n",
" hidden_layers: int = 3,\n",
" emb_size: int = 128,\n",
" **kwargs,\n",
" ):\n",
" super().__init__()\n",
" self.time_mlp = SinusoidalEmbedding(emb_size)\n",
" self.input_mlp1 = SinusoidalEmbedding(emb_size, scale=25.0)\n",
" self.input_mlp2 = SinusoidalEmbedding(emb_size, scale=25.0)\n",
"\n",
" concat_size = len(self.time_mlp) + len(self.input_mlp1) + len(self.input_mlp2)\n",
" layers = [nn.Linear(concat_size, hidden_size), nn.GELU()]\n",
" for _ in range(hidden_layers):\n",
" layers.append(Block(hidden_size))\n",
" layers.append(nn.Linear(hidden_size, 2))\n",
" self.joint_mlp = nn.Sequential(*layers)\n",
"\n",
" def forward(self, x, t):\n",
" # TODO\n",
" return x"
]
},
{
"cell_type": "markdown",
"id": "76121122",
"metadata": {},
"source": [
"## 4. Training the score network\n",
"\n",
"We train the network with **Denoising Score Matching** (DSM).\n",
"At each iteration we sample:\n",
"- a clean point $x_0\\sim p_0$,\n",
"- a time $t\\sim \\mathcal U[0,1]$,\n",
"- Gaussian noise $\\varepsilon\\sim\\mathcal N(0,I)$,\n",
"and we construct\n",
"$$\n",
"x_t = \\sqrt{\\bar\\alpha(t)}\\,x_0 + \\sqrt{1-\\bar\\alpha(t)}\\,\\varepsilon.\n",
"$$\n",
"\n",
"For Gaussian corruption, the conditional score is available in closed form:\n",
"$$\n",
"\\nabla_{x_t}\\log p(x_t\\mid x_0)\n",
"= -\\frac{x_t - \\sqrt{\\bar\\alpha(t)}\\,x_0}{1-\\bar\\alpha(t)}.\n",
"$$\n",
"DSM trains $s_\\theta(x_t,t)$ to match this target in expectation:\n",
"$$\n",
"\\mathcal L(\\theta)\n",
"= \\mathbb E_{t,x_0,\\varepsilon}\\Big[\n",
"\\lambda(t)\\,\\big\\| s_\\theta(x_t,t) - \\nabla_{x_t}\\log p(x_t\\mid x_0)\\big\\|_2^2\n",
"\\Big],\n",
"$$\n",
"where $\\lambda(t)$ is a weighting function. Here, we use \n",
"$$\n",
"\\lambda(t) = \\beta(t) \\quad \\text{or} \\quad \\lambda(t) = \\frac{1}{1-\\bar\\alpha(t)}.\n",
"$$\n",
"\n",
"**TODO:**\n",
"- Implement the DSM loss and verify that it decreases during training.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "413b1592",
"metadata": {},
"outputs": [],
"source": [
"score_model = MLP(hidden_size=128, hidden_layers=3, emb_size=128, dim_out=2).to(device)\n",
"optimizer = torch.optim.Adam(score_model.parameters(), lr=cfg[\"lr\"])\n",
"\n",
"\n",
"def score_matching_loss(model, x0, t=None, reduction=\"mean\"):\n",
" if t is None:\n",
" t = torch.rand(x0.shape[0], device=x0.device) * (1.0 - cfg[\"eps\"]) + cfg[\"eps\"]\n",
" z = torch.randn_like(x0)\n",
" # TODO\n",
"\n",
" # TODO\n",
" if reduction == \"mean\":\n",
" return per_sample.mean()\n",
" return per_sample\n",
"\n",
"\n",
"def train_score_model(model, loader, epochs=cfg[\"epochs\"]):\n",
" model.train()\n",
" loss_history = []\n",
" pbar = tqdm.tqdm(range(epochs), leave=True)\n",
" for epoch in pbar:\n",
" average_loss = 0.0\n",
" for batch in loader:\n",
" x0 = batch[0].to(device)\n",
" optimizer.zero_grad()\n",
" loss = score_matching_loss(model, x0)\n",
" average_loss += loss.item() * x0.size(0)\n",
" loss.backward()\n",
" optimizer.step()\n",
" average_loss /= len(loader.dataset)\n",
" loss_history.append(average_loss)\n",
" pbar.set_description(f\"epoch {epoch + 1}/{epochs} | loss={average_loss:4.4f}\")\n",
" return loss_history\n",
"\n",
"\n",
"loss_history = train_score_model(score_model, train_loader)"
]
},
{
"cell_type": "markdown",
"id": "954aed83",
"metadata": {},
"source": [
"### Diagnostics: loss curves\n",
"\n",
"We track two simple diagnostics:\n",
"1. **Training loss vs epoch**: it should decrease and then plateau.\n",
"2. **Loss as a function of time $t$**: some noise levels are harder (often intermediate times).\n",
"\n",
"\n",
"**TODO:**\n",
"- Run the next cell to make sure that the model has reached convergence.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9267b4a2",
"metadata": {},
"outputs": [],
"source": [
"def plot_loss_history(loss_history):\n",
" plt.figure(figsize=(5, 3))\n",
" plt.plot(loss_history)\n",
" plt.xlabel(\"epoch\")\n",
" plt.ylabel(\"loss\")\n",
" plt.title(\"Training loss\")\n",
" plt.grid(True, alpha=0.3)\n",
" plt.show()\n",
"\n",
"\n",
"plot_loss_history(loss_history)"
]
},
{
"cell_type": "markdown",
"id": "dcc60714",
"metadata": {},
"source": [
"Why look at loss-by-$t$?\n",
"- For large $t$, the distribution is close to $\\mathcal N(0,I)$ and the score is comparatively simple.\n",
"- For small $t$, the distribution is highly multi-modal and the score is more complex (but scaling issues via $1-\\bar\\alpha(t)$ also matter).\n",
"\n",
"**TODO:**\n",
"- Complete the `loss_by_t` function.\n",
"- Relate this diagnostic to the choice of weighting $\\lambda(t)$ and to the noise schedule.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "49bbcf38",
"metadata": {},
"outputs": [],
"source": [
"def loss_by_t(model, dataset, t_values, batch_size=1024):\n",
" model.eval()\n",
" if isinstance(dataset, TensorDataset):\n",
" X = dataset.tensors[0]\n",
" else:\n",
" X = dataset[0]\n",
" loader = DataLoader(TensorDataset(X), batch_size=batch_size, shuffle=False)\n",
" losses = []\n",
" with torch.no_grad():\n",
" for t in t_values:\n",
" total = 0.0\n",
" count = 0\n",
" t_batch = torch.full((batch_size,), t, device=device)\n",
" for batch in loader:\n",
" # TODO\n",
" \n",
" # TODO\n",
" losses.append(total / max(count, 1))\n",
" return losses\n",
"\n",
"\n",
"t_grid = torch.linspace(cfg[\"eps\"], 1.0, 21, device=device)\n",
"t_losses = loss_by_t(score_model, train_dataset, t_grid)\n",
"plt.figure(figsize=(5, 3))\n",
"plt.plot(t_grid.cpu().numpy(), t_losses, marker=\"o\")\n",
"plt.xlabel(\"t (noise level)\")\n",
"plt.ylabel(\"loss\")\n",
"plt.title(\"Loss vs noise level\")\n",
"plt.grid(True, alpha=0.3)\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "d8b74263",
"metadata": {},
"source": [
"## 5. Reverse sampling: SDE vs ODE\n",
"\n",
"Once the score is learned, we can **generate samples** by reversing the diffusion.\n",
"We start from $x_1\\sim\\mathcal N(0,I)$ and integrate from $t=1$ down to $t=0$.\n",
"\n",
"1) **Reverse-time SDE** (stochastic, Euler–Maruyama):\n",
"$$\n",
"\\mathrm d x_t\n",
"= \\Big[-\\tfrac12\\beta(t)x_t - \\beta(t)s_\\theta(x_t,t)\\Big]\\,\\mathrm dt\n",
"+ \\sqrt{\\beta(t)}\\,\\mathrm d\\bar W_t,\n",
"$$\n",
"where $\\bar W_t$ is a Brownian motion in reverse time.\n",
"The noise term injects randomness and often improves **diversity**. This SDE discretization is equivalent to samples a noise $\\varepsilon\\sim\\mathcal N(0,I)$ at each step and performing:\n",
"$$\n",
"x_{t-\\Delta t}\n",
"= x_t + \\Big[-\\tfrac12\\beta(t)x_t - \\beta(t)s_\\theta(x_t,t)\\Big]\\,\\Delta t\n",
"+ \\sqrt{\\beta(t)\\,\\Delta t}\\,\\varepsilon.\n",
"$$\n",
"\n",
"2) **Probability-flow ODE** (deterministic):\n",
"$$\n",
"\\frac{\\mathrm d x_t}{\\mathrm dt}\n",
"= -\\tfrac12\\beta(t)x_t - \\tfrac12\\beta(t)s_\\theta(x_t,t).\n",
"$$\n",
"This ODE has the same marginals $p_t$ as the SDE, but without injected noise.\n",
"With a sufficiently accurate numerical solver, it yields smoother trajectories and often better **quality** for a fixed number of steps. We can discretize it with Euler or higher-order methods; here we use Euler:\n",
"$$\n",
"x_{t-\\Delta t}\n",
"= x_t + \\Big[-\\tfrac12\\beta(t)x_t - \\tfrac12\\beta(t)s_\\theta(x_t,t)\\Big]\\,\\Delta t.\n",
"$$\n",
"and the second order Heun's method:\n",
"1. Compute an Euler step to get a provisional value $x^*$:\n",
" $$\n",
" x^* = x_t + \\Big[-\\tfrac12\\beta(t)x_t - \\tfrac12\\beta(t)s_\\theta(x_t,t)\\Big]\\,\\Delta t.\n",
" $$\n",
"2. Use $x^*$ to compute the slope at $t-\\Delta t$ and average it with the slope at $t$:\n",
" $$\n",
" x_{t-\\Delta t}\n",
" = x_t + \\tfrac12\\Big(-\\tfrac12\\beta(t)x_t - \\tfrac12\\beta(t)s_\\theta(x_t,t)\n",
" -\\tfrac12\\beta(t-\\Delta t)x^* - \\tfrac12\\beta(t-\\Delta t)s_\\theta(x^*,t-\\Delta t)\n",
" \\Big)\\,\\Delta t.\n",
" $$\n",
"\n",
"Discretization:\n",
"we use $N$ steps with $\\Delta t = 1/N$, and integrate in practice with decreasing $t$.\n",
"The number of steps matters: too few steps $\\Rightarrow$ visible discretization error (trajectories missing modes).\n",
"\n",
"**TODO:**\n",
"- Implement the three samplers (reverse SDE, ODE Euler, ODE Heun).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8ad6d98e",
"metadata": {},
"outputs": [],
"source": [
"@torch.no_grad()\n",
"def reverse_sde_step(x, t, dt, model):\n",
" batch_t = torch.full((x.shape[0],), t, device=x.device)\n",
"\n",
" return x\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def sample_reverse_sde(\n",
" model, num_samples=1000, num_steps=cfg[\"num_steps\"], return_traj=False\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" x = reverse_sde_step(x, t, dt, model)\n",
" if return_traj:\n",
" traj.append(x.cpu().numpy())\n",
" if return_traj:\n",
" return traj\n",
" return x\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def reverse_ode_step(x, t, dt, model):\n",
" batch_t = torch.full((x.shape[0],), t, device=x.device)\n",
"\n",
" return x\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def sample_reverse_ode(\n",
" model, num_samples=1000, num_steps=cfg[\"num_steps\"], return_traj=False\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" x = reverse_ode_step(x, t, dt, model)\n",
" if return_traj:\n",
" traj.append(x.cpu().numpy())\n",
" if return_traj:\n",
" return traj\n",
" return x\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def reverse_ode_step_heun(x, t, dt, model):\n",
" batch_t = torch.full((x.shape[0],), t, device=x.device)\n",
"\n",
" return x\n",
"\n",
"\n",
"@torch.no_grad()\n",
"def sample_reverse_ode_heun(\n",
" model, num_samples=1000, num_steps=cfg[\"num_steps\"], return_traj=False\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" x = reverse_ode_step_heun(x, t, dt, model)\n",
" if return_traj:\n",
" traj.append(x.cpu().numpy())\n",
" if return_traj:\n",
" return traj\n",
" return x"
]
},
{
"cell_type": "markdown",
"id": "ea10341b",
"metadata": {},
"source": [
"### Visualization: samples and trajectories\n",
"\n",
"We visualize two complementary objects:\n",
"\n",
"- **Final sample clouds** $x_0$: shows whether the modes are correctly covered.\n",
"- **Time trajectories** $x_t$: helps understand *how* the algorithm moves through space.\n",
"\n",
"How to read trajectories:\n",
"- At the beginning (near $t=1$), points are essentially noise around the origin.\n",
"- As $t\\to 0$, they “commit” to a mode and converge toward one of the Gaussian clusters.\n",
"- The SDE introduces additional dispersion; the ODE produces deterministic paths (given the initial seed).\n",
"\n",
"**TODO:**\n",
"- Check whether trajectories concentrate on the 8 modes as expected.\n",
"- Identify typical failures: missed modes, oscillatory paths, collapse onto a subset of modes.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fb0fa1c1",
"metadata": {},
"outputs": [],
"source": [
"def plot_generated_samples(samples, title, data=None, means=None, figsize=(5, 5)):\n",
" if isinstance(samples, torch.Tensor):\n",
" x = samples.detach().cpu().numpy()\n",
" else:\n",
" x = samples\n",
" plt.figure(figsize=figsize)\n",
" if data is not None:\n",
" if isinstance(data, TensorDataset):\n",
" X = data.tensors[0].numpy()\n",
" else:\n",
" X = data\n",
" plt.scatter(X[:, 0], X[:, 1], s=4, alpha=0.25, color=\"#9e9e9e\")\n",
" plt.scatter(x[:, 0], x[:, 1], s=6, alpha=0.6, color=\"#1f77b4\")\n",
" if means is not None:\n",
" plt.scatter(means[:, 0], means[:, 1], s=80, c=\"black\", marker=\"x\")\n",
" plt.title(title)\n",
" plt.axis(\"equal\")\n",
" plt.xlim(-3.5, 3.5)\n",
" plt.ylim(-3.5, 3.5)\n",
" plt.show()\n",
"\n",
"\n",
"def plot_trajectories(trajs, ax, title):\n",
" for mean_i in range(cfg[\"num_classes\"]):\n",
" for r_i in range(4):\n",
" ax.add_patch(\n",
" plt.Circle(\n",
" (\n",
" means[mean_i, 0],\n",
" means[mean_i, 1],\n",
" ),\n",
" data_std * (r_i + 1),\n",
" color=\"gray\",\n",
" fill=False,\n",
" linestyle=\"--\",\n",
" alpha=0.3,\n",
" )\n",
" )\n",
" trajs = np.array(trajs)\n",
" for traj_i in range(trajs.shape[1]):\n",
" traj = trajs[:, traj_i, :]\n",
" ax.scatter(traj[0, 0], traj[0, 1], s=20, edgecolor=\"black\")\n",
" ax.plot(traj[:, 0], traj[:, 1], alpha=0.4)\n",
" ax.set_title(title)\n",
" ax.set_aspect(\"equal\")\n",
" ax.set_xlim(-3.5, 3.5)\n",
" ax.set_ylim(-3.5, 3.5)\n",
"\n",
"\n",
"def compare_trajectories(model, num_traj=10, num_steps=200):\n",
" sde_traj = sample_reverse_sde(\n",
" model, return_traj=True, num_samples=num_traj, num_steps=num_steps\n",
" )\n",
" ode_traj = sample_reverse_ode(\n",
" model, return_traj=True, num_samples=num_traj, num_steps=num_steps\n",
" )\n",
" ode2_traj = sample_reverse_ode_heun(\n",
" model, return_traj=True, num_samples=num_traj, num_steps=num_steps\n",
" )\n",
" fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
" plot_trajectories(sde_traj, axes[0], \"SDE trajectories\")\n",
" plot_trajectories(ode_traj, axes[1], \"ODE trajectories\")\n",
" plot_trajectories(ode2_traj, axes[2], \"ODE (2nd order) trajectories\")\n",
" plt.show()\n",
"\n",
"\n",
"sde_samples = sample_reverse_sde(score_model, num_samples=2000)\n",
"ode_samples = sample_reverse_ode(score_model, num_samples=2000)\n",
"ode2_samples = sample_reverse_ode_heun(score_model, num_samples=2000)\n",
"\n",
"plot_generated_samples(sde_samples, \"Samples - SDE\", data=train_dataset, means=means)\n",
"plot_generated_samples(ode_samples, \"Samples - ODE\", data=train_dataset, means=means)\n",
"plot_generated_samples(\n",
" ode2_samples, \"Samples - ODE (2nd order)\", data=train_dataset, means=means\n",
")\n",
"\n",
"compare_trajectories(score_model, num_traj=20, num_steps=200)"
]
},
{
"cell_type": "markdown",
"id": "0f99b32d",
"metadata": {},
"source": [
"## 6. Metrics: quality \n",
"\n",
"Since we are in 2D with known modes, we can define simple evaluation metrics.\n",
"\n",
"**Quality**:\n",
"fraction of samples that are “close” to at least one mode.\n",
"For instance, we can call a point correct if it lies within distance $\\le 3\\sigma_0$ of some mean $\\mu_k$.\n",
"\n",
"This metric is not universal, but it are very useful here to compare samplers:\n",
"\n",
"- the ODE tends to increase quality (stable trajectories),\n",
"- the SDE tends to increase diversity (randomness during sampling).\n",
"\n",
"**TODO:**\n",
"\n",
"- Compare quality between SDE and ODE for several step counts.\n",
"- Identify the trade-off: does higher quality ?\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25fb6408",
"metadata": {},
"outputs": [],
"source": [
"def quality_metrics(samples, means, std):\n",
" if isinstance(samples, torch.Tensor):\n",
" x = samples.detach().cpu()\n",
" else:\n",
" x = torch.tensor(samples)\n",
" means_t = torch.tensor(means)\n",
" dists = torch.cdist(x, means_t)\n",
" # TODO\n",
" return quality\n",
"\n",
"\n",
"def evaluate_steps(model, steps_list):\n",
" rows = []\n",
" for n in steps_list:\n",
" sde_samples = sample_reverse_sde(model, num_steps=n)\n",
" ode_samples = sample_reverse_ode(model, num_steps=n)\n",
" ode2_samples = sample_reverse_ode_heun(model, num_steps=n)\n",
" q_sde = quality_metrics(sde_samples, means, data_std)\n",
" q_ode = quality_metrics(ode_samples, means, data_std)\n",
" q_ode2 = quality_metrics(ode2_samples, means, data_std)\n",
" rows.append((n, q_sde, q_ode, q_ode2))\n",
"\n",
" quality_metric = np.array(rows)\n",
" plt.figure(figsize=(8, 4))\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 1],\n",
" marker=\"o\",\n",
" label=\"SDE\",\n",
" )\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 2],\n",
" marker=\"o\",\n",
" label=\"ODE\",\n",
" )\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 3],\n",
" marker=\"o\",\n",
" label=\"ODE (2nd order)\",\n",
" )\n",
" plt.xlabel(\"Number of steps\")\n",
" plt.ylabel(\"Quality\")\n",
" plt.xscale(\"log\")\n",
" plt.title(\"Quality vs number of steps\")\n",
" plt.grid(True, alpha=0.3)\n",
" plt.legend()\n",
" plt.show()\n",
"\n",
"\n",
"evaluate_steps(score_model, steps_list=[1, 2, 5, 10, 20, 50, 100, 200, 500])"
]
},
{
"cell_type": "markdown",
"id": "e275da9f",
"metadata": {},
"source": [
"**TODO:**\n",
"\n",
"- Adapt this function to compute quality metrics as a function of the NFE, ie number of function evaluations, the number of times the score network is called during sampling."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ff2b43a8",
"metadata": {},
"outputs": [],
"source": [
"def evaluate_NFE(model, steps_list):\n",
" rows = []\n",
" for n in steps_list:\n",
" sde_samples = sample_reverse_sde(model, num_steps=n)\n",
" ode_samples = sample_reverse_ode(model, num_steps=n)\n",
" ode2_samples = sample_reverse_ode_heun(model, num_steps=n)\n",
" q_sde = quality_metrics(sde_samples, means, data_std)\n",
" q_ode = quality_metrics(ode_samples, means, data_std)\n",
" q_ode2 = quality_metrics(ode2_samples, means, data_std)\n",
" rows.append((n, q_sde, q_ode, q_ode2))\n",
"\n",
" quality_metric = np.array(rows)\n",
" plt.figure(figsize=(8, 4))\n",
" # TODO\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 1],\n",
" marker=\"o\",\n",
" label=\"SDE\",\n",
" )\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 2],\n",
" marker=\"o\",\n",
" label=\"ODE\",\n",
" )\n",
" plt.plot(\n",
" (quality_metric[:, 0] * 2) - 1,\n",
" quality_metric[:, 3],\n",
" marker=\"o\",\n",
" label=\"ODE (2nd order)\",\n",
" )\n",
" # TODO\n",
" plt.xlabel(\"Number of steps\")\n",
" plt.ylabel(\"Quality\")\n",
" plt.xscale(\"log\")\n",
" plt.title(\"Quality vs number of steps\")\n",
" plt.grid(True, alpha=0.3)\n",
" plt.legend()\n",
" plt.show()\n",
"\n",
"\n",
"evaluate_NFE(score_model, steps_list=[1, 2, 5, 10, 20, 50, 100, 200, 500])"
]
},
{
"cell_type": "markdown",
"id": "ee172301",
"metadata": {},
"source": [
"## 7. A classifier on noisy points\n",
"\n",
"We train a classifier $c_\\phi(x_t,t)$ that predicts the class $y$ of a noisy point $x_t$.\n",
"Why include $t$?\n",
"- For small $t$, classification is easy (modes are well separated).\n",
"- For large $t$, everything looks like noise: class information disappears.\n",
"\n",
"Training pipeline:\n",
"1. Sample $(x_0,y)$ from the dataset.\n",
"2. Sample $t\\sim\\mathcal U[0,1]$ and corrupt $x_0\\to x_t$.\n",
"3. Predict logits $\\ell = c_\\phi(x_t,t)$ and minimize cross-entropy $\\mathrm{CE}(\\ell,y)$.\n",
"\n",
"We will then analyze accuracy as a function of $t$:\n",
"it should drop as the noise level increases.\n",
"\n",
"**TODO:**\n",
"- Adapte the model given in the class `MLP` to build the classifier `ClassifierMLP`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "56c5ed04",
"metadata": {},
"outputs": [],
"source": [
"class ClassifierMLP(nn.Module):\n",
" raise NotImplementedError"
]
},
{
"cell_type": "markdown",
"id": "003ec979",
"metadata": {},
"source": [
"The training loop is similar to the score network. At each iteration we sample a clean point, a time, corrupt the point, and compute the cross-entropy loss.\n",
"\n",
"**TODO:**\n",
"- Complete the training loop for the noisy classifier.\n",
"- Run the training function. "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "44d1392e",
"metadata": {},
"outputs": [],
"source": [
"classifier = ClassifierMLP().to(device)\n",
"clf_opt = torch.optim.Adam(classifier.parameters(), lr=1e-3)\n",
"\n",
"\n",
"def classifier_loss(model, x0, y):\n",
" t = torch.rand(x0.shape[0], device=x0.device) * (1.0 - cfg[\"eps\"]) + cfg[\"eps\"]\n",
" z = torch.randn_like(x0)\n",
" # TODO\n",
" raise NotImplementedError\n",
"\n",
"\n",
"def train_classifier(model, loader, epochs=100):\n",
" model.train()\n",
" pbar = tqdm.tqdm(range(epochs), leave=True)\n",
" for epoch in pbar:\n",
" average_loss = 0.0\n",
" for batch in loader:\n",
" x0, y = batch[0].to(device), batch[1].to(device)\n",
" clf_opt.zero_grad()\n",
" loss = classifier_loss(model, x0, y)\n",
" average_loss += loss.item() * x0.size(0)\n",
" loss.backward()\n",
" clf_opt.step()\n",
" average_loss /= len(loader.dataset)\n",
" pbar.set_description(\n",
" f\"clf epoch {epoch + 1}/{epochs} | loss={average_loss:.4f}\"\n",
" )\n",
"\n",
"\n",
"train_classifier(classifier, train_loader)"
]
},
{
"cell_type": "markdown",
"id": "76aae74b",
"metadata": {},
"source": [
"Similarly to the score network, we can track accuracy as a function of time $t$.\n",
"\n",
"**TODO:**\n",
"- Complete the `accuracy_vs_noise` function.\n",
"- Observe how the accuracy behaves as noise increases."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6d5a7de1",
"metadata": {},
"outputs": [],
"source": [
"def accuracy_vs_noise(model, dataset, t_values):\n",
" model.eval()\n",
" if isinstance(dataset, TensorDataset):\n",
" X, y = dataset.tensors\n",
" else:\n",
" X, y = dataset\n",
" X = X.to(device)\n",
" y = y.to(device)\n",
" acc_list = []\n",
" with torch.no_grad():\n",
" for t in t_values:\n",
" t_batch = torch.full((X.shape[0],), t, device=device)\n",
"\n",
" acc_list.append(acc)\n",
" return acc_list\n",
"\n",
"\n",
"t_grid = torch.linspace(0.0, 1.0, 11)\n",
"acc = accuracy_vs_noise(classifier, train_dataset, t_grid)\n",
"plt.plot(t_grid.cpu().numpy(), acc, marker=\"o\")\n",
"plt.plot([min(t_grid), max(t_grid)], [1.0 / cfg[\"num_classes\"]] * 2, \"k--\")\n",
"plt.xlabel(\"t (noise level)\")\n",
"plt.ylabel(\"accuracy\")\n",
"plt.title(\"Classifier accuracy vs noise\")\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "0e2edfc1",
"metadata": {},
"source": [
"## 8. *Classifier guidance*\n",
"\n",
"Guidance **biases** the reverse dynamics toward a target class $y$.\n",
"We use the identity:\n",
"$$\n",
"\\nabla_x \\log p_t(x\\mid y)\n",
"= \\nabla_x \\log p_t(x) + \\nabla_x \\log p(y\\mid x).\n",
"$$\n",
"We already have an approximation of $\\nabla_x\\log p_t(x)$ from the score network,\n",
"and we obtain $\\nabla_x\\log p(y\\mid x_t)$ by backpropagating through the classifier.\n",
"\n",
"In practice, we introduce a guidance strength $w\\ge 0$:\n",
"$$\n",
"s^{\\text{guided}}(x_t,t)\n",
"= s_\\theta(x_t,t) + w\\,\\nabla_{x_t}\\log p_\\phi(y\\mid x_t,t).\n",
"$$\n",
"- $w=0$: unconditional generation.\n",
"- Larger $w$: samples more aligned with the target class, but potentially reduced diversity and more artifacts.\n",
"\n",
"We can apply this guidance inside either the reverse-time SDE or the probability-flow ODE.\n",
"\n",
"In Pytorch, we can compute the gradient of a function with respect to its input using `torch.autograd.grad`. Example:\n",
"```python\n",
"x = x.clone().detach().requires_grad_(True)\n",
"logits = classifier(x, t)\n",
"log_probs = torch.log_softmax(logits, dim=1)\n",
"selected = log_probs[range(x.shape[0]), y]\n",
"grad = torch.autograd.grad(selected.sum(), x)[0]\n",
"```\n",
"We have to ensure that `x` has `requires_grad=True` to compute gradients with respect to it.\n",
"\n",
"**TODO:**\n",
"- Compette the `guided_score` function.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2dcdf74e",
"metadata": {},
"outputs": [],
"source": [
"def guided_score(score_model, classifier, x, t, y, guidance_scale=1.0):\n",
" x = x.clone().detach().requires_grad_(True)\n",
" raise NotImplementedError"
]
},
{
"cell_type": "markdown",
"id": "c8219e19",
"metadata": {},
"source": [
"We can now generate conditional samples for several classes and plot the generated points using the following function."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5a0ffe54",
"metadata": {},
"outputs": [],
"source": [
"def plot_generated_samples_labeled(\n",
" samples, title, target, data=None, means=None, figsize=(5, 5)\n",
"):\n",
" if isinstance(samples, torch.Tensor):\n",
" x = samples.detach().cpu().numpy()\n",
" else:\n",
" x = samples\n",
" plt.figure(figsize=figsize)\n",
" if data is not None:\n",
" if isinstance(data, TensorDataset):\n",
" X = data.tensors[0].numpy()\n",
" else:\n",
" X = data\n",
" plt.scatter(X[:, 0], X[:, 1], s=4, alpha=0.25, color=\"#9e9e9e\")\n",
" for label in range(cfg[\"num_classes\"]):\n",
" pts = x[target == label]\n",
" plt.scatter(pts[:, 0], pts[:, 1], s=6, alpha=0.6, color=COLORS[label])\n",
" if means is not None:\n",
" plt.scatter(means[:, 0], means[:, 1], s=80, c=\"black\", marker=\"x\")\n",
" plt.title(title)\n",
" plt.axis(\"equal\")\n",
" plt.xlim(-3.5, 3.5)\n",
" plt.ylim(-3.5, 3.5)\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "6f2bd18d",
"metadata": {},
"source": [
"**TODO:**\n",
"- Run the next cell to generate and visualize samples to make sure everything works as expected."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3a463b19",
"metadata": {},
"outputs": [],
"source": [
"def sample_guided_sde(\n",
" score_model,\n",
" classifier,\n",
" guidance_scale=1.0,\n",
" num_samples=200,\n",
" num_steps=cfg[\"num_steps\"],\n",
" return_traj=False,\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" y_target = torch.randint(0, cfg[\"num_classes\"], (num_samples,), device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" t_batch = torch.full((x.shape[0],), t, device=device)\n",
" g = torch.sqrt(beta_t(t))\n",
" drift = -0.5 * beta_t(t) * x\n",
" score = guided_score(\n",
" score_model, classifier, x, t_batch, y_target, guidance_scale\n",
" )\n",
" noise = torch.randn_like(x)\n",
" x = x + (drift - g**2 * score) * dt + g * torch.sqrt(-dt) * noise\n",
" # if return_traj and (i % (num_steps // 100) == 0):\n",
" if return_traj:\n",
" traj.append(x.detach().cpu().numpy())\n",
" return x, traj, y_target.cpu().numpy()\n",
"\n",
"\n",
"sde_guided_samples, _, target = sample_guided_sde(\n",
" score_model, classifier, guidance_scale=1.0, num_samples=1000\n",
")\n",
"plot_generated_samples_labeled(\n",
" sde_guided_samples,\n",
" \"Guided Samples - SDE\",\n",
" target=target,\n",
" data=train_dataset,\n",
" means=means,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b9115088",
"metadata": {},
"outputs": [],
"source": [
"def sample_guided_ode(\n",
" score_model,\n",
" classifier,\n",
" guidance_scale=1.0,\n",
" num_samples=200,\n",
" num_steps=cfg[\"num_steps\"],\n",
" return_traj=False,\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" y_target = torch.randint(0, cfg[\"num_classes\"], (num_samples,), device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" t_batch = torch.full((x.shape[0],), t, device=device)\n",
" g = torch.sqrt(beta_t(t))\n",
" drift = -0.5 * beta_t(t) * x\n",
" score = guided_score(\n",
" score_model, classifier, x, t_batch, y_target, guidance_scale\n",
" )\n",
" x = x + (drift - 0.5 * g**2 * score) * dt\n",
" # if return_traj and (i % (num_steps // 100) == 0):\n",
" if return_traj:\n",
" traj.append(x.detach().cpu().numpy())\n",
" return x, traj, y_target.cpu().numpy()\n",
"\n",
"\n",
"samples_ode_guided, _, target = sample_guided_ode(\n",
" score_model, classifier, guidance_scale=1.0, num_samples=1000\n",
")\n",
"plot_generated_samples_labeled(\n",
" samples_ode_guided,\n",
" \"Guided Samples - ODE\",\n",
" target=target,\n",
" data=train_dataset,\n",
" means=means,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "93c46c0e",
"metadata": {},
"outputs": [],
"source": [
"def sample_guided_ode_heun(\n",
" score_model,\n",
" classifier,\n",
" guidance_scale=1.0,\n",
" num_samples=200,\n",
" num_steps=cfg[\"num_steps\"],\n",
" return_traj=False,\n",
"):\n",
" x = torch.randn(num_samples, 2, device=device)\n",
" ts = torch.linspace(1.0, cfg[\"eps\"], num_steps, device=device)\n",
" y_target = torch.randint(0, cfg[\"num_classes\"], (num_samples,), device=device)\n",
" traj = []\n",
" for i in range(len(ts) - 1):\n",
" t = ts[i]\n",
" dt = ts[i + 1] - ts[i]\n",
" t_batch = torch.full((x.shape[0],), t, device=device)\n",
" g = torch.sqrt(beta_t(t))\n",
" drift = -0.5 * beta_t(t) * x\n",
" score = guided_score(\n",
" score_model, classifier, x, t_batch, y_target, guidance_scale\n",
" )\n",
" f = drift - 0.5 * g**2 * score\n",
" x_pred = x + f * dt\n",
" t_next = t + dt\n",
" batch_t_next = torch.full((x.shape[0],), t_next, device=device)\n",
" g_next = torch.sqrt(beta_t(t_next))\n",
" drift_next = -0.5 * beta_t(t_next) * x_pred\n",
" score_next = guided_score(\n",
" score_model, classifier, x_pred, batch_t_next, y_target, guidance_scale\n",
" )\n",
" f_next = drift_next - 0.5 * g_next**2 * score_next\n",
" x = x + 0.5 * (f + f_next) * dt\n",
" # if return_traj and (i % (num_steps // 100) == 0):\n",
" if return_traj:\n",
" traj.append(x.detach().cpu().numpy())\n",
" return x, traj, y_target.cpu().numpy()\n",
"\n",
"\n",
"samples_ode2_guided, _, target = sample_guided_ode_heun(\n",
" score_model, classifier, guidance_scale=1.0, num_samples=1000\n",
")\n",
"plot_generated_samples_labeled(\n",
" samples_ode2_guided,\n",
" \"Guided Samples - ODE (2nd order)\",\n",
" target=target,\n",
" data=train_dataset,\n",
" means=means,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "26d788a8",
"metadata": {},
"source": [
"We can then visualize the effect of guidance strength on the generated samples and trajectories and observe the trajectories.\n",
"\n",
"**TODO:**\n",
"- Run the next cell to generate samples with different guidance strengths and observe the effect on the trajectories."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "91dc7083",
"metadata": {},
"outputs": [],
"source": [
"def plot_trajectories_labeled(trajs, y_target, title):\n",
" plt.figure(figsize=(5, 5))\n",
" trajs = np.array(trajs)\n",
" # scatter the first points\n",
" for mean_i in range(cfg[\"num_classes\"]):\n",
" for r_i in range(4):\n",
" plt.gca().add_patch(\n",
" plt.Circle(\n",
" (\n",
" means[mean_i, 0],\n",
" means[mean_i, 1],\n",
" ),\n",
" data_std * (r_i + 1),\n",
" color=\"gray\",\n",
" fill=False,\n",
" linestyle=\"--\",\n",
" alpha=0.3,\n",
" )\n",
" )\n",
" for traj_i in range(trajs.shape[1]):\n",
" traj = trajs[:, traj_i, :]\n",
" color_i = y_target[traj_i]\n",
" plt.scatter(\n",
" traj[0, 0], traj[0, 1], color=COLORS[color_i], s=20, edgecolor=\"black\"\n",
" )\n",
" plt.plot(traj[:, 0], traj[:, 1], alpha=0.4, color=COLORS[color_i])\n",
" plt.title(title)\n",
" plt.axis(\"equal\")\n",
" plt.xlim(-3.5, 3.5)\n",
" plt.ylim(-3.5, 3.5)\n",
" plt.show()\n",
"\n",
"\n",
"# Example usage after training\n",
"for scale in [0.0, 0.5, 1.0, 3.0, 5.0]:\n",
" samples, traj, y_target = sample_guided_ode(\n",
" score_model,\n",
" classifier,\n",
" guidance_scale=scale,\n",
" return_traj=True,\n",
" num_samples=100,\n",
" )\n",
" plot_trajectories_labeled(traj, y_target=y_target, title=f\"Guidance scale={scale}\")"
]
},
{
"cell_type": "markdown",
"id": "a9fab273",
"metadata": {},
"source": [
"Now that the model are conditional, we have to make sure that the final samples correspond to the target class. We have to adapt the quality metric function accordingly.\n",
"\n",
"**TODO:**\n",
"- Adapt the quality metric to account for class labels in the conditional generation setting.\n",
"- Run the adapted quality metric on samples generated with different guidance strengths."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f2123ea0",
"metadata": {},
"outputs": [],
"source": [
"def quality_metrics_labeled(samples, means, std, y_target):\n",
" if isinstance(samples, torch.Tensor):\n",
" x = samples.detach().cpu()\n",
" else:\n",
" x = torch.tensor(samples)\n",
" means_t = torch.tensor(means)\n",
" dists = torch.cdist(x, means_t)\n",
"\n",
" return quality\n",
"\n",
"\n",
"def evaluate_guided_guidance(\n",
" model, classifier, guidance_scales=[0.0, 0.5, 1.0, 3.0, 5.0], num_steps=100\n",
"):\n",
" rows = []\n",
" for scale in guidance_scales:\n",
" sde_samples, _, target = sample_guided_sde(\n",
" model,\n",
" classifier,\n",
" guidance_scale=scale,\n",
" num_steps=num_steps,\n",
" )\n",
" ode_samples, _, target = sample_guided_ode(\n",
" model,\n",
" classifier,\n",
" guidance_scale=scale,\n",
" num_steps=num_steps,\n",
" )\n",
" ode2_samples, _, target = sample_guided_ode_heun(\n",
" model,\n",
" classifier,\n",
" guidance_scale=scale,\n",
" num_steps=num_steps,\n",
" )\n",
" q_sde = quality_metrics_labeled(sde_samples, means, data_std, target)\n",
" q_ode = quality_metrics_labeled(ode_samples, means, data_std, target)\n",
" q_ode2 = quality_metrics_labeled(ode2_samples, means, data_std, target)\n",
" rows.append((scale, q_sde, q_ode, q_ode2))\n",
"\n",
" quality_metric = np.array(rows)\n",
" plt.figure(figsize=(8, 4))\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 1],\n",
" marker=\"o\",\n",
" label=\"SDE\",\n",
" )\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 2],\n",
" marker=\"o\",\n",
" label=\"ODE\",\n",
" )\n",
" plt.plot(\n",
" quality_metric[:, 0],\n",
" quality_metric[:, 3],\n",
" marker=\"o\",\n",
" label=\"ODE (2nd order)\",\n",
" )\n",
" plt.xlabel(\"Guidance scale\")\n",
" plt.ylabel(\"Quality\")\n",
" plt.title(f\"Guided Quality vs guidance scale (steps={num_steps})\")\n",
" plt.grid(True, alpha=0.3)\n",
" plt.legend()\n",
" plt.show()\n",
"\n",
"\n",
"evaluate_guided_guidance(\n",
" score_model,\n",
" classifier,\n",
" guidance_scales=[0.0, 0.5, 1.0, 3.0, 5.0],\n",
" num_steps=cfg[\"num_steps\"] // 2,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "438c7ba3",
"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
}