扩散模型（Diffusion Model）原理与图像生成实战

反向扩散过程与模型训练

数学原理

反向扩散是前向扩散的逆过程。目标是训练一个神经网络 ε_θ，从 x_t 中预测出添加的噪声 ε，然后逐步去除噪声，还原出 x₀。

核心公式为：

p_θ(x_{t-1}|x_t) = N(x_{t-1}; μ_θ(x_t, t), Σ_θ(x_t, t))

均值 μ_θ 由神经网络预测的噪声计算得到，方差 Σ_θ 通常设置为固定值。

训练目标是最小化预测噪声和真实噪声之间的均方误差：

L(θ) = E_{x_0, ε, t}[||ε - ε_θ(x_t, t)||²]

构建噪声预测网络

噪声预测网络是一个卷积神经网络。输入是 x_t 和步数 t 的嵌入，输出是预测的噪声 ε_θ。

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

# 位置嵌入层：将步数 t 转换为高维向量
class PositionalEncoding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = np.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None] * emb[None, :]
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
        return emb

# 残差块
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, time_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.time_mlp = nn.Linear(time_dim, out_channels)
        self.skip = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()

    def forward(self, x, t):
        h = F.relu(self.bn1(self.conv1(x)))
        h += self.time_mlp(t)[:, :, None, None]
        h = F.relu(self.bn2(self.conv2(h)))
        return h + self.skip(x)

# UNet 网络结构
class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, time_dim=256):
        super().__init__()
        self.time_dim = time_dim
        self.pos_encoding = PositionalEncoding(time_dim)

        # 下采样路径
        self.down1 = ResidualBlock(in_channels, 64, time_dim)
        self.down2 = ResidualBlock(64, 128, time_dim)
        self.down3 = ResidualBlock(128, 256, time_dim)
        self.pool = nn.MaxPool2d(2)

        # 瓶颈层
        self.bottleneck = ResidualBlock(256, 256, time_dim)

        # 上采样路径
        self.up1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.res_up1 = ResidualBlock(256, 128, time_dim)
        self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.res_up2 = ResidualBlock(128, 64, time_dim)
        self.up3 = nn.ConvTranspose2d(64, 32, 2, stride=2)  # 补充缺失的上采样层
        self.res_up3 = ResidualBlock(64, 32, time_dim)

        # 输出层
        self.out = nn.Conv2d(32, out_channels, 1)

    def forward(self, x, t):
        t = self.pos_encoding(t)

        # 下采样
        h1 = self.down1(x, t)
        h2 = self.down2(self.pool(h1), t)
        h3 = self.down3(self.pool(h2), t)

        # 瓶颈层
        bottleneck = self.bottleneck(self.pool(h3), t)

        # 上采样
        up1 = self.up1(bottleneck)
        up1 = torch.cat([up1, h3], dim=1)
        up1 = self.res_up1(up1, t)

        up2 = self.up2(up1)
        up2 = torch.cat([up2, h2], dim=1)
        up2 = self.res_up2(up2, t)

        up3 = self.up3(up2)
        up3 = torch.cat([up3, h1], dim=1)
        up3 = self.res_up3(up3, t)

        return self.out(up3)

model = UNet().to(device)
print(model)

模型训练流程

from torch.utils.data import DataLoader
from torch.optim import Adam

dataset = MNIST(root='./data', train=True, download=True, transform=ToTensor())
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

optimizer = Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()

def train_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    for x0, _ in dataloader:
        x0 = x0.to(device)
        batch_size = x0.shape[0]
        
        # 随机采样步数 t
        t = torch.randint(0, T, (batch_size,), device=device)
        
        # 前向扩散生成 xt 和真实噪声
        xt, eps_true = forward_diffusion(x0, t, device)
        
        # 模型预测噪声
        eps_pred = model(xt, t)
        
        # 计算损失
        loss = criterion(eps_pred, eps_true)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    
    return total_loss / len(dataloader)

epochs = 50
for epoch in range(epochs):
    loss = train_epoch(model, dataloader, optimizer, criterion, device)
    print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss:.4f}")

torch.save(model.state_dict(), 'ddpm_mnist.pth')

实战：基于 DDPM 的图像生成

反向扩散采样过程

训练完成后，从随机噪声出发，通过反向扩散过程逐步去除噪声，生成新图像。

def sample(model, batch_size, device):
    model.eval()
    xt = torch.randn((batch_size, 1, 28, 28)).to(device)
    
    with torch.no_grad():
        for t in reversed(range(1, T)):
            t_tensor = torch.tensor([t], device=device).repeat(batch_size)
            eps_pred = model(xt, t_tensor)
            
            alpha_t = alpha[t].to(device)
            alpha_bar_t = alpha_bar[t].to(device)
            alpha_bar_t_1 = alpha_bar[t-1].to(device)
            beta_t = beta[t].to(device)
            
            mean = (1 / torch.sqrt(alpha_t)) * (xt - (beta_t / torch.sqrt(1 - alpha_bar_t)) * eps_pred)
            
            if t == 1:
                variance = 0
            else:
                variance = beta_t
            
            z = torch.randn_like(xt).to(device) if t > 1 else torch.zeros_like(xt).to(device)
            xt = mean + torch.sqrt(variance) * z
    
    xt = torch.clamp(xt, 0, 1)
    return xt

model.load_state_dict(torch.load('ddpm_mnist.pth'))
generated_images = sample(model, batch_size=16, device=device)

plt.figure(figsize=(8, 8))
for i in range(16):
    img = generated_images[i].squeeze().cpu().detach().numpy()
    plt.subplot(4, 4, i + 1)
    plt.imshow(img, cmap='gray')
    plt.axis('off')
plt.show()

优化技巧

1. 余弦噪声调度：将 beta 的变化从线性改为余弦，可提升生成质量。
2. 分类引导：在生成过程中加入类别信息，实现指定类别的图像生成。
3. 更大网络架构：如 UNet++、Attention UNet，提升特征捕捉能力。

发展与展望

经典变体包括 DDIM（加速生成）、Stable Diffusion（潜在空间扩散）等。应用场景涵盖艺术创作、图像修复、文生图（如 Midjourney）及语音合成等领域。

总结

扩散模型通过前向加噪和反向去噪实现了稳定训练，解决了 GAN 的不稳定性问题。前向扩散固定，反向扩散训练网络预测噪声。基于 DDPM 可实现简单图像生成，优化调度和架构能进一步提升效果。Stable Diffusion 等已成为主流工具，广泛应用于创意设计。