JKO方案中的一阶最优性条件与生成框架

在上一篇文章什么是Jordan–Kinderlehrer–Otto (JKO)方案中，我们介绍了Jordan–Kinderlehrer–Otto (JKO)方案的核心公式及其在生成模型中的应用。JKO方案通过最小化传输成本、势能和负熵，描述了概率分布从噪声到数据的平滑演化。本文将进一步探讨JKO方案中的一阶最优性条件，并介绍一个结合最优传输（Optimal Transport, OT）和能量基础模型（Energy-Based Models, EBMs）的生成框架。我们将详细解析公式、提供直观解释，并通过例子说明其工作原理。

paper ：https://arxiv.org/pdf/2504.10612

一阶最优性条件

为了深入理解JKO方案的动态行为，我们需要分析其优化问题的一阶最优性条件。在JKO方案中，概率分布 ( ${\rho }_{t+\Delta t}$ ) 是通过最小化以下目标函数得到的：

$${\rho }_{t+\Delta t}=\arg \underset{\rho }{\min }\left\{\frac{W_2^2(\rho ,{\rho }_t)}{2\Delta t}+\int V_{\theta }(x)d\rho (x)+\epsilon (t)\int \rho (x)\log \rho (x)dx\right\}$$

根据Terpin等人（2024年）的研究，我们可以通过一阶最优性条件分析这一优化问题在每个时间步 ( $t$ ) 的行为。假设 ( ${\gamma }_t$ ) 是 ( ${\rho }_t$ ) 和 ( ${\rho }_{t+\Delta t}$ ) 之间的最优传输计划（OT plan），其支持集为 ( $\text{supp}({\gamma }_t)$ )，一阶最优性条件为：

$$\frac{1}{\Delta t}(x-y)+{\nabla }_x{V}_{\theta }(x)+\epsilon (t){\nabla }_x\log {\rho }_t(x)=0,(x,y)\in \text{supp}({\gamma }_t)$$

公式解析

这个公式描述了在最优传输计划 ( ${\gamma }_t$ ) 中，所有配对点 ( $(x,y)$ )（即从 ( ${\rho }_t$ ) 的点 ( $y$ ) 移动到 ( ${\rho }_{t+\Delta t}$ ) 的点 ( $x$ )）需要满足的平衡条件。公式包含三个项：

1. 传输项：($\frac{1}{\Delta t}(x-y)$)

   • 表示从 ( $y$ ) 到 ( $x$) 的位移向量，乘以 ($\frac{1}{\Delta t}$) 表示单位时间内的移动速度。
   • 这一项反映了Wasserstein距离中的传输成本，倾向于使 ( $x$ ) 和 ( $y$ ) 尽可能接近，以减少移动代价。

2. 势能梯度：(${\nabla }_x{V}_{\theta }(x)$)

   • ( $V_{\theta }(x)$ ) 是标量势函数，其梯度 (${\nabla }_x{V}_{\theta }(x)$) 指向能量增加的方向。
   • 这一项引导点 ( $x$ ) 向低能量区域移动，通常对应于数据密度高的区域。

3. 熵梯度：($\epsilon (t){\nabla }_x\log {\rho }_t(x)$)

   • (${\nabla }_x\log {\rho }_t(x)=\frac{{\nabla }_x{\rho }_t(x)}{{\rho }_t(x)}$) 是概率密度 ( ${\rho }_t(x)$ ) 的对数梯度，指向密度增加的方向。
   • ($\epsilon (t)$) 控制熵的影响力，类似于温度参数。当 ($\epsilon (t)$) 较大时，这一项使分布更均匀，防止过度集中。

直观解释

可以将这个公式想象为一个物理系统中粒子的运动：

• 传输项 像是一个“惯性”力，试图保持粒子位置不变（即 ( $x\approx y$ )）。
• 势能梯度 像是一个“引力”，拉动粒子向低能量区域（数据分布）。
• 熵梯度 像是一个“扩散”力，鼓励粒子分散，防止过于集中。

这三个力的平衡决定了粒子（即概率分布中的点）在每一步的移动方向和速度。

两个关键洞察

通过分析一阶最优性条件，我们可以得出两个关键洞察，分别对应于演化过程中的不同阶段：

1. 远离数据流形（早期阶段，( $t<{\tau }^{\ast }$ )）

在早期阶段，假设温度参数 ($\epsilon (t)=0$)（根据线性调度），公式简化为：

$$\frac{1}{\Delta t}(x-y)+{\nabla }_x{V}_{\theta }(x)=0,(x,y)\in \text{supp}({\gamma }_t)$$

即：
$$x-y=-\Delta t{\nabla }_x{V}_{\theta }(x)$$

意义：

• 此时，熵项消失，系统的演化完全由传输成本和势能驱动，呈现出最优传输（OT）流的特性。
• 点 ( $y$ )（在 ( ${\rho }_t$ ) 中）会沿着势能函数 ( $V_{\theta }(x)$ ) 的负梯度方向移动到 ( $x$ )，类似于一个确定性流（deterministic flow）。
• 这种行为适合在样本空间中进行长距离移动，例如从初始的噪声分布快速移向数据分布的大致区域。

直观类比：想象一群粒子从随机散布的位置开始，受到一个引力场（势能 ( $V_{\theta }(x)$ )）的牵引，快速向目标区域（数据流形）靠拢。

2. 接近数据流形（后期阶段，( $t\gg 1$ )）

当时间 ( $t$ ) 很大时，分布 ( ${\rho }_t$ ) 接近平衡分布 ( ${\rho }_{\text{eq}}$ )，并且 ( $\epsilon (t)\approx {\epsilon }_{\text{max}}$ )。此时，( ${\rho }_{t+\Delta t}\approx {\rho }_t\approx {\rho }_{\text{eq}}$ )，因此对于 ( $(x,y)\in \text{supp}({\gamma }_t)$ )，有 ( $x\approx y$ )。这使得传输项 ($\frac{1}{\Delta t}(x-y)\approx 0$)，公式简化为：

$${\nabla }_x{V}_{\theta }(x)+{\epsilon }_{\text{max}}{\nabla }_x\log {\rho }_{\text{eq}}(x)=0$$

即：

$${\nabla }_x\log {\rho }_{\text{eq}}(x)=-\frac{{\nabla }_x{V}_{\theta }(x)}{{\epsilon }_{\text{max}}}$$

由于 (${\nabla }_x\log {\rho }_{\text{eq}}(x)={\nabla }_x\left(\log {\rho }_{\text{eq}}(x)\right)$)，我们可以推导出：

$$\log {\rho }_{\text{eq}}(x)=-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}+C$$

因此：

$${\rho }_{\text{eq}}(x)\propto \exp \left(-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right)$$

意义：

• 平衡分布 ( ${\rho }_{\text{eq}}(x)$ ) 是一个Boltzmann分布，其中能量函数 ( $E(x)=\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}$ )。
• 这一分布由能量基础模型（EBM）定义，能够精确捕捉数据的未归一化似然。
• 在这一阶段，样本通过扩散（diffusion）过程在数据流形上稳定分布，覆盖所有数据模式。

直观类比：粒子已经接近目标区域（数据流形），现在在一个低能量区域内“扩散”，形成一个稳定的分布，类似于气体分子在平衡状态下的分布。

结合OT和EBM的生成框架

基于以上洞察，我们可以提出一个结合最优传输（OT）和能量基础模型（EBM）的生成框架，核心思想是学习一个时间无关的标量势函数 ( $V_{\theta }(x)$ )，其Boltzmann分布：

$${\rho }_{\text{eq}}(x)\propto \exp \left(-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right)$$

能够匹配目标数据分布 ( ${\rho }_{\text{data}}$ )。生成过程分为两个阶段：

1. 远离数据流形（($\epsilon \approx 0$)）

• 在这一阶段，($\epsilon (t)\approx 0$)，演化过程由最优传输流主导。
• 样本从初始噪声分布 ( ${\rho }_0$ )（如高斯分布）开始，沿着 ( $V_{\theta }(x)$ ) 的负梯度快速移动到数据流形附近。
• 这一阶段的特点是确定性和高效性，适合跨越样本空间中的大距离。

例子：假设数据分布是一个二维平面上的圆环，初始分布是均匀的高斯噪声。早期阶段的OT流会将噪声点快速“拉”向圆环区域，类似于粒子在引力场中的运动。

2. 接近数据流形（($\epsilon \approx {\epsilon }_{\text{max}}$)）

• 当样本接近数据流形时，($\epsilon (t)\approx {\epsilon }_{\text{max}}$)，熵项开始主导，样本进入一个扩散过程。
• 样本在数据流形上形成一个稳定的Boltzmann分布，覆盖所有数据模式（如圆环上的均匀分布）。
• 这一阶段能够精确建模数据的局部密度，同时保留显式的似然信息（通过 ( $V_{\theta }(x)$ )）。

例子：在圆环区域，样本不再快速移动，而是通过扩散均匀分布在圆环上，形成一个稳定的概率分布。

框架优势

通过结合OT和EBM，这个框架兼具两者的优点：

• 长距离传输：OT流的确定性使得样本可以高效地从噪声分布移动到数据流形。
• 局部密度建模：EBM的Boltzmann分布能够精确捕捉数据流形的局部结构，并提供显式的似然信息。
• 可控采样：通过调整 ($\epsilon (t)$) 的调度，可以平衡确定性流和扩散过程，实现高效且高质量的采样。

例子：生成圆环分布

为了更直观地说明，我们继续以二维圆环分布为例：

1. 初始分布：高斯噪声分布，点随机散布在平面。
2. 势能函数：设计 ( $V_{\theta }(x)\propto (\Vert x\Vert -r{)}^2$ )，使得圆环半径 ( $r$ ) 处的点具有最低能量。
3. 演化过程：
   • 早期（($\epsilon (t)=0$)）：点沿着 ( $V_{\theta }(x)$ ) 的负梯度快速向圆环移动，传输项确保移动路径尽量短。
   • 后期（($\epsilon (t)={\epsilon }_{\text{max}}$)）：点在圆环上扩散，形成均匀分布，熵项防止点过于集中。
4. 结果：最终分布 ( ${\rho }_{\text{eq}}(x)\propto \exp \left(-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right)$ ) 精确匹配圆环分布。

总结

JKO方案的一阶最优性条件揭示了概率分布演化的动态行为，分为远离数据流形的OT流阶段和接近数据流形的EBM扩散阶段。基于此，我们提出了一个结合OT和EBM的生成框架，通过学习时间无关的势能函数 ( $V_{\theta }(x)$ )，实现从噪声到数据的平滑过渡。这个框架不仅能够高效生成高质量样本，还能保留数据的显式似然信息，在生成模型领域具有广泛的应用前景。通过物理类比和具体例子，我们希望读者能够更直观地理解JKO方案及其生成框架的核心思想。

概念补充

在理解Jordan–Kinderlehrer–Otto (JKO)方案及其一阶最优性条件时，传输计划（OT plan） 和支持集（support） 是两个关键概念，特别是在最优传输（Optimal Transport, OT）理论的背景下。以下是对这两个概念的详细解释，包括定义、直观说明和在上下文中的作用。

---

什么是传输计划（OT Plan）？

定义

在最优传输理论中，传输计划（也称为耦合，英文为 coupling 或 transport plan）是一个联合概率分布，用于描述如何将一个概率分布的“质量”重新分配到另一个概率分布。具体来说，给定两个概率分布 (${\rho }_t$) 和 (${\rho }_{t+\Delta t}$)，定义在空间 (${\mathbb{R}}^d$) 上，传输计划 (${\gamma }_t$) 是一个联合概率分布，定义在 (${\mathbb{R}}^d\times {\mathbb{R}}^d$) 上，满足以下边缘条件：

• (${\gamma }_t$) 的第一个边缘分布等于 (${\rho }_{t+\Delta t}$)，即对于任意可测集 ($A\subset {\mathbb{R}}^d$)：
  $${\int }_{{\mathbb{R}}^d}{\gamma }_t(x,y)dy={\rho }_{t+\Delta t}(x)$$
• (${\gamma }_t$) 的第二个边缘分布等于 (${\rho }_t$)，即：
  $${\int }_{{\mathbb{R}}^d}{\gamma }_t(x,y)dx={\rho }_t(y)$$

集合 ($\Gamma ({\rho }_t,{\rho }_{t+\Delta t})$) 表示所有满足上述边缘条件的联合分布 (${\gamma }_t$)。在JKO方案中，最优传输计划 (${\gamma }_t$) 是通过最小化Wasserstein距离定义的特定耦合：

$$W_2^2({\rho }_{t+\Delta t},{\rho }_t)=\underset{\gamma \in \Gamma ({\rho }_{t+\Delta t},{\rho }_t)}{\min }{\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\Vert x-y{\Vert }^{2}d\gamma (x,y)$$

这里的 ($\Vert x-y{\Vert }^2$) 是移动点 ($y$)（来自 (${\rho }_t$)）到点 ($x$)（在 (${\rho }_{t+\Delta t}$)）的成本。最优传输计划 (${\gamma }_t$) 是实现这一最小值的联合分布。

直观解释

可以将传输计划想象为一个“搬运方案”：

• 假设 (${\rho }_t$) 是一堆沙子（概率分布），分布在空间 (${\mathbb{R}}^d$) 上，(${\rho }_{t+\Delta t}$) 是目标形状的沙堆。
• 传输计划 (${\gamma }_t(x,y)$) 描述了如何将位置 ($y$) 处的沙子（质量）搬运到位置 ($x$) 处。
• (${\gamma }_t(x,y)$) 的值表示从 ($y$) 移动到 ($x$) 的质量比例。例如，如果 (${\gamma }_t(x,y)=0$)，则没有质量从 ($y$) 移动到 ($x$)。
• 边缘条件确保所有沙子都被搬运（没有遗漏或多余），并且目标分布 (${\rho }_{t+\Delta t}$) 被完全填满。

最优传输计划则是成本最低的搬运方案，其中“成本”由移动距离的平方 ($\Vert x-y{\Vert }^2$) 衡量。换句话说，它是最节省能量的搬运方式。

在JKO方案中的作用

在JKO方案中，传输计划 (${\gamma }_t$) 定义了从当前分布 (${\rho }_t$) 到下一时间步分布 (${\rho }_{t+\Delta t}$) 的质量移动方式。具体来说：

• (${\gamma }_t$) 是Wasserstein距离 ($W_2^2({\rho }_{t+\Delta t},{\rho }_t)$) 的最小化器，决定了传输成本项 ($\frac{W_2^2({\rho }_{t+\Delta t},{\rho }_t)}{2\Delta t}$)。
• 在一阶最优性条件中，(${\gamma }_t$) 的支持集 ($\text{supp}({\gamma }_t)$) 指定了哪些点对 ($(x,y)$) 参与了质量搬运，从而约束了条件的适用范围。

---

什么是支持集（Support）？

定义

在概率论和测度论中，一个概率分布（或测度）的支持集（support）是该分布非零的部分所在的最小闭集。具体来说，对于联合分布 (${\gamma }_t$)（定义在 (${\mathbb{R}}^d\times {\mathbb{R}}^d$) 上），其支持集 ($\text{supp}({\gamma }_t)$) 是满足以下条件的集合：

• 对于任意 ($(x,y)\in \text{supp}({\gamma }_t)$)，(${\gamma }_t$) 在 ($(x,y)$) 的任意邻域内有非零概率质量。
• 在 ($\text{supp}({\gamma }_t)$) 之外的区域，(${\gamma }_t(x,y)=0$)。

在最优传输的背景下，($\text{supp}({\gamma }_t)$) 包含所有实际发生质量搬运的点对 ($(x,y)$)，即从 (${\rho }_t$) 的点 ($y$) 移动到 (${\rho }_{t+\Delta t}$) 的点 ($x$) 的配对。

直观解释

支持集可以看作是传输计划 (${\gamma }_t$) 的“活动区域”：

• 如果 ($(x,y)\in \text{supp}({\gamma }_t)$)，则表示存在从 ($y$) 到 ($x$) 的质量搬运。
• 如果 ($(x,y)\notin \text{supp}({\gamma }_t)$)，则没有任何质量从 ($y$) 移动到 ($x$)。
• 支持集通常是一个低维子集，例如在确定性传输映射（deterministic transport map）的情况下，支持集可能是 (${\mathbb{R}}^d\times {\mathbb{R}}^d$) 中的一个子流形。

类比：继续以沙子搬运为例，支持集就像是搬运过程中实际使用的“运输路线”。只有在这些路线上，沙子才会从一个位置移动到另一个位置，而其他路线（不在支持集中）没有被使用。

在JKO方案中的作用

在一阶最优性条件中：

$$\frac{1}{\Delta t}(x-y)+{\nabla }_x{V}_{\theta }(x)+\epsilon (t){\nabla }_x\log {\rho }_t(x)=0,(x,y)\in \text{supp}({\gamma }_t)$$

支持集 ($\text{supp}({\gamma }_t)$) 指定了条件适用的点对 ($(x,y)$)。换句话说，这一等式只需要对那些实际参与质量搬运的点对成立。具体来说：

• ($\text{supp}({\gamma }_t)$) 中的 ($(x,y)$) 表示从 (${\rho }_t$) 的点 ($y$) 移动到 (${\rho }_{t+\Delta t}$) 的点 ($x$)。
• 条件描述了这些点对在传输过程中的动态平衡，涉及传输速度（($\frac{x-y}{\Delta t}$)）、势能梯度（(${\nabla }_x{V}_{\theta }(x)$)）和熵梯度（($\epsilon (t){\nabla }_x\log {\rho }_t(x)$)）。

---

举例说明

为了更直观地理解传输计划和支持集，考虑一个简单的二维例子：

场景

• 假设 (${\rho }_t$) 是一个在二维平面上的离散分布，包含两个点 ($y_1=(0,0)$) 和 ($y_2=(1,0)$)，每个点有等概率 ($0.5$)。
• 假设 (${\rho }_{t+\Delta t}$) 也是一个离散分布，包含两个点 ($x_1=(0.5,0.5)$) 和 ($x_2=(1.5,0.5)$)，同样各有概率 (0.5)。
• 我们需要找到一个传输计划 (${\gamma }_t\in \Gamma ({\rho }_{t+\Delta t},{\rho }_t)$)，并确定其支持集。

传输计划

传输计划 (${\gamma }_t$) 是一个联合分布，描述如何将 (${\rho }_t$) 的质量分配到 (${\rho }_{t+\Delta t}$)。可能的传输计划之一是：

• 将 ($y_1=(0,0)$) 的全部质量（0.5）搬运到 ($x_1=(0.5,0.5)$)。
• 将 ($y_2=(1,0)$) 的全部质量（0.5）搬运到 ($x_2=(1.5,0.5)$)。

这个传输计划可以表示为：
$${\gamma }_t((x_1,y_1))=0.5,{\gamma }_t((x_2,y_2))=0.5,\text{其他点对}=0$$

支持集

该传输计划的支持集 ($\text{supp}({\gamma }_t)$) 是：
$$\text{supp}({\gamma }_t)=\{(x_1,y_1),(x_2,y_2)\}=\{((0.5,0.5),(0,0)),((1.5,0.5),(1,0))\}$$

这表示只有这两对点参与了质量搬运。

最优传输计划

如果我们计算Wasserstein距离，目标是找到最小化 ($\int \Vert x-y{\Vert }^{2}d\gamma (x,y)$) 的 (${\gamma }_t$)。上述传输计划可能是最优的，因为它将每个点直接映射到最近的目标点（取决于具体距离）。在这种情况下，支持集 ($\text{supp}({\gamma }_t)$) 仍然是上述点对集合。

在一阶最优性条件中的应用

在一阶最优性条件中，我们只需要对 ($\text{supp}({\gamma }_t)$) 中的点对 ($(x,y)$) 检查：

• 对于 ($(x_1,y_1)=((0.5,0.5),(0,0))$)，代入公式：
  $$\frac{1}{\Delta t}((0.5,0.5)-(0,0))+{\nabla }_x{V}_{\theta }(0.5,0.5)+\epsilon (t){\nabla }_x\log {\rho }_t(0,0)=0$$
• 对于 ($(x_2,y_2)=((1.5,0.5),(1,0))$)，类似地检查。

这确保了传输过程中的动态平衡仅在实际搬运的路径上成立。

---

总结

• 传输计划（OT Plan）：一个联合概率分布 (${\gamma }_t\in \Gamma ({\rho }_t,{\rho }_{t+\Delta t})$)，描述如何将 (${\rho }_t$) 的质量重新分配到 (${\rho }_{t+\Delta t}$)。最优传输计划是最小化Wasserstein距离的耦合，决定了质量搬运的最节省方式。
• 支持集（Support）：传输计划 (${\gamma }_t$) 非零的区域，包含所有实际参与质量搬运的点对 ($(x,y)$)。在一阶最优性条件中，公式仅对 ($\text{supp}({\gamma }_t)$) 中的点对成立。
• 在JKO方案中的作用：传输计划定义了分布演化的搬运方案，支持集指定了动态平衡适用的点对，共同支撑了从噪声到数据的平滑过渡。

通过这些概念，我们可以更清晰地理解JKO方案如何通过最优传输理论实现概率分布的演化，并在生成模型中发挥作用。

---

后记

2025年4月17日于上海，在grok 3大模型辅助下完成。