共计 5077 个字符，预计需要花费 13 分钟才能阅读完成。
AI 扩散与反向扩散模型入门指南：从数学基础到 PyTorch 实战
背景痛点：为什么选择扩散模型？
传统 GAN（生成对抗网络）虽然在图像生成领域取得了显著成果，但其训练过程存在两个主要问题：
- 训练不稳定：生成器和判别器的对抗训练容易导致模式崩溃（mode collapse）和训练振荡
- 生成质量难以控制：缺乏明确的概率框架，生成结果往往带有不可预测的伪影
扩散模型通过引入 渐进式噪声添加 和反向去噪过程，解决了这些问题：
- 训练目标明确：最小化噪声预测误差
- 生成过程稳定：基于马尔可夫链的逐步去噪
- 数学基础坚实：建立在变分推断的理论框架上
数学原理：扩散模型的核心思想
前向扩散过程
前向过程是一个逐步添加高斯噪声的马尔可夫链：
$$q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I})$$
其中 $\beta_t$ 是噪声调度参数。通过重参数化技巧，我们可以直接计算任意时刻 $t$ 的噪声图像：
$$x_t = \sqrt{\alpha_t}x_0 + \sqrt{1-\alpha_t}\epsilon$$
其中 $\alpha_t = \prod_{s=1}^t(1-\beta_s)$，$\epsilon \sim \mathcal{N}(0, \mathbf{I})$。
反向去噪过程
反向过程学习逐步去除噪声：
$$p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t,t), \Sigma_\theta(x_t,t))$$
目标函数是变分下界（ELBO）：
$$\mathcal{L} = \mathbb{E}{q(x|x_t))]$$}|x_0)}[-\log p_\theta(x_0|x_1) + \sum_{t=2}^T D_{KL}(q(x_{t-1}|x_t,x_0)||p_\theta(x_{t-1
PyTorch 实战：从零实现扩散模型
1. 噪声调度器实现
import torch
import math
class NoiseScheduler:
def __init__(self, num_timesteps: int, beta_start: float = 1e-4, beta_end: float = 0.02):
self.num_timesteps = num_timesteps
self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
self.alphas = 1. - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, x0: torch.Tensor, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
参数:
x0: 原始图像 [B, C, H, W]
t: 时间步 [B,]
返回:
noisy_image: 加噪后的图像
noise: 实际添加的噪声
"""
noise = torch.randn_like(x0)
sqrt_alpha_cumprod = torch.sqrt(self.alphas_cumprod[t])[:, None, None, None]
sqrt_one_minus_alpha = torch.sqrt(1. - self.alphas_cumprod[t])[:, None, None, None]
# 重参数化计算
noisy_image = sqrt_alpha_cumprod * x0 + sqrt_one_minus_alpha * noise
return noisy_image, noise2. UNet 条件化实现
import torch.nn as nn
import torch.nn.functional as F
class AttentionBlock(nn.Module):
def __init__(self, channels: int):
super().__init__()
self.norm = nn.GroupNorm(32, channels)
self.q = nn.Conv2d(channels, channels, 1)
self.k = nn.Conv2d(channels, channels, 1)
self.v = nn.Conv2d(channels, channels, 1)
self.proj_out = nn.Conv2d(channels, channels, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, C, H, W = x.shape
h = self.norm(x)
q = self.q(h).view(B, C, -1) # [B, C, H*W]
k = self.k(h).view(B, C, -1) # [B, C, H*W]
v = self.v(h).view(B, C, -1) # [B, C, H*W]
attn = torch.bmm(q.permute(0, 2, 1), k) # [B, H*W, H*W]
attn = F.softmax(attn * (C ** -0.5), dim=2)
out = torch.bmm(v, attn.permute(0, 2, 1)) # [B, C, H*W]
out = out.view(B, C, H, W)
return x + self.proj_out(out)
class UNet(nn.Module):
def __init__(self, in_channels: int = 3):
super().__init__()
# 下采样路径
self.down1 = nn.Sequential(nn.Conv2d(in_channels, 64, 3, padding=1),
nn.GroupNorm(32, 64),
nn.SiLU(),
AttentionBlock(64)
)
# 中间层
self.mid = nn.Sequential(nn.Conv2d(64, 128, 3, padding=1),
nn.GroupNorm(32, 128),
nn.SiLU(),
AttentionBlock(128)
)
# 上采样路径 (包含 skip connection)
self.up1 = nn.Sequential(nn.Conv2d(192, 64, 3, padding=1), # 128+64=192
nn.GroupNorm(32, 64),
nn.SiLU(),
AttentionBlock(64)
)
# 最终输出层
self.out = nn.Conv2d(64, in_channels, 3, padding=1)
def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
# t_emb 是时间步的嵌入表示
h1 = self.down1(x)
h2 = self.mid(h1)
h = self.up1(torch.cat([h2, h1], dim=1)) # skip connection
return self.out(h)3. DDPM 采样过程
class DiffusionModel(nn.Module):
def __init__(self, unet: nn.Module, scheduler: NoiseScheduler):
super().__init__()
self.unet = unet
self.scheduler = scheduler
def sample(self, shape: tuple[int], temp: float = 1.0) -> torch.Tensor:
"""
参数:
shape: 生成图像的形状 [B, C, H, W]
temp: 温度系数，控制采样随机性
返回:
生成的图像
"""
device = next(self.parameters()).device
x_t = torch.randn(shape, device=device)
for t in reversed(range(self.scheduler.num_timesteps)):
# 预测噪声
t_batch = torch.full((shape[0],), t, device=device)
pred_noise = self.unet(x_t, t_batch)
# 计算均值方差
alpha_t = self.scheduler.alphas[t]
alpha_cumprod_t = self.scheduler.alphas_cumprod[t]
beta_t = self.scheduler.betas[t]
if t > 0:
noise = torch.randn_like(x_t) * temp # 应用温度系数
else:
noise = 0
# 反向过程更新
x_t = (x_t - (1 - alpha_t)/torch.sqrt(1 - alpha_cumprod_t) * pred_noise) / torch.sqrt(alpha_t)
x_t = x_t + torch.sqrt(beta_t) * noise
return x_t性能优化技巧
显存管理策略
- 梯度累积：当显存不足时，可以通过多次前向传播积累梯度再更新
batch_size = 32
accum_steps = 4 # 实际 batch_size=128
optimizer.zero_grad()
for i, (x, _) in enumerate(dataloader):
# 计算损失
t = torch.randint(0, num_timesteps, (x.shape[0],))
noisy_x, noise = scheduler.add_noise(x, t)
pred_noise = model(noisy_x, t)
loss = F.mse_loss(pred_noise, noise)
# 梯度累积
loss = loss / accum_steps
loss.backward()
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()- 混合精度训练：使用 AMP 自动混合精度
scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
t = torch.randint(0, num_timesteps, (x.shape[0],))
noisy_x, noise = scheduler.add_noise(x, t)
pred_noise = model(noisy_x, t)
loss = F.mse_loss(pred_noise, noise)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()常见错误及解决方案
- 训练发散问题
- 现象：损失值突然变为 NaN
- 原因：噪声调度参数 $\beta_t$ 设置不当，导致数值不稳定
-
解决：使用 cosine 调度代替线性调度
-
生成质量差
- 现象：生成图像模糊或颜色异常
- 原因：UNet 容量不足或训练 epoch 不够
-
解决：增加模型通道数或延长训练时间
-
采样速度慢
- 现象：生成一张图需要几分钟
- 原因：采样步数过多
- 解决：使用 DDIM 加速采样或减少步数
延伸思考：文本条件生成
要将模型扩展到文本条件生成，可以：
- 添加文本编码器（如 CLIP 或 BERT）
- 在 UNet 中引入交叉注意力层
- 使用 Classifier-Free Guidance 技术增强文本对齐
# 在 UNet 中添加文本条件
class TextConditionedUNet(UNet):
def __init__(self, text_dim: int):
super().__init__()
self.text_proj = nn.Linear(text_dim, 64)
def forward(self, x: torch.Tensor, t_emb: torch.Tensor, text_emb: torch.Tensor):
text_emb = self.text_proj(text_emb)[..., None, None] # [B, 64, 1, 1]
return super().forward(x + text_emb, t_emb)总结
本文详细介绍了扩散模型的数学原理和 PyTorch 实现，包括：
- 前向扩散和反向去噪过程的数学推导
- 完整的噪声调度器和 UNet 实现
- 采样过程中的温度系数控制
- 显存优化和训练技巧
- 常见问题的解决方案
扩散模型作为当前最先进的生成模型，在图像生成、音频合成等领域展现出强大能力。希望通过本文，初学者能够掌握其核心思想并实现自己的扩散模型。
正文完
