TracingFlow: 一个基于二阶动力学的无模拟轨迹推断框架

# 基于二阶动力学的无模拟轨迹推断框架
来源: https://arxiv.org/html/2608.21070

## TracingFlow：基于二阶动力学的无模拟轨迹推断框架

孙宇浩¹ ²
吴泽坤¹ ¹ 通信作者
¹ 北京大学数学科学学院
黄梓勋
北京大学数学科学学院
周培杰
谢谢: 通信作者: pjzhou@pku\.edu\.cn
北京大学机器学习研究中心
北京大学定量生物学中心
北京大数据分析与应用国家工程实验室
北京科学智能研究院

###### 摘要
从稀疏的时间快照中推断连续系统演化是生成模型和单细胞组学中的一个关键挑战。虽然最优传输（OT）很流行，但现有框架大多局限于一阶动力学，假设速度场无记忆。这限制了表达能力，因为一阶系统无法解释细胞分化等过程中固有的调节动量和时间延迟响应。在此，我们介绍了TracingFlow，一个推广到二阶动力学的无模拟流匹配框架。通过使用神经网络来回归加速度场，TracingFlow为动态最优加速传输（DOAT）问题提供了一个精确、高效的解决方案。与产生过度平滑轨迹的一阶方法不同，我们的二阶公式通过学习底层力场来捕捉高曲率转变和非线性演化。在复杂合成数据集和大规模scRNA-seq数据集上的评估表明，TracingFlow在分布重建和轨迹保真度方面取得了卓越的准确性。此外，通过整合谱系追踪先验，它恢复了在数学上最优且生物学上合理的动态结构。

††脚注文本: 预印本。2026年8月21日。

## 1引言
从离散观测中恢复底层动力学是单细胞组学中的一项关键任务，被称为轨迹推断（TI）[9 (https://arxiv.org/html/2608.21070#bib.bib9), 72 (https://arxiv.org/html/2608.21070#bib.bib72)]。为了识别分布之间的最小代价演化，最优传输（OT）框架应运而生[7 (https://arxiv.org/html/2608.21070#bib.bib7), 71 (https://arxiv.org/html/2608.21070#bib.bib71)]。这些框架通常分为静态OT方法[51 (https://arxiv.org/html/2608.21070#bib.bib51), 28 (https://arxiv.org/html/2608.21070#bib.bib28), 22 (https://arxiv.org/html/2608.21070#bib.bib22)]（学习直接映射）和动态OT方法[57 (https://arxiv.org/html/2608.21070#bib.bib57), 24 (https://arxiv.org/html/2608.21070#bib.bib24), 67 (https://arxiv.org/html/2608.21070#bib.bib67), 44 (https://arxiv.org/html/2608.21070#bib.bib44)]（通过流映射捕捉连续演化）。早期的动态OT框架依赖于神经ODE[9 (https://arxiv.org/html/2608.21070#bib.bib9)]，由于数值积分而产生了高计算成本。为解决此问题，流匹配[35 (https://arxiv.org/html/2608.21070#bib.bib35), 36 (https://arxiv.org/html/2608.21070#bib.bib36), 48 (https://arxiv.org/html/2608.21070#bib.bib48)]被提出作为一种无模拟的替代方案，通过回归速度场将源分布推向目标分布。通过将传输成本分解为单粒子轨迹，流匹配有效地解决了动态OT问题[58 (https://arxiv.org/html/2608.21070#bib.bib58), 27 (https://arxiv.org/html/2608.21070#bib.bib27), 50 (https://arxiv.org/html/2608.21070#bib.bib50)]，及其各种变体[59 (https://arxiv.org/html/2608.21070#bib.bib59), 15 (https://arxiv.org/html/2608.21070#bib.bib15), 8 (https://arxiv.org/html/2608.21070#bib.bib8), 45 (https://arxiv.org/html/2608.21070#bib.bib45)]。然而，大多数TI框架假设一阶动力学（$\dot{\bm{x}}=\bm{v}(\bm{x},t)$），固有地将$\bm{v}$限制为相对于$\bm{x}$的单值函数。当建模复杂的生物先验（如谱系追踪）[30 (https://arxiv.org/html/2608.21070#bib.bib30), 38 (https://arxiv.org/html/2608.21070#bib.bib38)]时，这限制了表达能力，可能产生与附加生物先验信息相矛盾的动力学。此外，[19 (https://arxiv.org/html/2608.21070#bib.bib19), 34 (https://arxiv.org/html/2608.21070#bib.bib34), 23 (https://arxiv.org/html/2608.21070#bib.bib23)]指出，如果考虑更多模态，如染色质可及性或蛋白质组学，生物动力学可能遵循更高阶的动力学，或者可以在增广空间中建模。虽然3MSBM[56 (https://arxiv.org/html/2608.21070#bib.bib56)]最近引入了二阶动量薛定谔桥以获得更平滑的轨迹，但它依赖于类似于Rectified Flow[36 (https://arxiv.org/html/2608.21070#bib.bib36)]的迭代重训练加速度场来确定最优耦合。此外，它采用启发式方法估计初始速度，而不是将此估计作为最优控制问题的一部分纳入。为克服这些限制，我们引入了**TracingFlow**，一个利用二阶动力学系统的无模拟框架。我们提出了动态最优加速传输（DOAT）问题，它通过最小化加速度成本来确定演化。TracingFlow直接回归加速度场和初始速度，避免了ODE模拟。我们的实验表明，TracingFlow与现有的无模拟框架[58 (https://arxiv.org/html/2608.21070#bib.bib58), 59 (https://arxiv.org/html/2608.21070#bib.bib59), 69 (https://arxiv.org/html/2608.21070#bib.bib69), 42 (https://arxiv.org/html/2608.21070#bib.bib42)]相比，取得了有竞争力且通常更好的重建精度，同时有效地纳入了生物先验。我们的贡献包括：
- • 我们提出了TracingFlow，一个基于二阶动力学的无模拟TI框架。它通过回归加速度场解决了所提出的DOAT问题，与基于ODE的方法相比降低了计算成本。
- • 我们通过解耦最优传输计划和单粒子控制的设计提供了理论保证。我们进一步引入了一种迭代策略，将实际问题中的速度缺失-DOAT（VM-DOAT）问题转化为标准DOAT问题，从而能够获得近似解。
- • 我们在多时间点的真实世界数据集上证明了TracingFlow的有效性。结果表明，与现有最先进方法相比，其具有优越的分布重建精度和增强的保留生物先验的能力。

参见标题图1: TracingFlow示意图

## 2相关工作

#### 通过流匹配求解最优传输
流匹配[35 (https://arxiv.org/html/2608.21070#bib.bib35), 36 (https://arxiv.org/html/2608.21070#bib.bib36), 48 (https://arxiv.org/html/2608.21070#bib.bib48), 1 (https://arxiv.org/html/2608.21070#bib.bib1)]通过学习到的流映射传输概率分布，提供了可扩展的无模拟训练。作为一个最优控制问题，最优传输（OT）[47 (https://arxiv.org/html/2608.21070#bib.bib47)]可以分解为最优耦合和单粒子测地线问题，从而便于通过流匹配解决[58 (https://arxiv.org/html/2608.21070#bib.bib58), 27 (https://arxiv.org/html/2608.21070#bib.bib27), 50 (https://arxiv.org/html/2608.21070#bib.bib50), 15 (https://arxiv.org/html/2608.21070#bib.bib15)]。这种方法已扩展到各种变体和推广[59 (https://arxiv.org/html/2608.21070#bib.bib59), 15 (https://arxiv.org/html/2608.21070#bib.bib15), 8 (https://arxiv.org/html/2608.21070#bib.bib8), 13 (https://arxiv.org/html/2608.21070#bib.bib13), 61 (https://arxiv.org/html/2608.21070#bib.bib61), 45 (https://arxiv.org/html/2608.21070#bib.bib45), 26 (https://arxiv.org/html/2608.21070#bib.bib26), 68 (https://arxiv.org/html/2608.21070#bib.bib68), 4 (https://arxiv.org/html/2608.21070#bib.bib4), 46 (https://arxiv.org/html/2608.21070#bib.bib46)]。然而，这些通常采用一阶动力学。虽然二阶方法已经出现[56 (https://arxiv.org/html/2608.21070#bib.bib56)]，但它们需要迭代训练加速度场以实现最优耦合。此外，它们依赖于启发式方法进行初始速度估计，而不是将其作为最优控制问题的组成部分。TracingFlow通过将动态加速度最优传输（DOAT）问题分解为最优耦合和单粒子最优轨迹来解决此问题，实现了DOAT的真正无模拟过程。

#### 克服流匹配中的单值速度约束
在标准流匹配中，速度场被建模为仅依赖于$\bm{x}$和$t$的函数。这种单值依赖性阻止模型表示具有不同速度的相交轨迹，通常导致弯曲的推断路径和增加的计算成本。为了将速度与当前位置解耦，速度网络可以引入额外的输入，如类别标签[73 (https://arxiv.org/html/2608.21070#bib.bib73), 54 (https://arxiv.org/html/2608.21070#bib.bib54)]、初始位置[14 (https://arxiv.org/html/2608.21070#bib.bib14)]或隐藏状态[20 (https://arxiv.org/html/2608.21070#bib.bib20)]。此外，[69 (https://arxiv.org/html/2608.21070#bib.bib69)]采用分层生成来缓解此问题。然而，这些方法并未扩展到二阶动力学。虽然[42 (https://arxiv.org/html/2608.21070#bib.bib42)]利用了二阶系统，但它仅限于恒定加速度路径。相比之下，TracingFlow提供了一个灵活的二阶设置。通过设计最小成本路径，它有效地解决了单值速度场的限制。

#### 单细胞轨迹推断与谱系整合
最优传输（OT）是从scRNA-seq数据推断细胞动力学的强大框架[51 (https://arxiv.org/html/2608.21070#bib.bib51), 28 (https://arxiv.org/html/2608.21070#bib.bib28), 57 (https://arxiv.org/html/2608.21070#bib.bib57), 24 (https://arxiv.org/html/2608.21070#bib.bib24), 40 (https://arxiv.org/html/2608.21070#bib.bib40), 41 (https://arxiv.org/html/2608.21070#bib.bib41), 2 (https://arxiv.org/html/2608.21070#bib.bib2), 74 (https://arxiv.org/html/2608.21070#bib.bib74), 37 (https://arxiv.org/html/2608.21070#bib.bib37), 66 (https://arxiv.org/html/2608.21070#bib.bib66), 12 (https://arxiv.org/html/2608.21070#bib.bib12), 60 (https://arxiv.org/html/2608.21070#bib.bib60), 33 (https://arxiv.org/html/2608.21070#bib.bib33), 53 (https://arxiv.org/html/2608.21070#bib.bib53), 29 (https://arxiv.org/html/2608.21070#bib.bib29), 6 (https://arxiv.org/html/2608.21070#bib.bib6), 10 (https://arxiv.org/html/2608.21070#bib.bib10), 25 (https://arxiv.org/html/2608.21070#bib.bib25), 70 (https://arxiv.org/html/2608.21070#bib.bib70), 55 (https://arxiv.org/html/2608.21070#bib.bib55), 44 (https://arxiv.org/html/2608.21070#bib.bib44), 65 (https://arxiv.org/html/2608.21070#bib.bib65), 52 (https://arxiv.org/html/2608.21070#bib.bib52), 11 (https://arxiv.org/html/2608.21070#bib.bib11)]。然而，仅凭转录组相似性通常无法解析复杂的轨迹。为缓解此问题，研究引入了谱系追踪先验，例如克隆条形码信息。现有方法通过正则化[17 (https://arxiv.org/html/2608.21070#bib.bib17), 49 (https://arxiv.org/html/2608.21070#bib.bib49)]、结构对齐[32 (https://arxiv.org/html/2608.21070#bib.bib32)]、速度映射[62 (https://arxiv.org/html/2608.21070#bib.bib62)]或稀疏转移建模[63 (https://arxiv.org/html/2608.21070#bib.bib63), 21 (https://arxiv.org/html/2608.21070#bib.bib21), 18 (https://arxiv.org/html/2608.21070#bib.bib18)]整合此类信息。虽然提高了准确性，但这些通常关注离散耦合或静态矩阵。TracingFlow通过将先验嵌入基于连续OT的二阶流匹配框架中，捕捉与谱系信息一致的复杂动态，从而推进了这一领域。

## 3数学背景
本节提供TracingFlow所解决问题的数学表述。

#### 动态最优加速传输（DOAT）问题
设$\mathcal{X}\subset\mathbb{R}^{d}$和$\mathcal{V}\subset\mathbb{R}^{d}$分别表示*位置*空间和*速度*空间。我们将*增广空间*定义为$\mathcal{S}\coloneqq\mathcal{X}\times\mathcal{V}$，这样完整的系统状态就是$\bm{s}=(\bm{x},\bm{v})\in\mathcal{S}$。考虑二阶动态系统：
$$
\dot{\bm{x}}=\bm{v}\quad\dot{\bm{v}}=\bm{a}(\bm{x},\bm{v},t),\quad t\in[t_{0},t_{K}].
\tag{1}
$$
通过增广空间$\mathcal{S}$中的概率密度$\rho_{t}$建模大量此类粒子，我们遵循[5 (https://arxiv.org/html/2608.21070#bib.bib5), 11 (https://arxiv.org/html/2608.21070#bib.bib11)]制定以下最优控制问题：
$$
\small{\min_{\bm{a},\rho}\mathcal{J}_{\text{DOAT}}(\bm{a},\rho)=\int_{t_{0}}^{t_{K}}\int_{\mathcal{S}}\frac{1}{2}\|\bm{a}(\bm{x},\bm{v},t)\|^{2}\rho_{t}(\bm{x},\bm{v})\,\mathrm{d}\bm{x}\mathrm{d}\bm{v}\mathrm{d}t}
\tag{2}
$$
$$
\text{s.t.}\quad\displaystyle\partial_{t}\rho_{t}+\nabla_{\bm{x}}\cdot(\bm{v}\rho_{t})+\nabla_{\bm{v}}\cdot(\bm{a}\rho_{t})=0,
\tag{3}
$$
$$
\displaystyle\rho_{t_{i}}(\bm{x},\bm{v})=\mu_{i}(\bm{x},\bm{v}),\quad i=0,\dots,K.
\tag{4}
$$
这里，在$K+1$个时间戳$t_{i}$处施加约束，给定密度$\mu_{i}$。方程3 (https://arxiv.org/html/2608.21070#S3.E3)代表增广空间中的连续性方程。类似于标准的动态最优传输[47 (https://arxiv.org/html/2608.21070#bib.bib47)]，我们称之为*动态最优加速传输（DOAT）*问题。为了便于后续通过粒子方法进行数值求解，我们假设存在一个最优流映射。

#### 速度缺失动态最优加速传输（VM-DOAT）问题
然而，在单细胞数据集等实际应用中，只有粒子位置是可观察的，速度是未知量。这与DOAT公式存在差异。因此，TracingFlow解决的是该问题的一个松弛形式：
$$
\small{\min_{\bm{a},\rho}\mathcal{J}_{\text{VM-DOAT}}(\bm{a},\rho)=\int_{t_{0}}^{t_{K}}\int_{\mathcal{S}}\frac{1}{2}\|\bm{a}(\bm{x},\bm{v},t)\|^{2}\rho_{t}(\bm{x},\bm{v})\,\mathrm{d}\bm{x}\mathrm{d}\bm{v}\mathrm{d}t}
\tag{5}
$$
$$
\text{s.t.}\quad\displaystyle\partial_{t}\rho_{t}(\bm{x},\bm{v})+\nabla_{\bm{x}}\cdot(\bm{v}\rho_{t})+\nabla_{\bm{v}}\cdot(\bm{a}\rho_{t})=0,
\tag{6}
$$
$$
\displaystyle\int\rho_{t_{i}}(\bm{x},\bm{v})\,\mathrm{d}\bm{v}=\mu_{i}^{(\text{pos})}(\bm{x}),\ i=0,1,\dots,K.
\tag{7}
$$
这里，约束仅施加在位置的边际分布$\int\rho_{t_{j}}(\bm{x},\bm{v})\,\mathrm{d}\bm{v}$上。