扩散模型统一理论：测度论视角

概述

扩散模型、Score Matching和Flow Matching是当前生成式AI的核心技术¹²。最近的理论工作提出了一个统一的测度论框架，将这三类方法纳入同一理论体系³。本章从测度论角度深入分析这些方法的数学基础和内在联系。

测度论基础

概率测度与随机过程

设 $(\Omega ,\mathcal{F},\mathbb{P})$ 为概率空间，$X:\Omega \to {\mathbb{R}}^d$ 为随机变量。

定义（概率分布）：随机变量 $X$ 的概率分布是测度空间 $({\mathbb{R}}^d,\mathcal{B}({\mathbb{R}}^d))$ 上的概率测度 $\mu$：

$$\mu (A)=\mathbb{P}(X\in A),\forall A\in \mathcal{B}({\mathbb{R}}^d)$$

Kullback-Leibler散度

两个概率分布 $\mu$ 和 $\nu$ 之间的KL散度：

$$D_{\text{KL}}(\mu \Vert \nu )=\int \log \left(\frac{d\mu }{d\nu }\right)d\mu$$

前提是 $\mu \ll \nu$（$\mu$ 关于 $\nu$ 绝对连续）。

Wasserstein距离

${\mathcal{P}}_p({\mathbb{R}}^d)$ 表示具有 $p$ 阶矩的概率分布空间：

$$W_p(\mu ,\nu )={\left(\underset{\gamma \in \Gamma (\mu ,\nu )}{\inf }\int \Vert x-y{\Vert }^{p}d\gamma (x,y)\right)}^{1/p}$$

其中 $\Gamma (\mu ,\nu )$ 是边际为 $\mu ,\nu$ 的联合分布集合。

随机微分方程视角

前向扩散过程

设 $\mathbf{x}(t)\in {\mathbb{R}}^d$ 为 $t\in [0,T]$ 上的随机过程：

$$d\mathbf{x}(t)=\mathbf{f}(\mathbf{x}(t),t)dt+g(t)d\mathbf{w}(t)$$

其中：

• $\mathbf{f}:{\mathbb{R}}^d\times [0,T]\to {\mathbb{R}}^d$ 是漂移函数（drift）
• $g:[0,T]\to \mathbb{R}$ 是扩散系数（diffusion coefficient）
• $\mathbf{w}(t)$ 是标准布朗运动

特殊情况：DDPM使用常系数：

$$d\mathbf{x}(t)=-\frac{1}{2}\mathbf{x}(t)dt+d\mathbf{w}(t)$$

SDE的解与转移概率

SDE的解给出转移概率核 $p_{0t}({\mathbf{x}}_0|{\mathbf{x}}_t)$：

$$\mathbf{x}(t)|\mathbf{x}(0)\sim \mathcal{N}\left(e^{-t/2}\mathbf{x}(0),(1-e^{-t})\mathbf{I}\right)$$

Fokker-Planck方程

SDE的前向过程对应Fokker-Planck方程：

$$\frac{\partial p_t}{\partial t}=-\nabla \cdot (\mathbf{f}{p}_t)+\frac{1}{2}\Delta (g^2{p}_t)$$

扩散模型的理论框架

证据下界（ELBO）

设 $q({\mathbf{x}}_{0:T})$ 为数据生成的前向过程，$p_{\theta }({\mathbf{x}}_{0:T})$ 为模型。ELBO：

$$\log p_{\theta }({\mathbf{x}}_0)\ge {\mathbb{E}}_q\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\right]$$

展开后：

$$\mathcal{L}={\mathbb{E}}_q\left[\sum _{t=1}^T{D}_{\text{KL}}(q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})\Vert p_{\theta }({\mathbf{x}}_t|{\mathbf{x}}_{t-1}))\right]+\text{常数}$$

Score Matching目标

Stein分数：${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})$

score matching损失⁴：

$${\mathcal{L}}_{\text{SM}}(\theta )={\mathbb{E}}_{t,{\mathbf{x}}_t}\left[{\Vert {\mathbf{s}}_{\theta }({\mathbf{x}}_t,t)-{\nabla }_{{\mathbf{x}}_t}\log p_{0t}({\mathbf{x}}_t)\Vert }^2\right]$$

去噪扩散模型

DDPM的训练目标简化形式：

$${\mathcal{L}}_{\text{simple}}={\mathbb{E}}_{t,{\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]$$

Flow Matching

条件概率路径

Flow Matching通过定义确定性的概率路径 $p_t$ 从 $\pi$（噪声）到 $p_0$（数据）：

$$p_t=(1-t)\cdot \pi +t\cdot p_0,t\in [0,1]$$

这不是合理的概率密度混合，而是插值分布。

最优传输Flow Matching

最优传输条件Flow Matching (OT-CFM)⁵：

$${\mathbf{v}}_t(\mathbf{x})=\frac{(1-t)\sigma \mathbf{x}-\nabla \log p_0(\mathbf{x})}{1-t+t{\sigma }^2}$$

这定义了一个向量场 ${\mathbf{v}}_t$，其积分曲线从噪声流向数据。

速度预测模型

Flow Matching训练速度预测模型 $v_{\theta }$：

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

其中 ${\mathbf{x}}_t=(1-t){\mathbf{x}}_1+t{\mathbf{x}}_0$（线性插值）。

统一测度论框架

核心定理

定理（统一框架³）：扩散模型、Score Matching和Flow Matching都可以统一为**Wasserstein梯度流（Wasserstein Gradient Flow）**的学习问题。

Wasserstein梯度流

概率分布空间上的梯度流定义为：

$$\frac{\partial {\mu }_t}{\partial t}=\nabla \cdot \left({\mu }_t\nabla \frac{\delta \mathcal{F}}{\delta \mu }[{\mu }_t]\right)$$

其中 $\mathcal{F}$ 是自由能泛函：

$$\mathcal{F}(\mu )=\int f(\rho )d\mu +\frac{1}{2}\int Vd\mu +\int W\ast \mu d\mu$$

三种方法的统一视角

方法	梯度流形式	目标泛函
Score Matching	${\partial }_t{\mu }_t=\nabla \cdot ({\mu }_t\nabla \log p_t)$	匹配真实分布
Flow Matching	${\partial }_t{\mu }_t=-\nabla \cdot ({\mu }_t{\mathbf{v}}_t)$	沿向量场流动
DDPM	${\partial }_t{\mu }_t=-\frac{1}{2}\nabla \cdot ({\mu }_t\nabla \log {\sigma }_t^2)$	最小化KL散度

统一损失函数

在统一框架下，所有方法都可以表示为：

$$\mathcal{L}={\mathbb{E}}_{t,{\mathbf{x}}_t\sim {\mu }_t}\left[{\Vert {\mathbf{u}}_{\theta }({\mathbf{x}}_t,t)-{\mathbf{u}}_t^{\ast }({\mathbf{x}}_t)\Vert }^2\right]$$

其中：

• ${\mathbf{u}}_{\theta }$ 是模型预测的向量场/分数
• ${\mathbf{u}}_t^{\ast }$ 是目标向量场（具体形式取决于方法）

训练动态分析

分数空间的收敛性

设 $s^{\ast }(\mathbf{x},t)=\nabla \log p_t(\mathbf{x})$ 为真实分数，$s_{\theta }(\mathbf{x},t)$ 为预测分数。

收敛定理：在适当条件下，当 $\theta$ 优化至最优时：

$$\Vert s_{\theta }-s^{\ast }{\Vert }_{L^2(p_t)}\to 0$$

分布收敛

定理：如果训练收敛，则生成分布 $\hat{p}$ 满足：

$$W_2(\hat{p},p_{\text{data}})\le ϵ$$

其中 $ϵ$ 与训练误差相关。

实现细节

统一训练框架

import torch
import torch.nn as nn
from typing import Callable, Tuple
 
class UnifiedDiffusionFramework:
    """
    统一的扩散/Flow Matching训练框架
    支持：DDPM, Score Matching, Flow Matching
    """
    
    def __init__(
        self,
        model: nn.Module,
        method: str = 'ddpm',  # 'ddpm', 'score_matching', 'flow_matching'
        sigma_min: float = 0.01,
        sigma_max: float = 50.0,
    ):
        self.model = model
        self.method = method
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
    
    def get_noise_schedule(self, t: torch.Tensor) -> torch.Tensor:
        """噪声调度"""
        if self.method == 'ddpm':
            # DDPM调度
            return self.sigma_min * (self.sigma_max / self.sigma_min) ** t
        elif self.method == 'flow_matching':
            # Flow Matching：线性调度
            return self.sigma_min + (self.sigma_max - self.sigma_min) * t
        else:
            # Score Matching
            return torch.exp(t * torch.log(self.sigma_max) + 
                           (1-t) * torch.log(self.sigma_min))
    
    def compute_target(
        self,
        x0: torch.Tensor,
        xt: torch.Tensor,
        t: torch.Tensor,
        noise: torch.Tensor,
    ) -> torch.Tensor:
        """
        计算目标向量（根据方法不同）
        
        Returns:
            目标向量：分数、速度或噪声
        """
        sigma = self.get_noise_schedule(t)
        
        if self.method == 'ddpm':
            # DDPM：预测噪声
            return noise
        
        elif self.method == 'score_matching':
            # Score Matching：预测分数
            return -noise / sigma.unsqueeze(-1)
        
        elif self.method == 'flow_matching':
            # Flow Matching：预测速度
            return x0 - xt  # 从噪声到数据的向量
    
    def training_loss(
        self,
        x0: torch.Tensor,
    ) -> Tuple[torch.Tensor, dict]:
        """
        计算训练损失
        
        Args:
            x0: 真实数据 [batch, dim]
        
        Returns:
            loss: 标量损失
            info: 诊断信息
        """
        batch_size = x0.shape[0]
        device = x0.device
        
        # 采样时间步
        t = torch.rand(batch_size, device=device)
        
        # 采样噪声
        noise = torch.randn_like(x0)
        
        # 获取噪声调度
        sigma = self.get_noise_schedule(t)
        
        # 加噪
        xt = x0 + sigma.unsqueeze(-1) * noise
        
        # 模型预测
        pred = self.model(xt, t)
        
        # 计算目标
        target = self.compute_target(x0, xt, t, noise)
        
        # 计算损失
        loss = 0.5 * ((pred - target) ** 2).sum(dim=-1).mean()
        
        info = {
            't_mean': t.mean().item(),
            'sigma_mean': sigma.mean().item(),
            'pred_norm': pred.norm().item() / batch_size,
        }
        
        return loss, info
    
    @torch.no_grad()
    def sampling(
        self,
        shape: Tuple[int, ...],
        num_steps: int = 100,
        ode_solver: str = 'euler',
    ) -> torch.Tensor:
        """
        采样生成
        
        Args:
            shape: 输出形状
            num_steps: 采样步数
            ode_solver: ODE求解器 ('euler', 'heun')
        
        Returns:
            x0: 生成样本
        """
        x = torch.randn(shape, device=next(self.model.parameters()).device)
        dt = 1.0 / num_steps
        
        for i in range(num_steps):
            t = torch.full((shape[0],), i / num_steps, device=x.device)
            
            # 预测向量场
            v = self.model(x, t)
            
            if self.method == 'ddpm':
                # DDPM采样（离散）
                alpha_bar = 1 - (self.sigma_min ** 2) * \
                           (self.sigma_max ** 2 / self.sigma_min ** 2) ** t
                pred_x0 = (x - (1 - alpha_bar).sqrt() * v) / alpha_bar.sqrt()
                x = x - (1 - alpha_bar) * pred_x0 + v * dt
            else:
                # Flow Matching / Score Matching
                x = x + v * dt
        
        return x

概率路径的可视化

import matplotlib.pyplot as plt
import numpy as np
 
def visualize_probability_paths(
    framework: UnifiedDiffusionFramework,
    data_samples: np.ndarray,
    num_times: int = 5,
):
    """
    可视化不同方法的概率路径
    
    Args:
        data_samples: 真实数据样本 [n_samples, dim]
        num_times: 要可视化的时间步数
    """
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    methods = ['ddpm', 'score_matching', 'flow_matching']
    titles = ['DDPM', 'Score Matching', 'Flow Matching']
    
    for ax, method, title in zip(axes, methods, titles):
        framework.method = method
        
        for i, t_val in enumerate(np.linspace(0, 1, num_times)):
            t = torch.full((len(data_samples),), t_val)
            sigma = framework.get_noise_schedule(t).numpy()
            
            # 加噪
            noise = np.random.randn(*data_samples.shape)
            xt = data_samples + sigma[:, None] * noise
            
            ax.scatter(xt[:, 0], xt[:, 1], 
                      alpha=0.5, s=10,
                      label=f't={t_val:.1f}' if i == 0 else None)
        
        ax.set_title(title)
        ax.legend()
        ax.set_xlabel('x₁')
        ax.set_ylabel('x₂')
    
    plt.tight_layout()
    plt.savefig('probability_paths.png')
    plt.show()

实践注意事项

方法选择指南

场景	推荐方法	理由
高维图像生成	DDPM / Flow Matching	理论基础扎实
快速采样	Flow Matching / Consistency	收敛快
连续时间建模	Score Matching	SDE框架自然
密度估计	Score Matching	分数易计算

训练技巧

1. 噪声调度：DDPM推荐 $\beta$ 调度，Flow Matching推荐线性调度
2. 时间嵌入：使用Transformer式的时间编码
3. 网络架构：U-Net + Self-Attention是图像生成的标准架构
4. 正则化：避免数值不稳定，特别是大时间步

统一理论的意义

理论价值

1. 统一视角：将三类方法纳入同一数学框架
2. 理论保证：提供收敛性和最优性保证
3. 新方法发现：为设计新方法提供指导

实践价值

1. 灵活切换：不同场景可选择最优方法
2. 性能预测：理论指导实践参数选择
3. 组合优化：混合不同方法的优点

总结

从测度论视角，扩散模型、Score Matching和Flow Matching都统一于Wasserstein梯度流框架：

1. 数学基础：SDE、Fokker-Planck方程、Wasserstein距离
2. 统一目标：学习目标向量场，使分布沿梯度流演化
3. 方法差异：体现在目标向量场 ${\mathbf{u}}_t^{\ast }$ 的定义方式
4. 实践选择：根据任务特性选择合适的方法和调度策略

这一统一理论不仅深化了我们对生成模型的理解，也为未来更强大的生成方法奠定了理论基础。

参考文献

Footnotes

1. Ho et al. (2020). “Denoising Diffusion Probabilistic Models.” NeurIPS 2020 ↩

2. Song et al. (2021). “Score-Based Generative Modeling through Stochastic Differential Equations.” ICLR 2021 ↩

3. [arXiv 2605.06829] “A Unified Measure-Theoretic View of Diffusion, Score-Based, and Flow Matching Models” ↩ ↩²

4. Hyvarinen (2005). “Estimation of Non-Normalized Statistical Models by Score Matching.” JMLR 2005 ↩

5. Lipman et al. (2022). “Flow Matching for Generative Modeling.” ICLR 2022 ↩