技术摘要：Wasserstein 残差：从群体动力学中学习梯度流

问题陈述

本文研究了从稀疏观测快照中重建群体动力学的逆问题。在计算生物学和人群动力学等许多科学领域，群体演化被建模为 Wasserstein 梯度流 (WGF)：即由能量泛函 $F$ 的最速下降驱动的概率分布 $\rho_t$ 曲线。目标是恢复底层的能量泛效 $F$，使得其梯度流能够拟合在离散观测时间 $t \in \mathcal{T}_{obs}$ 下的观测边缘分布 $\{q_t\}$。

现有的主流方法依赖于 Jordan–Kinderlehrer–Otto (JKO) 方案，该方案将 WGF 定义为通过最小化能量加上最优传输 (OT) 惩罚项来进行的一系列近端步（proximal steps）。然而，基于 JKO 的方法存在两个主要局限性：

1. 对时间离散化的不灵活性： 它们需要在连续快照之间求解昂贵的最优传输问题。
2. 对间隙的敏感性： 它们难以处理观测之间较大的时间间隔，因为它们依赖于快照之间的线性弦插值（linear chord interpolation），这导致无法捕捉到弯曲的轨迹（例如，粒子沿正弦谷底运动）。

方法论：Wasserstein 残差框架

作者提出了一种基于残差的方法，绕过了 JKO 方案和最优传输耦合。他们不再强制执行 JKO 条件，而是直接通过一个非负损失函数来强制执行刻画梯度流的基本微分方程。

1. 残差公式化

论文利用了梯度流的 切向量 (Tangent) 和 能量耗散等式 (EDE) 公式。具体而言，它侧重于速度约束：
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
其中 $v_t$ 是速度场，$\frac{\delta F}{\delta \rho_t}$ 是能量泛函的一阶变分。

作者定义了一个速度残差 $R_{vel}$，当且仅当 $\rho$ 是 $F$ 的梯度流时，$R_{vel}$ 为零：
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

全局目标函数将此残差与观测时间 $t$ 处的数据拟合散度 $D$ 相结合：
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
该框架将现有的方法（如 Path-Finding 和 Action Matching）统一为单一残差原理的实例化。

2. “缝合”算法 (The "Stitching" Algorithm)

为了最小化该目标函数，作者引入了 Stitching，这是一种无模拟（simulation-free）、基于粒子的方法，具有以下特征：

• 可学习的轨迹： 与 JKO 方法中 $\rho$ 是隐式定义的不同，Stitching 将曲线 $\rho_\theta$ 显式参数化为一组 $N$ 个可学习粒子轨迹 $x_{t,k}^\theta$ 的核密度估计 (KDE)。
• 无模拟： 该方法不需要求解神经常微分方程 (Neural ODEs) 或积分随机微分方程 (SDEs)。粒子位置和速度是直接的可学习参数。
• 经验近似： 为了提高计算效率，KDE 通过粒子中心处的经验测度进行近似。这简化了速度残差为对粒子的求和：
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• 泛函参数化： 能量泛函 $F_\theta$ 被参数化为势能项 $V_\theta$、熵项（对数密度）和相互作用项 $W_\theta$ 的总和，这些项均由神经网络表示。

核心贡献

1. 残差框架： 作者将 WGF 学习重新表述为最小化非负残差项，提供了一个能够涵盖 Path-Finding 和 Action Matching 的统一视角。
2. Stitching 算法： 他们引入了一种新颖的、无模拟的粒子方法，将轨迹曲线 $\rho$ 视为一等公民的可学习变量。这使得该方法能够容忍观测之间的巨大间隙，并恢复 JKO 基于弦预测器会丢失的弯曲动力学。
3. 最先进的性能： 该方法在轨迹推理基准测试中取得了卓越的结果，特别是在稀疏或非配对数据场景下。

实验结果

论文在三个主要基准测试上评估了 Stitching 的性能：

1. 说明性示例（连续流 vs. JKO）

在一个合成的“波浪谷”势能场中，粒子遵循正弦轨迹。

• 结果： JKO 方法（特别是 JKOnet*）无法捕捉稀疏快照之间（$t=10$ 到 $t=20$）的曲率，产生直线插值。Stitching 正确地恢复了弯曲轨迹，从而实现了更好的底层势能恢复（$R^2 = 0.62$ 对比 $0.51$）。

2. 合成势能恢复

该方法在配对（相关轨迹）和非配对（去相关快照）机制下对 15 个二维势能进行了测试。

• 结果： Stitching 在两种机制下都保持了稳定的性能。相比之下，JKO 方法在非配对设置下性能显著下降或崩溃，因为它们依赖于相邻快照之间的 OT 耦合。Stitching 的误差在两种机制之间保持在约 $\sim 2\times$ 的因子内，而 JKO 方法在非配对情况下往往无法恢复有信息的梯度场。

3. 单细胞轨迹推理（胚胎体数据集）

应用于人类胚胎干细胞分化数据（5 个快照，跨度 27 天）。

• 结果： 与包括 Neural SDEs、TrajectoryNet 和各种 JKO 方法在内的基准模型相比，Stitching 实现了最佳性能（最低的 $W_1$ 和 $W_2$ 距离）。它在全数据训练和“留两个出”的时间插值协议中均优于竞争对手。

4. 相互作用动力学恢复

该方法用于解构 McKean–Vlasov 系统中的环境势能 ($V$) 和相互作用核 ($W$)。

• 结果： Stitching 成功恢复了约束势能和吸引相互作用核，实现了高模式 $R^2$ 分数（$V$ 为 $0.84$，$W$为$0.81$）。此外，它在每轮迭代中比需要求解 OT 问题的 JKO 方法快得多。

5. 非梯度流

作者展示了通过将梯度项替换为一般向量场，可以将残差框架扩展到梯度流之外，从而成功追踪手性（非保守）动力学。

重要性与主张

论文声称，Wasserstein 残差视角为学习群体动力学提供了一种比 JKO 方案更灵活、更鲁棒的替代方案。通过将能量泛函的学习与最优传输耦合的约束解耦，Stitching 实现了：

• 对稀疏性的鲁棒性： 能够从间隔较宽的观测中学习动力学，避免了 JKO 线性插值固有的“曲率损失”。
• 非配对数据的兼容性： 能够从粒子身份未随时间追踪的数据集（非配对快照）中学习，这是生物成像中的常见情况。
• 无模拟的高效性： 避免了在训练期间求解 ODE 或 OT 问题的计算开销。

作者将这项工作定位为现有范式的统一以及科学发现领域的实际进展，涵盖了从单细胞生物学到集体行为的广泛领域，同时也承认了诸如相互作用项的 $O(N^2)$ 计算复杂度以及缺乏 $N \to \infty$ 时理论收敛保证等局限性。