""" diffusion_lab.py —— 第 7–8 讲的扩散模型工具箱：噪声调度、前向加噪、精确后验、三种参数化、训练与采样 ===================================================================== 【这个文件要回答的问题】 DDPM（Ho et al., 2020）只需要四样东西： 1. 一个固定的噪声调度 α_1..α_T，决定前向过程每一步加多少噪声； 2. 前向过程的两个闭式解：q(x_t | x_0)（slides 的性质 1）和 q(x_{t−1} | x_t, x_0)（性质 2）； 3. 一个去噪网络，输入 (x_t, t)，输出三种东西之一（slides 的选项 A / B / C）： A：直接输出均值 μ_θ(x_t, t) B：输出对 x_0 的估计 C：输出对噪声 ε 的估计 4. 训练循环（随机挑一个 t、加噪、回归）和采样循环（从纯噪声出发，一步步去噪）。 这个文件把它们写成几十行能运行的代码，记号和 slides 一一对应。 【记号（和 slides 相同）】 前向一步：q(x_t | x_{t−1}) = N(√α_t · x_{t−1}, (1 − α_t) I)，t = 1..T 累乘： ᾱ_t = α_1 α_2 ⋯ α_t，约定 ᾱ_0 = 1 性质 1： q(x_t | x_0) = N(√ᾱ_t · x_0, (1 − ᾱ_t) I)，即 x_t = √ᾱ_t · x_0 + √(1 − ᾱ_t) · ε 性质 2： q(x_{t−1} | x_t, x_0) = N(μ̃_q, σ_t² I)， μ̃_q = α_t^(0) · x_0 + α_t^(t) · x_t， α_t^(0) = √ᾱ_{t−1} (1 − α_t) / (1 − ᾱ_t)， α_t^(t) = √α_t (1 − ᾱ_{t−1}) / (1 − ᾱ_t)， σ_t² = (1 − ᾱ_{t−1})(1 − α_t) / (1 − ᾱ_t) 性质 3： 把 x_0 = (x_t − √(1 − ᾱ_t) ε)/√ᾱ_t 代进 μ̃_q：μ̃_q = (x_t − (1 − α_t)/√(1 − ᾱ_t) · ε) / √α_t 很多论文（包括 Ho et al.）用 β_t = 1 − α_t 表示"第 t 步加的噪声方差"，两种写法等价。 【本文件提供】 Schedule 噪声调度（linear / cosine）以及上面所有系数；可以"重新取样"成更少的步数（第 8 讲） time_embedding 把整数时间步 t 变成正弦 / 余弦特征（和第 2 讲的正弦位置编码同一个公式） MLPDenoiser 二维（或任意维向量）数据的去噪网络 diffusion_loss / train_diffusion 三种参数化的训练目标与训练循环（可选 EMA） predict_mean / sample 按三种参数化从网络输出算出反向均值；祖先采样（slides 的 Algorithm "Sampling"） swiss_data / ring_data / spiral_stats 二维玩具数据与评估 【约定】 - 数组下标就是时间步：alpha[t]、alpha_bar[t] 对 t = 0..T 都有定义（alpha[0] 不用，alpha_bar[0] = 1）。 - 所有系数先用 float64 算，再存成 float32（训练用）；随机性全部来自显式的 torch.Generator。 """ import math import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) for _d in ("prereq2", "lecture06"): _p = os.path.join(HERE, "..", _d) if _p not in sys.path: sys.path.insert(0, _p) from prob_lab import swiss_roll # noqa: E402 from gan_lab import RING, mode_stats # noqa: E402,F401 （mode_stats 在这里再导出一次，给 03、04 号脚本用） # ===================================================================== # 1. 噪声调度 # ===================================================================== class Schedule: """ 噪声调度与 slides 里的全部系数。 【参数】 T 步数（DDPM 用 1000） kind "linear"：β_t = 1 − α_t 从 beta_start 线性增长到 beta_end（Ho et al. 2020：1e-4 → 0.02） "cosine"：直接规定 ᾱ_t = f(t)/f(0)，f(t) = cos²((t/T + s)/(1 + s) · π/2)，s = 0.008 （Nichol & Dhariwal 2021；第 8 讲比较两种调度） alpha_bar 直接给定 ᾱ_1..ᾱ_T（长度 T 的序列）；给了它就忽略 kind，用于"重新取样"出的短调度 【属性】都是长度 T+1 的 float32 张量，下标 = 时间步： alpha, alpha_bar, c0（α_t^(0)）, ct（α_t^(t)）, sigma2（σ_t²） 【边界】t = 1 时 ᾱ_0 = 1，所以 σ_1² = 0、α_1^(0) = 1、α_1^(t) = 0：最后一步是确定性的，直接输出 x_0 的估计。 """ def __init__(self, T=1000, kind="linear", beta_start=1e-4, beta_end=0.02, alpha_bar=None): self.T = T if alpha_bar is not None: ab = np.concatenate([[1.0], np.asarray(alpha_bar, dtype=np.float64)]) elif kind == "linear": beta = np.linspace(beta_start, beta_end, T, dtype=np.float64) ab = np.concatenate([[1.0], np.cumprod(1 - beta)]) elif kind == "cosine": s = 0.008 f = np.cos((np.arange(T + 1) / T + s) / (1 + s) * math.pi / 2) ** 2 ab = f / f[0] # 相邻两步之比就是 α_t；按 Nichol & Dhariwal 的做法把 β_t = 1 − α_t 截断在 0.999 以内，避免最后一步 α_t ≈ 0 alpha = np.clip(ab[1:] / ab[:-1], 1 - 0.999, 1.0) ab = np.concatenate([[1.0], np.cumprod(alpha)]) else: raise ValueError(kind) alpha = np.concatenate([[1.0], ab[1:] / ab[:-1]]) ab_prev = np.concatenate([[1.0], ab[:-1]]) with np.errstate(divide="ignore", invalid="ignore"): c0 = np.where(np.arange(T + 1) > 0, np.sqrt(ab_prev) * (1 - alpha) / (1 - ab), 0.0) ct = np.where(np.arange(T + 1) > 0, np.sqrt(alpha) * (1 - ab_prev) / (1 - ab), 0.0) sigma2 = np.where(np.arange(T + 1) > 0, (1 - ab_prev) * (1 - alpha) / (1 - ab), 0.0) self.alpha_bar64 = ab # float64 版本：分析用（第 7 讲 01、02） self.alpha64 = alpha self.alpha = torch.tensor(alpha, dtype=torch.float32) self.alpha_bar = torch.tensor(ab, dtype=torch.float32) self.c0 = torch.tensor(c0, dtype=torch.float32) self.ct = torch.tensor(ct, dtype=torch.float32) self.sigma2 = torch.tensor(sigma2, dtype=torch.float32) def _b(self, v, t, x): """把按时间步取出的系数 v[t]（形状 (B,)）变成能和 x（形状 (B, ...)）逐元素相乘的形状。""" return v[t].view(-1, *([1] * (x.dim() - 1))) def q_sample(self, x0, t, eps): """性质 1：x_t = √ᾱ_t · x_0 + √(1 − ᾱ_t) · ε。t 是形状 (B,) 的整数张量（取值 1..T）。""" ab = self._b(self.alpha_bar, t, x0) return ab.sqrt() * x0 + (1 - ab).sqrt() * eps def posterior_mean(self, x0, xt, t): """性质 2：μ̃_q(x_t, x_0) = α_t^(0) · x_0 + α_t^(t) · x_t。""" return self._b(self.c0, t, x0) * x0 + self._b(self.ct, t, xt) * xt def x0_from_eps(self, xt, t, eps): """性质 3 的第一步：x_0 = (x_t − √(1 − ᾱ_t) · ε) / √ᾱ_t。""" ab = self._b(self.alpha_bar, t, xt) return (xt - (1 - ab).sqrt() * eps) / ab.sqrt() def respace(self, timesteps): """ "重新取样"：只保留原调度里的一部分时间步 τ_1 < τ_2 < ⋯ < τ_S（第 8 讲用更少的步数采样）。 新调度的第 i 步对应原来的第 τ_i 步：ᾱ'_i = ᾱ_{τ_i}，于是 α'_i = ᾱ_{τ_i} / ᾱ_{τ_{i−1}}—— 也就是把原来连续的几小步噪声合并成一大步。网络还是用原来的时间步 τ_i 去调用。 【返回】(新的 Schedule, τ 的张量（长度 S+1，τ_0 = 0）) """ ts = sorted(int(t) for t in timesteps) sch = Schedule(T=len(ts), alpha_bar=[self.alpha_bar64[t] for t in ts]) return sch, torch.tensor([0] + ts) # ===================================================================== # 2. 去噪网络 # ===================================================================== def time_embedding(t, dim=32, max_period=10000.0): """ 时间步的正弦 / 余弦嵌入：和第 2 讲 Transformer 的正弦位置编码同一个公式。 【为什么不直接把 t 当一个数输入】网络很难从一个标量里分辨"第 3 步"和"第 4 步"的细微区别； 一组不同频率的正弦把 t 展开成一个向量，低频分量分辨"大概在哪个阶段"，高频分量分辨相邻的步。 【参数】t：(B,) 整数或浮点张量；dim：嵌入维数（偶数） 【返回】(B, dim) """ half = dim // 2 freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32) / half) args = t.float()[:, None] * freqs[None] return torch.cat([torch.sin(args), torch.cos(args)], dim=1) class MLPDenoiser(nn.Module): """ 向量数据（比如二维点）的去噪网络：输入 (x_t, t)，输出和 x_t 同样维数的向量。 输出的含义由训练时选的参数化决定：ε（选项 C）、x_0（选项 B）或均值 μ（选项 A）。 h = Linear(x_t) + MLP(time_embedding(t)) h = h + Linear(SiLU(h)) × n_layers （残差块） 输出 = Linear(SiLU(h)) """ def __init__(self, dim=2, hidden=128, n_layers=3, t_dim=32): super().__init__() self.t_dim = t_dim self.inp = nn.Linear(dim, hidden) self.t_proj = nn.Sequential(nn.Linear(t_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden)) self.blocks = nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(n_layers)]) self.out = nn.Linear(hidden, dim) def forward(self, x, t): h = self.inp(x) + self.t_proj(time_embedding(t, self.t_dim)) for blk in self.blocks: h = h + blk(F.silu(h)) return self.out(F.silu(h)) # ===================================================================== # 3. 训练：三种参数化 # ===================================================================== def diffusion_loss(model, sched, x0, gen, param="eps", t=None): """ 一个小批量的训练损失（slides 的 Algorithm "Training"，第 4–7 行）： t ~ Uniform{1..T}；ε ~ N(0, I)；x_t = √ᾱ_t x_0 + √(1 − ᾱ_t) ε； 选项 C（param="eps"）：‖ε − ε_θ(x_t, t)‖² 选项 B（param="x0"） ：‖x_0 − x_θ^(0)(x_t, t)‖² 选项 A（param="mu"） ：‖μ̃_q(x_t, x_0) − μ_θ(x_t, t)‖² 【参数】t：可选，指定每个样本的时间步（第 8 讲按 t 统计损失时用）；None 表示均匀随机 【返回】标量损失（对样本和维度取平均） 【边界】三种损失的数值大小不可直接比较：它们回归的目标量级不同（见 04 号脚本）。 """ B = x0.shape[0] if t is None: t = torch.randint(1, sched.T + 1, (B,), generator=gen) eps = torch.randn(x0.shape, generator=gen) xt = sched.q_sample(x0, t, eps) out = model(xt, t) if param == "eps": target = eps elif param == "x0": target = x0 elif param == "mu": target = sched.posterior_mean(x0, xt, t) else: raise ValueError(param) return F.mse_loss(out, target) def train_diffusion(model, sample_x0, sched, steps, batch=256, lr=1e-3, param="eps", seed=0, ema_decay=None, callback=None, every=0, warmup=0): """ 训练循环。 【参数】 sample_x0 函数 sample_x0(batch, gen) → 一批真实数据 x_0 ema_decay 不为 None 时同时维护参数的指数滑动平均（EMA）：θ_ema ← d·θ_ema + (1−d)·θ。 DDPM 及之后几乎所有扩散模型都用 EMA 的权重来采样——它平滑掉了 SGD 最后阶段的抖动。 warmup 学习率从 0 线性升到 lr 的步数（大网络更稳定；二维玩具不需要） callback 每 every 步调用 callback(step, model, ema_model, losses) 【返回】(model, ema_model 或 None, 每步损失的列表) """ import copy gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) ema = copy.deepcopy(model) if ema_decay is not None else None losses = [] model.train() for step in range(1, steps + 1): if warmup: for g in opt.param_groups: g["lr"] = lr * min(1.0, step / warmup) loss = diffusion_loss(model, sched, sample_x0(batch, gen), gen, param) opt.zero_grad() loss.backward() opt.step() if ema is not None: with torch.no_grad(): for pe, pm in zip(ema.parameters(), model.parameters()): pe.mul_(ema_decay).add_(pm, alpha=1 - ema_decay) losses.append(loss.item()) if callback is not None and every and step % every == 0: callback(step, model, ema, losses) model.eval() if ema is not None: ema.eval() return model, ema, losses # ===================================================================== # 4. 采样 # ===================================================================== def predict_mean(model, sched, xt, t, param="eps", net_t=None): """ 从网络输出算出反向过程的均值 μ_θ(x_t, t)（slides 的三种选项）： A（"mu"） ：网络输出就是 μ_θ B（"x0"） ：μ_θ = α_t^(0) · x̂_0 + α_t^(t) · x_t，x̂_0 = 网络输出 C（"eps"）：x̂_0 = (x_t − √(1 − ᾱ_t) ε̂)/√ᾱ_t，再同 B 【参数】net_t：调用网络时用的时间步（重新取样的调度里，它和 sched 的下标 t 不同）；None 表示就用 t 【返回】(μ_θ, x̂_0)；选项 A 没有 x̂_0，返回 None """ out = model(xt, t if net_t is None else net_t) if param == "mu": return out, None x0_hat = out if param == "x0" else sched.x0_from_eps(xt, t, out) return sched.posterior_mean(x0_hat, xt, t), x0_hat @torch.no_grad() def sample(model, sched, n, shape, gen, param="eps", keep=(), clip=None, net_steps=None): """ 祖先采样（slides 的 Algorithm "Sampling"）： x_T ~ N(0, I)；对 t = T, …, 1：x_{t−1} = μ_θ(x_t, t) + σ_t · z，z ~ N(0, I)（t = 1 时 σ_1 = 0） 【参数】n：样本数；shape：单个样本的形状（二维点是 (2,)，图像是 (1, 28, 28)）； keep：要记录中间状态的时间步（比如 (1000, 500, 100)），结果里 key 是 t，value 是 x_t； clip：不为 None 时把 x̂_0 截断到 [−clip, clip]（图像在 [−1, 1] 里，截断能减少越界的像素；只对 B、C 有效）； net_steps：重新取样的调度里每一步对应的原始时间步 τ（Schedule.respace 的第二个返回值） 【返回】(x_0, {t: x_t}) """ model.eval() x = torch.randn((n,) + tuple(shape), generator=gen) traj = {sched.T: x.clone()} if sched.T in keep else {} for t in range(sched.T, 0, -1): tt = torch.full((n,), t, dtype=torch.long) nt = None if net_steps is None else torch.full((n,), int(net_steps[t]), dtype=torch.long) if clip is not None and param != "mu": out = model(x, tt if nt is None else nt) x0_hat = out if param == "x0" else sched.x0_from_eps(x, tt, out) mean = sched.posterior_mean(x0_hat.clamp(-clip, clip), x, tt) else: mean, _ = predict_mean(model, sched, x, tt, param, nt) if t > 1: x = mean + sched.sigma2[t].sqrt() * torch.randn(x.shape, generator=gen) else: x = mean if t - 1 in keep: traj[t - 1] = x.clone() return x, traj # ===================================================================== # 5. 二维玩具数据与评估 # ===================================================================== SWISS_NOISE = 0.05 def swiss_data(n, rng): """前置讲 P2 的二维"瑞士卷"（沿螺线的点，参数 s 在 [1.5π, 4.5π] 上均匀），返回 float32 张量 (n, 2)。""" return torch.from_numpy(swiss_roll(n, rng, noise=SWISS_NOISE)).float() def ring_data(n, rng): """第 6 讲的八高斯圆环，返回 float32 张量 (n, 2)。""" return torch.from_numpy(RING.sample(n, rng)).float() _S_GRID = np.linspace(1.5 * np.pi, 4.5 * np.pi, 6001) _CURVE = np.stack([_S_GRID * np.cos(_S_GRID), _S_GRID * np.sin(_S_GRID)], axis=1) / 5.0 def spiral_stats(x, n_bins=6): """ 评估二维样本有多像瑞士卷。 【思路】对每个点找螺线（无噪声的中心线）上最近的点：距离说明"在不在线上"，最近点的参数 s 说明"落在螺线的哪一段"。 【返回】字典 {median_dist: 到螺线的距离中位数, on: 距离 < 3×噪声 的比例, bins: 落在 n_bins 段螺线上的比例（列表）, tv: 和"均匀分布在各段"之间的总变差距离（0 = 覆盖得完全均匀）} """ x = np.asarray(x, dtype=np.float64) d2 = ((x[:, None, :] - _CURVE[None, ::4]) ** 2).sum(-1) # 先粗找（每 4 个格点取一个），够用且省内存 j = d2.argmin(1) * 4 dist = np.sqrt(d2.min(1)) seg = np.minimum((j / len(_S_GRID) * n_bins).astype(int), n_bins - 1) on = dist < 3 * SWISS_NOISE bins = np.bincount(seg[on], minlength=n_bins) / max(on.sum(), 1) return {"median_dist": float(np.median(dist)), "on": float(on.mean()), "bins": bins.tolist(), "tv": float(0.5 * np.abs(bins - 1.0 / n_bins).sum())} def count_parameters(model): return sum(p.numel() for p in model.parameters())