从零实现DDPM扩散模型

从零实现DDPM扩散模型

本节摘要:Ho、Jain、Abbeel(2020)给了这个领域一个戒不掉的配方——用一千个小步把数据摧毁成噪声,训一个神经网络去预测噪声,推理时再反向去噪。今天每一个主流的图像、视频、3D、音乐模型都跑在这个循环上(可能外加流匹配或一致性技巧)。本节是整章的数学核心:我们从零推导前向过程的闭式解 q(x_t|x_0) = N(√ᾱ_t·x_0, (1−ᾱ_t)·I)、反向去噪步骤的来历、为什么「预测噪声」等价于「学分数 ∇log p(x)」,以及为什么完整的变分 ELBO 在 DDPM 参数化下塌缩成一行 MSE(「简单损失」)。读完本节,你能在 40 行纯 Python 里手写一个能学到双峰分布的一维 DDPM,并理解后续所有潜在扩散、ControlNet、流匹配的数学地基。

对应原课程:Phase 08 · Lesson 06 · diffusion-ddpm-from-scratch(原英文 phases/08-generative-ai/06-diffusion-ddpm-from-scratch/docs/en.md)。

学习目标

阅读完本节,你应当能够:

1. 推导前向过程的闭式解 q(x_t|x_0),解释 α̅_t = ∏(1−β_s) 为何让「一步到位加噪」成为可能。
2. 写出反向去噪步骤的公式,并说清它是从后验 q(x_{t−1}|x_t, x_0) 解出来的代数结果。
3. 推导简单损失 L_simple = E||ε − ε_θ(√ᾱ_t·x_0 + √(1−ᾱ_t)·ε, t)||²,并解释为什么完整 ELBO 在 DDPM 参数化下塌缩成这一行 MSE。
4. 用三条直觉解释扩散为什么有效:去噪易、生成难、分数匹配伪装、ELBO 化简。
5. 在一维数据上手写一个 DDPM:前向调度、单步采样、一步训练、反向采样循环。

一、问题与直觉

你想要一个 p_data(x) 的采样器。GAN 玩 minimax 博弈常常发散;VAE 因高斯解码器产出模糊样本。你真正想要的是一个训练目标,它(a)是单一稳定的损失(无鞍点、无 minimax),(b)是 log p(x) 的下界(于是你有似然),(c)样本质量达到 SOTA。

Sohl-Dickstein 等(2015)给出理论答案:定义一个马尔可夫链 q(x_t|x_{t−1}) 逐步加高斯噪声,训一个反向链 p_θ(x_{t−1}|x_t) 去去噪。Ho、Jain、Abbeel(2020)证明损失可以化简成一行——预测噪声——并理顺了数学。2020 年这是个新奇事物;2021 年它产出了 SOTA 样本;2022 年它变成了 Stable Diffusion;2026 年它是底料。

前向过程:为什么有闭式解?

逐加噪 T 步:

q(x_t | x_{t−1}) = N( √(1−β_t) · x_{t−1},  β_t · I )
​

β_t 是第 t 步的噪声方差(从 1e-4 线性到 0.02,共 T=1000 步)。关键:高斯的线性组合仍是高斯,所以累积多步后的分布有闭式解:

q(x_t | x_0) = N( √(α̅_t) · x_0,  (1 − α̅_t) · I )
​

其中 α̅_t = ∏_{s=1..t} (1 − β_s)。这个闭式解是整个数学可处理的基石——它让你能在训练时一步从 x_0 跳到任意 x_t,无需逐步前向。到 t=T 时,α̅_T ≈ 0,x_T ≈ N(0, I),数据被完全摧毁成纯噪声。

反向过程:去噪一步的公式

学一个神经网络 ε_θ(x_t, t) 预测当时加的噪声。给定 x_t,去噪一步:

x_{t−1} = (1/√(α_t)) · ( x_t − (β_t/√(1−α̅_t)) · ε_θ(x_t, t) )  +  σ_t · z
​

其中 σ_t 是 √(β_t) 或学到的方差,z ~ N(0, I)。这公式看着丑,但只是代数——它是从后验 q(x_{t−1}|x_t, x_0) 解出 x_{t−1},再把 x_0 用「从噪声倒推的估计」代入得到的。

训练损失:一行 MSE

L_simple = E_{x_0, t, ε} [ || ε − ε_θ( √(α̅_t)·x_0 + √(1−α̅_t)·ε,  t ) ||² ]
​

从数据采样 x_0,随机选 t,采样 ε ~ N(0, I),用闭式一步算出带噪的 x_t,然后对噪声做回归。一个损失、无 minimax、无 KL、无重参数化。推理时从 x_T ~ N(0, I) 起,从 t=T 迭代到 t=1 即可。

二、为什么扩散能工作:三条直觉

1. 去噪易,生成难。 在 t=T,数据是纯噪声,网络要解的是个平凡问题;在 t=0,网络只需清理几个像素;在中间的 t,问题难,但同一组权重从每一个噪声层级都流过大量梯度,互相强化。

2. 分数匹配的伪装。 Vincent(2011)证明:预测噪声等价于估计 ∇_x log q(x_t|x_0),即分数(score)。反向 SDE 用这个分数沿密度梯度「往上走」——一次朝高概率区域的有引导的随机游走。这正是为什么扩散与第 13 节的分数/流匹配是一家人。

3. ELBO 塌缩成简单 MSE。 完整的变分下界在每个时间步都有一个 KL 项。在 DDPM 的参数化下,这些 KL 项化简成「对噪声预测的 MSE」配上特定系数;Ho 把系数全扔掉(叫它「简单损失」),质量反而提升。

💡 预测噪声 = 学分数的数学:从 x_t = √ᾱ·x_0 + √(1−ᾱ)·ε 解出 x_0 = (x_t − √(1−ᾱ)·ε)/√ᾱ。而 ∇_{x_t} log q(x_t|x_0) = −ε/√(1−ᾱ)。所以预测 ε 与预测分数只差一个已知常数因子。这一等价让扩散模型与 Song & Ermon 的分数匹配、第 13 节的流匹配在数学上统一为「连续时间生成模型」。

三、从零实现:一维 DDPM

code/main.py 实现一个一维 DDPM。数据是双峰混合高斯;「网络」是一个吃 (x_t, t) 输出预测噪声的小 MLP;训练用那一行损失;采样迭代反向链。

步骤 1:前向调度(闭式)

betas = [1e-4 + (0.02 - 1e-4) * t / (T - 1) for t in range(T)]
alphas = [1 - b for b in betas]
alpha_bars = []
cum = 1.0
for a in alphas:
    cum *= a
    alpha_bars.append(cum)        # α̅_t = ∏_{s≤t} α_s
​

步骤 2:一步采样 x_t

def forward_sample(x0, t, alpha_bars, rng):
    a_bar = alpha_bars[t]
    eps = rng.gauss(0, 1)
    x_t = math.sqrt(a_bar) * x0 + math.sqrt(1 - a_bar) * eps   # 闭式!
    return x_t, eps
​

步骤 3:一步训练

def train_step(x0, model, alpha_bars, rng):
    t = rng.randrange(T)                       # 随机时间步
    x_t, eps = forward_sample(x0, t, alpha_bars, rng)
    eps_hat = model_forward(model, x_t, t)
    loss = (eps - eps_hat) ** 2                # 一行 MSE
    return loss, gradient_step(model, ...)
​

步骤 4:反向采样

def sample(model, alpha_bars, T, rng):
    x = rng.gauss(0, 1)                        # x_T ~ N(0,I)
    for t in range(T - 1, -1, -1):
        eps_hat = model_forward(model, x, t)
        beta_t = 1 - alphas[t]
        x = (x - beta_t / math.sqrt(1 - alpha_bars[t]) * eps_hat) / math.sqrt(alphas[t])
        if t > 0:
            x += math.sqrt(beta_t) * rng.gauss(0, 1)   # σ_t·z
    return x
​

对一个 40 时间步、24 单元 MLP 的一维问题,这套代码约 200 个 epoch 就能学到双峰分布。

时间步条件注入

网络必须知道自己在去噪第几步。两种标准做法:

• 正弦嵌入(Sinusoidal Embedding):像 Transformer 的位置编码。embed(t) = [sin(t/ω_0), cos(t/ω_0), sin(t/ω_1), ...],过一个 MLP,广播进网络。
• FiLM / 组归一化条件化:把嵌入投影成每通道的缩放/偏移(FiLM),注入每个块。

玩具代码用「正弦嵌入 → 拼接」;生产 U-Net 用 FiLM。

四、典型陷阱与修复

• 调度很关键:线性 β 是 DDPM 默认,但 cosine 调度(Nichol & Dhariwal,2021)在同等算力下 FID 更好。质量停滞时换调度。

• 时间步嵌入脆弱:把原始 t 当 float 传进去,在一维玩具上能跑,在图像上会失败——必须用正规嵌入。

• V-prediction vs ε-prediction:在极端 t(很小或很大)处,ε 信噪比差。V-prediction(v = α·ε − σ·x)更稳;SDXL、SD3、Flux 都用它。

• 无分类器引导(CFG):推理时同时算条件与无条件 ε,然后 ε_cfg = (1+w)·ε_cond − w·ε_uncond,w ≈ 3-7。详见第 08 节。

• 1000 步太多:生产用 DDIM(20-50 步)、DPM-Solver(10-20 步)或蒸馏(1-4 步)。

💡 为什么 ε-prediction 在中间 t 工作得最好? 训练时每个 t 的损失权重,在完整 ELBO 推导里是不同的系数。Ho 发现「全部当作 1」反而最好——这意味着中间 t 的去噪任务对最终质量贡献最大(那里既不平凡也非不可能),而网络恰好在那里得到最多的梯度信号。这也是为什么蒸馏/调度优化都聚焦在「如何更聪明地采样中间 t」。

五、框架对比:扩散在 2026 年的角色

扩散是通用生成骨干。流匹配(第 13 节)是 2024—2026 的竞争者,在同等质量下推理速度上通常胜出。

六、生产推理:扩散推理是「步数问题」

DDPM 论文跑 T=1000 反向步。生产里没人这么交付。每个真实推理栈从三种策略里选一种——每种都干净地对应「延迟从哪来」:

1. 更快采样器、同模型:DDIM(20-50 步)、DPM-Solver++(10-20)、UniPC(8-16)。直接替换反向循环,训练好的 ε_θ 权重不动。砍延迟 20-50 倍。

2. 蒸馏:训一个学生用更少步匹配老师:渐进蒸馏(2→1)、一致性模型(任意→1-4)、LCM、SDXL-Turbo、SD3-Turbo。再砍 5-10 倍,但需重训。

3. 缓存与编译:torch.compile(unet, mode="reduce-overhead")、TensorRT-LLM 的扩散后端、xformers/SDPA 注意力、bf16 权重。每步延迟砍约 2 倍。可与(1)(2)叠加。

对生产扩散服务器,预算对话与生产文献描述的 LLM 一致:延迟 = 步数 × 单步成本 + VAE_解码,吞吐 = 批大小 × (步数 × 单步成本)^{-1}。TTFT 很小(一步);TPOT 等价物是完整响应时间,因为图像生成对用户而言是「一次性」的。

七、可复用产物

本节产出一个训练技能文件(位于原课程 outputs/skill-diffusion-trainer.md)。

• skill-diffusion-trainer.md:输入「数据集 + 算力预算」,输出:调度(linear/cosine/sigmoid)、预测目标(ε/v/x)、步数、引导比例、采样器家族、评估协议。

八、练习

1. 简单。 把 code/main.py 的 T 从 40 改成 10。样本质量(输出的可视化直方图)如何退化?在哪个 T 下双峰结构塌掉?

2. 中等。 从 ε-prediction 切换到 v-prediction。重新推导反向步。比较最终样本质量。

3. 困难。 加入无分类器引导。条件是类别标签 c ∈ {0, 1},训练时 10% 概率丢弃条件,采样时用 ε = (1+w)·ε_cond − w·ε_uncond。测量 w = 0, 1, 3, 7 时的「命中正确条件模式」率。

本节要点回顾

1. 扩散 = 加噪 + 学去噪:前向固定马尔可夫链把数据摧毁成 N(0,I),反向学习链逐步还原。
2. 闭式前向:q(x_t|x_0) = N(√ᾱ_t·x_0, (1−ᾱ_t)·I),α̅_t = ∏(1−β_s),让训练时一步从 x_0 跳到任意 x_t。
3. 反向去噪:x_{t−1} = (1/√α_t)(x_t − β_t/√(1−ᾱ_t)·ε_θ) + σ_t·z,是从后验 q(x_{t−1}|x_t,x_0) 解出的代数结果。
4. 简单损失:L_simple = E||ε − ε_θ(√ᾱ·x_0 + √(1−ᾱ)·ε, t)||²,完整 ELBO 在 DDPM 参数化下塌缩成这一行 MSE。
5. 三条直觉:去噪易生成难、预测噪声=学分数(∇log p)、ELBO 化简成 MSE。
6. 时间步条件:正弦嵌入 / FiLM,玩具用拼接、生产用 FiLM。
7. 生产推理三策略:更快采样器(DDIM/DPM-Solver,20-50×)、蒸馏(1-4 步)、缓存与编译(2×);延迟 = 步数 × 单步成本 + VAE 解码。

下一节,我们进入扩散家族的工程化里程碑——潜在扩散与 Stable Diffusion:把 DDPM 从像素空间搬到 VAE 的低维潜空间,再注入文本条件,让扩散能在消费级 GPU 上跑起来,并彻底改变了文生图的格局。