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

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

**来源：** https://arxiv.org/html/2606.24157

###### 摘要

具有有限二阶矩的概率测度空间 \(\mathcal{P}_2(\mathbb{R}^d)\) 承载着一种自然的几何结构：二次 Wasserstein 距离 \(W_2\) 使其成为一个完备度量空间，并且，遵循 Otto 的观点，成为一个（形式上的）黎曼流形，其测地线正是最优传输插值。在该流形上，自由能 \(\mathcal{F}(\rho) = \mathrm{KL}(\rho\|\pi)\) 的梯度流恰好是 Fokker-Planck 方程，其隐式欧拉离散化即为 JKO 格式。这正是扩散模型背后的几何结构：前向过程沿自由能下降，而每个去噪步骤实现了 JKO 的一步，从而统一了 DDPM、DDIM、NCSN/SMLD 和 Energy Matching——同一格式，而非分离的理论。同一流形还支持第二个变分原理。其测地线——Benamou–Brenier 公式的最小作用量曲线——正是流匹配所学习的最优传输路径。固定两个端点并沿测地线进行，生成过程变为沿直线的确定性常微分方程，因此所需采样步数大幅减少。将这两类模型置于同一流形上，它们的关系变得精确：扩散遵循自由能梯度流，是一个初值问题；最优传输流匹配遵循 Wasserstein 测地线，是一个边值问题。两者沿着不同路径到达相同端点。

Wasserstein 距离 → 赋予几何结构 → 概率空间成为黎曼流形，承载两个互补的变分原理：
- (A) 扩散——自由能*梯度流*（*初值*问题：给定 \(\rho_k\)，在 \(W_2\) 中保持接近的同时下降 \(\mathcal{F}\)）  
  自由能 → Fokker-Planck → 时间离散化 → JKO 格式
- (B) 流匹配——Benamou–Brenier*最小作用量*（*边值*问题：给定 \(\rho_0\) 和 \(\rho_1\)，最小化动能）  
  最小作用量 → Wasserstein 测地线 → OT 传输路径

本质区别：*相同端点* \(\rho_1 = p_{\text{data}}\)，*不同路径*。

*   大纲
    -   **序言**从扩散模型和流匹配中出现的连续性方程和 Fokker–Planck 方程开始，并揭示这些处理隐含的联系。
    -   **第 1 节**介绍 Wasserstein 距离——两个分布之间的传输代价。
    -   **第 2 节**推导连续性方程——质量守恒的局部形式。
    -   **第 3 节**建立黎曼结构：Benamou–Brenier 公式将 \(W_2\) 识别为测地距离，其测地线即为流匹配学习的最优传输路径。
    -   **第 4 节**从随机微分方程推导 Fokker–Planck 方程，并将其重写为连续性方程。
    -   **第 5 节**证明核心恒等式：Fokker–Planck 方程是自由能的 Wasserstein 梯度流。
    -   **第 6 节**将流离散化为 JKO 格式，证明收敛性，并从单个变分模板恢复 DDPM、DDIM、NCSN/SMLD、流匹配和 Energy Matching。

**预备知识。** 我们假设读者熟悉多元微积分、线性代数和初等概率论；测度论、微分几何和凸分析将在自成体系的附录中展开。

###### 目录

1. 符号与约定 (https://arxiv.org/html/2606.24157#Sx1)
2. 序言：你可能已经知道的事情 (https://arxiv.org/html/2606.24157#Sx2)
    1. 场景 1：在扩散模型中遇到 Fokker-Planck 方程 (https://arxiv.org/html/2606.24157#Sx2.SSx1)
    2. 场景 2：在流匹配中遇到连续性方程 (https://arxiv.org/html/2606.24157#Sx2.SSx2)
    3. 场景 3：概率流常微分方程——两个方程交汇之处 (https://arxiv.org/html/2606.24157#Sx2.SSx3)
    4. 这些联系有多深刻？ (https://arxiv.org/html/2606.24157#Sx2.SSx4)
    5. 本文的使命 (https://arxiv.org/html/2606.24157#Sx2.SSx5)
    6. 相关视角 (https://arxiv.org/html/2606.24157#Sx2.SSx6)
3. 1 Wasserstein 距离：运输土壤的成本 (https://arxiv.org/html/2606.24157#S1)
    1. 1.1 为什么我们需要一个新的距离？ (https://arxiv.org/html/2606.24157#S1.SS1)
    2. 1.2 Kantorovich 形式 (https://arxiv.org/html/2606.24157#S1.SS2)
    3. 1.3 Monge 形式和 Brenier 定理 (https://arxiv.org/html/2606.24157#S1.SS3)
    4. 1.4 一个具体例子：高斯分布 (https://arxiv.org/html/2606.24157#S1.SS4)
    5. 1.5 关键性质 (https://arxiv.org/html/2606.24157#S1.SS5)
4. 2 连续性方程：质量守恒 (https://arxiv.org/html/2606.24157#S2)
    1. 2.1 物理图像 (https://arxiv.org/html/2606.24157#S2.SS1)
    2. 2.2 从第一原理推导 (https://arxiv.org/html/2606.24157#S2.SS2)
    3. 2.3 弱形式 (https://arxiv.org/html/2606.24157#S2.SS3)
    4. 2.4 与粒子常微分方程的联系 (https://arxiv.org/html/2606.24157#S2.SS4)
5. 3 Wasserstein 空间的黎曼结构 (https://arxiv.org/html/2606.24157#S3)
    1. 3.1 Benamou–Brenier 公式 (https://arxiv.org/html/2606.24157#S3.SS1)
    2. 3.2 为什么连续性方程 + 最小动能给出最优传输路径 (https://arxiv.org/html/2606.24157#S3.SS2)
    3. 3.3 Otto 的黎曼解释 (https://arxiv.org/html/2606.24157#S3.SS3)
    4. 3.4 测地线 = 位移插值 (https://arxiv.org/html/2606.24157#S3.SS4)
6. 4 Fokker–Planck 方程 (https://arxiv.org/html/2606.24157#S4)
    1. 4.1 概述 (https://arxiv.org/html/2606.24157#S4.SS1)
    2. 4.2 动机：物理设定 (https://arxiv.org/html/2606.24157#S4.SS2)
    3. 4.3 预备知识：Itô 公式 (https://arxiv.org/html/2606.24157#S4.SS3)
    4. 4.4 Fokker–Planck 的推导：完整细节 (https://arxiv.org/html/2606.24157#S4.SS4)
    5. 4.5 重写为连续性方程 (https://arxiv.org/html/2606.24157#S4.SS5)
    6. 4.6 稳态分布 (https://arxiv.org/html/2606.24157#S4.SS6)
    7. 4.7 与现代生成模型的联系 (https://arxiv.org/html/2606.24157#S4.SS7)
    8. 4.8 概念澄清：什么需要能量，什么不需要？ (https://arxiv.org/html/2606.24157#S4.SS8)
    9. 4.9 扩散中的建模选择在哪里？ (https://arxiv.org/html/2606.24157#S4.SS9)
7. 5 自由能与 Wasserstein 梯度 (https://arxiv.org/html/2606.24157#S5)
    1. 5.1 自由能泛函 (https://arxiv.org/html/2606.24157#S5.SS1)
    2. 5.2 一阶变分（泛函导数） (https://arxiv.org/html/2606.24157#S5.SS2)
    3. 5.3 计算 Wasserstein 梯度 (https://arxiv.org/html/2606.24157#S5.SS3)
    4. 5.4 主要结论：Fokker–Planck 作为 Wasserstein 梯度流 (https://arxiv.org/html/2606.24157#S5.SS4)
    5. 5.5 节末比较：欧几里得 vs Wasserstein 梯度流逻辑 (https://arxiv.org/html/2606.24157#S5.SS5)
8. 6 JKO 格式 (https://arxiv.org/html/2606.24157#S6)
    1. 6.1 隐式欧拉方法：从梯度流到优化问题 (https://arxiv.org/html/2606.24157#S6.SS1)
    2. 6.2 JKO 格式：Wasserstein 空间中的隐式欧拉 (https://arxiv.org/html/2606.24157#S6.SS2)
    3. 6.3 从梯度流推导 JKO (https://arxiv.org/html/2606.24157#S6.SS3)
    4. 6.4 JKO 的一阶最优性条件 (https://arxiv.org/html/2606.24157#S6.SS4)
    5. 6.5 严格收敛定理 (https://arxiv.org/html/2606.24157#S6.SS5)
    6. 6.6 JKO 的结构优势 (https://arxiv.org/html/2606.24157#S6.SS6)
    7. 6.7 例子：热方程 (https://arxiv.org/html/2606.24157#S6.SS7)
    8. 6.8 应用：Energy Matching (Balcerak et al., 2025) (https://arxiv.org/html/2606.24157#S6.SS8)
    9. 6.9 从 JKO 到主流生成算法的统一推导 (https://arxiv.org/html/2606.24157#S6.SS9)
9. 7 全局图景 (https://arxiv.org/html/2606.24157#S7)
10. A 测度与耦合 (https://arxiv.org/html/2606.24157#A1)
    1. A.1 什么是测度？ (https://arxiv.org/html/2606.24157#A1.SS1)
    2. A.2 绝对连续：真正含义是什么？ (https://arxiv.org/html/2606.24157#A1.SS2)
    3. A.3 关于测度的积分 (https://arxiv.org/html/2606.24157#A1.SS3)
    4. A.4 前推测度 (https://arxiv.org/html/2606.24157#A1.SS4)
    5. A.5 耦合：最优传输的关键概念 (https://arxiv.org/html/2606.24157#A1.SS5)
11. B 微分几何预备知识 (https://arxiv.org/html/2606.24157#A2)
12. C 分析工具：分部积分与 du Bois-Reymond 引理 (https://arxiv.org/html/2606.24157#A3)
    1. C.1 分部积分 (https://arxiv.org/html/2606.24157#A3.SS1)
        1. C.1.1 一维情形 (https://arxiv.org/html/2606.24157#A3.SS1.SSS1)
        2. C.1.2 高维情形：散度定理版本 (https://arxiv.org/html/2606.24157#A3.SS1.SSS2)
        3. C.1.3 传递拉普拉斯算子（两次分部积分） (https://arxiv.org/html/2606.24157#A3.SS1.SSS3)
        4. C.1.4 时间上的分部积分 (https://arxiv.org/html/2606.24157#A3.SS1.SSS4)
        5. C.1.5 公式总结表 (https://arxiv.org/html/2606.24157#A3.SS1.SSS5)
    2. C.2 Du Bois-Reymond 引理（变分法基本引理） (https://arxiv.org/html/2606.24157#A3.SS2)
        1. C.2.1 直观理解 (https://arxiv.org/html/2606.24157#A3.SS2.SSS1)
        2. C.2.2 证明概要 (https://arxiv.org/html/2606.24157#A3.SS2.SSS2)
        3. C.2.3 约束版本（用于正文） (https://arxiv.org/html/2606.24157#A3.SS2.SSS3)
        4. C.2.4 在正文中的使用位置 (https://arxiv.org/html/2606.24157#A3.SS2.SSS4)
13. D Gibbs–Boltzmann 分布 (https://arxiv.org/html/2606.24157#A4)
    1. D.1 核心直觉 (https://arxiv.org/html/2606.24157#A4.SS1)
    2. D.2 温度 \(T\) 的作用——单一控制旋钮 (https://arxiv.org/html/2606.24157#A4.SS2)
    3. D.3 跨学科应用 (https://arxiv.org/html/2606.24157#A4.SS3)
    4. D.4 与正文的核心联系 (https://arxiv.org/html/2606.24157#A4.SS4)
    5. D.5 为什么是指数形式？ (https://arxiv.org/html/2606.24157#A4.SS5)
14. E Fisher 信息 (https://arxiv.org/html/2606.24157#A5)
    1. E.1 统计学中的 Fisher 信息（参数版本） (https://arxiv.org/html/2606.24157#A5.SS1)
    2. E.2 信息论/偏微分方程中的 Fisher 信息（分布版本） (https://arxiv.org/html/2606.24157#A5.SS2)
    3. E.3 两个版本的统一 (https://arxiv.org/html/2606.24157#A5.SS3)
    4. E.4 相对 Fisher 信息 (https://arxiv.org/html/2606.24157#A5.SS4)
    5. E.5 三个恒等式 (https://arxiv.org/html/2606.24157#A5.SS5)
    6. E.6 信息不等式 (https://arxiv.org/html/2606.24157#A5.SS6)
    7. E.7 例子 (https://arxiv.org/html/2606.24157#A5.SS7)
15. F 生成模型：分数、速度和扩散 (https://arxiv.org/html/2606.24157#A6)
    1. F.1 分数、速度和漂移之间的完整关系 (https://arxiv.org/html/2606.24157#A6.SS1)
    2. F.2 分数的几何意义：熵的 Wasserstein 梯度 (https://arxiv.org/html/2606.24157#A6.SS2)
    3. F.3 主流模型的统一框架 (https://arxiv.org/html/2606.24157#A6.SS3)
        1. F.3.1 DDPM（去噪扩散概率模型） (https://arxiv.org/html/2606.24157#A6.SS3.SSS1)
        2. F.3.2 NCSN / SMLD (https://arxiv.org/html/2606.24157#A6.SS3.SSS2)
        3. F.3.3 VE-SDE（方差爆炸） (https://arxiv.org/html/2606.24157#A6.SS3.SSS3)
        4. F.3.4 VP-SDE（方差保持） (https://arxiv.org/html/2606.24157#A6.SS3.SSS4)
        5. F.3.5 流匹配 (https://arxiv.org/html/2606.24157#A6.SS3.SSS5)
    4. F.4 统一对照表 (https://arxiv.org/html/2606.24157#A6.SS4)
    5. F.5 从 Fokker–Planck/JKO 视角的统一理解 (https://arxiv.org/html/2606.24157#A6.SS5)
16. G 凸分析快速回顾 (https://arxiv.org/html/2606.24157#A7)
    1. G.1 凸函数 (https://arxiv.org/html/2606.24157#A7.SS1)
    2. G.2 海森矩阵 (https://arxiv.org/html/2606.24157#A7.SS2)
    3. G.3 正定性与半正定性 (https://arxiv.org/html/2606.24157#A7.SS3)
    4. G.4 强凸性与 \(\lambda\)-凸性 (https://arxiv.org/html/2606.24157#A7.SS4)
17. H 补充概念 (https://arxiv.org/html/2606.24157#A8)
    1. H.1 Langevin 动力学 (https://arxiv.org/html/2606.24157#A8.SS1)
    2. H.2 Kolmogorov 前向方程 (https://arxiv.org/html/2606.24157#A8.SS2)
    3. H.3 Ornstein-Uhlenbeck (OU) 过程 (https://arxiv.org/html/2606.24157#A8.SS3)
    4. H.4 Otto 微积分 (https://arxiv.org/html/2606.24157#A8.SS4)
    5. H.5 测地线 (https://arxiv.org/html/2606.24157#A8.SS5)
    6. H.6 Riesz 表示定理 (https://arxiv.org/html/2606.24157#A8.SS6)
    7. H.7 Lyapunov 函数 (https://arxiv.org/html/2606.24157#A8.SS7)
    8. H.8 KKT 条件 (https://arxiv.org/html/2606.24157#A8.SS8)

## 符号与约定

以下符号在本文中频繁使用，现集中于此以便参考。

### 集合与空间

-   \(\mathbb{R}^d\)：\(d\) 维欧几里得空间。
-   Borel 集：从 \(\mathbb{R}^d\) 中的所有开集出发，对可数并、交和补运算封闭，得到 Borel \(\sigma\)-代数，记为 \(\mathcal{B}(\mathbb{R}^d)\)。其元素称为 Borel 集。直观上，所有“合理”的几何集合（开集、闭集、可数并/交）都是 Borel 集；你在实践中遇到的几乎所有集合都是 Borel 集。
-   \(\mathcal{P}_2(\mathbb{R}^d)\)：所有定义在 \(\mathbb{R}^d\) 上且具有有限二阶矩的概率测度的集合，即满足 \(\int_{\mathbb{R}^d} \|x\|^2 d\mu(x) < \infty\) 的概率测度 \(\mu\)。
-   \(L^2(\mathbb{R}^d)\)（或 \(L^2(\rho)\)）：平方可积函数空间。\(L^2(\mathbb{R}^d)\) 由所有满足 \(\int_{\mathbb{R}^d} |f(x)|^2 dx < \infty\) 的函数 \(f\) 组成。\(L^2(\rho)\) 是带权 \(\rho\) 的版本：\(\int |f|^2 \rho dx < \infty\)。内积为 \(\langle f,g \rangle_{L^2(\rho)} = \int f(x) g(x) \rho(x) dx\)。这是一个**希尔伯特空间**（完备内积空间）。
-   \(C_c^\infty(\mathbb{R}^d)\)：紧支撑光滑函数空间。它包含所有无穷可微且在某个有界区域外为零的函数。这里：
    -   光滑 (\(C^\infty\)) = 所有阶导数存在且连续。
    -   紧支撑 (\(c\) 代表紧凑支撑) = 存在一个有界闭集 \(K\)，使得对所有 \(x \notin K\)，有 \(f(x) = 0\)。直观上，函数具有“有限的活动范围”，在远处为零。

### 算子与符号

-   \(\inf A\)（下确界）：集合 \(A\) 的最大下界。例如，\(\inf\{1/n : n \in \mathbb{N}\} = 0\)。类似于 \(\min\)，但允许最小值未达到（极限情况）。
-   \(\sup A\)（上确界）：集合 \(A\) 的最小上界，是 \(\max\) 的推广。
-   \(\arg\min_x f(x)\)：使 \(f(x)\) 达到最小值的 \(x\) 的值（不是最小值本身，而是达到最小值的“位置”）。
-   \(\nabla f\)：梯度（\(f\) 关于空间变量的一阶偏导数向量）。
-   \(\nabla \cdot v\)：散度（向量场 \(v\) 各分量偏导数之和）。
-   \(\Delta f = \nabla \cdot (\nabla f)\)：拉普拉斯算子。

### 关键术语

-   a.e.（几乎处处）：一个性质“几乎处处成立”意味着它可能在某些点失效，但失效点的集合测度为零（“体积为零”）。例如，函数 \(f\) 和 \(g\) 在勒贝格测度下几乎处处相等意味着集合 \(\{x: f(x) \neq g(x)\}\) 的体积为零。
-   测试函数：\(C_c^\infty(\mathbb{R}^d)\) 中的函数。