扩散模型(Diffusion Models)图解：从数学原理到PyTorch实战

1次阅读

共计 2250 个字符，预计需要花费 6 分钟才能阅读完成。

扩散模型的核心思想是通过逐步添加噪声（前向过程）和逐步去噪（反向过程）来实现数据生成。从信息论的角度来看，前向过程是一个熵增的过程，数据逐渐失去结构信息；而反向过程则是熵减的过程，模型学习如何从噪声中恢复出有意义的数据结构。

与 GAN 和 VAE 相比，扩散模型有几个显著差异：

GAN 通过对抗训练直接学习数据分布，但训练不稳定且容易模式崩溃
VAE 通过变分推断近似数据分布，但生成的样本往往模糊
扩散模型通过固定的前向过程和可学习的反向过程，既能稳定训练，又能生成高质量的样本

前向过程定义为逐步添加高斯噪声的马尔可夫链：

q(x_t|x_{t-1}) = N(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_tI)

其中 β_t 是噪声调度器，控制着每一步添加的噪声量。通过重参数化技巧，我们可以直接计算任意时间步 t 的 x_t：

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

其中 α_t=1-β_t，而̄α_t=∏_{s=1}^tα_s。

反向过程使用 UNet 结构来预测噪声。UNet 的关键组件包括：

下采样路径：逐步提取高层语义特征
上采样路径：逐步恢复空间分辨率
skip connection：将低级特征与高级特征融合，帮助梯度传播

skip connection 的作用尤为重要：

缓解梯度消失问题，使深层网络也能有效训练
保留空间细节信息，提高生成质量
实现多尺度特征融合，增强模型表达能力

扩散模型的训练目标是最小化变分下界：

L_{ELBO} = E_q[-log p_θ(x_0|x_1)] + \sum_{t=2}^T D_{KL}(q(x_{t-1}|x_t,x_0)||p_θ(x_{t-1}|x_t))

实践中，我们通常使用简化的目标函数：

L_{simple} = E_{t,x_0,\epsilon}[||\epsilon - \epsilon_θ(x_t,t)||^2]

以下是完整的 DDPM 类实现：

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

class DDPM(nn.Module):
    def __init__(self, model, T=1000, beta_start=1e-4, beta_end=0.02):
        super().__init__()
        self.model = model
        self.T = T
        # 线性调度器
        self.register_buffer('betas', torch.linspace(beta_start, beta_end, T))
        self.register_buffer('alphas', 1. - self.betas)
        self.register_buffer('alphas_bar', torch.cumprod(self.alphas, dim=0))

    @torch.compile
    def forward(self, x, t):
        # 混合精度训练
        with torch.cuda.amp.autocast():
            return self.model(x, t)

    def loss(self, x0):
        # 随机选择时间步
        t = torch.randint(0, self.T, (x0.shape[0],), device=x0.device)
        # 生成噪声
        epsilon = torch.randn_like(x0)
        # 计算加噪后的样本
        x_t = torch.sqrt(self.alphas_bar[t])[:,None,None,None] * x0 + \
              torch.sqrt(1 - self.alphas_bar[t])[:,None,None,None] * epsilon
        # 预测噪声
        epsilon_pred = self(x_t, t)
        # 计算损失
        return F.mse_loss(epsilon_pred, epsilon)