0基础学扩散模型：从数学原理到PyTorch实战

1次阅读

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

扩散模型的核心思想其实很有趣——想象一杯清水滴入墨水，墨水会逐渐扩散直到整杯水变得均匀。这个过程反过来，就是从混沌中恢复出清晰图像的神奇能力。让我们拆解两个关键概念：

马尔可夫链假设：扩散模型认为每一步的噪声添加只依赖前一步的状态，就像多米诺骨牌，当前状态只被前一块影响。数学表达为：
$$q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I})$$
其中 $\beta_t$ 是预设的噪声计划表（noise schedule）。
变分下界(ELBO)：模型通过最大化这个下界来学习逆向过程。简单说，就是让模型预测的 ” 去噪路径 ” 尽可能接近真实的扩散路径：
$$\text{ELBO} = \mathbb{E}q[\log p\theta(x_0|T)] – \sum_{t=1}^T D_{KL}(q(x_{t-1}|x_t,x_0)||p_\theta(x_{t-1}|x_t))$$

DDPM：
严格遵循马尔可夫过程
需要完整 T 步采样（通常 T =1000）
训练稳定但采样慢
DDIM：
非马尔可夫过程，允许跳步采样
采样步数可缩减至 20-50 步
需要更精细的超参调节

选择建议：优先用 DDPM 打基础，理解原理后再尝试 DDIM 加速。

import matplotlib.pyplot as plt
import torch

def visualize_diffusion(image, betas, steps=5):
    alphas = 1 - betas
    alpha_bars = torch.cumprod(alphas, dim=0)

    plt.figure(figsize=(15,3))
    for t in range(0, len(betas), len(betas)//steps):
        noise = torch.randn_like(image)
        noisy_img = torch.sqrt(alpha_bars[t])*image + torch.sqrt(1-alpha_bars[t])*noise
        plt.subplot(1, steps+1, t//(len(betas)//steps)+1)
        plt.imshow(noisy_img.squeeze().cpu().numpy(), cmap='gray')

class UNetBlock(nn.Module):
    def __init__(self, in_c, out_c, time_emb_dim):
        super().__init__()
        self.time_mlp = nn.Sequential(nn.Linear(time_emb_dim, out_c),
            nn.SiLU(),
            nn.Linear(out_c, out_c)
        )
        self.conv = nn.Sequential(nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.SiLU())

    def forward(self, x, t):
        h = self.conv(x)
        time_emb = self.time_mlp(t)[:,:,None,None]
        return h + time_emb

# 使用 EMA 包装器
class EMA:
    def __init__(self, beta=0.9999):
        self.beta = beta
        self.shadow = {}

    def register(self, module):
        for name, param in module.named_parameters():
            self.shadow[name] = param.data.clone()

    def update(self, module):
        for name, param in module.named_parameters():
            self.shadow[name] = self.beta * self.shadow[name] + (1-self.beta) * param.data

def cosine_beta_schedule(timesteps, s=0.008):
    """https://arxiv.org/abs/2102.09672"""
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0, 0.999)

现象：深层网络参数更新幅度趋近于 0
解决方案：
使用残差连接（如 U -Net 的 skip connection）
采用梯度裁剪（torch.nn.utils.clip_grad_norm_）

混合精度训练（AMP）：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    loss = model(x, t)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()