Diffusion Models from Scratch

从零实现 DDPM，用代码理解每一个组件。目标不是复现 SOTA，而是把扩散模型的损失推导里的公式一行行落地。

参考来源

Bishop & Bishop (2024)《Deep Learning Foundations and Concepts》第 20 章 — 公式与算法伪代码的权威来源（Algorithm 20.1 训练 / 20.2 采样）
Ho et al. (2020) DDPM — 原始算法
Ronneberger et al. (2015) U-Net — 网络架构
本文符号与 DDPM、diffusion process 笔记一致： $x_t$ 为带噪数据， $\epsilon_\theta$ 为噪声预测网络， $\alpha_t=1-\beta_t$ ， $\bar\alpha_t=\prod_{s=1}^t\alpha_s$ 。

1. 噪声调度（Noise Schedule）

$\beta_t$ 是每步加噪强度，DDPM 采用线性调度（ $\beta_1=10^{-4}$ 到 $\beta_T=0.02$ ）。据此预计算 $\bar\alpha_t$ ：

import torch

T = 1000
betas = torch.linspace(1e-4, 0.02, T)      # β_t：线性调度
alphas = 1.0 - betas                        # α_t = 1 - β_t
alpha_bar = torch.cumprod(alphas, dim=0)    # ᾱ_t = ∏ α_s

对应 diffusion process 的 $q(x_t\mid x_0)=\mathcal N(x_t;\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I)$ 。

2. 前向加噪（Diffusion Kernel）

利用闭式解，一步从 $x_0$ 采样任意 $x_t$ ，无需逐步加噪：

def q_sample(x0, t, eps):
    """Bishop (20.8)：z_t = √ᾱ_t·x + √(1-ᾱ_t)·ε"""
    a_bar = alpha_bar[t][:, None, None, None]
    return torch.sqrt(a_bar) * x0 + torch.sqrt(1 - a_bar) * eps

3. U-Net 噪声预测网络

网络 $\epsilon_\theta(x_t, t)$ 预测「加到 $x_0$ 上的总噪声」。U-Net 由编码器、中间块、解码器组成，靠 skip connection 保留细节，时间步 $t$ 通过正弦嵌入注入每一层。

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

def sinusoidal_embedding(t, dim):
    """时间步正弦位置编码"""
    half = dim // 2
    freqs = torch.exp(-torch.arange(half).float() * (torch.log(torch.tensor(10000.0)) / half))
    args = t[:, None].float() * freqs[None, :]
    return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

class Block(nn.Module):
    """残差块，含时间嵌入注入"""
    def __init__(self, in_ch, out_ch, t_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_ch)
        self.t_proj = nn.Linear(t_dim, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.res = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x, t_emb):
        h = F.silu(self.norm1(self.conv1(x)))
        h = h + self.t_proj(F.silu(t_emb))[:, :, None, None]   # 时间注入
        h = F.silu(self.norm2(self.conv2(h)))
        return h + self.res(x)

class UNet(nn.Module):
    def __init__(self, in_ch=1, base_ch=64, t_dim=128):
        super().__init__()
        self.t_emb = nn.Sequential(
            nn.Linear(t_dim, t_dim * 4), nn.SiLU(), nn.Linear(t_dim * 4, t_dim)
        )
        # 编码器（下采样）
        self.enc1 = Block(in_ch, base_ch, t_dim)
        self.enc2 = Block(base_ch, base_ch * 2, t_dim)
        self.enc3 = Block(base_ch * 2, base_ch * 4, t_dim)
        # 中间块
        self.mid = Block(base_ch * 4, base_ch * 4, t_dim)
        # 解码器（上采样 + skip connection）
        self.dec3 = Block(base_ch * 4 + base_ch * 4, base_ch * 2, t_dim)
        self.dec2 = Block(base_ch * 2 + base_ch * 2, base_ch, t_dim)
        self.dec1 = Block(base_ch + base_ch, base_ch, t_dim)
        self.out = nn.Conv2d(base_ch, in_ch, 3, padding=1)
        self.down = nn.MaxPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)

    def forward(self, x, t):
        t_emb = self.t_emb(sinusoidal_embedding(t, 128))
        e1 = self.enc1(x, t_emb)
        e2 = self.enc2(self.down(e1), t_emb)
        e3 = self.enc3(self.down(e2), t_emb)
        m = self.mid(self.down(e3), t_emb)
        d3 = self.dec3(torch.cat([self.up(m), e3], dim=1), t_emb)
        d2 = self.dec2(torch.cat([self.up(d3), e2], dim=1), t_emb)
        d1 = self.dec1(torch.cat([self.up(d2), e1], dim=1), t_emb)
        return self.out(d1)

4. 训练循环（Bishop Algorithm 20.1）

每步：采样数据点 $x$ 、采样时间步 $t$ 、采样噪声 $\epsilon$ ，构造带噪 $x_t$ ，让网络预测 $\epsilon$ ：

model = UNet()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for x0 in dataloader:                       # x0 ~ 训练数据
    t = torch.randint(0, T, (x0.shape[0],)) # 采样时间步
    eps = torch.randn_like(x0)              # 采样噪声 ε ~ N(0,I)
    xt = q_sample(x0, t, eps)               # 构造带噪样本 x_t
    pred = model(xt, t)                     # 网络预测噪声
    loss = F.mse_loss(pred, eps)            # L_simple = ‖ε - ε_θ‖²
    opt.zero_grad(); loss.backward(); opt.step()

对应扩散模型的损失推导的 $L_{simple}=\mathbb E_{t,x_0,\epsilon}\|\epsilon-\epsilon_\theta(x_t,t)\|^2$ ，即 Bishop (20.40)。

5. 采样（Bishop Algorithm 20.2）

从纯噪声 $x_T\sim\mathcal N(0,I)$ 出发，逐步去噪 $T$ 步：

@torch.no_grad()
def sample(model, shape, T=T):
    x = torch.randn(shape)                  # x_T ~ N(0,I)
    for t in reversed(range(T)):
        t_batch = torch.full((shape[0],), t)
        pred_eps = model(x, t_batch)
        # Bishop (20.36)：μ = (1/√α_t)(x_t - (β_t/√(1-ᾱ_t))·ε_θ)
        mu = (1 / torch.sqrt(alphas[t])) * (x - (betas[t] / torch.sqrt(1 - alpha_bar[t])) * pred_eps)
        if t > 0:
            x = mu + torch.sqrt(betas[t]) * torch.randn_like(x)   # Bishop (20.41)
        else:
            x = mu                          # 最后一步不加噪，输出干净样本
    return x

[!note] 采样方差的两种选择这里噪声系数用 $\sqrt{\beta_t}$ （Bishop 教材的简化版）。DDPM 原论文用后验方差 $\tilde\beta_t=\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t$ ，二者效果接近。当 $t$ 很小或 $\beta_t$ 很小时差异可忽略。

6. 加速与条件生成

加速：把采样换成 DDIM 的确定性跳步，可 50 步代替 1000 步；
条件引导：采样时用 CFG 放大条件信号；
高分辨率：把 $x$ 换成 latent，即 Latent Diffusion 路线。