Flow Matching Guide and Code(3)

Flow Matching Guide and Code(3)

Code: flow_matching library at https://github.com/facebookresearch/flow_matching

4. Flow Matching

给定一个源分布 $p$ 和一个目标分布 $q$，流匹配（Flow Matching, FM） (Lipman et al., 2022; Liu et al., 2022; Albergo and Vanden-Eijnden, 2022) 是一种可扩展的方法，用于训练一个由可学习的速度场 $u_t^{\theta }$ 定义的流模型，并解决以下流匹配问题：
$$\mathrm{Find}{u}_t^{\theta }\text{ generating }{p}_t,\mathrm{with}{p}_0=p\mathrm{and}{p}_1=q.\text{\hspace{1em}(4.1)}$$

上式中的“生成”是依据方程 (3.24) 的含义。回顾图2中的流匹配蓝图，FM框架 (a) 确定一个已知的源分布 $p$ 和一个未知的数据目标分布 $q$；(b) 规定一条从 $p_0=p$ 插值到 $p_1=q$ 的概率路径 $p_t$；© 学习一个用神经网络实现的、生成路径 $p_t$ 的速度场 $u_t^{\theta }$；(d) 通过求解带有 $u_t^{\theta }$ 的ODE来从学习到的模型中采样。

为了在步骤 © 中学习速度场 $u_t^{\theta }$，FM通过最小化以下回归损失来实现：

$${\mathcal{L}}_{\mathrm{FM}}(\theta )={\mathbb{E}}_{X_t\sim p_t}D\left(u_t(X_t),u_t^{\theta }(X_t)\right),\text{\hspace{1em}(4.2)}$$

其中 $D$ 是向量间的差异度量，例如平方 $l_2$ 范数 $D(u,v)={\Vert u-v\Vert }^2$ 。直观地说，FM损失函数鼓励我们可学习的速度场 $u_t^{\theta }$ 去匹配已知能生成所需概率路径 $p_t$ 的真实速度场 $u_t$。图9描绘了流匹配框架中的主要对象及它们之间的依赖关系。让我们通过描述如何构建 $p_t$ 和 $u_t$，以及损失函数 (4.2) 的一个实际实现，来开始我们对流匹配的阐述。

Figure 9 流匹配（Flow Matching）框架的主要对象及其相互关系。一个流（Flow） 由 速度场（Velocity field） 表示，该速度场定义了一个生成 概率路径（Probability path） 的随机过程。流匹配的核心思想是将构建一个满足期望边界条件（Boundary conditions） 的复杂流（顶行）这一问题，分解（break down） 为构建多个满足更简单边界条件的条件流（conditional flows）（中间行）的问题，后者因此更容易解决。箭头表示了不同对象之间的依赖关系（dependencies）；蓝色箭头标示了流匹配框架所采用的（employed）核心关系。“损失（Loss）” 列列出了用于学习速度场的损失函数，其中 CFM 损失（条件流匹配损失）（中间行和底行）是在实践中实际使用的损失。底行列出了在第2节中描述的最简单的流匹配算法实例（instantiation）。

4.1 Data

重申一下，令源样本为随机变量 $X_0\sim p$，目标样本为随机变量 $X_1\sim q$。通常，源样本遵循一个已知的、易于采样的分布（如高斯分布），而目标样本则以一个有限大小的数据集的形式提供给我们。根据应用的不同，目标样本可以构成图像、视频、音频片段或其他类型的高维、富含结构的数据。源样本和目标样本可以是独立的，或者源自一个称为耦合（coupling） 的通用联合分布：

$$(X_0,X_1)\sim {\pi }_{0,1}(X_0,X_1),\text{\hspace{1em}(4.3)}$$
其中，如果没有已知的耦合，则源-目标样本遵循独立耦合（independent coupling） ${\pi }_{0,1}(X_0,X_1)=p(X_0)q(X_1)$。独立源-目标分布的一个常见例子是：考虑从随机高斯噪声向量 $X_0\sim \mathcal{N}(0,I)$ 生成图像 $X_1$。作为依赖耦合（dependent coupling） 的一个例子，可以考虑从其低分辨率版本 $X_0$ 生成高分辨率图像 $X_1$，或从其灰度版本 $X_0$ 生成彩色视频 $X_1$ 的情况。

4.2 Building probability paths

流匹配（Flow Matching）通过采用一种条件策略（conditional strategy），极大地简化了设计概率路径 $p_t$ 及其对应速度场 $u_t$ 的问题。作为第一个例子，考虑将 $p_t$ 的设计条件于一个单独的目标样本 $X_1=x_1$，这就产生了如图3a所示的条件概率路径 $p_{t|1}(x|x_1)$。然后，我们可以通过聚合（aggregating） 这些条件概率路径 $p_{t|1}$ 来构建整体的、边际的概率路径 $p_t$：

$$p_t(x)=\int p_{t|1}(x|x_1)q(x_1)d{x}_1\text{\hspace{1em}(4.4)}$$

如图3b所示。为了解决流匹配问题，我们希望 $p_t$ 满足以下边界条件（boundary conditions）：
$$p_0=p,p_1=q,\text{\hspace{1em}(4.5)}$$
也就是说，边际概率路径 $p_t$ 从时间 $t=0$ 的源分布 $p$ 插值（interpolates） 到时间 $t=1$ 的目标分布 $q$。这些边界条件可以通过要求条件概率路径满足以下条件来强制执行：
$$p_{0|1}(x|x_1)={\pi }_{0|1}(x|x_1),\text{and}{p}_{1|1}(x|x_1)={\delta }_{x_1}(x),\text{\hspace{1em}(4.6)}$$
其中条件耦合（conditional coupling）$ \pi_{0|1} (x_0 | x_1) = \pi_{0,1} (x_0,x_1) / q(x_1) $，而 ${\delta }_{x_1}$ 是集中于 $x_1$ 的狄拉克测度（delta measure）。对于独立耦合（independent coupling） ${\pi }_{0,1}(x_0,x_1)=p(x_0)q(x_1)$，上述第一个约束简化为 $p_{0|1}(x|x_1)=p(x)$。由于狄拉克测度没有密度函数，第二个约束应理解为：对于连续函数 $f$，当 $t\to 1$ 时，有 $\int p_{t|1}(x|y)f(y)\mathrm{d}y\to f(x)$。请注意，将 (4.6) 代入 (4.4) 即可验证边界条件 (4.5)。

一个满足 (4.6) 中条件的流行条件概率路径的例子在 (2.2) 式中给出：

$$\mathcal{N}(\cdot |t{x}_1,(1-t{)}^{2}I)\to {\delta }_{x_1}(\cdot )\text{ as }t\to 1.$$

（即：一个均值向 $t{x}_1$ 移动、方差逐渐缩小的正态分布，在 $t=1$ 时坍缩到点 $x_1$ 上的狄拉克分布。）

4.3 Deriving generating velocity fields

在具备了边际概率路径 $p_t$ 之后，我们现在来构建一个生成 $p_t$ 的速度场 $u_t$。这个生成性的边际速度场 $u_t$ 是多个条件速度场 $u_t(x|x_1)$（图示见图3c）的平均，其中每个条件速度场满足：

$$u_t(\cdot |x_1)\text{ generates }{p}_{t|1}(\cdot |x_1).\text{\hspace{1em}(4.7)}$$
那么，生成边际路径 $p_t(x)$ 的边际速度场 $u_t(x)$（图示见图3d）由以下公式给出，即在所有目标样本上对条件速度场 $u_t(x|x_1)$ 进行平均：

$$u_t(x)=\int u_t(x|x_1)p_{1|t}(x_1|x)\mathrm{d}{x}_1.\text{\hspace{1em}(4.8)}$$

为了用已知项来表达上述方程，我们回顾贝叶斯规则：

$$p_{1|t}(x_1|x)=\frac{p_{t|1}(x|x_1)q(x_1)}{p_t(x)},\text{\hspace{1em}(4.9)}$$

该式对所有满足 $p_t(x)>0$ 的 $x$ 有定义。方程 (4.8) 可以解释为条件速度 $u_t(x|x_1)$ 的加权平均，其权重 $p_{1|t}(x_1|x)$ 代表了在给定当前样本 $x$ 的情况下，目标样本 $x_1$ 的后验概率。方程 (4.8) 的另一种解释可以通过条件期望（见第3.2节）给出。即，如果存在一个随机变量 $X_t$，使得 $X_t\sim p_{t|1}(\cdot |X_1)$，或者等价地，$(X_t,X_1)$ 的联合分布具有密度 $p_{t,1}(x,x_1)=p_{t|1}(x|x_1)q(x_1)$，那么使用 (3.12) 将 (4.8) 写成一个条件期望，我们得到：

$$u_t(x)=\mathbb{E}\left[u_t(X_t|X_1)|X_t=x\right],\text{\hspace{1em}(4.10)}$$

这得到了一个对 $u_t(x)$ 有用的解释：它是 $u_t(X_t|X_1)$ 在给定 $X_t=x$ 时的最小二乘近似（least-squares approximation）（见第3.2节）。请注意，(4.10) 中的 $X_t$ 通常是一个与由最终流模型 (3.16) 定义的 $X_t$ 不同的随机变量，尽管它们共享相同的边际概率 $p_t(x)$。

4.4 General conditioning and the Marginalization Trick

为了证明上述构造的合理性，我们需要证明在温和的假设下，由方程 (4.8) 和 (4.10) 定义的边际速度场 $u_t$ 生成了由方程 (4.4) 定义的边际概率路径 $p_t$。证明所需的数学工具是质量守恒定理（定理 2）。为此，让我们考虑一个稍后在本手册中会有用的、更一般的设定。特别地，通过条件于 $X_1=x_1$ 来构建条件概率路径和速度场并没有什么特殊之处。正如 Tong et al. (2023) 所指出的，前一节的分析可以推广到条件于任意随机变量 $Z\in {\mathbb{R}}^m$（其概率密度函数为 $p_Z$） 的情况。这就产生了边际概率路径：
$$p_t(x)=\int p_{t|Z}(x|z)p_Z(z)dz,\text{\hspace{1em}(4.11)}$$
而该路径又由边际速度场生成：

$$u_t(x)=\int u_t(x|z)p_{Z|t}(z|x)dz=\mathbb{E}\left[u_t(X_t|Z)\mid X_t=x\right],\text{\hspace{1em}(4.12)}$$

其中：

• $u_t(\cdot |z)$ 生成 $p_{t|Z}(\cdot |z)$,
• $p_{Z|t}(z|x)=\frac{p_{t|Z}(x|z)p_Z(z)}{p_t(x)}$ （由贝叶斯规则得出，给定 $p_t(x)>0$）,
• 且 $X_t\sim p_{t|Z}(\cdot |Z)$。

自然，通过设置 $Z=X_1$，我们可以恢复前面章节中的构造。

在证明主要结果之前，我们需要一些正则性假设，总结如下：

Assumption 1：
$p_{t|Z}(x|z)$ 是 $(t,x)$ 的 $C^1([0,1)\times {\mathbb{R}}^d)$ 函数，且 $u_t(x|z)$ 是 $(t,x)$ 的 $C^1([0,1)\times {\mathbb{R}}^d,{\mathbb{R}}^d)$ 函数。此外，$p_Z$ 具有有界支撑，即 $p_Z(x)=0$ 在 ${\mathbb{R}}^m$ 的某个有界集之外成立。最后，对于所有 $x\in {\mathbb{R}}^d$ 和 $t\in [0,1)$，有 $p_t(x)>0$。

这些是温和的假设。例如，通过找到一个满足 $p_Z(z)>0$ 且 $p_{t|Z}(\cdot |z)>0$ 的条件 $z$，可以证明 $p_t(x)>0$。在实践中，可以通过考虑 $(1-(1-t)\epsilon )p_{t|Z}+(1-t)\epsilon \mathcal{N}(0,I)$（其中 $\epsilon >0$ 是任意小的数）来满足此条件。满足此假设的 $p_{t|Z}(\cdot |z)$ 的一个例子是 (2.2) 式中的路径，其中我们令 $Z=X_1$。

我们现在准备陈述主要结果：

Theorem 3 (边际化技巧, Marginalization Trick)：
在假设 1 下，如果 $u_t(x|z)$ 是条件可积的且生成条件概率路径 $p_t(\cdot |z)$，那么边际速度场 $u_t$ 生成边际概率路径 $p_t$，对于所有 $t\in [0,1)$ 成立。

在上述定理中，“条件可积”指的是质量守恒定理（3.26）中可积性条件的条件版本，即：
$${\int }_0^1\int \int \Vert u_t(x|z)\Vert p_{t|Z}(x|z)p_Z(x)dz\mathrm{d}x\mathrm{d}t<\infty .\text{\hspace{1em}(4.13)}$$

证明：
该结果通过验证质量守恒定理（定理 2）的两个条件得出。首先，让我们验证组合 $(u_t,p_t)$ 满足连续性方程 (3.25)。因为 $u_t(\cdot |z)$ 生成 $p_t(\cdot |z)$，我们有：

$$\begin{array}{rl}\frac{\mathrm{d}}{\mathrm{d}t}{p}_t(x)& \stackrel{(i)}{=}\int \frac{\mathrm{d}}{\mathrm{d}t}{p}_{t|Z}(x|z)p_Z(x)dz\\ & \stackrel{(ii)}{=}-\int {\mathrm{div}}_x\left[u_t(x|z)p_{t|Z}(x|z)\right]p_Z(z)dz\\ & \stackrel{(i)}{=}-{\mathrm{div}}_x\int u_t(x|z)p_{t|Z}(x|z)p_Z(z)dz\\ & \stackrel{(iii)}{=}-{\mathrm{div}}_x\left[u_t(x)p_t(x)\right].\end{array}\text{\hspace{1em}(4.14)(4.15)(4.16)(4.17)}$$

等号 $(i)$ 成立是根据莱布尼茨法则，交换了微分（分别是 $\frac{d}{dt}$ 和 ${\text{div}}_x$）与积分的顺序，这由以下事实证明是合理的：$p_{t|Z}(x|z)$ 和 $u_t(x|z)$ 关于 $t,x$ 是 $C^1$ 的，并且 $p_Z$ 具有有界支撑（因此所有被积函数作为有界集上的连续函数都是可积的）。等号 $(ii)$ 成立是因为 $u_t(\cdot |z)$ 生成 $p_{t|Z}(\cdot |z)$ 以及定理 2。等号 $(iii)$ 成立是通过乘以并除以 $p_t(x)$（由假设严格大于零）并使用了 $u_t$ 的公式 (4.12)。

为了验证定理 2 的第二个也是最后一个条件，我们应证明 $u_t$ 是可积且局部利普希茨连续的。因为 $C^1$ 函数是局部利普希茨连续的，所以只需检查 $u_t(x)$ 对所有 $(t,x)$ 是 $C^1$ 的。这将由 $u_t(x|z)$ 和 $p_{t|Z}(x|z)$ 是 $C^1$ 且 $p_t(x)>0$（由假设成立）来保证。此外，$u_t(x)$ 是可积的，因为 $u_t(x|z)$ 是条件可积的：

$${\int }_0^1\int \Vert u_t(x)\Vert p_t(x)\mathrm{d}x\mathrm{d}t\le {\int }_0^1\int \int \Vert u_t(x|z)\Vert p_{t|Z}(x|z)p_Z(z)dz\mathrm{d}x\mathrm{d}t<\infty ,\text{\hspace{1em}(4.18)}$$

其中第一个不等式源于向量詹森不等式 (vector Jensen’s inequality)。

4.5 Flow Matching loss

在确立了目标速度场 $u_t$ 生成从 $p$ 到 $q$ 的指定概率路径 $p_t$ 之后，缺失的部分是一个易于处理的损失函数，用于学习一个尽可能接近目标 $u_t$ 的速度场模型 $u_t^{\theta }$。

直接使用此损失函数的一个主要障碍是计算目标 $u_t$ 是不可行的，因为它需要对整个训练集进行边际化（即，对方程 (4.8) 中的 $x_1$ 或方程 (4.12) 中的 $z$ 进行积分）。

幸运的是，有一族称为布雷格曼散度 (Bregman divergences) 的损失函数，可以仅根据条件速度 $u_t(x|z)$ 来为学习 $u_t^{\theta }(x)$ 提供无偏梯度。

Figure 10 布雷格曼散度

布雷格曼散度测量两个向量 $u,v\in {\mathbb{R}}^d$ 之间的差异性，定义为：

$$D(u,v):=\Phi (u)-\left[\Phi (v)+⟨u-v,\nabla \Phi (v)⟩\right],\text{\hspace{1em}(4.19)}$$

其中 $\Phi :{\mathbb{R}}^d\to \mathbb{R}$ 是一个在某个凸集 $\Omega \subset {\mathbb{R}}^d$ 上定义的严格凸函数。如图10所示，布雷格曼散度测量了 $\Phi (u)$ 与在 $v$ 处展开并 evaluated at $u$ 的 $\Phi$ 的线性近似之间的差异。因为线性近似是凸函数的全局下界，所以有 $D(u,v)\ge 0$。此外，由于 $\Phi$ 是严格凸的，当且仅当 $u=v$ 时 $D(u,v)=0$。

最基础的布雷格曼散度是平方欧几里得距离 $D(u,v)=|u-v{|}^2$，这源于选择 $\Phi (u)=|u{|}^2$。

使得布雷格曼散度对 Flow Matching 有用的关键性质是它们关于第二个参数的梯度是仿射不变的 (affine invariant) (Holderrieth et al., 2024)：

$${\nabla }_{v}D(a{u}_1+b{u}_2,v)=a{\nabla }_{v}D(u_1,v)+b{\nabla }_{v}D(u_2,v),\text{ for any }a+b=1,\text{\hspace{1em}(4.20)}$$

这可以从方程 (4.19) 验证。仿射不变性允许我们如下交换期望和梯度：

$${\nabla }_{v}D(\mathbb{E}[Y],v)=\mathbb{E}[{\nabla }_{v}D(Y,v)]\text{for any RV}Y\in {\mathbb{R}}^d.\text{\hspace{1em}(4.21)}$$

Flow Matching 损失采用一个布雷格曼散度，将我们的可学习速度 $u_t^{\theta }(x)$ 回归 (regress) 到沿着概率路径 $p_t$ 的目标速度 $u_t(x)$ 上：

$${\mathcal{L}}_{\mathrm{FM}}(\theta )={\mathbb{E}}_{t,X_t\sim p_t}D(u_t(X_t),u_t^{\theta }(X_t)),\text{\hspace{1em}(4.22)}$$

其中时间 $t\sim \mathcal{U}[0,1]$。然而，如上所述，目标速度 $u_t$ 是难以处理的，因此上面的损失无法按原样计算。相反，我们考虑更简单且易于处理的条件流匹配 (Conditional Flow Matching, CFM) 损失：

$${\mathcal{L}}_{\mathrm{CFM}}(\theta )={\mathbb{E}}_{t,Z,X_t\sim p_{t|Z}(\cdot |Z)}D(u_t(X_t|Z),u_t^{\theta }(X_t)).\text{\hspace{1em}(4.23)}$$

这两个损失对于学习目的是等价的，因为它们的梯度重合 (Holderrieth et al., 2024)：

定理 4. Flow Matching 损失和 Conditional Flow Matching 损失的梯度重合：
$${\nabla }_{\theta }{\mathcal{L}}_{FM}(\theta )={\nabla }_{\theta }{\mathcal{L}}_{CFM}(\theta ).\text{\hspace{1em}(4.24)}$$
特别地，Conditional Flow Matching 损失的最小化器就是边际速度场 $u_t(x)$。

证明. 证明遵循直接计算：

$$\begin{array}{rl}{\nabla }_{\theta }{\mathcal{L}}_{\mathrm{FM}}(\theta )& ={\nabla }_{\theta }{\mathbb{E}}_{t,X_t\sim p_t}D(u_t(X_t),u_t^{\theta }(X_t))\\ & ={\mathbb{E}}_{t,X_t\sim p_t}{\nabla }_{\theta }D(u_t(X_t),u_t^{\theta }(X_t))\\ & \stackrel{(i)}{=}{\mathbb{E}}_{t,X_t\sim p_t}{\nabla }_{v}D(u_t(X_t),u_t^{\theta }(X_t)){\nabla }_{\theta }{u}_t^{\theta }(X_t)\\ & \stackrel{(4.12)}{=}{\mathbb{E}}_{t,X_t\sim p_t}{\nabla }_{v}D({\mathbb{E}}_{Z\sim p_{Z|t}(\cdot |X_t)}[u_t(X_t|Z)],u_t^{\theta }(X_t)){\nabla }_{\theta }{u}_t^{\theta }(X_t)\\ & \stackrel{(ii)}{=}{\mathbb{E}}_{t,X_t\sim p_t}{\mathbb{E}}_{Z\sim p_{Z|t}(\cdot |X_t)}\left[{\nabla }_{v}D(u_t(X_t|Z),u_t^{\theta }(X_t)){\nabla }_{\theta }{u}_t^{\theta }(X_t)\right]\\ & \stackrel{(iii)}{=}{\mathbb{E}}_{t,X_t\sim p_t}{\mathbb{E}}_{Z\sim p_{Z|t}(\cdot |X_t)}[{\nabla }_{\theta }D(u_t(X_t|Z),u_t^{\theta }(X_t))]\\ & \stackrel{(iv)}{=}{\nabla }_{\theta }{\mathbb{E}}_{t,Z\sim q,X_t\sim p_{t|Z}(\cdot |Z)}[D(u_t(X_t|Z),u_t^{\theta }(X_t))]\\ & ={\nabla }_{\theta }{\mathcal{L}}_{\mathrm{CFM}}(\theta )\end{array}$$

其中在 $(i)$ 和 $(iii)$ 中我们使用了链式法则；$(ii)$ 源于将方程 (4.21) 条件于 $X_t$ 应用；在 $(iv)$ 中我们使用了贝叶斯规则。

用于学习条件期望的布雷格曼散度. 定理 4 是利用布雷格曼散度学习条件期望的一个更一般结果的特定实例，如下所述。它将在本手册中全程使用，并为 Flow Matching 背后所有可扩展的损失提供基础：

Proposition(命题) 1 (用于学习条件期望的布雷格曼散度). 设 $X\in S_X$, $Y\in S_Y$ 是状态空间 $S_X$, $S_Y$ 上的随机变量，且 $g:{\mathbb{R}}^p\times S_X\to {\mathbb{R}}^n$, $(\theta ,x)\mapsto g_{\theta }(x)$，其中 $\theta \in {\mathbb{R}}^p$ 表示可学习参数。设 $D_x(u,v)$, $x\in S_X$ 是一个在凸集 $\Omega \subset {\mathbb{R}}^n$ 上的布雷格曼散度，该凸集包含 $f$ 的像。那么，$${\nabla }_{\theta }{\mathbb{E}}_{X,Y}{D}_X\left(Y,g^{\theta }(X)\right)={\nabla }_{\theta }{\mathbb{E}}_X{D}_X\left(\mathbb{E}\left[Y\mid X\right],g^{\theta }(X)\right).\text{\hspace{1em}(4.25)}$$

特别地，对于所有满足 $p_X(x)>0$ 的 $x$，$g_{\theta }(x)$ 关于 $\theta$ 的全局最小值满足

$$g^{\theta }(x)=\mathbb{E}\left[Y\mid X=x\right].\text{\hspace{1em}(4.26)}$$

证明. 我们假设 $g_{\theta }$ 关于 $\theta$ 可微，并且 $X$ 和 $Y$ 的分布，以及 $D_x$ 和 $g$ 允许交换微分和积分，展开：

$$\begin{array}{rl}{\nabla }_{\theta }{\mathbb{E}}_{X,Y}{D}_X\left(Y,g^{\theta }(X)\right)& \stackrel{(i)}{=}{\mathbb{E}}_X\left[\mathbb{E}\left[{\nabla }_v{D}_X\left(Y,g^{\theta }(X)\right){\nabla }_{\theta }{g}^{\theta }(X)\mid X\right]\right]\\ & \stackrel{(ii)}{=}{\mathbb{E}}_X\left[{\nabla }_v{D}_X\left(\mathbb{E}\left[Y\mid X\right],g^{\theta }(X)\right){\nabla }_{\theta }{g}^{\theta }(X)\right]\\ & \stackrel{(iii)}{=}{\mathbb{E}}_X\left[{\nabla }_{\theta }{D}_X\left(\mathbb{E}\left[Y\mid X\right],g^{\theta }(X)\right)\right]\\ & ={\nabla }_{\theta }{\mathbb{E}}_X{D}_X\left(\mathbb{E}\left[Y\mid X\right],g^{\theta }(X)\right),\end{array}$$

其中 $(i)$ 遵循链式法则和期望的塔性质 (3.11)。等号 $(ii)$ 遵循 (4.21)。等号 $(iii)$ 再次使用了链式法则。最后，对于每个满足 $p_X(x)>0$ 的 $x\in S_X$，我们可以选择 $g_{\theta }(x)=\mathbb{E}[Y|X=x]$，得到 $\mathbb{E}\ast X\left[D_X(\mathbb{E}[Y|X],g\ast \theta (X))\right]=0$，这必然是关于 $\theta$ 的全局最小值。

通过选择 $X=X_t$, $Y=u_t(X_t|Z)$, $g_{\theta }(x)=u_t^{\theta }(x)$，并对 $t\sim \mathcal{U}[0,1]$ 取期望，可以立即从命题 1 证明定理 4。

一般时间分布. FM 损失的一个有用变体是从均匀分布以外的分布中采样时间 $t$。具体来说，考虑 $t\sim \omega (t)$，其中 $\omega$ 是 $[0,1]$ 上的一个概率密度函数。这导致以下加权目标：

$${\mathcal{L}}_{\mathrm{CFM}}(\theta )={\mathbb{E}}_{t\sim \omega ,Z,X_t}D(u_t(X_t|Z),u_t^{\theta }(X_t))={\mathbb{E}}_{t\sim U,Z,X_t\omega }(t)D(u_t(X_t|Z),u_t^{\theta }(X_t)).\text{\hspace{1em}(4.27)}$$

尽管在数学上是等价的，但在大规模图像生成任务中，采样 $t\sim \omega$ 比使用权重 $\omega (t)$ 能带来更好的性能 (Esser et al., 2024)。

4.6 Solving conditional generation with conditional flows

到目前为止，我们已将训练流模型 $u_t^{\theta }$ 的问题简化为以下三个步骤：
(i) 寻找能够产生满足 (4.5) 中边界条件的边际概率路径 $p_t(x)$ 的条件概率路径 $p_{t|Z}(x|z)$；
(ii) 寻找能够生成该条件概率路径的条件速度场 $u_t(x|z)$；
(iii) 使用条件流匹配损失（参见方程 (4.23)）进行训练。

我们现在讨论如何具体实现步骤 (i) 和 (ii)，即如何设计此类条件概率路径和速度场。我们将提出一种灵活的方法，通过条件流（conditional flows） 的特定构造来设计此类条件概率路径和速度场。其核心思想如下：定义一个满足边界条件 (4.6) 的流模型 $X_{t|1}$（类似于 (3.16)），然后通过微分 (3.20) 从中提取速度场。此过程同时定义了 $p_{t|1}(x|x_1)$ 和 $u_t(x|x_1)$。更详细地，定义条件流模型：

$$X_{t|1}={\psi }_t(X_0\mid x_1),\text{where }{X}_0\sim {\pi }_{0|1}(\cdot \mid x_1),\text{\hspace{1em}(4.28)}$$

其中 $\psi :[0,1)\times {\mathbb{R}}^d\times {\mathbb{R}}^d\to {\mathbb{R}}^d$ 是一个条件流，由下式定义：

$$\begin{array}{r}{\psi }_t(x|x_1)=\{\begin{array}{ll}x& t=0\\ x_1& t=1\end{array},\end{array}\text{\hspace{1em}(4.29)}$$

该函数在 $(t,x)$ 上光滑，并且在 $x$ 上是微分同胚。（此处的“光滑”是指 ${\psi }_t(x|x_1)$ 关于 $t$ 和 $x$ 的所有导数均存在且连续：即 $C^{\infty }([0,1)\times {\mathbb{R}}^d,{\mathbb{R}}^d)$。这些条件可以放宽至 $C^2([0,1)\times {\mathbb{R}}^d,{\mathbb{R}}^d)$，但会牺牲一定的简洁性。）

推送映射公式 (3.15) 将 $X_{t|1}$ 的概率密度定义为：

$$p_{t|1}(x|x_1):=\left[\begin{array}{c}{\psi }_t(\cdot |x_1{)}_{♯}{\pi }_{0|1}(\cdot |x_1)\end{array}\right](x),\text{\hspace{1em}(4.30)}$$

尽管在 CFM 损失的实际优化中我们不需要此表达式，但它在理论上用于证明 $p_{t|1}$ 满足两个边界条件 (4.6)。首先，根据 (4.29)，${\psi }_0(\cdot |x_1)$ 是恒等映射，在时间 $t=0$ 时保持 ${\pi }_{0|1}(\cdot |x_1)$ 不变。其次，${\psi }_1(\cdot |x_1)=x_1$ 是常值映射，随着 $t\to 1$ 将所有概率质量集中于 $x_1$。此外，请注意 ${\psi }_t(\cdot |x_1)$ 对于 $t\in [0,1)$ 是光滑微分同胚。因此，根据流与速度场的等价性（第 3.4.1 节），存在一个唯一的光滑条件速度场（参见方程 (3.20)），其形式为：
$$u_t(x|x_1)={\dot{\psi }}_t({\psi }_t^{-1}(x|x_1)|x_1).\text{\hspace{1em}(4.31)}$$

总结而言：我们进一步将寻找条件路径及相应生成速度场的任务，简化为只需构建一个满足 (4.29) 的条件流 ${\psi }_t(\cdot |x_1)$。在第 4.7 节中，我们将选择一个特别简单且具有某些理想性质（条件最优传输流）的 ${\psi }_t(x|x_1)$，它将引致第 1 节中看到的标准 Flow Matching 算法。在第 4.8 节中，我们将讨论一个特定且众所周知的条件流族，即仿射流（affine flows），其中包含扩散模型文献中的一些已知示例。在第 5 节中，我们将使用条件流在流形（manifold） 上定义 Flow Matching，从而展示此方法的灵活性。