2D扩散模型原理解析与实战：从数学基础到高效实现

1次阅读

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

扩散模型（Diffusion Models）是近年来在图像生成领域崭露头角的一类生成模型，其核心思想是通过逐步添加和去除噪声来学习数据分布。相比于 GAN 和 VAE，扩散模型具有训练稳定性高、生成质量好等优势，但也面临着计算成本高、生成速度慢等挑战。

扩散模型的主要优势包括：

训练过程稳定，不易出现模式崩溃（mode collapse）问题
生成图像质量高，细节丰富
理论框架清晰，数学基础扎实

然而，扩散模型也存在一些挑战：

采样过程需要多步迭代，生成速度较慢
计算资源消耗大，尤其是高分辨率图像生成
超参数选择对模型性能影响显著

前向扩散过程是一个逐步向数据添加高斯噪声的过程，定义为一个马尔可夫链：

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

其中 $\beta_t$ 是噪声调度参数，控制着噪声添加的速度。

通过重参数化技巧，我们可以直接计算任意时刻 $t$ 的 $x_t$：

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

其中 $\bar{\alpha}t = \prod^t(1-\beta_s)$。

反向扩散过程的目标是从噪声中逐步恢复出原始图像，其核心是学习一个网络来预测每一步的噪声：

p_\theta(x_{t-1}|x_t) = N(x_{t-1}; \mu_\theta(x_t,t), \Sigma_\theta(x_t,t))

其中 $\mu_\theta$ 和 $\Sigma_\theta$ 是由神经网络参数化的均值和方差。

噪声调度器控制着噪声添加的节奏，常见的有线性调度和余弦调度：

# 线性调度
def linear_beta_schedule(timesteps):
    beta_start = 0.0001
    beta_end = 0.02
    return torch.linspace(beta_start, beta_end, timesteps)

# 余弦调度
def cosine_beta_schedule(timesteps, s=0.008):
    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)

我们采用改进的 U -Net 架构作为噪声预测网络：

class UNet(nn.Module):
    def __init__(self, dim=64, dim_mults=(1, 2, 4, 8)):
        super().__init__()

        # 时间嵌入
        self.time_mlp = nn.Sequential(SinusoidalPositionEmbeddings(dim),
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim * 4)
        )

        # 下采样路径
        self.downs = nn.ModuleList([])
        dims = [3] + [dim * m for m in dim_mults]
        for i in range(len(dims) - 1):
            self.downs.append(Block(dims[i], dims[i+1]))

        # 上采样路径
        self.ups = nn.ModuleList([])
        dims = [dim * m for m in dim_mults][::-1]
        for i in range(len(dims) - 1):
            self.ups.append(Block(dims[i] * 2, dims[i+1]))

        self.mid_block = Block(dims[-1], dims[-1])
        self.final_conv = nn.Conv2d(dim, 3, kernel_size=1)

    def forward(self, x, t):
        # 时间嵌入
        t = self.time_mlp(t)

        # 保存跳连
        h = []

        # 下采样
        for block in self.downs:
            x = block(x, t)
            h.append(x)
            x = F.avg_pool2d(x, 2)

        # 中间层
        x = self.mid_block(x, t)

        # 上采样
        for block in self.ups:
            x = F.interpolate(x, scale_factor=2, mode='nearest')
            x = torch.cat([x, h.pop()], dim=1)
            x = block(x, t)

        return self.final_conv(x)

训练过程的核心是最小化噪声预测误差：

# 采样时间步
batch_size = x_start.shape[0]

# 均匀采样时间步
t = torch.randint(0, timesteps, (batch_size,), device=device).long()

# 添加噪声
noise = torch.randn_like(x_start)
x_noisy = q_sample(x_start, t, noise)

# 预测噪声
predicted_noise = model(x_noisy, t)

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

混合精度训练可以显著减少显存占用并加速训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    predicted_noise = model(x_noisy, t)
    loss = F.mse_loss(predicted_noise, noise)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

使用 PyTorch 的分布式数据并行 (DDP) 进行多 GPU 训练：

# 初始化分布式环境
torch.distributed.init_process_group('nccl')
local_rank = int(os.environ['LOCAL_RANK'])

# 包装模型
model = UNet().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 数据加载器
train_sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=64, sampler=train_sampler)