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

1次阅读

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

扩散模型（Diffusion Models）近年来在生成式 AI 领域大放异彩，尤其是图像生成任务中表现优异。然而，其背后的数学原理（如马尔可夫链、变分推断）常常让开发者感到困惑。许多人在初次接触时会被其复杂的数学推导吓退，或者在实际实现时遇到训练不稳定、计算开销大等问题。

本文将用图解的方式拆解扩散模型的核心机制，并提供一个 PyTorch 实现，帮助开发者从原理到实践全面掌握这一技术。

前向扩散过程可以理解为逐步向数据添加噪声的过程。假设我们有一张清晰的图片，每一步都对其添加少量高斯噪声，经过足够多的步骤后，图片会变成纯噪声。这一过程可以用马尔可夫链来描述：

初始状态：清晰图片 $x_0$
第 $t$ 步的状态 $x_t$ 由前一步的状态 $x_{t-1}$ 决定，满足 $q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I})$
其中 $\beta_t$ 是噪声调度参数，控制每一步添加的噪声量

反向去噪过程是前向过程的逆过程，目的是从噪声中逐步恢复出原始数据。这一过程通过训练一个神经网络来预测每一步的噪声：

给定噪声图片 $x_t$，网络预测噪声 $\epsilon_\theta(x_t, t)$
利用预测的噪声，可以逐步去除噪声，得到更清晰的图片
反向过程的每一步可以表示为 $p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$

GAN：训练不稳定，容易出现模式崩溃
VAE：生成质量通常不如扩散模型
扩散模型 ：训练稳定，生成质量高，但计算开销较大

下面是一个基础的扩散模型实现，包含噪声预测网络和训练循环。

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

class NoisePredictor(nn.Module):
    """噪声预测网络"""
    def __init__(self):
        super().__init__()
        # 网络结构定义
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.time_embed = nn.Linear(1, 128)  # 时间步嵌入
        self.out = nn.Conv2d(128, 3, kernel_size=3, padding=1)

    def forward(self, x, t):
        # 前向传播
        h = F.relu(self.conv1(x))
        t_emb = self.time_embed(t.view(-1, 1))
        h = h + t_emb.view(-1, 128, 1, 1)
        h = F.relu(self.conv2(h))
        return self.out(h)

# 训练循环
def train_diffusion(model, dataloader, optimizer, num_steps=1000):
    model.train()
    for epoch in range(num_epochs):
        for x0 in dataloader:
            # 随机采样时间步
            t = torch.randint(0, num_steps, (x0.size(0),))

            # 前向扩散过程
            noise = torch.randn_like(x0)
            xt = sqrt_alphas_cumprod[t] * x0 + sqrt_one_minus_alphas_cumprod[t] * noise

            # 预测噪声
            pred_noise = model(xt, t)

            # 计算损失
            loss = F.mse_loss(pred_noise, noise)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

扩散模型的主要性能瓶颈在于需要多次迭代（通常 50-1000 步）才能生成一张图片。以下是一些优化建议：