扩散模型与流匹配背后的几何学：Wasserstein 空间中的梯度流与测地线

背景

在生成式人工智能领域，扩散模型（Diffusion Models）和流匹配（Flow Matching）是两大主流技术范式。尽管它们在工程实现上各有侧重，且常被分别讨论，但二者在数学本质上共享着深厚的几何根基。

概率测度空间 $\mathcal{P}_2(\mathbb{R}^d)$（即具有有限二阶矩的概率测度空间）并非仅仅是一个统计集合，它承载着一种自然的几何结构。通过二次 Wasserstein 距离 $W_2$，该空间构成了一个完备度量空间。更为深刻的是，根据 Felix Otto 的理论，这个空间可以被形式化地视为一个黎曼流形，而该流形上的测地线（Geodesics）恰好对应于最优传输（Optimal Transport, OT）的插值路径。

理解这一几何结构，是统一理解扩散模型与流匹配的关键。扩散模型本质上是自由能（Free Energy）在 Wasserstein 空间上的梯度流，而流匹配则是在同一空间上寻找连接起点和终点的测地线。

核心内容

1. Wasserstein 空间的几何结构

概率测度空间 $\mathcal{P}_2(\mathbb{R}^d)$ 拥有独特的几何性质。二次 Wasserstein 距离 $W_2$ 不仅定义了空间中的距离，还赋予了其类似黎曼流形的结构。在这种几何视角下，概率分布之间的转换不再是简单的向量加减，而是沿着流形上的最短路径（测地线）或能量下降最快的方向进行。

2. 扩散模型的几何本质：自由能的梯度流

扩散模型的前向过程（加入噪声）和反向过程（去噪）可以被精确地映射到这一几何框架中：

• 自由能泛函：定义自由能泛函 $F(\rho) = KL(\rho || \pi)$，其中 $\rho$ 是当前分布，$\pi$ 是目标数据分布，$KL$ 为 Kullback-Leibler 散度。
• Fokker-Planck 方程：在该流形上，自由能 $F(\rho)$ 的梯度流（Gradient Flow）恰好对应于 Fokker-Planck 方程。这意味着概率分布随时间的演化遵循能量最小化的方向。
• JKO 方案：Fokker-Planck 方程的隐式欧拉离散化（Implicit-Euler discretization）即为 JKO 方案（Jordan-Kinderlehrer-Otto scheme）。
• 统一视角：扩散模型的前向过程是在自由能面上“下坡”，而每一步去噪操作实际上执行了一次 JKO 步骤。这一几何解释统一了 DDPM、DDIM、NCSN/SMLD 以及 Energy Matching 等模型。它们并非独立的理论，而是同一几何框架下的不同实现方案。

3. 流匹配的几何本质：Wasserstein 测地线

Wasserstein 空间支持另一种变分原理，这与扩散模型形成了鲜明对比：

• Benamou-Brenier 公式：该公式定义了最优传输问题的最小作用量曲线，这些曲线正是 Wasserstein 空间中的测地线。
• 边界值问题：流匹配（Flow Matching）学习的正是这些最优传输路径。与扩散模型不同，流匹配固定了起点（噪声分布）和终点（数据分布），并沿着连接这两点的测地线进行生成。
• 确定性 ODE：由于沿着测地线（直线）移动，生成过程可以转化为一个确定性的常微分方程（ODE）。这种路径的直线性使得生成过程所需的采样步数远少于扩散模型。

4. 两种范式的统一与对比

将扩散模型和流匹配置于同一个 Wasserstein 流形上，可以清晰地揭示它们的关系：

• 扩散模型：遵循自由能梯度流。这是一个初值问题（Initial-value problem），从噪声出发，逐步沿能量下降方向逼近数据分布。
• 流匹配：遵循 Wasserstein 测地线。这是一个边界值问题（Boundary-value problem），直接学习连接噪声和数据的最短路径。

尽管路径不同（一个是能量下降的曲线，一个是最优传输的直线），两者最终都到达了相同的终点（目标数据分布）。

关键要点

• 几何统一性：扩散模型和流匹配共享同一个几何基础——具有二次 Wasserstein 距离 $W_2$ 的概率测度空间 $\mathcal{P}_2(\mathbb{R}^d)$。
• 扩散模型的机制：扩散过程是自由能 $F(\rho) = KL(\rho || \pi)$ 的梯度流。前向过程是能量下降，反向去噪步骤等价于 JKO 方案的离散化步骤。这解释了为何 DDPM、DDIM 等模型在几何上是同源的。
• 流匹配的机制：流匹配学习的是 Benamou-Brenier 公式定义的最优传输路径，即 Wasserstein 空间中的测地线。
• 问题类型的差异：
  • 扩散模型处理的是初值问题（从噪声开始演化）。
  • 流匹配处理的是边界值问题（固定起点和终点，寻找中间路径）。
• 采样效率：由于流匹配沿着测地线（近似直线）生成，其生成轨迹更直接，因此通常比扩散模型需要更少的采样步数即可达到高质量结果。
• 殊途同归：两种模型虽然遵循不同的动力学路径（梯度流 vs. 测地线），但目标一致，最终都收敛到相同的数据分布。

意义与影响

这篇论文（及 arXiv 预印本所传达的观点）的核心价值在于去神秘化和统一化。

1. 理论澄清：长期以来，扩散模型和流匹配在文献中常被作为独立的技术分支进行介绍。本文通过引入 Otto 几何和最优传输理论，证明了二者本质上是同一几何空间上的两种不同动力学系统。这消除了社区中存在的理论割裂感。
2. 算法设计的指导：理解扩散模型是“梯度流”有助于解释为何需要大量采样步（因为能量景观复杂，路径曲折）；理解流匹配是“测地线”则解释了为何其采样效率高（路径最短）。这种几何洞察可以指导未来模型的设计，例如结合梯度流的稳定性与测地线的高效性。
3. 简化模型认知：对于从业者而言，无需将 DDPM、DDIM、NCSN 等视为截然不同的算法，而应视其为同一 JKO 离散化框架下的不同变体。这有助于更系统地掌握生成模型的核心原理。
4. 推动最优传输在生成模型中的应用：强调 Wasserstein 几何的重要性，鼓励研究者利用最优传输理论来改进生成模型的训练稳定性、收敛速度和生成质量。

总之，该文章提供了一个强有力的几何视角，将看似不同的生成模型家族统一在 Wasserstein 空间的框架下，为理解和发展下一代生成式 AI 奠定了坚实的理论基础。