Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

技术摘要：基于对抗学习插值器的多边际流匹配

问题陈述
本文解决了在仅给定离散时间点采样观测值（快照）而无法获取地面真值轨迹的情况下，学习系统随时间变化的动力学过程的挑战。这是生物学（如单细胞 RNA 测序、细胞追踪）和医学等领域的一个核心任务。形式上，目标是推断一个控制常微分方程 (ODE) $dx_t = v_t(x_t)dt$ 的随时间变化的向量场 $v_t$，使得积分映射将初始边缘分布 $q_0$ 推向目标 $q_1$，同时满足中间边缘约束 $p_{t_i} = q_{t_i}$，其中时间步为 $0 = t_1 < \dots < t_K = 1$。

现有的方法（例如将标准流匹配 (Flow Matching, FM) 应用于连续的边缘对，或使用分段线性/三次样条插值器）往往会导致非平滑的插值曲线。这种非平滑性问题会导致学习目标中的梯度方差过高，并在高维设置或数据存在噪声时导致向量场归一化困难。

方法论：ALI-CFM
作者提出了 ALI-CFM（对抗学习插值器结合条件流匹配），这是一种将插值器学习与向量场学习解耦的两阶段方法。

1. 插值器的对抗学习 (ALI)：
   不同于强制插值器通过特定数据点（逐点匹配），ALI 学习一个神经插值器 $G_\phi(x_0, x_1, t)$ 来匹配中间边缘的分布。该插值器被参数化为线性插值加上一个神经修正项：
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   为了确保插值器的前推分布与观测到的中间边缘 $q_{t_i}$ 相匹配，作者采用了受 GAN 启发的对抗目标。判别器 $D_\gamma$ 用于区分来自真实边缘 $q_{t_i}$ 的样本与由插值器 $G_\phi$ 生成的样本。生成器（插值器）通过最小化这些分布之间的 Jensen-Shannon 散度来进行训练。

   为了确保插值器的唯一性并防止出现任意弯曲的路径，作者引入了正则项：

   • 线性参考 (Linear Reference)： 惩罚偏离耦合端点之间直线路径的行为。
   • 分段线性参考 (Piecewise Linear Reference)： 将插值器回归到由马尔可夫链最优传输 (OT) 耦合定义的分段线性路径。
   • 二阶导数范数 (Second Derivative Norm)： 通过惩罚插值器的曲率来确保平滑性。

2. 通过条件流匹配进行归一化 (Marginalization via CFM)：
   一旦训练好插值器 $G_\phi$，将其固定。随后训练第二个神经网络 $u_\theta$ 来逼近生成这些插值器的向量场。这是通过标准的条件流匹配目标实现的，其中目标向量场是学习到的插值器的对时间导数：
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   $q_0$ 与 $q_1$ 之间的耦合 $\pi$ 可以是独立的，也可以是 Optimal Transport 平面。

核心贡献

1. 新颖的插值技术： 本文引入了一种利用对抗目标来学习插值器的方法，该方法匹配的是分布而非点。这使得生成的插值器在温和假设下是平滑且唯一的（在附录中进行了证明）。
2. 对噪声和复杂性的鲁棒性： 该方法旨在处理噪声数据集和复杂的随时间变化的几何结构，在这些场景下，逐点插值方法（如样条或分段线性）会失效或变得不稳定。
3. 可扩展性与通用性： 作者展示了 ALI-CFM 在高维数据和大量边缘分布数量 ($K$) 下的可扩展性，在特定复杂场景下优于现有基准方法，同时在其他场景中保持竞争力。

实验结果
作者在四个不同的任务上评估了 ALI-CFM：

• 合成“结” (Knot) 分布： 在一个具有 1,200 个边缘分布形成“结”结构的合成数据集上，ALI-CFM 是唯一能够准确捕捉随时间变化的几何形状的方法。竞争方法（OT-MFM、三次样条、分段线性）产生了具有高梯度方差的非平滑插值器。
• 细胞追踪 (胶质母细胞瘤)： 使用移动细胞的显微镜数据，ALI-CFM 成功从噪声、欠采样的数据中学习了时空动力学。相比之下，OT-CFM 轨迹发生了发散，而不随时间变化的 OT-MFM 则未能捕捉到时间依赖性。
• 单细胞轨迹推断 (scRNA-seq)： 在 Embryoid Body (EB)、Cite-seq 和 Multiome 数据集上，ALI-CFM 在地球移动距离 (EMD) 指标上与最先进的基准方法（OT-MFM, OT-CFM）表现相当。作者指出，虽然对抗训练使得在这些特定指标上难以严格超越过度拟合于逐点数据的算法，但该方法仍然保持了极强的竞争力。
• 肿瘤坐标推断 (空间转录组学)： 在一个具有空间有序组织切片的乳腺癌数据集上，ALI-CFM 在预测留出部分 (held-out sections) 的肿瘤区域分布方面显著优于所有基准方法（包括 OT-CFM, OT-MMFM 和 OT-MFM）。作者将其成功归功于该方法处理多模态分布以及应对组织切片对齐挑战的能力。

意义与主张
本文声称，ALI-CFM 通过从严格的逐点分配转向分布匹配，为多边际流匹配提供了一种概念上新颖的方法。这种转变使得学习平滑、唯一的插值器成为可能，使其在面对噪声时更具鲁棒性，并能模拟以往方法难以处理的复杂动力学。作者强调，在拥有数百或数千个边缘分布以及噪声数据的场景下，该方法特别有效，而在这些领域，目前尚无其他流匹配算法展现出同等的性能。虽然在标准的单细胞基准测试中，该方法的表现与现有最先进技术持平，但其在空间转录组学和合成复杂几何任务上的卓越表现，凸显了其通用性和可扩展性。