扩散模型理论基础

扩散概率模型（Diffusion Probabilistic Models）是一类基于马尔可夫链的生成模型，通过逐步加噪和去噪实现数据生成。本文档从第一性原理推导其数学基础。¹

核心思想

扩散模型包含两个过程：

1. 前向过程（Forward/Diffusion）：逐步向数据添加高斯噪声，直到接近纯噪声
2. 反向过程（Reverse/Denoising）：学习逆转前向过程，从噪声逐步恢复数据

┌─────────────────────────────────────────────────────────────────┐
│                      扩散模型工作流程                              │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  数据 x₀ ──加噪──▶ x₁ ──加噪──▶ x₂ ──...──▶ xT ──▶ 纯噪声     │
│   ↓                                          (前向过程 q)       │
│   │                                                             │
│   │                                                             │
│   │                     ┌─────────────────────┐               │
│   │                     │    学习反向过程      │               │
│   │                     │  pθ(x_{t-1}|x_t)   │               │
│   │                     └─────────────────────┘               │
│   │                                                             │
│   ▼                                                             │
│  数据 x₀ ←──去噪── x₁ ←──去噪── x₂ ←──...←── xT              │
│                              (反向过程 pθ)                       │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

前向过程

高斯噪声调度

前向过程 $q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)$ 是预先定义的超参马尔可夫链：

$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$

其中 $\{{\beta }_t{\}}_{t=1}^T$ 是噪声调度（noise schedule），通常满足：

$${\beta }_1<{\beta }_2<...<{\beta }_T,{\beta }_t\in (0,1)$$

封闭形式的噪声注入

利用重参数化技巧，可以直接计算任意时间步 $t$ 的噪声状态：

$${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$$

其中：

• ${\alpha }_t=1-{\beta }_t$
• ${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i={\prod }_{i=1}^t(1-{\beta }_i)$

关键性质：任意时刻的分布 $q({\mathbf{x}}_t|{\mathbf{x}}_0)$ 可以直接计算，无需迭代！

import torch
import numpy as np
 
def make_beta_schedule(schedule='linear', n_timestep=1000, linear_start=1e-4, linear_end=2e-2):
    """
    生成噪声调度表
    
    参数:
        schedule: 调度类型 ('linear', 'cosine', 'quadratic')
        n_timestep: 总时间步数 T
        linear_start/end: 线性调度的起止值
    """
    if schedule == 'linear':
        betas = np.linspace(linear_start, linear_end, n_timestep)
    elif schedule == 'cosine':
        # 余弦调度，DDPM论文推荐
        steps = n_timestep + 1
        x = np.linspace(0, n_timestep, steps)
        alphas_cumprod = np.cos(((x / n_timestep) + 0.008) / 1.008 * np.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        betas = np.clip(betas, 0, 0.999)
    elif schedule == 'quadratic':
        betas = np.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep) ** 2
    else:
        raise NotImplementedError(schedule)
    
    return betas
 
# 示例
betas = make_beta_schedule('cosine', n_timestep=1000)
alphas = 1 - betas
alphas_cumprod = np.cumprod(alphas)
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1 - alphas_cumprod)

常用噪声调度对比

调度方式	公式	特点
Linear	${\beta }_t=\text{lerp}({\beta }_1,{\beta }_T,t/T)$	简单，DDPM原始方案
Cosine	${\bar{\alpha }}_t={\cos }^2\left(\frac{t/T+s}{1+s}\cdot \frac{\pi }{2}\right)$	训练更稳定，推荐方案
Quadratic	${\beta }_t=\text{lerp}({\beta }_1^2,{\beta }_T^2,t/T)$	中间段变化更慢
Sigmoid	${\beta }_t=\text{sigmoid}(\text{lerp}(-5,5,t/T))$	两端平滑过渡

反向过程

学习目标

反向过程 $p_{\theta }({\mathbf{x}}_{0:T})$ 是参数化的马尔可夫链：

$$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$$

其中 $p({\mathbf{x}}_T)=\mathcal{N}({\mathbf{x}}_T;0,\mathbf{I})$ 是先验分布。

每个去噪转移定义为：

$$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$$

变分下界（ELBO）

通过变分推断，训练目标转化为最大化证据下界（ELBO）：

$$\mathcal{L}={\mathbb{E}}_q[\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]-{\mathbb{E}}_q[D_{\text{KL}}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\Vert p({\mathbf{x}}_T))]$$ $$-\sum _{t=2}^T{\mathbb{E}}_q\left[D_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))\right]$$

三项的物理意义：

1. 重建项 $\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)$：从 ${\mathbf{x}}_1$ 重建 ${\mathbf{x}}_0$
2. 先验匹配项：$q({\mathbf{x}}_T|{\mathbf{x}}_0)\approx p({\mathbf{x}}_T)$，当 $T\to \infty$ 时自动满足
3. 去噪匹配项：学习 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$ 的近似

简化的损失函数

经过推导（详见下方数学推导），损失函数可以简化为：

$${\mathcal{L}}_t\approx {\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}\left[{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t)\Vert }^2\right]$$

即：预测注入的噪声！

三种参数化方式

1. ε-预测（Noise Prediction）

直接预测注入的噪声 $\mathbf{ϵ}$：

$${\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\right)$$

class DiffusionModel_eps(torch.nn.Module):
    def __init__(self, denoiser, T=1000):
        super().__init__()
        self.denoiser = denoiser  # UNet等去噪网络
        self.T = T
    
    def training_losses(self, x0):
        """计算训练损失"""
        batch_size = x0.shape[0]
        t = torch.randint(0, self.T, (batch_size,), device=x0.device)
        
        # 采样噪声
        eps = torch.randn_like(x0)
        
        # 计算带噪图像
        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
        
        xt = sqrt_alphas_cumprod_t * x0 + sqrt_one_minus_alphas_cumprod_t * eps
        
        # 预测噪声
        eps_theta = self.denoiser(xt, t)
        
        # MSE损失
        loss = F.mse_loss(eps_theta, eps, reduction='mean')
        return loss

2. x₀-预测（Direct Prediction）

直接预测原始图像 ${\mathbf{x}}_0$：

$${\mathbf{x}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t))$$

对应的均值：

$${\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t$$

3. v-预测（Velocity Prediction）

预测速度向量 $\mathbf{v}=\sqrt{1-{\bar{\alpha }}_t}{\mathbf{x}}_0-\sqrt{{\bar{\alpha }}_t}\mathbf{ϵ}$：

这是 Stable Diffusion 2.0+ 使用的方案，实验表明效果更好。

$${\mathcal{L}}_v={\mathbb{E}}_t\left[{\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{v}}\left[{\Vert \mathbf{v}-{\mathbf{v}}_{\theta }({\mathbf{x}}_t,t)\Vert }^2\right]\right]$$

class VelocityPrediction(torch.nn.Module):
    def predict_velocity(self, x0, eps, t):
        """计算velocity"""
        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
        
        # v = sqrt(1-ᾱₜ)·x₀ - √(ᾱₜ)·ε
        return sqrt_one_minus_alphas_cumprod_t * x0 - sqrt_alphas_cumprod_t * eps
    
    def training_losses(self, x0):
        t = torch.randint(0, self.T, (x0.shape[0],))
        eps = torch.randn_like(x0)
        v = self.predict_velocity(x0, eps, t)
        
        xt = torch.sqrt(self.alphas_cumprod[t].view(-1, 1, 1, 1)) * x0 + \
             torch.sqrt(1 - self.alphas_cumprod[t].view(-1, 1, 1, 1)) * eps
        
        v_theta = self.denoiser(xt, t)
        return F.mse_loss(v_theta, v)

参数化对比

参数化	输出	优点	缺点
ε-预测	噪声	简单，最常用	中间步骤重建质量差
x₀-预测	原始图像	直观	训练初期不稳定
v-预测	速度向量	收敛快，质量好	需要调整网络结构

数学推导

后验分布 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$

这是推导的核心，利用贝叶斯定理：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}$$

由于所有分布都是高斯分布，乘积仍然是高斯分布，其均值和方差可以直接计算：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}\left({\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I}\right)$$

其中：

$${\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t$$ $${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$$

损失函数的推导

将 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$ 和 $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 的KL散度展开：

$$D_{\text{KL}}(q\Vert p_{\theta })=\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mathbf{\mu }}}_t-{\mathbf{\mu }}_{\theta }\Vert }^2$$

将 ${\tilde{\mathbf{\mu }}}_t$ 用 ${\mathbf{x}}_t$ 和 $\mathbf{ϵ}$ 表示：

$${\tilde{\mathbf{\mu }}}_t=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\mathbf{ϵ}\right)$$

假设 ${\mathbf{\mu }}_{\theta }$ 也是 ${\mathbf{x}}_t$ 和 ${\mathbf{ϵ}}_{\theta }$ 的函数：

$${\mathbf{\mu }}_{\theta }=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }\right)$$

代入KL散度：

$$D_{\text{KL}}(q\Vert p_{\theta })=\frac{{\beta }_t^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)}{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }\Vert }^2$$

由于常数项与 $\theta$ 无关，最小化KL散度等价于最小化噪声预测的MSE！

采样过程

DDPM采样

@torch.no_grad()
def ddpm_sampling(model, xt, betas, T):
    """DDPM反向采样过程"""
    alphas = 1 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    
    for t in reversed(range(T)):
        # 预测噪声
        eps_theta = model(xt, t)
        
        # 计算均值
        alpha_t = alphas[t]
        alpha_bar_t = alphas_cumprod[t]
        beta_t = betas[t]
        
        # 系数
        coef1 = beta_t / torch.sqrt(1 - alpha_bar_t)
        coef2 = torch.sqrt(alpha_t) * (1 - alphas_cumprod[t-1] if t > 0 else 1)
        coef3 = 1 - alphas_cumprod[t-1] if t > 0 else 1
        
        mean = (xt - coef1 * eps_theta) / torch.sqrt(alpha_t)
        
        # 添加噪声（最后一步不加噪声）
        if t > 0:
            noise = torch.randn_like(xt)
            xt = mean + torch.sqrt(beta_t) * noise
        else:
            xt = mean
    
    return xt

与Score Matching的关系

扩散模型的损失函数与Score Matching有深刻联系。

Score Matching的目标是学习对数密度的梯度（score function）：

$$s_{\theta }(\mathbf{x})={\nabla }_{\mathbf{x}}\log p(\mathbf{x})$$

Song Yang等人的工作（SDE框架）证明了：

$${\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)=-{\sigma }_t{\mathbf{s}}_{\theta }({\mathbf{x}}_t,t)$$

其中 ${\sigma }_t$ 是时间 $t$ 对应的噪声水平。这意味着扩散模型的去噪器本质上学习的是不同噪声水平下的score函数。

详见 Score Matching与SDE。

网络架构

UNet骨干网络

大多数扩散模型使用 U-Net 作为去噪网络：

class UNet(torch.nn.Module):
    def __init__(self, in_channels=3, out_channels=3, base_channels=128):
        super().__init__()
        # 编码器（下采样）
        self.enc1 = ResidualBlock(in_channels, base_channels)
        self.enc2 = ResidualBlock(base_channels, base_channels*2, downsample=True)
        self.enc3 = ResidualBlock(base_channels*2, base_channels*4, downsample=True)
        self.enc4 = ResidualBlock(base_channels*4, base_channels*4, downsample=True)
        
        # 瓶颈
        self.bottleneck = ResidualBlock(base_channels*4, base_channels*4)
        
        # 解码器（上采样）
        self.dec4 = ResidualBlock(base_channels*8, base_channels*4)
        self.dec3 = ResidualBlock(base_channels*6, base_channels*2)
        self.dec2 = ResidualBlock(base_channels*3, base_channels)
        self.dec1 = ResidualBlock(base_channels*2, out_channels)
        
        self.time_mlp = TimeEmbedding(base_channels)
        self.downsample = Downsample()
        self.upsample = Upsample()
    
    def forward(self, x, t):
        # 时间嵌入
        t_emb = self.time_mlp(t)
        
        # 编码路径
        e1 = self.enc1(x, t_emb)
        e2 = self.enc2(e1, t_emb)
        e3 = self.enc3(e2, t_emb)
        e4 = self.enc4(e3, t_emb)
        
        # 瓶颈
        b = self.bottleneck(e4, t_emb)
        
        # 解码路径（包含跳跃连接）
        d4 = self.dec4(torch.cat([b, e4], dim=1), t_emb)
        d3 = self.dec3(torch.cat([d4, e3], dim=1), t_emb)
        d2 = self.dec2(torch.cat([d3, e2], dim=1), t_emb)
        d1 = self.dec1(torch.cat([d2, e1], dim=1), t_emb)
        
        return d1

时间条件机制

扩散模型需要将时间步信息注入网络，常用方法：

1. 正弦位置编码：将 $t$ 映射到高频振荡基函数
2. AdaIN (Adaptive Instance Normalization)：将时间嵌入注入归一化层
3. Adaptive Layer Norm：$\text{AdaLN}(x,\gamma ,\beta )=\gamma \cdot \frac{x-\mu }{\sigma }+\beta$

参考

Footnotes

1. Ho et al., “Denoising Diffusion Probabilistic Models”, NeurIPS 2020. https://arxiv.org/abs/2006.11239 ↩