4.2 DDPM 原理与从零实现

本章目标：把扩散模型的”加噪—去噪”从玄学变成你能手推、能手写的算法。学完标准：能从零推导闭式加噪公式，能解释”预测噪声”这个训练目标从哪来，并在 MNIST 上跑通一个 50 行的迷你扩散模型。

1. 前向过程：主动把图搅成噪声（不需要学习）

定义一个马尔可夫链，每步加一点高斯噪声，T 步后图像变成纯噪声：

$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t;\ \sqrt{1-\beta_t}\, x_{t-1},\ \beta_t I)$

$\beta_t$ 是噪声强度表（schedule，从 1e-4 涨到 0.02 之类）。关键推导（必会）：令 $\alpha_t = 1-\beta_t$、$\bar{\alpha}_t = \prod_{s\le t}\alpha_s$，反复代入高斯叠加性，得到任意时刻 t 的闭式解——不用逐步迭代，一步采样：

$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\, \epsilon, \quad \epsilon\sim\mathcal{N}(0,I)$

读法：$x_t$ 是”原图按 $\bar\alpha_t$ 衰减 + 噪声按 $1-\bar\alpha_t$ 注入”的混合。t 小偏原图，t 大近纯噪声——训练时对每个样本随机抽一个 t，立刻得到带标签的”输入-答案”对（输入 $x_t$，答案 ε），这就是自监督信号。

2. 反向过程：学”去噪方向”（模型登场）

真实反向 $q(x_{t-1}|x_t)$ 不可得（需要整张图的数据分布），所以训练神经网络 $p_\theta(x_{t-1}|x_t) = \mathcal{N}(\mu_\theta(x_t,t), \sigma_t^2 I)$ 去逼近它。贝叶斯代入后可把 $\mu_\theta$ 重参数化——模型实际要预测的是混在 $x_t$ 里的噪声 ε：

$L_{\text{simple}} = \mathbb{E}_{x_0, \epsilon, t}\Big[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \Big]$

这个 MSE 就是 DDPM 的全部训练目标（DDPM 论文证明它是 ELBO 的简化等价——严格推导选修，直觉必修）。

直觉版理解（推荐记忆这个）：预测噪声 = 预测”去噪方向”。把 $\| \epsilon - \epsilon_\theta(x_t,t)\|$ 改写一下，它近似于在学 $\nabla_{x_t} \log p(x_t)$——数据对数密度的梯度（score）。也就是说，模型学的是”在当前位置，往哪个方向走概率密度更高”。去噪生成的本质是沿着学到的密度梯度，从噪声一路走到数据流形。这句话打通了扩散和”压缩即智能”（2.1）——都在学数据的结构。

3. 训练与采样算法（各 4 行，背下来）

训练：① 取 $x_0$；② 随机 t、随机 ε；③ 闭式合成 $x_t$；④ 梯度下降 $\|\epsilon - \epsilon_\theta(x_t,t)\|^2$。

采样：从 $x_T \sim \mathcal{N}(0,I)$ 开始，循环 $t=T\to1$：

$x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\Big(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t)\Big) + \sigma_t z$

（减去预测噪声的贡献，再加一点随机性 $z$；t=1 时停止加噪。）1000 步走完，噪声里”长”出一张图。

4. 从零实现（MNIST，50 行核心）

import torch, torch.nn as nn, torch.nn.functional as F

betas = torch.linspace(1e-4, 0.02, 500)                    # T=500（MNIST 够了）
alphas, abar = 1-betas, torch.cumprod(1-betas, dim=0)      # α 与 ᾱ 的闭式表

def add_noise(x0, t):                                      # 闭式加噪
    a = abar[t].view(-1,1,1,1)
    eps = torch.randn_like(x0)
    return a.sqrt()*x0 + (1-a).sqrt()*eps, eps

class Denoiser(nn.Module):                                 # 简化网络：真实用 U-Net
    def __init__(self):
        super().__init__()
        self.t_emb = nn.Sequential(nn.Embedding(500, 64), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 64))
        self.net = nn.Sequential(                          # 仅示意，U-Net 见 4.4
            nn.Conv2d(1, 64, 3, 1, 1), nn.SiLU(),
            nn.Conv2d(64, 64, 3, 1, 1), nn.SiLU(), nn.Conv2d(64, 1, 3, 1, 1))
    def forward(self, xt, t):
        h = self.net(xt)                                   # 真实实现要把 t_emb 调制进每层
        return h

model = Denoiser()
for step in range(20000):                                  # 训练
    x0 = next(loader)                                      # MNIST batch, 归一化到 [-1,1]
    t = torch.randint(0, 500, (x0.size(0),))
    xt, eps = add_noise(x0, t)
    loss = F.mse_loss(model(xt, t), eps)                   # 全部训练目标就这一行
    ...

@torch.no_grad()
def sample(n=8):
    x = torch.randn(n, 1, 28, 28)
    for t in reversed(range(500)):
        eps = model(x, torch.full((n,), t))
        a, ab = alphas[t], abar[t]
        x = (x - (1-a)/ (1-ab).sqrt() * eps) / a.sqrt()    # 去噪一步
        if t > 0: x = x + betas[t].sqrt() * torch.randn_like(x)
    return x

（网络按理应把 t 编码注入每层——教学版从简，作业里补全；U-Net 结构与注入方式见 4.4。）跑通后你会看到：采样从一团雪花里逐渐浮现数字——这是你亲手实现的”从噪声中创造”。

5. 预测目标的三胞胎（为 4.5 埋线）

同一个网络可以学三种”答案”：ε-pred（预测噪声，DDPM）、x0-pred（直接预测原图）、v-pred（预测二者的组合 $v = \sqrt{\bar\alpha}\,\epsilon - \sqrt{1-\bar\alpha}\,x_0$）。数学上可互相换算，但训练动态和数值稳定性不同——v-pred 在低噪声端更稳，是渐进式蒸馏和流模型（4.5 的速度场正是它的推广）的桥。知道存在即可，别背公式。

6. 自测题

1. 闭式加噪公式的两个分量各是什么？它让训练获得了什么？（$\sqrt{\bar\alpha_t}x_0$ 衰减信号 + $\sqrt{1-\bar\alpha_t}\epsilon$ 注入噪声；任意 t 一步采样 → 海量自监督对）
2. 为什么反向过程必须”学”而前向不用？（前向是人为定义的马尔可夫加噪；反向依赖未知数据分布）
3. 训练 loss 的直觉含义是什么？（学数据对数密度的梯度方向——score matching）
4. 采样循环里 σ_t z 的作用？去掉会怎样？（随机性，保分布多样性；去掉=确定性轨迹，多样性降——正是 DDIM 的思路）
5. T=1000 而采样也 1000 步，慢在哪？预感怎么破？（串行迭代；见 4.3 跳步采样）
6. ε/x0/v 三种预测目标的关系？（可互相换算的三种参数化；稳定性与下游蒸馏的适配不同）

7. 动手作业

1. 跑通 MNIST 迷你版（补全 t 注入：把 t_emb 过 MLP 后 reshape 加到特征图上），存 8 张采样过程的中间帧 gif；
2. 手推闭式公式：从 $q(x_t|x_{t-1})$ 出发两步归纳到 $q(x_t|x_0)$（高斯叠加：均值方差各自相加）；
3. 把 β schedule 换成 cosine（Nichol & Dhariwal）对比采样质量——体会 schedule 不是小事；
4. 挑战：给迷你模型加条件（0–9 数字标签，用 4.3 的 CFG 思路），实现”指定数字生成”。

图解与案例 · 加噪-去噪的一来一回

flowchart LR
  subgraph F["前向（不需要学习）"]
    X0["x₀ 原图"] -->|"√ᾱ_t·x₀ + √(1−ᾱ_t)·ε"| XT["x_t 加噪图"]
  end
  XT --> NET["网络 ε_θ(x_t, t)<br>任务：指认出噪声 ε"]
  subgraph R["反向（迭代采样）"]
    NET --> X1["x_{t−1} = 减去预测的噪声贡献"]
  end
  X1 -. "重复 t−1 步直到 x₀" .-> NET

案例：给一个像素加噪（数字全算出）

取单像素 $x_0 = 0.6$，某时刻 $\bar\alpha_t = 0.25$（信号只剩 1/4），采样噪声 $\epsilon = 0.8$：

$x_t = \sqrt{0.25}\times 0.6 + \sqrt{0.75}\times 0.8 = 0.30 + 0.69 = 0.99$

训练时：把 $x_t=0.99$ 和 $t$ 喂给网络，要求它指认出当初那个 0.8，loss = $(\epsilon_\theta - 0.8)^2$。

读出三点：① 信号被 √ᾱ 衰减、噪声被 √(1−ᾱ) 放大——t 就是信噪比的旋钮；② 模型的任务永远不变：从混合物里指认噪声；③ 因为闭式公式，任意 t 的训练样本可以一步合成——这是扩散训练可以大规模并行的根基。

8. 延伸资源

• DDPM 论文（Ho et al., 2020）+ Lilian Weng 博客——精读组合拳
• 李沐《论文精读》DDPM 一期——推导全程带画
• Hugging Face《Annotated Diffusion》——带注释的完整实现

附录 · 完整可运行脚本：ddpm_mnist.py（含小 U-Net）

依赖：torch torchvision。4.2 §4 教学版的完整化：真 U-Net（含时间注入与跳跃连接）、余弦可换 schedule、采样出网格图。

# ddpm_mnist.py —— MNIST 上的完整 DDPM（T=400）
import torch, torch.nn as nn, torch.nn.functional as F, math
from torchvision import datasets, transforms, utils
from torch.utils.data import DataLoader

T = 400
betas = torch.linspace(1e-4, 0.02, T)
alphas, abar = 1 - betas, torch.cumprod(1 - betas, 0)

def time_emb(t, dim=64):                                   # 正弦时间嵌入（1.3 位置编码同款）
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half) / half)
    ang = t[:, None].float() * freqs[None]
    return torch.cat([torch.sin(ang), torch.cos(ang)], 1)

class Block(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.c1, self.c2 = nn.Conv2d(cin, cout, 3, 1, 1), nn.Conv2d(cout, cout, 3, 1, 1)
        self.n1, self.n2 = nn.GroupNorm(8, cout), nn.GroupNorm(8, cout)
        self.tmlp = nn.Linear(64, cout)                    # 时间注入：加到特征上
        self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity()
    def forward(self, x, t):
        h = F.silu(self.n1(self.c1(x)))
        h = h + self.tmlp(t)[:, :, None, None]
        h = F.silu(self.n2(self.c2(h)))
        return h + self.skip(x)

class UNet(nn.Module):
    def __init__(self, ch=64):
        super().__init__()
        self.inp = nn.Conv2d(1, ch, 3, 1, 1)
        self.d1, self.d2 = Block(ch, ch), Block(ch, ch * 2)
        self.pool = nn.Conv2d(ch, ch, 4, 2, 1)             # 步长卷积下采样
        self.mid = Block(ch * 2, ch * 2)
        self.u2 = Block(ch * 4, ch), self.u1 = Block(ch * 2, ch)
        self.out = nn.Conv2d(ch, 1, 3, 1, 1)
    def forward(self, x, t):
        te = time_emb(t)
        x0 = self.inp(x)                                    # 28
        x1 = self.d1(x0, te)                                # 28  (skip1)
        x2 = self.pool(x1); x2 = self.d2(x2, te)            # 14  (skip2)
        m = self.mid(F.max_pool2d(x2, 2), te)               # 7
        u = F.interpolate(m, scale_factor=2)                # 14
        u = self.u2[0](torch.cat([u, x2], 1), te)
        u = F.interpolate(u, scale_factor=2)                # 28
        u = self.u1[0](torch.cat([u, x1], 1), te)
        return self.out(u)

def add_noise(x0, t):
    a = abar[t].view(-1, 1, 1, 1).to(x0.device)
    eps = torch.randn_like(x0)
    return a.sqrt() * x0 + (1 - a).sqrt() * eps, eps

def train():
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    ds = datasets.MNIST("data", train=True, download=True,
                        transform=transforms.Compose([transforms.Pad(2),
                            transforms.ToTensor(), transforms.Lambda(lambda x: x * 2 - 1)]))
    dl = DataLoader(ds, batch_size=128, shuffle=True, num_workers=2, drop_last=True)
    net = UNet().to(dev); opt = torch.optim.AdamW(net.parameters(), 2e-4)
    for ep in range(8):
        for x0, _ in dl:
            x0 = x0.to(dev)
            t = torch.randint(0, T, (x0.size(0),), device=dev)
            xt, eps = add_noise(x0, t)
            loss = F.mse_loss(net(xt, t), eps)
            opt.zero_grad(set_to_none=True); loss.backward()
            torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0); opt.step()
        print(f"ep{ep} loss={loss.item():.4f}")
        torch.save(net.state_dict(), "ddpm.pt")

@torch.no_grad()
def sample(n=64, ckpt="ddpm.pt"):
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    net = UNet().to(dev); net.load_state_dict(torch.load(ckpt, map_location=dev)); net.eval()
    x = torch.randn(n, 1, 32, 32, device=dev)
    for t in reversed(range(T)):
        tt = torch.full((n,), t, device=dev)
        eps = net(x, tt)
        a, ab, b = alphas[t].to(dev), abar[t].to(dev), betas[t].to(dev)
        x = (x - (1 - a) / (1 - ab).sqrt() * eps) / a.sqrt()
        if t > 0: x = x + b.sqrt() * torch.randn_like(x)
    utils.save_image((x.clamp(-1, 1) + 1) / 2, "samples.png", nrow=8)
    print("saved samples.png")

if __name__ == "__main__":
    import sys
    train() if len(sys.argv) < 2 or sys.argv[1] == "train" else sample()