FM × scDynamics(2): Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling
Published:
第二篇文章来讲一下NIPS 2025最新的一篇paper,Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling,文章探讨了在做single cell trajectory inference的时候,之前的方法一般都不考虑质量变化,也就是细胞数量的变化。但实际上,细胞在生长发育的时候会发生凋亡和增殖,也就会发生细胞数量上的变化。之前的很多方法,并不考虑这一点,比如之前提到过的OT-CFM,他就是从起点终点数据集中随机抽取一个batch然后做OT匹配,再进行训练,并不考虑前后整体数量的差距。本文提出的VGFM在问题建模的时候,就考虑到了这一点,它通过流匹配(Flow Matching)技术,联合学习描述细胞状态变化的“速度场”和描述细胞数量变化的“生长函数”。
- 论文原文: arXiv:2505.13413v1
- 代码仓库: VGFM
1. 背景介绍
从稀疏和带噪声的快照数据中推断复杂系统的潜在动力学是科学与工程中的一个基本挑战,在单细胞 RNA 测序(scRNA-seq)领域尤为突出。Lähnemann et al., 2020 概括了当前单细胞数据科学面临的关键难题,而破坏性测序流程只能在离散时间点上提供群体级快照,Macosko et al., 2015 等工作展示了这一常态。由于细胞在发育或响应过程中会持续增殖与凋亡,不同时间点的细胞总数往往不守恒,因此研究目标从学习归一化概率密度转变为刻画随时间演化、允许质量增减的非归一化密度函数。
近年来出现了一系列以深度学习为核心的动力学推断方法,它们多采用数值求解常微分或随机微分方程来拟合速度场。Chen et al., 2018 提出的 Neural ODE、Tong et al., 2020 的 TrajectoryNet 以及 Yeo et al., 2021 的 PRESCIENT 都属于这种基于模拟(simulation-based)范式:模型在训练时需要反复调用数值求解器生成轨迹,并与观测快照对齐,这既带来了高昂的计算代价,也在高维空间放大了参数搜索的难度。
为缓解这一问题,研究者提出了免模拟(simulation-free)的训练框架,直接对构造好的概率路径回归速度场。Lipman et al., 2023 的 Flow Matching、Liu et al., 2023 的 Rectified Flow 以及 Tong et al., 2024 的 OT-CFM 展示了这一思路在效率和稳定性上的优势。
然而,上述免模拟方法大多仅关注速度场 $v_t$,忽略了观测分布中的质量不平衡,这可能导致对真实动力学的偏差重构。Chizat et al., 2018 提出的 Wasserstein–Fisher–Rao (WFR) 框架和 Pariset et al., 2023 的 Unbalanced Diffusion Schrödinger Bridge 试图用统一的度量或联合训练策略来同时建模速度与质量,但前者从数学上要求 $v_t = \nabla g_t$,后者仍需要复杂的优化过程。针对单细胞数据的后续工作(如 Sha et al., 2024)虽然放宽了这一约束,却往往牺牲了可扩展性或可解释性。
本文介绍的 VGFM(Joint Velocity-Growth Flow Matching) 正是在这一背景下提出。它基于对静态半松弛最优传输的动态理解,先求得同时包含状态转移与质量变化的理想“运输蓝图”,再用神经网络联合回归速度场和生长函数,并辅以基于 Wasserstein 距离的分布拟合损失,从而在训练效率、可扩展性与生物学解释之间取得更好的平衡。
2. 方法(Methods)
这一节主要是详细分析一下论文的数学背景和方法设计的动机,从单细胞动力学建模中的根本挑战——“不平衡”问题出发,展示为何传统方法失效,然后引入最优传输(OT)作为寻找“最优蓝图”的工具,并最终推导出一种高效、可学习的流匹配训练方案。
2.1 挑战:对不平衡动态系统建模
在不考虑细胞增殖或凋亡的理想情况下,一个细胞群体的动态可以被一个简单的常微分方程(ODE)描述。每个细胞的状态 $\mathbf{x}_t$ 随着一个速度场 $v_t(\mathbf{x}_t)$ 演化: \(\frac{d\mathbf{x}_t}{dt} = v_t(\mathbf{x}_t) \tag{1}\) 这对应于一个标准的连续性方程,它描述了细胞密度 $\rho_t$ 如何在速度场 $v_t$ 的“推动”下随时间变化,但总质量(细胞总数)保持不变: \(\partial_t \rho_t = -\nabla \cdot (\rho_t v_t) \tag{2}\) 然而,在真实的生物系统中,细胞会增殖和死亡,导致细胞总数发生变化。这是一个“不平衡”系统。为了对此建模,论文引入了一个关键的补充项:生长函数 $g_t(\mathbf{x})$。系统的 ODE 被扩展为包含两个部分: \(\begin{cases} \frac{d\mathbf{x}_t}{dt} = v_t(\mathbf{x}_t) & \text{(状态变化)} \\ \frac{d \log w_t(\mathbf{x}_t)}{dt} = g_t(\mathbf{x}_t) & \text{(质量/数量变化)} \end{cases} \tag{3}\) 这里的 $w_t$ 可以理解为跟随每个细胞轨迹的一个动态权重。当 $g_t > 0$ 时,权重增加,代表细胞增殖;当 $g_t < 0$ 时,权重减小,代表细胞凋亡。这个扩展后的系统对应的连续性方程也增加了一个“源/汇”项: \(\partial_t \rho_t = \underbrace{-\nabla \cdot (\rho_t v_t)}_{\text{状态转移(搬运)}} + \underbrace{g_t \rho_t}_{\text{质量增减(源/汇)}} \tag{4}\) 其中,$g_t \rho_t$ 这一项直接描述了在每个位置 $\mathbf{x}$ 上,细胞密度因增殖或死亡而发生的净变化率。
至此,核心问题被清晰地定义:给定两个时间点的细胞快照(分布 $p_0$ 和 $p_1$),我们如何找到一对最“合理”的速度场 $v_t$ 和增长函数 $g_t$ 来描述这个演化过程?
2.2 困境:无法简单构造演化路径
直接求解 $v_t$ 和 $g_t$ 是一个不适定问题(ill-posed problem),因为存在无数种可能的解。常规的流匹配(Flow Matching)方法通过一个巧妙的简化来解决这个问题:它们假设一个已知的路径。例如,通过从起点和终点分布中随机配对 $(\mathbf{x}_0, \mathbf{x}_1)$,可以定义一条简单的直线路径 $\mathbf{x}_t = (1-t)\mathbf{x}_0 + t\mathbf{x}_1$。对这条路径求导,就能立刻得到目标速度 $v_t = \mathbf{x}_1 - \mathbf{x}_0$,从而将复杂的动力学问题简化为一个直接的回归问题。
但在我们的不平衡场景下,这个方法失效了。由于 $p_0$ 和 $p_1$ 的细胞总数不同,我们无法在两者之间建立稳定的一一对应关系。既然无法配对,就无法定义一条简单的“直线路径”,也就无法预先知道目标速度场 $v_t$ 和目标增长函数 $g_t$ 应该是什么。我们需要一个更有原则性的方法来确定唯一的、最优的演化路径。
2.3 蓝图:用半松弛最优传输(SROT)寻找最优路径
为了找到那条“最优”的演化路径,论文从半松弛最优传输 (Semi-Relaxed Optimal Transport, SROT) 理论中寻找灵感。SROT 旨在找到两个总质量不同的分布之间的最优匹配方案。论文中用以下公式(5)定义了一个综合成本函数: \(\min_{\pi \ge 0} \underbrace{\int_{\Omega^2} c(\mathbf{x}_0, \mathbf{x}_1) d\pi(\mathbf{x}_0, \mathbf{x}_1)}_{\text{移动成本}} + \underbrace{\mathrm{KL}(P^0_{\#}\pi \,\|\, p_0)}_{\text{质量变化成本}} \quad \text{s.t.} \quad P^1_{\#}\pi = p_1 \tag{5}\) 其中:
- 移动成本: 将一个质量微元从 $\mathbf{x}_0$ 移动到 $\mathbf{x}_1$ 的代价,通常使用二次成本 $c(\mathbf{x}_0, \mathbf{x}_1) = |\mathbf{x}_0-\mathbf{x}_1|^2$。
- 质量变化成本: 这是一个KL散度项,它惩罚传输方案 $\pi$ 的起始边缘分布 \(P^0_{\#}\pi\) 与原始分布 $p_0$ 之间的差异。直观上,这个差异就代表了在每个位置上需要发生的质量增减。
可以发现这是一个静态优化问题,因为它不包含任何关于时间 $t$ 或速度 $v_t$ 的信息。它输入两个静态分布 $p_0$ 和 $p_1$,输出一个静态的“传输方案” $\pi^$。这个方案 $\pi^$ 是一个最优蓝图,它精确地描述了 $p_0$ 中的每个点应该与 $p_1$ 中的哪些点建立对应关系,以及在此过程中质量应该如何变化。
这个SROT为我们提供了一个静态的、最优的“蓝图” $\pi^*$,但它本身并不是一个动态的过程。它只告诉你“起点和终点的最优匹配”,不告诉你中间的路径。
2.4 动态化:从静态蓝图到两阶段模型
论文的核心insight就在于,可以将上述静态的 SROT 过程创造性地理解为一个分为两步走的动态过程。并且,通过文中的Proposition 1 在数学上严格证明了这个动态过程的最优解与原始静态 SROT 问题的最优解是等价的。这个动态优化问题如下公式(6)所示: \(\min_{(v_t,g_t) \in \mathcal{C}_\lambda(p_0, p_1)} \mathcal{J}^\lambda_{\rm tpt}(v_t,g_t) := (1 - \lambda) \int_\Omega \int_\lambda^1 p_t(\mathbf{x}) \|v_t(\mathbf{x})\|^2 \mathrm{d}t \mathrm{d}\mathbf{x} + \mathcal{H}(v_t,g_t,p_t) \tag{6}\) 其中:
- $\mathcal{J}^\lambda_{\rm tpt}(v_t,g_t)$ 是总的动态成本。
- 第一项是运输成本:在第二阶段($t \in [\lambda, 1]$),对速度场 $v_t$ 的能量(平方范数)进行积分,衡量状态转移的效率。
- 第二项 $\mathcal{H}(v_t,g_t,p_t)$ 是增长成本:它与第一阶段($t \in [0, \lambda]$)的增长函数 $g_t$ 相关,本质上是衡量从 $p_0$ 变为 $p_\lambda$ 所需的“代价”,其形式与KL散度紧密相关。
- $\mathcal{C}_\lambda(p_0, p_1)$ 代表所有可能的路径 $(v_t, g_t)$ 的集合,这些路径必须满足两阶段的演化约束。
从公式中就可以看出,这个动态模型将复杂的演化过程解耦为两个独立的阶段:
- 第一阶段:纯质量增长 (t ∈ [0, λ])
- 在此阶段,细胞状态不移动 ($v_t = 0$),只有细胞数量发生变化。
- 该过程完全由增长函数 $g_t$ 控制,它将初始分布 $p_0$ 演化为一个中间分布 $p_\lambda$,这个中间分布的总质量与最终分布 $p_1$ 完全相同。
- 第二阶段:纯状态转移 (t ∈ (λ, 1])
- 在此阶段,细胞数量不再变化 ($g_t = 0$),只有细胞状态发生移动。
- 该过程完全由速度场 $v_t$ 控制,它通过“搬运”将中间分布 $p_\lambda$ 演化为最终分布 $p_1$。
这是一个动态优化问题,因为它是在所有可能的路径 $(v_t, g_t)$ 空间中寻找最优解,其优化目标里明确包含了对速度 $|v_t|^2$ 的积分。这个模型虽然在理论上很完美,但它与生物学现实并不完全相符,因为细胞的状态变化和数量变化通常是同时发生的。
2.5 现实化:通过重参数化实现联合动力学
为了使模型更贴近生物学现实,作者进行了一次巧妙的数学变换(重参数化),将两阶段模型中解耦的 $v_t$ 和 $g_t$ 融合为一组新的、在整个时间区间 [0, 1] 上同时起作用的联合速度场 $\tilde{v}_t$ 和联合增长函数 $\tilde{g}_t$: \(\begin{aligned} \tilde{v}_t(\mathbf{x}) &= (1 - \lambda) \cdot v_{(1 - \lambda)t + \lambda}(\mathbf{x}) \\ \tilde{g}_t(\mathbf{x}) &= \lambda \cdot g_{\lambda t}\left(\psi_{\tilde{v},t}^{-1}(\mathbf{x}) \right) \end{aligned} \tag{7}\) 更重要的是,通过文章中的Theorem 1证明了这个更真实的联合动态过程,其最终结果与理论上最优的两阶段过程完全相同。这样,我们就获得了一个既有坚实理论基础(源于SROT)又在形式上更合理(状态与数量同步变化)的理想动态模型 $(\tilde{v}_t, \tilde{g}_t)$。
2.6 发现规律:最优路径的简单形式
尽管我们有了理想的动态模型 \((\tilde{v}_t, \tilde{g}_t)\),但其数学表达式(公式(7))依然很复杂。特别是 \(\tilde{g}_t\) 的定义中包含了对流映射 $\psi_{\tilde{v},t}^{-1}(\mathbf{x})$ 的依赖,这意味着在计算 $t$ 时刻的增长率时,我们需要知道粒子在该时刻的位置 \(\mathbf{x}_t = \psi_{\tilde{v},t}(\mathbf{x}_0)\),而这个位置本身就是常微分方程 $\frac{d\mathbf{x}}{dt} = \tilde{v}_t(\mathbf{x})$ 的解。这种依赖关系使得直接从数据中学习它变得非常困难。
因此,我们需要找到更简化的形式。这里的关键突破在于分析并找到最优形式的速度场和生长函数,再将其代入我们的联合动态模型中。
- 分析速度 $\tilde{v}_t$: 论文的 Proposition 1 已经为我们指明了方向。它给出了在两阶段模型中第二阶段(纯状态转移)的最优速度场形式: \(v^*_t\left(x+\frac{t-\lambda}{1-\lambda}(T^*(x)-x)\right) = \frac{T^*(x)-x}{1-\lambda}\) 这个公式的物理意义是,从位置 $\mathbf{x}$ 到达其最优终点 \(T^*(\mathbf{x})\) 的路径是一条直线,并且在这条路径上的速度是恒定的。这背后的原理是经典的布雷尼耶定理 (Brenier’s Theorem),该定理指出,在二次成本下,最优传输的路径是由一簇互不相交的直线构成的。每个粒子都从其起点 $\mathbf{x}_0$ 匀速直线运动到其唯一的终点 $T^*(\mathbf{x}_0)$。 将这个最优的、恒定的速度代入我们的重参数化公式(公式(7)),就可以得到联合模型中同样简洁的目标速度:
- 目标速度: \(\tilde{v}_t(\psi_{\tilde{v},t}(\mathbf{x}_0)) = T^*(\mathbf{x}_0) - \mathbf{x}_0\)
- 分析增长 $\tilde{g}_t$: 对于第一阶段的纯质量增长,满足 \(p_\lambda(\mathbf{x}) = p_0(\mathbf{x}) \exp\left( \int_0^\lambda g_s(\mathbf{x}) ds \right)\) 的生长函数 \(g_t\) 其实有很多种。为了便于实现,论文做了一个简化,即假设最优生长率 $g^*_t$ 不随时间变化。在这个假设下,可以得到: \(g^*_t(\mathbf{x}) = \frac{\log {\rm P}_{\#}^0\pi^*(\mathbf{x}) - \log p_0(\mathbf{x})}{\lambda}\) 同样地,将这个时间无关的生长率代入重参数化公式(公式(7)),我们得到了一个同样不依赖于 $t$ 的目标增长率:
- 目标增长率: \(\tilde{g}_t(\psi_{\tilde{v},t}(\mathbf{x}_0)) = \log\big( ({\rm P}^0_{\#}\pi^*)(\mathbf{x}_0) \big) - \log\big( p_0(\mathbf{x}_0) \big)\)
经过这一系列推导,我们把一个极其复杂的动力学问题,转化为了寻找两个非常简单的目标:一个常数向量(目标速度)和一个常数标量(目标增长率)。这为使用流匹配进行训练铺平了道路。
2.7 Flow Matching 训练
现在,我们可以定义一个简单直接的回归损失来训练我们的神经网络 $v_\theta$ 和 $g_\omega$。然而,在真实的单细胞实验中,我们无法得到连续的概率分布 $p_0$ 和 $p_1$,也无法直接得到最优传输映射 $T^*$。我们拥有的只是一堆离散的、像沙子一样的样本点:第一个时间点的细胞样本 ${\mathbf{x}_0^i}$ 和第二个时间点的细胞样本 ${\mathbf{x}_1^j}$。
于是,VGFM学习OT-CFM,先使用Unbalanced OT进行一次匹配,然后再进行训练。因此,论文采用了基于熵正则的 Sinkhorn 算法来求解一个离散化的 SROT 问题,其优化目标如下: \(\pi^{0\to1} = \arg\min_{\pi \geq 0} \sum_{i,j} c_{ij} \pi_{ij} + \epsilon H(\pi) + \tau \mathrm{KL}(\pi\mathbf{1}_m||\mathbf{1}_n),\quad \text{subject to} \quad \pi^\top\mathbf{1}_n = \mathbf{1}_m\) 其中 $c_{ij}=\Vert x_0^i-x_1^j\Vert^2$,$H(\pi)$ 是负熵项,$\epsilon$ 和 $\tau$ 是超参数。和传统的OT-CFM的区别就在于,VGFM面向Unbalanced OT,所以要添加一个KL散度项来$\mathrm{KL}(\pi\mathbf{1}_m||\mathbf{1}_n)$ 来处理非平衡最优传输的质量变化。
求解这个优化问题后,我们得到一个离散的传输方案(一个矩阵)$\pi^{0\to1}$,它告诉我们每个起始细胞 $\mathbf{x}_0^i$ 与每个终点细胞 $\mathbf{x}_1^j$ 之间的关联强度。基于这个方案,我们可以近似出训练所需的目标:
近似 \(T^*(\mathbf{x}_0^i)\):细胞 $i$ 的理想终点是什么?一个自然的想法是将其所有可能的终点 \(\mathbf{x}_1^j\) 按关联强度 $\pi_{ij}$ 进行加权平均,即“质心映射”(Barycentric Mapping): \(T^*(\mathbf{x}_0^i) \approx \frac{\sum_j \pi_{ij}^{0\to1} \mathbf{x}_1^j}{\sum_j \pi_{ij}^{0\to1}}\) 在论文的实现中,为了更简单,直接从这个权重分布中随机采样一个 $\mathbf{x}_1^j$ 作为近似终点。也就是说,对于每个起始细胞 $i$,我们将运输方案 $\pi$ 的第 $i$ 行看作一个概率分布,从中抽取一个终点 $j$,并认定在这次计算中,细胞 $i$ 的归宿就是细胞 $j$。
近似细胞路径:一旦确定了起点 $\mathbf{x}_0^i$ 和近似的终点 $\mathbf{x}_1^j$,最简单的路径就是一条直线。在时间 $t$ 上的位置就是: \(\mathbf{x}_t \approx \mathbf{x}_0^i + t(\mathbf{x}_1^j - \mathbf{x}_0^i)\)
有了这些近似,我们就可以定义 VGFM 损失函数,让神经网络的预测去逼近这些简单的目标:
\[\mathcal{L}_{VGFM}(\theta,\omega)=\mathbb{E}_{(\mathbf{x}_0^i,\mathbf{x}_1^j)\sim\pi^{0\to1}}\,\mathbb{E}_{t\sim U(0,1)}\big[\,\|v_\theta(\mathbf{x}_t,t)-(\mathbf{x}_1^j-\mathbf{x}_0^i)\|^2+|g_\omega(\mathbf{x}_t,t)-\log([\pi^{0\to1}\mathbf{1}_m]_i)|^2\,\big] \tag{8}\]这个损失函数可以直观地理解为比较“学生答案”和“标准答案”:
- 速度部分: \(\|v_\theta(\mathbf{x}_t, t) - (\mathbf{x}_1^j-\mathbf{x}_0^i)\|^2\)
- $v_\theta(\mathbf{x}_t, t)$ 是我们神经网络模型给出的“学生答案”,它预测在 $t$ 时刻 $\mathbf{x}_t$ 位置的速度。
- $(\mathbf{x}_1^j-\mathbf{x}_0^i)$ 是通过“直线近似”得到的“标准答案”。如果一个物体从 $\mathbf{x}_0^i$ 匀速直线运动到 $\mathbf{x}_1^j$,它的速度向量就是恒定的 $\mathbf{x}_1^j-\mathbf{x}_0^i$。
- 生长部分: \(\|g_\omega(\mathbf{x}_t, t) - \log([\pi^{0\to1}\mathbf{1}_m]_i)\|^2\)
- $g_\omega(\mathbf{x}_t, t)$ 是生长网络的“学生答案”。
- $\log([\pi^{0\to1}\mathbf{1}_m]_i)$ 是生长率的“标准答案”,代表了起始细胞 $i$ 最终的总“质量”变化(取对数)。这里的 $[\pi^{0\to1}\mathbf{1}_m]_i$ 是传输矩阵 $\pi$ 第 $i$ 行的和,表示从 $\mathbf{x}_0^i$ 出发的总传出质量。
- 一个细节: 为什么标准答案里没有
log(p₀(x₀ⁱ))这一项了?因为在处理离散样本时,我们通常假设每个起始样本点的初始“质量”都是1,而 $\log(1)=0$,因此可以省略。
在实际训练中,我们并不会一次性计算所有匹配对的损失,而是采用小批量梯度下降:
- 根据运输方案 $\pi$ 的概率,随机抽取一批匹配对 $(i, j)$。
- 为这批匹配对随机抽取一些时间点 $t$。
- 计算这批样本的损失,并据此更新神经网络。
- 不断重复此过程。 此外,为了增强模型的鲁棒性,在计算 $\mathbf{x}_t$ 时还会加入少量高斯噪声,这是一种在流匹配方法中常用的技巧。
不过,作者发现仅依靠 \(\mathcal{L}_{VGFM}\) 不足以完美地拟合数据。因此,论文还引入了一个基于模拟的分布拟合损失 \(\mathcal{L}_{OT}\)。具体做法是,在通过 \(\mathcal{L}_{VGFM}\) 进行一段时间的训练后,用训练好的速度网络 $v_\theta$ 和生长网络 $g_\omega$ 从初始时间点出发,通过数值求解器(模拟)生成中间时刻的细胞分布,然后将这个预测分布与真实的观测数据快照用 1-Wasserstein距离 进行比较。
最终,VGFM采用了两阶段的训练策略:
- 热身阶段: 只使用 $\mathcal{L}_{VGFM}$ 损失函数,让模型快速学到路径的基本方向和节奏。
- 联合训练阶段: 启用完整的损失函数 \(\mathcal{L}(\theta,\omega) = \mathcal{L}_{VGFM}(\theta,\omega) + \mathcal{L}_{OT}(\theta,\omega)\) 进行微调,以确保模型在整体分布上也与真实数据对齐,达到最佳性能。
3. 实验与讨论
本节将重组并深入探讨 VGFM 的实验表现,重点解读其核心优势、设计选择的合理性,并结合局限性提供实践建议。
3.1 核心评价指标
为了全面评估模型性能,尤其是在处理细胞数量变化的场景下,作者主要采用了两个关键指标:
分布匹配度 (Wasserstein-1 距离, W1): 这是单细胞动力学研究中常用的度量,用于衡量模型预测的细胞群体分布与真实的细胞群体分布有多接近。这里的“分布”可以理解为所有细胞在基因表达空间中的整体位置和形状。W1 距离越小,说明模型对细胞状态演化的轨迹预测得越准。
质量匹配度 (相对质量误差, RME): 该指标专门用于衡量模型预测的细胞总数与真实的细胞总数之间的差距。在细胞会发生增殖和凋亡的系统中,总质量(细胞数)会发生变化。RME 分数越低,说明模型对细胞增殖/死亡的速率和位置判断得越准,这是 VGFM 相比于传统方法的核心优势之一。
3.2 综合实验表现
VGFM 在合成数据与真实单细胞数据(如 EB、CITE-seq)上与多种基线方法进行了比较。实验结果表明:
性能优势: 在分布拟合(W1)和质量匹配(RME)两项指标上,VGFM 普遍优于或不输于现有强基线模型,包括那些专门为不平衡任务设计的方法。这证明了 VGFM “速度+生长”的联合学习框架,能够在精确刻画质量变化的同时,不牺牲对细胞状态演化路径的拟合精度。
高维可扩展性: 在 1000 维的高斯混合模型上进行了压力测试。结果显示,若干基于模拟(simulation-based)的方法在高维空间中训练不稳甚至难以收敛,而 VGFM 依然能够学到合理的轨迹与质量变化。这充分体现了其免模拟(simulation-free)训练框架在高维任务中的稳定性与可扩展性。
多场景验证: 在留一补空(leave-one-out)等多个插值实验中,VGFM 的效果同样好于其他方法,验证了其模型的鲁棒性和泛化能力。
3.3 消融研究:损失函数的重要性
VGFM 的最终损失函数由两部分构成:基于流匹配的回归损失 \(\mathcal{L}_{VGFM}\) 和基于模拟的分布拟合损失 \(\mathcal{L}_{OT}\)。消融实验旨在探究这两个部分各自的贡献:
仅使用 \(\mathcal{L}_{VGFM}\): 去掉 \(\mathcal{L}_{OT}\) 后,虽然模型训练速度有显著提升,因为不再需要数值求解器生成中间分布,但最终的分布拟合效果(W1 距离)会变差。这说明 \(\mathcal{L}_{VGFM}\) 虽能为模型提供正确的演化“方向”和“节奏”,但缺少了 \(\mathcal{L}_{OT}\) 进行全局校准,容易在中间时刻产生累积误差。
仅使用 \(\mathcal{L}_{OT}\): 去掉 \(\mathcal{L}_{VGFM}\) 后,模型效果差很多,并且训练时间会暴增。这说明如果完全依赖模拟和分布匹配,模型在训练初期会因为缺少明确的路径引导而难以优化,收敛缓慢且不稳定。
结论: 两个损失函数缺一不可。\(\mathcal{L}_{VGFM}\) 提供了一个高效的“热身”阶段,让模型快速学到从 SROT 蓝图导出的最优路径的大致形态;而 \(\mathcal{L}_{OT}\) 则在此基础上进行“微调”,确保模型生成的整体动力学过程与所有观测快照都精确对齐。二者结合,才实现了既快又准的训练。
3.4 训练效率的飞跃
VGFM 最大的亮点之一是其训练效率。在 50 维的 EB 数据集上,与完全依赖模拟的现有最优方法之一 DeepRUOT 相比:
- DeepRUOT: 完成一次训练需要约 90 分钟。
- VGFM: 完成一次训练仅需 13 分钟。
效率的大幅提升,根本原因在于其“先有蓝图,再回归”的免模拟训练范式。传统方法需要在每次训练迭代中反复调用 ODE/SDE 求解器来生成轨迹,计算量巨大。而 VGFM 首先通过 SROT 一次性计算出静态的“最优运输蓝图”,然后将复杂的动力学学习问题转化为一个简单的回归问题,让神经网络去拟合这个蓝图所蕴含的目标速度和目标生长率,从而避免了高昂的模拟开销。
4. 讨论与展望
综合来看,VGFM 的“静态蓝图 + 动态重构 + 免模拟回归”路线,为不平衡单细胞动力学建模提供了一个兼顾准确性、效率与可解释性的优秀方案。
优点: 其核心优势在于路径由 SROT 理论导出,避免了主观构造路径的偏差;免模拟训练带来了效率和高维稳定性;并且速度 $v$ 和生长 $g$ 两个分量分别对应细胞的“位移”和“增减”,具有良好的生物学可解释性。
局限与实践建议: 方法的效果在很大程度上依赖于 SROT 求解的质量以及数据的计数校正。实践中,SROT 的熵正则、松弛系数等超参数需要仔细调优。此外,对数据进行批次校正、文库大小归一化等预处理,以弱化技术噪声对生长函数 $g$ 学习的干扰,是十分必要的。
未来方向: 当前模型假设生长率 $g$ 是时间常数,未来可以探索时变的生长函数以适应更复杂的生物节律。同时,将多个时间点的快照联合纳入多边缘最优传输框架,或融合基因调控网络等生物先验知识作为正则项,都是值得探索的方向。
