内容阅读
扩散模型：从噪声到图像——DDPM、潜在扩散与 Stable Diffusion 的完整路径
本文迁移自 mindcarver/91ai · 原始位置 docs/machine learning/stage8 advanced/02 diffusion models.md · 由 @阿新聊ai 整理。 扩散模型：从噪声到图像——DDPM、潜在扩散与 Stable Diffusion 的完整路径 <! 调研来源...
本文迁移自 mindcarver/91ai · 原始位置
docs/machine-learning/stage8-advanced/02-diffusion-models.md· 由 @阿新聊ai 整理。
扩散模型：从噪声到图像——DDPM、潜在扩散与 Stable Diffusion 的完整路径
TL;DR： 扩散模型通过"加噪声→学习去噪声"的过程来生成图像。本文从数学基础（前向扩散、反向去噪的推导）开始，经过 DDPM 的训练与采样算法，到 Stable Diffusion 的潜在空间架构，再到 DiT 和 Flow Matching 的前沿发展，给出从公式到代码的完整路径。
为什么这很重要
2022年8月，Stable Diffusion 开源发布。任何人只要有一张消费级显卡，就能用文字描述生成高质量图像。这标志着图像生成从实验室走向了大众。
但扩散模型的影响远不止于此：
- 视频生成：OpenAI 的 Sora、Runway Gen-3 都基于扩散模型
- 音频生成：OpenAI 的 Jukebox、Google 的 AudioLDM
- 3D 生成：DreamFusion、Point-E 用扩散模型生成三维模型
- 蛋白质设计：RFdiffusion 用扩散模型设计新蛋白质结构
- 药物发现：基于扩散的分子生成正在加速新药研发
扩散模型为什么能取代 GAN 成为主流？因为 GAN 训练不稳定、模式崩塌（生成的图像多样性差）是固有问题，而扩散模型的理论框架保证了训练的稳定性和生成质量的可控性。
核心概念
扩散模型的两大过程
┌────────────────────────────────────────────────────────────┐
│ 扩散模型的全景图 │
│ │
│ 前向过程 (Forward Process) │
│ ┌────┐ ┌────┐ ┌────┐ ┌────┐ │
│ │ x_0 │───▶│ x_1 │───▶│ x_2 │──...──▶│ x_T │ │
│ │原图 │ │+噪声│ │+更多│ │纯噪声│ │
│ └────┘ └────┘ └────┘ └────┘ │
│ ▲ │ │
│ │ ▼ │
│ │ 反向过程 (Reverse Process) │
│ │ ┌────┐ ┌────┐ ┌────┐ │
│ └─────────│x̂_0 │◀───│x̂_1 │◀───│x̂_T │ │
│ 生成图像 │去噪 │ │去噪 │ │随机噪声│ │
│ └────┘ └────┘ └────┘ │
│ │
│ 关键：每一步由神经网络 ε_θ(x_t, t) 预测噪声并去除 │
│ │
│ 训练：L_simple = E[‖ε - ε_θ(x_t, t)‖²] │
│ 采样：从 x_T ~ N(0,I) 开始，逐步去噪得到 x_0 │
└────────────────────────────────────────────────────────────┘
DDPM 的关键数学概念
| 概念 | 符号 | 含义 |
|---|---|---|
| 噪声调度 | β₁, β₂, ..., β_T | 控制每一步加噪的强度 |
| 累积乘积 | α_t = 1 - β_t, ᾱ_t = ∏α_s | 简化任意时刻的采样 |
| 前向过程 | q(x_t | x_0) = N(√ᾱ_t·x_0, (1-ᾱ_t)·I) | 任意一步的闭式解 |
| 反向过程 | p_θ(x_{t-1} | x_t) | 由神经网络近似 |
| 训练目标 | L_simple = E[||ε - ε_θ(x_t, t)||²] | 预测噪声而非原始图像 |
工作原理（简化的心智模型）
用"修复古画"来理解扩散模型
想象你是一位古画修复师，面对一幅被严重损坏的画。
前向过程——画是如何变损坏的：一幅完好的画（x_0），经过风吹日晒（加噪声），逐渐变得模糊不清。经过 T 步损坏后，变成了一团看不出原样的斑点（纯高斯噪声）。
关键洞察：如果我们知道每一步损坏的规律（噪声调度 β_t），那么从完好的画到纯噪声的每一步都可以精确计算，不需要真的走 T 步——这就是"重参数化技巧"。
反向过程——修复画作：你训练了一个助手（神经网络），给它看部分损坏的画，让它预测"这幅画上哪些是损坏的痕迹（噪声）"。然后你把预测的噪声去掉，画就恢复了一点。重复 T 步，从纯噪声恢复出完整的画。
训练：你有很多完好的画（训练数据），随机选一幅，随机加一些噪声，让助手猜加的是什么噪声。猜对了就奖励（损失减小），猜错了就纠正。经过大量练习，助手变得非常擅长识别噪声。
生成新画：现在你不需要原有的画了。随机撒一团噪声给助手，让它一步步去除噪声——出来的就是一幅全新的、从未存在过的画。
潜在扩散（Stable Diffusion）的改进
如果画太大了（比如 1024×1024 像素），直接修复太慢。于是你先用压缩工具把画缩小到 64×64 的小缩略图（潜在空间），在小图上做修复，最后再放大还原。
- VAE Encoder：压缩工具，把大图压缩到小图
- U-Net/DiT：在小图上工作的修复助手
- VAE Decoder：放大还原工具，把小图变回大图
- CLIP Text Encoder：翻译官，把你的文字描述翻译成修复助手能理解的条件信号
工作原理（详细机制）
一、前向扩散过程的数学推导
前向过程是一个马尔可夫链，每一步给图像添加少量高斯噪声：
q(x_t | x_{t-1}) = N(x_t; √(1-β_t) · x_{t-1}, β_t · I)
其中 β_t 是噪声方差调度（noise schedule），通常从 β₁ = 10⁻⁴ 线性增长到 β_T = 0.02。
重参数化技巧——任意时刻的闭式解：
定义 α_t = 1 - β_t，ᾱ_t = ∏(s=1 to t) α_s，利用高斯分布的叠加性质，可以递归推导：
x_1 = √α₁ · x_0 + √(1-α₁) · ε₀
x_2 = √α₂ · x_1 + √(1-α₂) · ε₁
= √(α₂α₁) · x_0 + √(1-α₂α₁) · ε
...
x_t = √ᾱ_t · x_0 + √(1-ᾱ_t) · ε
其中 ε ~ N(0, I)。这意味着给定原始图像 x_0，我们可以在一步之内采样任意时刻 t 的噪声图像——不需要真的跑 t 步。
"""
前向扩散过程的实现
演示如何从原始图像一步加噪到任意时刻
"""
import torch
import torch.nn.functional as F
import numpy as np
from typing import Tuple
class ForwardDiffusion:
"""
前向扩散过程
核心公式：x_t = sqrt(α_bar_t) * x_0 + sqrt(1 - α_bar_t) * ε
其中 ε ~ N(0, I)
"""
def __init__(self, num_timesteps: int = 1000,
beta_start: float = 1e-4,
beta_end: float = 0.02,
schedule: str = 'linear'):
self.num_timesteps = num_timesteps
# 噪声调度
if schedule == 'linear':
betas = torch.linspace(beta_start, beta_end, num_timesteps)
elif schedule == 'cosine':
# cosine 调度（Nichol & Dhariwal 2021 推荐）
steps = torch.arange(num_timesteps + 1)
f = torch.cos((steps / num_timesteps + 0.008) / 1.008 * np.pi / 2) ** 2
f = f / f[0]
betas = torch.clip(1 - f[1:] / f[:-1], 0.0001, 0.9999)
else:
raise ValueError(f"Unknown schedule: {schedule}")
# 预计算所有需要的系数
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # ᾱ_t
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
# 注册为 buffer（不参与梯度计算，但会随模型移动到 GPU）
self.register_buffer = lambda name, val: setattr(self, name, val)
self.betas = betas
self.alphas = alphas
self.alphas_cumprod = alphas_cumprod
# q(x_t | x_0) 的参数
self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # √ᾱ_t
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # √(1-ᾱ_t)
# 反向过程 q(x_{t-1} | x_t, x_0) 的参数
self.sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # 1/√α_t
self.posterior_variance = (
betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
)
# 用于对数空间计算的系数
self.posterior_log_variance_clipped = torch.log(
self.posterior_variance.clamp(min=1e-20)
)
self.posterior_mean_coef1 = (
betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
)
self.posterior_mean_coef2 = (
(1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod)
)
def q_sample(self, x_0: torch.Tensor, t: torch.Tensor,
noise: torch.Tensor = None) -> torch.Tensor:
"""
前向采样：从 x_0 直接采样 x_t（重参数化技巧）
x_t = sqrt(ᾱ_t) * x_0 + sqrt(1 - ᾱ_t) * ε
参数：
x_0: 原始图像 [B, C, H, W]
t: 时间步 [B]（每个样本可以有不同的 t）
noise: 可选的预采样噪声
"""
if noise is None:
noise = torch.randn_like(x_0)
# 提取对应时间步的系数，reshape 为 [B, 1, 1, 1] 以便广播
sqrt_alpha = self._extract(self.sqrt_alphas_cumprod, t, x_0.shape)
sqrt_one_minus_alpha = self._extract(
self.sqrt_one_minus_alphas_cumprod, t, x_0.shape
)
return sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise
def _extract(self, a: torch.Tensor, t: torch.Tensor,
x_shape: tuple) -> torch.Tensor:
"""从预计算数组中提取对应时间步的值"""
batch_size = t.shape[0]
out = a.gather(-1, t.cpu()).to(t.device)
return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
def q_posterior_mean_variance(self, x_0: torch.Tensor,
x_t: torch.Tensor,
t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
计算后验 q(x_{t-1} | x_t, x_0) 的均值和方差
μ̃_t = (√ᾱ_{t-1} * β_t) / (1 - ᾱ_t) * x_0
+ (√α_t * (1 - ᾱ_{t-1})) / (1 - ᾱ_t) * x_t
"""
posterior_mean = (
self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_0 +
self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
)
posterior_variance = self._extract(self.posterior_variance, t, x_t.shape)
return posterior_mean, posterior_variance
# 可视化前向扩散过程
def visualize_forward_process():
"""展示前向扩散如何逐步将图像变成噪声"""
diffusion = ForwardDiffusion(num_timesteps=1000)
# 创建一个简单的测试图像（例如一个圆）
x_0 = torch.zeros(1, 1, 32, 32)
y, x = torch.meshgrid(torch.arange(32), torch.arange(32), indexing='ij')
center = 16
radius = 10
circle = ((y - center)**2 + (x - center)**2) <= radius**2
x_0[0, 0] = circle.float()
# 在不同时间步采样
timesteps = [0, 50, 100, 250, 500, 999]
print("前向扩散过程：从清晰图像到纯噪声")
print("=" * 50)
for t_val in timesteps:
t = torch.tensor([t_val])
x_t = diffusion.q_sample(x_0, t)
# x_t 中信号的能量比
signal_ratio = diffusion.sqrt_alphas_cumprod[t_val].item() ** 2
noise_ratio = 1 - signal_ratio
print(f"t={t_val:4d}: 信号占比={signal_ratio:.4f}, "
f"噪声占比={noise_ratio:.4f}")
if __name__ == "__main__":
visualize_forward_process()
二、反向去噪过程的推导
前向过程我们完全知道，但反向过程 q(x_{t-1} | x_t) 是不可直接计算的（需要整个数据集）。关键发现是：如果我们额外知道 x_0，反向过程就变得可处理了。
利用贝叶斯定理：
q(x_{t-1} | x_t, x_0) = N(x_{t-1}; μ̃_t(x_t, x_0), β̃_t · I)
其中：
- β̃_t = (1 - ᾱ_{t-1}) / (1 - ᾱ_t) · β_t
- μ̃_t(x_t, x_0) = (√ᾱ_{t-1} · β_t)/(1-ᾱ_t) · x_0 + (√α_t · (1-ᾱ_{t-1}))/(1-ᾱ_t) · x_t
用 x_0 = (x_t - √(1-ᾱ_t) · ε) / √ᾱ_t 代入，可以把均值写成只关于 x_t 和 ε 的函数：
μ̃_t(x_t) = (1/√α_t) · (x_t - β_t/√(1-ᾱ_t) · ε)
这意味着：如果我们能预测出噪声 ε，就能计算出反向过程的均值。这就是为什么 DDPM 让神经网络预测噪声——不是预测原始图像，不是预测均值，而是预测被添加的噪声。
三、训练目标——从 ELBO 到简化损失
扩散模型的训练基于变分推断，最大化数据的对数似然的证据下界（ELBO）：
log p(x_0) ≥ E_q[log p_θ(x_0|x_1)] - D_KL(q(x_T|x_0) || p(x_T)) - Σ E_q[D_KL(q(x_{t-1}|x_t,x_0) || p_θ(x_{t-1}|x_t))]
经过一系列化简，核心的训练损失可以表示为预测噪声与真实噪声之间的均方误差。Ho et al. (2020) 发现，忽略权重项的简化版本效果更好：
L_simple = E_{x_0, t, ε} [‖ε - ε_θ(x_t, t)‖²]
训练算法：
- 采样真实图像 x_0 ~ q(x_0)
- 随机选择时间步 t ~ Uniform(1, T)
- 采样噪声 ε ~ N(0, I)
- 计算噪声图像 x_t = √ᾱ_t · x_0 + √(1-ᾱ_t) · ε
- 让模型预测噪声 ε_θ(x_t, t)
- 计算损失 L = ‖ε - ε_θ(x_t, t)‖²
- 反向传播更新参数
四、采样算法——从噪声生成图像
训练完成后，从纯噪声开始逐步去噪：
"""
DDPM 完整的采样过程
"""
class DDPMSampler:
"""
DDPM 反向采样器
从 x_T ~ N(0, I) 开始，逐步去噪到 x_0
"""
def __init__(self, diffusion: ForwardDiffusion):
self.diffusion = diffusion
@torch.no_grad()
def sample(self, model: torch.nn.Module,
shape: tuple,
device: str = 'cpu') -> torch.Tensor:
"""
DDPM 采样算法
参数：
model: 噪声预测模型 ε_θ(x_t, t)
shape: 生成图像的形状 [B, C, H, W]
device: 计算设备
"""
# 步骤1：从纯高斯噪声开始
x = torch.randn(shape, device=device)
# 步骤2：逐步去噪，从 T 到 1
for t in reversed(range(self.diffusion.num_timesteps)):
t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
# 模型预测噪声
predicted_noise = model(x, t_batch)
# 计算去噪后的均值
# μ_θ = (1/√α_t) * (x_t - β_t/√(1-ᾱ_t) * ε_θ(x_t, t))
sqrt_recip_alpha = self.diffusion._extract(
self.diffusion.sqrt_recip_alphas, t_batch, x.shape
)
beta_over_sqrt_one_minus_alpha = self.diffusion.betas[t] / \
self.diffusion.sqrt_one_minus_alphas_cumprod[t]
mean = sqrt_recip_alpha * (
x - beta_over_sqrt_one_minus_alpha * predicted_noise
)
if t > 0:
# 添加随机噪声（除了最后一步）
noise = torch.randn_like(x)
sigma = torch.sqrt(self.diffusion.posterior_variance[t])
x = mean + sigma * noise
else:
x = mean
return x
class DDIMSampler:
"""
DDIM（Denoising Diffusion Implicit Models）采样器
Song et al. 2020 提出的确定性采样方法
可以用更少的步数生成高质量图像
核心改进：将随机采样变为确定性映射
η = 0 时为完全确定性（DDIM）
η = 1 时退化为 DDPM
"""
def __init__(self, diffusion: ForwardDiffusion):
self.diffusion = diffusion
@torch.no_grad()
def sample(self, model: torch.nn.Module,
shape: tuple,
num_inference_steps: int = 50,
eta: float = 0.0,
device: str = 'cpu') -> torch.Tensor:
"""
DDIM 采样
参数：
model: 噪声预测模型
shape: 生成图像的形状
num_inference_steps: 实际采样步数（远小于训练步数 T）
eta: 随机性控制，0=确定性，1=DDPM
"""
T = self.diffusion.num_timesteps
# 创建子采样时间表
step_size = T // num_inference_steps
timesteps = list(range(0, T, step_size))
timesteps = list(reversed(timesteps))
x = torch.randn(shape, device=device)
alphas_cumprod = self.diffusion.alphas_cumprod
for i in range(len(timesteps)):
t = timesteps[i]
t_prev = timesteps[i + 1] if i < len(timesteps) - 1 else 0
t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
# 预测噪声
predicted_noise = model(x, t_batch)
# DDIM 的核心公式
alpha_t = alphas_cumprod[t]
alpha_t_prev = alphas_cumprod[t_prev] if t_prev > 0 else torch.tensor(1.0)
# 预测 x_0
x0_pred = (x - torch.sqrt(1 - alpha_t) * predicted_noise) / torch.sqrt(alpha_t)
# 计算方向指向的噪声
sigma = eta * torch.sqrt(
(1 - alpha_t_prev) / (1 - alpha_t) * (1 - alpha_t / alpha_t_prev)
)
# 预测 x_{t-1} 的方向
pred_dir = torch.sqrt(1 - alpha_t_prev - sigma**2) * predicted_noise
# DDIM 更新
x = torch.sqrt(alpha_t_prev) * x0_pred + pred_dir
if sigma > 0:
x = x + sigma * torch.randn_like(x)
return x
五、U-Net 噪声预测网络
DDPM 使用 U-Net 作为噪声预测网络 ε_θ(x_t, t)。U-Net 的编码器-解码器结构配合跳跃连接，非常适合保留空间信息。
"""
简化版 U-Net 噪声预测网络
用于 DDPM 训练和采样
"""
import torch
import torch.nn as nn
import math
class SinusoidalPositionEmbeddings(nn.Module):
"""
正弦位置编码——将时间步 t 编码为向量
与 Transformer 中的位置编码类似
让模型知道"当前处于哪个噪声水平"
"""
def __init__(self, dim: int):
super().__init__()
self.dim = dim
def forward(self, t: torch.Tensor) -> torch.Tensor:
device = t.device
half_dim = self.dim // 2
embeddings = math.log(10000) / (half_dim - 1)
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
embeddings = t[:, None].float() * embeddings[None, :]
embeddings = torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
return embeddings
class ResidualBlock(nn.Module):
"""残差块，融合时间步嵌入"""
def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int):
super().__init__()
self.conv1 = nn.Sequential(
nn.GroupNorm(8, in_channels),
nn.SiLU(),
nn.Conv2d(in_channels, out_channels, 3, padding=1)
)
self.time_mlp = nn.Sequential(
nn.SiLU(),
nn.Linear(time_emb_dim, out_channels)
)
self.conv2 = nn.Sequential(
nn.GroupNorm(8, out_channels),
nn.SiLU(),
nn.Conv2d(out_channels, out_channels, 3, padding=1)
)
self.shortcut = (
nn.Conv2d(in_channels, out_channels, 1)
if in_channels != out_channels else nn.Identity()
)
def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
h = self.conv1(x)
# 将时间嵌入加到特征图上
t_emb = self.time_mlp(t_emb)
h = h + t_emb[:, :, None, None]
h = self.conv2(h)
return h + self.shortcut(x)
class AttentionBlock(nn.Module):
"""自注意力块——在低分辨率特征图上捕获全局依赖"""
def __init__(self, channels: int):
super().__init__()
self.norm = nn.GroupNorm(8, channels)
self.qkv = nn.Conv2d(channels, channels * 3, 1)
self.proj = nn.Conv2d(channels, channels, 1)
self.scale = channels ** -0.5
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, C, H, W = x.shape
h = self.norm(x)
qkv = self.qkv(h).reshape(B, 3, C, H * W)
q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
attn = torch.softmax(torch.bmm(q.transpose(1, 2), k) * self.scale, dim=-1)
h = torch.bmm(v, attn.transpose(1, 2)).reshape(B, C, H, W)
return self.proj(h) + x
class SimpleUNet(nn.Module):
"""
简化版 U-Net
架构：
- 编码器：逐步下采样
- 瓶颈层：自注意力
- 解码器：逐步上采样 + 跳跃连接
输入：噪声图像 x_t 和时间步 t
输出：预测的噪声 ε_θ(x_t, t)
"""
def __init__(self, in_channels: int = 1, out_channels: int = 1,
base_channels: int = 64):
super().__init__()
time_dim = base_channels * 4
# 时间步嵌入
self.time_mlp = nn.Sequential(
SinusoidalPositionEmbeddings(base_channels),
nn.Linear(base_channels, time_dim),
nn.SiLU(),
nn.Linear(time_dim, time_dim)
)
# 编码器
self.enc1 = ResidualBlock(in_channels, base_channels, time_dim)
self.enc2 = ResidualBlock(base_channels, base_channels * 2, time_dim)
self.enc3 = ResidualBlock(base_channels * 2, base_channels * 4, time_dim)
self.down = nn.MaxPool2d(2)
# 瓶颈层
self.bottleneck = nn.Sequential(
ResidualBlock(base_channels * 4, base_channels * 8, time_dim),
AttentionBlock(base_channels * 8),
ResidualBlock(base_channels * 8, base_channels * 4, time_dim)
)
# 解码器
self.up3 = nn.ConvTranspose2d(base_channels * 4, base_channels * 4, 2, 2)
self.dec3 = ResidualBlock(base_channels * 8, base_channels * 2, time_dim)
self.up2 = nn.ConvTranspose2d(base_channels * 2, base_channels * 2, 2, 2)
self.dec2 = ResidualBlock(base_channels * 4, base_channels, time_dim)
self.up1 = nn.ConvTranspose2d(base_channels, base_channels, 2, 2)
self.dec1 = ResidualBlock(base_channels * 2, base_channels, time_dim)
# 输出
self.final = nn.Sequential(
nn.GroupNorm(8, base_channels),
nn.SiLU(),
nn.Conv2d(base_channels, out_channels, 3, padding=1)
)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
# 时间步嵌入
t_emb = self.time_mlp(t)
# 编码器
e1 = self.enc1(x, t_emb) # [B, 64, H, W]
e2 = self.enc2(self.down(e1), t_emb) # [B, 128, H/2, W/2]
e3 = self.enc3(self.down(e2), t_emb) # [B, 256, H/4, W/4]
# 瓶颈层
b = self.bottleneck[0](https://raw.githubusercontent.com/mindcarver/91ai/main/docs/machine-learning/stage8-advanced/self.down(e3), t_emb)
b = self.bottleneck[1](https://raw.githubusercontent.com/mindcarver/91ai/main/docs/machine-learning/stage8-advanced/b)
b = self.bottleneck[2](https://raw.githubusercontent.com/mindcarver/91ai/main/docs/machine-learning/stage8-advanced/b, t_emb)
# 解码器 + 跳跃连接
d3 = self.up3(b)
d3 = self.dec3(torch.cat([d3, e3], dim=1), t_emb)
d2 = self.up2(d3)
d2 = self.dec2(torch.cat([d2, e2], dim=1), t_emb)
d1 = self.up1(d2)
d1 = self.dec1(torch.cat([d1, e1], dim=1), t_emb)
return self.final(d1)
六、完整训练流程
"""
DDPM 完整训练流程
在 MNIST 数据集上训练一个无条件扩散模型
"""
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from tqdm import tqdm
class DDPMTrainer:
"""
DDPM 训练器
训练循环的核心：
1. 取真实图像 x_0
2. 随机选时间步 t
3. 随机采样噪声 ε
4. 计算噪声图像 x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε
5. 让模型预测噪声 ε_θ(x_t, t)
6. 损失 = ‖ε - ε_θ(x_t, t)‖²
"""
def __init__(self, model: nn.Module,
diffusion: ForwardDiffusion,
lr: float = 2e-4,
device: str = 'cuda' if torch.cuda.is_available() else 'cpu'):
self.model = model.to(device)
self.diffusion = diffusion
self.device = device
self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
# EMA（指数移动平均）提升生成质量
self.ema_model = copy.deepcopy(model)
self.ema_decay = 0.9999
def train_step(self, x_0: torch.Tensor) -> float:
"""单步训练"""
self.optimizer.zero_grad()
batch_size = x_0.shape[0]
x_0 = x_0.to(self.device)
# 随机采样时间步
t = torch.randint(0, self.diffusion.num_timesteps, (batch_size,),
device=self.device)
# 随机采样噪声
noise = torch.randn_like(x_0)
# 前向加噪
x_t = self.diffusion.q_sample(x_0, t, noise)
# 模型预测噪声
predicted_noise = self.model(x_t, t)
# 简化损失：预测噪声与真实噪声的 MSE
loss = nn.functional.mse_loss(predicted_noise, noise)
loss.backward()
# 梯度裁剪
nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
# 更新 EMA
self._update_ema()
return loss.item()
def _update_ema(self):
for ema_p, model_p in zip(self.ema_model.parameters(),
self.model.parameters()):
ema_p.data.mul_(self.ema_decay).add_(model_p.data, alpha=1 - self.ema_decay)
def train(self, num_epochs: int = 50, batch_size: int = 128):
"""完整训练循环"""
# 加载 MNIST
transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]) # 归一化到 [-1, 1]
])
dataset = datasets.MNIST('./data', train=True, download=True,
transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=4)
for epoch in range(num_epochs):
total_loss = 0
pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}")
for batch, _ in pbar:
loss = self.train_step(batch)
total_loss += loss
pbar.set_postfix(loss=f"{loss:.4f}")
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch+1}: 平均损失 = {avg_loss:.4f}")
# 每10个 epoch 生成样本看看效果
if (epoch + 1) % 10 == 0:
self.generate_samples(epoch + 1)
@torch.no_grad()
def generate_samples(self, epoch: int, num_images: int = 16):
"""用 EMA 模型生成样本"""
sampler = DDPMSampler(self.diffusion)
samples = sampler.sample(
self.ema_model,
shape=(num_images, 1, 32, 32),
device=self.device
)
# 从 [-1, 1] 转回 [0, 1]
samples = (samples + 1) / 2
samples = samples.clamp(0, 1)
print(f"生成了 {num_images} 张图像 (epoch {epoch})")
import copy # 用于 EMA
七、Stable Diffusion——潜在扩散模型
Stable Diffusion 的核心创新在于：不在像素空间做扩散，而是先通过 VAE 将图像压缩到低维潜在空间，在潜在空间中做扩散。
┌─────────────────────────────────────────────────────────────┐
│ Stable Diffusion 架构 │
│ │
│ 文本输入 ──── CLIP Text Encoder ──── 文本嵌入 [B, 77, 768] │
│ │ │
│ ▼ │
│ ┌───────────────────────────────────────────────────┐ │
│ │ 条件去噪 U-Net │ │
│ │ │ │
│ │ 噪声潜在 z_t ──▶ [ResBlock + CrossAttention] ──▶ │ │
│ │ ▲ ▲ │ │
│ │ │ 文本嵌入作为 │ │
│ │ 时间步 t Key 和 Value │ │
│ │ │ │ │ │
│ │ └──────────────────────────┘ │ │
│ │ │ │
│ │ 输出：预测的噪声 ε_θ(z_t, t, c) │ │
│ └───────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ 去噪后的潜在 z_0 │
│ │ │
│ ▼ │
│ VAE Decoder ──── 生成图像 [B, 3, 512, 512] │
│ │
│ 训练流程： │
│ 图像 x ──▶ VAE Encoder ──▶ z_0 ──▶ 加噪 ──▶ z_t │
│ 损失 = ‖ε - ε_θ(z_t, t, text_embedding)‖² │
└─────────────────────────────────────────────────────────────┘
为什么潜在空间有效？ 图像的绝大部分比特用于编码高频细节（纹理、噪声），而语义信息集中在低维空间。VAE 将 512×512×3 = 786432 维压缩到 64×64×4 = 16384 维（压缩比约 48:1），但保留了足够的语义信息用于扩散。
计算效率对比：
- 像素空间扩散：在 512×512×3 上做扩散，每步计算量巨大
- 潜在空间扩散：在 64×64×4 上做扩散，计算量减少约 48 倍
"""
Stable Diffusion 的核心组件实现
展示 VAE、CLIP 文本编码和交叉注意力的协同工作
"""
class CrossAttention(nn.Module):
"""
交叉注意力机制
在 Stable Diffusion 中，U-Net 通过交叉注意力接收文本条件
- Query: 来自图像特征
- Key, Value: 来自文本嵌入
这使得模型能根据文本描述调整去噪方向
"""
def __init__(self, query_dim: int, context_dim: int,
n_heads: int = 8, head_dim: int = 64):
super().__init__()
inner_dim = n_heads * head_dim
self.n_heads = n_heads
self.scale = head_dim ** -0.5
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Linear(inner_dim, query_dim)
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
"""
参数：
x: 图像特征 [B, H*W, query_dim]
context: 文本嵌入 [B, seq_len, context_dim]
"""
B = x.shape[0]
q = self.to_q(x) # [B, H*W, inner_dim]
k = self.to_k(context) # [B, seq_len, inner_dim]
v = self.to_v(context) # [B, seq_len, inner_dim]
# 多头注意力
q = q.reshape(B, -1, self.n_heads, q.shape[-1] // self.n_heads).transpose(1, 2)
k = k.reshape(B, -1, self.n_heads, k.shape[-1] // self.n_heads).transpose(1, 2)
v = v.reshape(B, -1, self.n_heads, v.shape[-1] // self.n_heads).transpose(1, 2)
attn = torch.softmax(torch.matmul(q, k.transpose(-1, -2)) * self.scale, dim=-1)
out = torch.matmul(attn, v)
out = out.transpose(1, 2).reshape(B, -1, out.shape[-1] * self.n_heads)
return self.to_out(out)
class ClassifierFreeGuidance:
"""
无分类器引导（Classifier-Free Guidance）
Ho & Salimans 2021 提出，Stable Diffusion 的核心采样技术
核心思想：
同时训练条件模型 ε_θ(x_t, t, c) 和无条件模型 ε_θ(x_t, t, ∅)
（通过训练时随机丢弃条件实现）
采样时：
ε_guided = ε_θ(x_t, t, ∅) + s * (ε_θ(x_t, t, c) - ε_θ(x_t, t, ∅))
s > 1 增强条件的影响，生成更符合文本描述的图像
但 s 过大可能导致图像质量下降（过饱和、伪影）
典型值：s = 7.5（Stable Diffusion 的默认值）
"""
def __init__(self, guidance_scale: float = 7.5):
self.guidance_scale = guidance_scale
@torch.no_grad()
def guided_denoise(self, model: nn.Module,
x_t: torch.Tensor,
t: torch.Tensor,
text_embedding: torch.Tensor,
uncond_embedding: torch.Tensor) -> torch.Tensor:
"""
引导去噪一步
同时预测条件和无条件噪声，然后线性外推
"""
# 无条件预测
noise_uncond = model(x_t, t, uncond_embedding)
# 条件预测
noise_cond = model(x_t, t, text_embedding)
# 引导：远离无条件，靠近条件
guided_noise = noise_uncond + self.guidance_scale * (
noise_cond - noise_uncond
)
return guided_noise
class StableDiffusionPipeline:
"""
简化版 Stable Diffusion Pipeline
整合所有组件：文本编码 + 潜在扩散 + VAE 解码
"""
def __init__(self, unet, vae, text_encoder, diffusion,
guidance_scale: float = 7.5):
self.unet = unet
self.vae = vae
self.text_encoder = text_encoder
self.diffusion = diffusion
self.cfg = ClassifierFreeGuidance(guidance_scale)
@torch.no_grad()
def __call__(self, prompt: str, num_inference_steps: int = 50,
height: int = 512, width: int = 512) -> torch.Tensor:
"""
文本到图像生成
参数：
prompt: 文本描述
num_inference_steps: 采样步数
height, width: 输出图像尺寸
"""
device = next(self.unet.parameters()).device
# 步骤1：编码文本
# 实际中使用 CLIP Text Encoder
# 这里用随机嵌入代替演示
text_embedding = self.text_encoder(prompt) # [1, 77, 768]
uncond_embedding = self.text_encoder("") # 无条件嵌入
# 步骤2：在潜在空间中采样
latent_height = height // 8 # VAE 下采样 8 倍
latent_width = width // 8
latent_shape = (1, 4, latent_height, latent_width) # 4 通道潜在空间
sampler = DDIMSampler(self.diffusion)
# 使用 CFG 的采样循环
latents = torch.randn(latent_shape, device=device)
T = self.diffusion.num_timesteps
step_size = T // num_inference_steps
timesteps = list(reversed(range(0, T, step_size)))
for t_val in timesteps:
t = torch.tensor([t_val], device=device)
# CFG 引导去噪
noise_pred = self.cfg.guided_denoise(
self.unet, latents, t, text_embedding, uncond_embedding
)
# DDIM 更新步骤（简化）
alpha_t = self.diffusion.alphas_cumprod[t_val]
x0_pred = (latents - torch.sqrt(1 - alpha_t) * noise_pred) / torch.sqrt(alpha_t)
t_prev = t_val - step_size if t_val > step_size else 0
alpha_t_prev = self.diffusion.alphas_cumprod[t_prev] if t_prev > 0 else torch.tensor(1.0)
latents = torch.sqrt(alpha_t_prev) * x0_pred + \
torch.sqrt(1 - alpha_t_prev) * noise_pred
# 步骤3：VAE 解码
image = self.vae.decode(latents)
return image
八、Diffusion Transformer（DiT）——架构革命
2023年 Peebles & Xie 提出的 DiT（Diffusion Transformer），用 Transformer 替代 U-Net 作为扩散模型骨干网络。这个改变看起来简单，却深刻影响了后续发展——Sora、Stable Diffusion 3、Flux 都采用了 DiT 或其变体。
"""
Diffusion Transformer (DiT) 的核心实现
参考：Peebles & Xie, "Scalable Diffusion Models with Transformers", ICCV 2023
"""
class DiTBlock(nn.Module):
"""
DiT 块
与标准 Transformer 块的区别：
1. 使用自适应层归一化（adaLN）注入时间步和条件信息
2. 使用 adaLN-Zero 初始化策略（残差连接最后一层初始化为0）
adaLN-Zero 让训练初期接近恒等映射，加速收敛
"""
def __init__(self, dim: int, num_heads: int, mlp_ratio: int = 4):
super().__init__()
self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
mlp_hidden = dim * mlp_ratio
self.mlp = nn.Sequential(
nn.Linear(dim, mlp_hidden),
nn.GELU(),
nn.Linear(mlp_hidden, dim)
)
# adaLN 调制参数：每个块有 6 个参数（scale, shift, gate × 2 个子层）
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
nn.Linear(dim, 6 * dim)
)
# Zero 初始化
nn.init.zeros_(self.adaLN_modulation[-1].weight)
nn.init.zeros_(self.adaLN_modulation[-1].bias)
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
"""
参数：
x: patch tokens [B, N, dim]
c: 条件向量（时间步 + 类别/文本嵌入） [B, dim]
"""
# 从条件向量计算调制参数
shift1, scale1, gate1, shift2, scale2, gate2 = \
self.adaLN_modulation(c).chunk(6, dim=-1)
# 自注意力 + adaLN + gate
h = self.norm1(x) * (1 + scale1[:, None, :]) + shift1[:, None, :]
h, _ = self.attn(h, h, h)
x = x + gate1[:, None, :] * h
# MLP + adaLN + gate
h = self.norm2(x) * (1 + scale2[:, None, :]) + shift2[:, None, :]
h = self.mlp(h)
x = x + gate2[:, None, :] * h
return x
class DiT(nn.Module):
"""
Diffusion Transformer
工作流程：
1. 将潜在表示切分为 patch（类似 ViT）
2. 通过 Transformer 块处理
3. 重组为原始空间形状，输出噪声预测
为什么 DiT 比 U-Net 好？
- 可扩展性：Transformer 架构更容易 scale up
- 训练效率：在相同计算量下 DiT 的 FID 更低
- 灵活性：条件注入方式更统一
"""
def __init__(self, input_size: int = 32, in_channels: int = 4,
patch_size: int = 2, dim: int = 384,
depth: int = 12, num_heads: int = 6):
super().__init__()
self.input_size = input_size
self.patch_size = patch_size
self.num_patches = (input_size // patch_size) ** 2
# Patch 嵌入
self.patch_embed = nn.Conv2d(
in_channels, dim, kernel_size=patch_size, stride=patch_size
)
# 位置编码
self.pos_embed = nn.Parameter(
torch.randn(1, self.num_patches, dim) * 0.02
)
# Transformer 块
self.blocks = nn.ModuleList([
DiTBlock(dim, num_heads) for _ in range(depth)
])
# 最终层归一化
self.final_norm = nn.LayerNorm(dim, elementwise_affine=False)
# 输出投影（预测噪声和对角方差）
self.output_proj = nn.Linear(dim, patch_size * patch_size * in_channels)
nn.init.zeros_(self.output_proj.weight)
nn.init.zeros_(self.output_proj.bias)
# 条件嵌入（时间步 + 类别/文本）
self.cond_embed = nn.Sequential(
SinusoidalPositionEmbeddings(dim),
nn.Linear(dim, dim),
nn.SiLU(),
nn.Linear(dim, dim)
)
def forward(self, x: torch.Tensor, t: torch.Tensor,
y: torch.Tensor = None) -> torch.Tensor:
"""
参数：
x: 噪声潜在 [B, in_channels, H, W]
t: 时间步 [B]
y: 条件嵌入 [B, cond_dim]（可选）
"""
B = x.shape[0]
# Patch 嵌入 + 位置编码
patches = self.patch_embed(x) # [B, dim, H/P, W/P]
patches = patches.flatten(2).transpose(1, 2) # [B, N, dim]
patches = patches + self.pos_embed
# 条件嵌入
c = self.cond_embed(t)
# Transformer 块
for block in self.blocks:
patches = block(patches, c)
# 输出
patches = self.final_norm(patches)
output = self.output_proj(patches) # [B, N, P*P*C]
# 重组为空间形状
P = self.patch_size
C = x.shape[1]
H = W = self.input_size // P
output = output.reshape(B, H, W, P, P, C)
output = output.permute(0, 5, 1, 3, 2, 4).reshape(B, C, H * P, W * P)
return output
九、Flow Matching——简化的训练范式
2022-2024年，Flow Matching 作为扩散模型的替代训练框架出现，被 Stable Diffusion 3、Flux 等最新模型采用。它用连续标准化流（CNF）替代 DDPM 的离散马尔可夫链。
"""
Flow Matching 核心概念
与 DDPM 的对比：
- DDPM：训练模型预测噪声 ε，采样时从噪声逐步去噪
- Flow Matching：训练模型学习向量场 v_t，采样时沿向量场积分
优势：
- 训练目标更简单（不需要噪声调度）
- 采样路径更直（理论上更少的步数）
- 与扩散模型等价但更灵活
"""
class FlowMatching:
"""
最简形式的 Flow Matching
核心公式：
x_t = (1 - t) * x_0 + t * x_1
其中 x_0 ~ 数据分布，x_1 ~ 噪声分布（标准高斯）
t ~ Uniform(0, 1)
目标：学习向量场 v_θ(x_t, t) 使得 dx/dt = v_θ
"""
def __init__(self, model: nn.Module, sigma_min: float = 1e-4):
self.model = model
self.sigma_min = sigma_min
def compute_loss(self, x_0: torch.Tensor) -> torch.Tensor:
"""
Flow Matching 训练损失
与 DDPM 的关键区别：
- 不需要预定义噪声调度
- 不需要 T 步训练，只需连续时间 t ∈ [0, 1]
- 损失直接回归向量场
参数：
x_0: 真实数据样本
"""
batch_size = x_0.shape[0]
device = x_0.device
# 从标准高斯采样噪声
x_1 = torch.randn_like(x_0)
# 随机采样时间
t = torch.rand(batch_size, device=device)
# 线性插值（最优传输路径）
t_shape = t.reshape(-1, *((1,) * (x_0.ndim - 1)))
x_t = (1 - t_shape) * x_0 + t_shape * x_1
# 目标向量场：从 x_t 指向 x_1 的方向
target_v = x_1 - x_0
# 模型预测向量场
predicted_v = self.model(x_t, t)
# 损失
loss = nn.functional.mse_loss(predicted_v, target_v)
return loss
@torch.no_grad()
def sample(self, shape: tuple, num_steps: int = 50,
device: str = 'cpu') -> torch.Tensor:
"""
Euler 方法沿向量场积分
从 x_1（噪声）到 x_0（数据）
"""
dt = 1.0 / num_steps
x = torch.randn(shape, device=device) # x_1
for i in range(num_steps):
t = torch.full((shape[0],), 1.0 - i * dt, device=device)
v = self.model(x, t)
x = x - dt * v # 注意是减，因为从 t=1 到 t=0
return x
完整可运行的 Python 代码示例
"""
扩散模型完整演示
从零实现 DDPM 训练和生成，在 MNIST 上运行
运行方法：
python diffusion_demo.py
依赖：pip install torch torchvision numpy tqdm matplotlib
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import numpy as np
import math
from tqdm import tqdm
from typing import Tuple, Optional
import copy
# ========== 组件1：噪声调度器 ==========
class NoiseScheduler:
"""管理所有噪声调度的预计算系数"""
def __init__(self, num_timesteps: int = 1000,
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.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
# 前向采样系数
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
# 反向采样系数
self.sqrt_recip_alphas = torch.sqrt(1.0 / self.alphas)
self.posterior_variance = (
self.betas * (1.0 - self.alphas_cumprod_prev) /
(1.0 - self.alphas_cumprod)
)
def add_noise(self, x_0: torch.Tensor, t: torch.Tensor,
noise: torch.Tensor = None) -> torch.Tensor:
"""前向加噪：x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε"""
if noise is None:
noise = torch.randn_like(x_0)
sqrt_a = self._gather(self.sqrt_alphas_cumprod, t, x_0.shape)
sqrt_1ma = self._gather(self.sqrt_one_minus_alphas_cumprod, t, x_0.shape)
return sqrt_a * x_0 + sqrt_1ma * noise
def _gather(self, a: torch.Tensor, t: torch.Tensor, shape: tuple) -> torch.Tensor:
b = t.shape[0]
out = a.gather(-1, t)
return out.reshape(b, *((1,) * (len(shape) - 1)))
# ========== 组件2：简化 U-Net ==========
class TimeEmbedding(nn.Module):
def __init__(self, dim: int):
super().__init__()
self.dim = dim
self.mlp = nn.Sequential(
nn.Linear(dim, dim * 4),
nn.SiLU(),
nn.Linear(dim * 4, dim * 4)
)
def forward(self, t: torch.Tensor) -> torch.Tensor:
half = self.dim // 2
freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
args = t[:, None].float() * freqs[None]
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
return self.mlp(emb)
class ConvBlock(nn.Module):
def __init__(self, in_ch: int, out_ch: int, time_dim: int):
super().__init__()
self.conv1 = nn.Sequential(
nn.GroupNorm(8, in_ch), nn.SiLU(),
nn.Conv2d(in_ch, out_ch, 3, padding=1)
)
self.time_proj = nn.Linear(time_dim, out_ch)
self.conv2 = nn.Sequential(
nn.GroupNorm(8, out_ch), nn.SiLU(),
nn.Conv2d(out_ch, out_ch, 3, padding=1)
)
self.shortcut = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
h = self.conv1(x)
h = h + self.time_proj(t)[:, :, None, None]
h = self.conv2(h)
return h + self.shortcut(x)
class MiniUNet(nn.Module):
"""精简版 U-Net，适合 MNIST 演示"""
def __init__(self, channels: int = 1, base: int = 32):
super().__init__()
td = base * 4
self.time_emb = TimeEmbedding(base)
# 编码器
self.enc1 = ConvBlock(channels, base, td)
self.enc2 = ConvBlock(base, base * 2, td)
self.enc3 = ConvBlock(base * 2, base * 4, td)
self.pool = nn.MaxPool2d(2)
# 瓶颈
self.bottleneck = ConvBlock(base * 4, base * 8, td)
# 解码器
self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, 2)
self.dec3 = ConvBlock(base * 8, base * 4, td)
self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, 2)
self.dec2 = ConvBlock(base * 4, base * 2, td)
self.up1 = nn.ConvTranspose2d(base * 2, base, 2, 2)
self.dec1 = ConvBlock(base * 2, base, td)
self.out = nn.Conv2d(base, channels, 1)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
t_emb = self.time_emb(t)
e1 = self.enc1(x, t_emb)
e2 = self.enc2(self.pool(e1), t_emb)
e3 = self.enc3(self.pool(e2), t_emb)
b = self.bottleneck(self.pool(e3), t_emb)
d3 = self.dec3(torch.cat([self.up3(b), e3], 1), t_emb)
d2 = self.dec2(torch.cat([self.up2(d3), e2], 1), t_emb)
d1 = self.dec1(torch.cat([self.up1(d2), e1], 1), t_emb)
return self.out(d1)
# ========== 组件3：训练 ==========
def train_ddpm(num_epochs: int = 30, batch_size: int = 128,
device: str = 'cpu'):
"""完整训练流程"""
scheduler = NoiseScheduler(num_timesteps=1000)
model = MiniUNet(channels=1, base=32).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)
ema_model = copy.deepcopy(model)
ema_decay = 0.999
for epoch in range(num_epochs):
model.train()
total_loss = 0
for batch, _ in tqdm(loader, desc=f"Epoch {epoch+1}/{num_epochs}"):
batch = batch.to(device)
bs = batch.shape[0]
t = torch.randint(0, 1000, (bs,), device=device)
noise = torch.randn_like(batch)
x_t = scheduler.add_noise(batch, t, noise)
pred_noise = model(x_t, t)
loss = F.mse_loss(pred_noise, noise)
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
# EMA 更新
with torch.no_grad():
for p_ema, p in zip(ema_model.parameters(), model.parameters()):
p_ema.data.mul_(ema_decay).add_(p.data, alpha=1 - ema_decay)
total_loss += loss.item()
avg_loss = total_loss / len(loader)
print(f"Epoch {epoch+1}: loss = {avg_loss:.4f}")
return ema_model, scheduler
# ========== 组件4：采样 ==========
@torch.no_grad()
def generate(model: nn.Module, scheduler: NoiseScheduler,
num_images: int = 16, device: str = 'cpu') -> torch.Tensor:
"""DDPM 采样：从纯噪声逐步生成图像"""
model.eval()
x = torch.randn(num_images, 1, 32, 32, device=device)
for t_val in tqdm(reversed(range(1000)), desc="Sampling"):
t = torch.full((num_images,), t_val, device=device, dtype=torch.long)
pred = model(x, t)
# 去噪均值
beta_t = scheduler.betas[t_val]
sqrt_1ma_t = scheduler.sqrt_one_minus_alphas_cumprod[t_val]
sqrt_ra_t = scheduler.sqrt_recip_alphas[t_val]
mean = sqrt_ra_t * (x - beta_t / sqrt_1ma_t * pred)
if t_val > 0:
sigma = torch.sqrt(scheduler.posterior_variance[t_val])
x = mean + sigma * torch.randn_like(x)
else:
x = mean
return (x + 1) / 2 # [-1,1] -> [0,1]
if __name__ == "__main__":
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"训练设备: {device}")
print("开始训练 DDPM...")
model, scheduler = train_ddpm(num_epochs=5, device=device)
print("生成样本...")
samples = generate(model, scheduler, num_images=4, device=device)
print(f"生成完成！输出形状: {samples.shape}")
print(f"像素值范围: [{samples.min():.3f}, {samples.max():.3f}]")
真实案例
案例1：Stable Diffusion 的工程实现
Stable Diffusion 1.5（2022年）的架构参数：
| 组件 | 参数 | 说明 |
|---|---|---|
| VAE Encoder | 34.2M 参数 | 将 512×512×3 压缩到 64×64×4 |
| VAE Decoder | 49.5M 参数 | 将 64×64×4 还原到 512×512×3 |
| CLIP Text Encoder | 123.7M 参数 | 将文本编码为 77×768 嵌入 |
| U-Net | 860.6M 参数 | 核心去噪网络，含交叉注意力层 |
| 总计 | ~1068M 参数 | ~1B 参数规模 |
关键设计决策：
- 潜在空间维度：4 通道（不是 3 或 8），平衡压缩率和重建质量
- 下采样倍率：8×（512→64），这是效率和质量的最优折中
- CFG Scale：默认 7.5，范围 1-20，越大越符合文本但多样性降低
案例2：Sora（OpenAI，2024）
Sora 使用 DiT（Diffusion Transformer）架构，核心改进：
- 时空 Patch：将视频切分为 3D 时空 Patch（类似 ViT 的 2D Patch 扩展到 3D）
- 可变分辨率：支持不同分辨率和时长，通过 Patch 化统一处理
- 规模效应：随着计算量增加，视频质量持续提升（没有明显的饱和点）
案例3：Stable Diffusion 3（Stability AI，2024-2025）
SD3 的三大改进：
- Flow Matching 替代 DDPM 的训练目标，训练更稳定
- MMDiT（Multimodal DiT）：文本和图像使用不同的权重，通过交叉注意力交互
- T5 文本编码器：配合 CLIP，增强对复杂文本的理解
案例4：Flux（Black Forest Labs，2024）
Flux 由 Stable Diffusion 原始团队创建：
- Rectified Flow：Flow Matching 的变体，采样路径更直
- 12B 参数的 DiT 架构
- 支持多种分辨率和宽高比
- 在 4 步采样下即可产生高质量图像
权衡
扩散模型 vs GAN
| 维度 | 扩散模型 | GAN |
|---|---|---|
| 训练稳定性 | 高（固定的前向过程，损失明确） | 低（对抗训练不稳定） |
| 生成质量 | 高（逐步精修） | 中-高（单次前向） |
| 多样性 | 高（覆盖数据分布的多个模式） | 低（容易模式崩塌） |
| 采样速度 | 慢（需要多步迭代，10-1000 步） | 快（单步前向） |
| 可控性 | 高（CFG、ControlNet 等方法） | 中（需要精心设计条件） |
| 理论基础 | 强（变分推断、SDE 有清晰推导） | 弱（训练动态缺乏理论保证） |
采样步数 vs 生成质量
| 步数 | 方法 | 速度 | 质量（FID↓） |
|---|---|---|---|
| 1000 | DDPM | 很慢 | 3.17（CIFAR-10） |
| 50 | DDIM | 中等 | 4.67 |
| 20 | DPM-Solver | 快 | 5.12 |
| 8 | LCM | 很快 | 7.84 |
| 4 | Consistency Model | 极快 | 8.56 |
| 1 | Consistency Model | 最快 | 12.43 |
模型规模 vs 训练成本
- SD 1.5（~1B 参数）：训练成本约 15,000 A100 GPU 天
- SDXL（~2.6B 参数）：训练成本约 50,000 A100 GPU 天
- SD3（~8B 参数）：训练成本估计 200,000+ A100 GPU 天
- Sora（未公开）：估计在数十万 GPU 天级别
要点总结
-
扩散模型的本质：通过"加噪声→学习去噪声"来建模数据分布。前向过程是固定的加噪过程，反向过程由神经网络学习。
-
从数学到代码的关键路径：
- 重参数化技巧 → 任意时刻的闭式采样
- 后验分布的推导 → 预测噪声（而非预测原始图像）
- 简化损失 L_simple = E[‖ε - ε_θ(x_t, t)‖²] → 训练目标
- 反向采样循环 → 从噪声生成图像
-
Stable Diffusion 的三大创新：
- VAE 压缩到潜在空间（降低 48 倍计算量）
- CLIP 文本条件 + 交叉注意力（文字控制生成）
- 无分类器引导 CFG（平衡多样性和条件遵循）
-
架构演进：U-Net → DiT（Transformer），DDPM → Flow Matching。核心趋势是向 Transformer 靠拢、训练目标简化。
-
前沿方向：单步/少步生成（Consistency Models、LCM）、视频生成（Sora）、3D 生成、Flow Matching 替代扩散框架。
延伸阅读
基础论文：
- Ho, Jain, & Abbeel. "Denoising Diffusion Probabilistic Models." NeurIPS 2020. — DDPM 原始论文，整个领域的奠基之作
- Song, Meng, & Ermon. "Denoising Diffusion Implicit Models." ICLR 2021. — DDIM，将采样步数从 1000 减到 50 以下
- Nichol & Dhariwal. "Improved Denoising Diffusion Probabilistic Models." ICML 2021. — cosine 调度、学习方差
潜在扩散与条件生成：
- Rombach et al. "High-Resolution Image Synthesis with Latent Diffusion Models." CVPR 2022. — Stable Diffusion 的理论基础
- Ho & Salimans. "Classifier-Free Diffusion Guidance." NeurIPS Workshop 2021. — CFG，最重要的条件控制技术
- Zhang et al. "Adding Conditional Control to Text-to-Image Diffusion Models." ICCV 2023. — ControlNet
架构演进：
- Peebles & Xie. "Scalable Diffusion Models with Transformers." ICCV 2023. — DiT，用 Transformer 替代 U-Net
- "From U-Nets to DiTs: The Architectural Evolution of Text-to-Image Diffusion Models." ICLR Blog 2026. — 架构演变的全面综述
前沿：
- Song et al. "Consistency Models." ICML 2023. — 单步生成的突破
- Lipman et al. "Flow Matching for Generative Modeling." ICLR 2023. — Flow Matching 框架
- MIT 6.S184. "Flow Matching and Diffusion Models." 2026. — MIT 最新课程讲义
教程与博客：
- Lilian Weng. "What Are Diffusion Models?" (2021, updated 2024). — 最全面的扩散模型综述博客
- Hugging Face. "The Annotated Diffusion Model." (2022). — 逐行代码解读 DDPM
- Sander Dieleman. "Noise Schedules Considered Harmful." (2024). — 对噪声调度的深入反思
- AI Summer. "Diffusion Models: Toward State-of-the-Art Image Generation." (2022). — 从零推导 DDPM 的数学
评论
登录后可以参与评论和讨论。
💬
还没有评论
欢迎留下第一条评论，帮助这篇内容更快形成讨论。