最优流匹配：学习单步直线轨迹.

paper：Optimal Flow Matching: Learning Straight Trajectories in Just One Step

0. TL; DR

在过去几年中，用于生成式建模的流匹配（Flow Matching, FM）方法取得了爆炸性的发展。其中一个社区所追求的迷人特性是学习具有直线轨迹的流，这种流能够实现最优传输（Optimal Transport, OT）位移。直线性对于学习到的流路径的快速积分（推理）至关重要。

然而，大多数现有的流拉直方法都基于非平凡的迭代FM程序，这些程序在训练过程中会累积误差，或者利用基于小批量OT的启发式方法。为了解决这些问题，作者开发并从理论上证明了一种新颖的最优流匹配（Optimal Flow Matching, OFM）方法，该方法能够在仅一个FM步骤中恢复出二次传输成本下的直线OT位移。该方法的核心思想是在FM中使用由凸函数参数化的向量场。

1. 背景介绍

近年来，流匹配（Flow Matching, FM）模型已成为驱动生成式建模成功的关键力量。这些模型通过常微分方程（ODE）来描述质量的运动，从而将一个已知的分布移动到一个目标分布。然而，这些过程通常具有弯曲的轨迹，导致采样时需要耗时较长的ODE积分。

为了解决这个问题，研究人员开发了几种改进的FM方法，旨在恢复更直的路径。

校正流 (Rectified Flow, RF)：该方法通过迭代地求解FM问题来逐步“拉直”轨迹。然而，每一次FM迭代都会累积误差，这可能会损害方法的性能。
最优传输条件流匹配 (OT-CFM)：该方法基于直线路径与最优传输（OT）之间的联系。它在FM中应用minibatch* OT解。然而，这种启发式方法由于小批量OT的偏差，并不能保证得到真正的直线路径。

本文旨在解决上述流拉直方法的根本问题。作者提出了一种新颖的最优流匹配（Optimal Flow Matching, OFM）方法。OFM在单次FM迭代后就能获得精确的直线轨迹，这些轨迹可以被精确模拟而无需ODE求解器。它恢复了二次代价函数下的OT流，即解决了Benamou–Brenier问题。

OFM的核心思想是在FM过程中，只考虑一类通过构造就能产生直线路径的特殊向量场。这些向量场是凸函数的梯度，在实践中可以通过输入凸神经网络（Input Convex Neural Networks, ICNNs）来参数化。在OFM中，可以选择性地使用小批量OT或任何其他传输方案作为输入，并且这种做法在理论上是完全合理的。

2. 最优流匹配 (Optimal Flow Matching, OFM)

2.1 理论：推导优化损失

作者的目标是设计一种通过精确的直线轨迹来移动分布$p_0$到$p_1$的方法，并且这个过程应该在单次FM迭代中完成。为此，作者提出了OFM，其核心思想是将FM损失的优化范围限制在一类特殊的“最优向量场”上。

最优向量场 (Optimal Vector Fields)

一个向量场$u_\Psi$被称为“最优的”，如果它生成的轨迹${z_t}_{t \in [0,1]}$是线性的，并且存在一个凸函数$\Psi: \mathbb{R}^D \to \mathbb{R}$，使得对于任何一条轨迹，其终点$z_1$都是起点$z_0$通过该凸函数的梯度映射得到的，即$z_1 = \nabla\Psi(z_0)$。

这样的轨迹可以被明确地写为：

\[z_t = (1-t)z_0 + t\nabla\Psi(z_0), \quad t \in [0, 1]\]

这个ODE的向量场可以表示为：

\[u_t^\Psi(x_t) = \nabla\Psi(z_0) - z_0 = \frac{x_t - z_0}{t}\]

其中，$z_0 = (\phi_t^\Psi)^{-1}(x_t)$是轨迹在$t=0$的初始点。

重要的是，二次代价下的动态最优传输解，其向量场就属于这类“最优向量场”，其中$\Psi$是布伦尼耶势（Brenier potential）。

最优流匹配 (OFM) 损失

作者将上述最优向量场$u_\Psi$代入标准的FM损失，得到了OFM损失：

\[\mathcal{L}_{\text{OFM}}^\pi (\Psi) := \int_0^1 \mathbb{E}_{(x_0, x_1)\sim\pi} \left[ \| u_t^\Psi(x_t) - (x_1 - x_0) \|^2 \right] dt, \quad x_t=(1-t)x_0+tx_1\]

作者证明了，对于任意初始传输方案$\pi$，最小化OFM损失等价于最小化对偶形式的最优传输损失。即：

\[\arg\min_{\text{convex }\Psi} \mathcal{L}_{\text{OFM}}^\pi (\Psi) = \arg\min_{\text{convex }\Psi} \mathcal{L}_{\text{OT}}(\Psi)\]

这意味着，无论我们从哪个初始耦合（甚至是最简单的独立耦合）开始，OFM都能在单次训练中精确地找到二次代价下的最优传输解（即布伦尼耶势$\Psi^*$），从而得到完美的直线OT流。

2.2 实际实现

在实践中，有几个关键的技术细节需要处理。

参数化$\Psi$：作者使用输入凸神经网络（Input Convex Neural Networks, ICNNs）来参数化凸函数$\Psi_\theta$。ICNN通过对网络权重施加非负约束和使用凸的激活函数来保证其输出相对于输入是凸的。
计算OFM损失梯度：直接对OFM损失求梯度是复杂的。作者推导出了一个显式的梯度计算公式，它将复杂的求导过程转化为一个更容易处理的形式，使得训练可行。
流图的反演：在计算向量场$u_t^\Psi(x_t)$时，需要找到当前点$x_t$对应的轨迹起始点$z_0 = (\phi_t^\Psi)^{-1}(x_t)$。这可以通过求解一个凸优化子问题来实现：

\[(\phi_t^\Psi)^{-1}(x_t) = \arg\min_{z_0 \in \mathbb{R}^D} \left[ \frac{(1-t)}{2}\|z_0\|^2 + t\Psi(z_0) - \langle x_t, z_0 \rangle \right]\]

这个强凸优化问题可以使用LBFGS等标准优化器高效求解。

2.3 与先前工作的关系

与OT-CFM的比较：OT-CFM使用小批量OT作为一种启发式方法来试图拉直路径，但其解是有偏的，且不保证路径的直线性。相比之下，OFM从理论上保证了对于任意初始耦合（包括独立耦合或小批量OT耦合），都能恢复出精确的直线OT流。
与Rectified Flow的比较：RF通过迭代来拉直路径，但这个过程会累积误差，并且在迭代后期需要模拟ODE，不再是无模拟训练。OFM则在单次、无模拟的训练中就达到了目标，概念更简洁，实现也更直接。

3. 实验分析

作者通过一系列实验，展示了OFM的有效性和优越性。

3.1 二维示例说明

作者在一个从标准高斯分布到八个高斯混合分布的二维生成任务上，测试了OFM对不同初始传输方案$\pi$的依赖性。展示了使用三种不同初始方案（独立耦合、小批量OT、反小批量OT）训练OFM的结果。

在所有情况下，OFM都学习到了几乎完全相同的直线轨迹和最优传输映射，有力地证明了其对初始传输方案的不依赖性。这与普通的FM形成鲜明对比，后者的轨迹会随着初始方案的不同而发生显著变化。

3.2 高维OT基准测试

作者在一个高维连续分布的OT基准测试上，将OFM与多种FM变体和OT求解器进行了定量比较。评估指标是L2-UVP（未解释方差百分比，越低越好）。

列出了在不同维度下的L2-UVP结果。在所有基于FM的方法中，OFM（无论是使用独立耦合还是小批量OT）的性能都是最好的，尤其是在高维情况下。与专门的OT求解器（如MMv1, Amortization, ICNN）相比，OFM的性能也极具竞争力。RF和c-RF的性能随着维度的增加而迅速恶化，而OT-CFM也逐渐落后于OFM。

3.3 无配对图像到图像翻译

作者将OFM应用于一个在FFHQ数据集上，进行“成人”到“儿童”面部特征翻译的任务。该任务在预训练的ALAE自编码器的512维潜空间中进行。

展示了不同FM方法生成的翻译结果。OFM的翻译结果在视觉上更自然、更合理。比较了不同方法的FID分数。OFM（使用小批量OT）取得了最低的FID分数（11.0），优于RF, c-RF, 和OT-CFM，证明了其在实际高维应用中的优越性。