第4章：去噪扩散概率模型（DDPM）

4.3 损失函数与训练细节

理论推导

DDPM的核心训练目标是最小化负对数似然的变分上界（ELBO），其损失函数可分解为以下三部分：

1. 正向过程损失（常数项，通常忽略）：

   $$L_T = D_{KL}(q(x_T|x_0) \parallel p(x_T))$$

2. 逆向过程损失（关键优化项）：

   $$L_{t-1} = \mathbb{E}_q\left[ D_{KL}(q(x_{t-1}|x_t,x_0) \parallel p_\theta(x_{t-1}|x_t)) \right]$$

3. 重构损失：

   $$L_0 = -\log p_\theta(x_0|x_1)$$

通过推导可得简化后的均方误差损失（实际实现形式）：

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

其中 $x_t = \sqrt{\alpha_t}x_0 + \sqrt{1-\alpha_t}\epsilon$，$\epsilon$为标准高斯噪声。

训练算法

def train_step(model, x0, optimizer):
    # 1. 随机采样时间步
    t = torch.randint(0, T, (x0.shape[0],))
    
    # 2. 生成带噪样本
    alpha_bar = compute_alpha_bar(t)  # 预计算的噪声调度系数
    eps = torch.randn_like(x0)
    xt = torch.sqrt(alpha_bar) * x0 + torch.sqrt(1-alpha_bar) * eps
    
    # 3. 预测噪声
    eps_pred = model(xt, t)
    
    # 4. 计算损失
    loss = F.mse_loss(eps_pred, eps)
    
    # 5. 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss

关键实现细节

1. 噪声调度策略：

   • 线性调度：$\beta_t$从$\beta_1=10^{-4}$到$\beta_T=0.02$线性增长
   • 余弦调度（改进版）：$\alpha_t = \frac{f(t)}{f(0)}$, $f(t)=\cos(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2})^2$

2. 模型架构选择：

   • U-Net结构（含残差连接）
   • 时间步$t$通过正弦位置编码嵌入
   • 自注意力机制（用于高分辨率生成）

3. 训练技巧：

   • 混合精度训练（FP16）
   • 梯度裁剪（防止梯度爆炸）
   • EMA（指数移动平均）模型参数

案例研究：CIFAR-10训练

训练曲线示例：

图：噪声预测损失随训练步数的变化趋势

数学补充

逆向过程真实后验分布：

$$q(x_{t-1}|x_t,x_0) = \mathcal{N}(\tilde{\mu}_t(x_t,x_0), \tilde{\beta}_t I)$$

其中：

$$\tilde{\mu}_t = \frac{\sqrt{\alpha_{t-1}}\beta_t}{1-\alpha_t}x_0 + \frac{\sqrt{\alpha_t}(1-\alpha_{t-1})}{1-\alpha_t}x_t$$

$$\tilde{\beta}_t = \frac{1-\alpha_{t-1}}{1-\alpha_t}\beta_t$$