TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics¶

会议: NeurIPS 2026
arXiv: 2608.21070
OpenReview: https://openreview.net/forum?id=WETDhvImnp
领域: 其他
关键词: 轨迹推断, 流匹配, 动力学最优传输, 二阶动力学, 谱系示踪

一句话总结¶

针对单细胞时序观测中一阶流匹配模型因单值速度场限制无法捕捉相交轨迹与调控惯性的缺陷，本文提出了首个基于二阶动力学的无仿真轨迹推断框架 TracingFlow，通过在增广相空间中定义动态最优加速度传输（DOAT）问题并交替求解离散三对角系统补全缺失速度，实现了高曲率非线性演化轨迹的精确重构与谱系先验的自然融合。

研究背景与动机¶

从离散时间快照中推断细胞分化与发育的连续演化轨迹（Trajectory Inference, TI）是单细胞组学与生成模型的核心挑战。为了在不同时间分布间建立能量最小的演化路径，最优传输（Optimal Transport, OT）理论被广泛应用。早期的动力学最优传输方法（如 TrajectoryNet）依赖神经常微分方程（Neural ODE），在训练期间必须执行高开销的数值积分求解，导致难以拓展至大规模单细胞数据集。流匹配（Flow Matching, FM）的出现通过直接回归速度场实现了无仿真（simulation-free）的高效训练，迅速成为连续动力学建模的标准范式。

然而，现有的流匹配轨迹推断框架绝大多数局限于一阶动力系统（$\dot{\bm{x}} = \bm{v}(\bm{x}, t)$），这意味着模型学习到的速度场 $\bm{v}$ 关于位置 $\bm{x}$ 是严格单值的。在单细胞生物学过程中，这一假设存在致命缺陷：细胞状态演化往往受到表观遗传、染色质可及性以及蛋白质组等多模态记忆调控的影响，具备明显的“调控动量”（regulatory momentum）与时间延迟效应；处于相似转录组状态的细胞可能因为历史动量不同而走向完全不同的分化命运，导致其演化轨迹在转录表达空间中发生交叉。一阶流匹配由于单值约束，强行规避路径交叉，往往产生过度平滑的虚假轨迹，并且无法与谱系示踪（lineage tracing）等离散条形码生物先验兼容。近期虽有二阶动量薛定谔桥（如 3MSBM）等探索，但其仍依赖类似 Rectified Flow 的多轮重训练迭代耦合，且初始速度主要通过启发式采样获取，缺乏统一的数学最优控制解。

本文针对这一根本矛盾，提出了拓展至二阶动力系统的无仿真流匹配框架。核心 idea：将轨迹推断形式化为在位置-速度增广空间中最小化总加速度能量的动力学最优加速度传输（DOAT）问题，利用三次样条插值的解析最优解构建非交叉静态耦合，并通过求解三对角稀疏线性系统实现数据驱动的缺失速度迭代补全与加速度场无仿真回归。

方法详解¶

整体框架¶

TracingFlow 将单细胞粒子状态从纯位置空间 $\mathcal{X} \subset \mathbb{R}^d$ 拓展至增广空间 $\mathcal{S} = \mathcal{X} \times \mathcal{V} \subset \mathbb{R}^{2d}$，系统演化遵循二阶动力学方程 $\dot{\bm{x}} = \bm{v}, \dot{\bm{v}} = \bm{a}(\bm{x}, \bm{v}, t)$。整个算法流程包含三个紧密衔接的核心阶段： 1. 缺失速度迭代补全：真实单细胞数据仅包含位置快照而缺失瞬时速度，TracingFlow 将速度缺失的 VM-DOAT 问题松弛为交替优化过程，利用多时间点三次样条轨迹的光滑性求解三对角线性方程组，迭代估计每个细胞的局部条件速度分布。 2. 静态最优加速度传输（SOAT）求解：基于补全的增广状态 $(\bm{x}, \bm{v})$，利用单粒子加速度积分的解析闭式代价矩阵，调用 Sinkhorn 算法求解相邻时间点间的联合耦合计划 $\pi^*_{i \to i+1}$；若存在谱系示踪条形码，则对跨谱系跳跃施加软惩罚矩阵。 3. 条件加速度流匹配（CAFM）与初速度回归：依据最优传输计划采样多时间步离散轨迹链作为条件变量，基于解析三次多项式均值路径构造高斯条件概率路径，分别通过 MSE 损失训练加速度网络 $\bm{a}_{\bm{\theta}}(\bm{x}, \bm{v}, t)$ 与初始速度推断网络 $\bm{v}_{0, \bm{\xi}}(\bm{x})$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["单细胞时序离散快照<br/>仅含表达位置分布 X_0...K"] --> B["三对角稀疏线性系统补全速度<br/>求解多快照平滑速度 V A = B"]
    B --> C["静态最优加速度传输求解<br/>Sinkhorn 计算增广空间耦合计划"]
    C -->|融入谱系示踪条形码| D["谱系先验软约束惩罚<br/>跨 Barcode 传输代价加权 p_0"]
    D --> E["条件加速度流匹配训练<br/>回归二阶加速度网络与初速度网络"]
    E --> F["二阶常微分方程数值推断<br/>重构连续细胞分化演化轨迹"]

关键设计¶

1. 动力学最优加速度传输与增广空间非交叉性：打破空间单值束缚 针对一阶流匹配中速度场关于位置单值导致轨迹不能交叉的限制，TracingFlow 定义了动力学最优加速度传输（DOAT）泛函： $$\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$$ 约束条件为增广空间中的连续性方程 $\partial_t \rho_t + \nabla_{\bm{x}} \cdot (\bm{v} \rho_t) + \nabla_{\bm{v}} \cdot (\bm{a} \rho_t) = 0$ 以及各时间快照的边缘分布 $\rho_{t_i} = \mu_i$。根据变分极值条件，单粒子在固定初末状态 $[(\bm{x}_0, \bm{v}_0), (\bm{x}_T, \bm{v}_T)]$ 下使总加速度能量最小的轨迹是唯一的三次多项式样条 $\bm{\gamma}(t) = \bm{c}_3 t^3 + \bm{c}_2 t^2 + \bm{c}_1 t + \bm{c}_0$。将其代入目标泛函，可解析得出两点间的精确输运代价： $$\mathcal{C}_{[0,T]} = \frac{2}{T^3} \Big\{ (\|\bm{v}_0\|^2 + \langle \bm{v}_0, \bm{v}_T \rangle + \|\bm{v}_T\|^2)T^2 + 3\langle \bm{v}_0 + \bm{v}_T, \bm{x}_0 - \bm{x}_T \rangle T + 3\|\bm{x}_0 - \bm{x}_T\|^2 \Big\}$$ 本文严格证明了增广空间中的非交叉性（Remark 4.3）：任意两条三次样条轨迹如果在增广状态中发生相交，其总加速度代价必然严格大于重新交换终点后的代价。这从理论上确保了最优加速度场 $\bm{a}(\bm{x}, \bm{v}, t)$ 在增广空间内处处单值，同时允许其投影在纯位置空间 $\mathcal{X}$ 中的路径自由相交，彻底解决了细胞发育分支与回流时的轨迹重叠问题。

2. 三对角稀疏线性系统补全速度：将 VM-DOAT 转化为标准 DOAT 真实生物测序实验只能观测到表达位置 $\mu_i^{(\text{pos})}(\bm{x})$ 而无法直接测量瞬时速度 $\bm{v}$。为解决该速度缺失问题（VM-DOAT），TracingFlow 提出基于全局轨迹一致性的确定性迭代逼近算法。给定整条轨迹在各个时间快照上的固定离散位置节点 $\bm{x}_{0:K}$，使整条分段三次多项式在全区间 $[t_0, t_K]$ 上二阶导数能量最小的节点速度矩阵 $\bm{V} = [\bm{v}_0, \bm{v}_1, \dots, \bm{v}_K] \in \mathbb{R}^{d \times (K+1)}$ 恰好满足由欧拉-拉格朗日边界条件导出的线性系统： $$\bm{V} \bm{A} = \bm{B}$$ 其中 $\bm{A} \in \mathbb{R}^{(K+1) \times (K+1)}$ 是一个严格对称的三对角稀疏矩阵，$\bm{B} \in \mathbb{R}^{d \times (K+1)}$ 由相邻节点的位移差确定。利用三对角矩阵的托马斯算法（Thomas algorithm），仅需 $\mathcal{O}(K)$ 的计算复杂度即可精确解出整条轨迹在所有快照上的全局平滑速度。算法从零速度初始化，交替执行“SOAT 计划求解 $\to$ 采样轨迹链并用线性方程组更新节点期望速度”，在保证总传输能量单调递减的前提下快速收敛至相容的确定性速度场 $\bm{v}_{\text{exp}}(\bm{x}_i)$。

3. 谱系先验软约束惩罚：物理最优与生物学合理性的统一 单纯依赖转录组表达相似性进行动态推断容易导致跨谱系混淆，难以应对细胞多能性分化。TracingFlow 在 SOAT 代价矩阵构建时引入离散谱系条形码（Lineage Barcode）先验信息。对于时间点 $t_i$ 的第 $l$ 个细胞与时间点 $t_{i+1}$ 的第 $k$ 个细胞，若其条形码记录不一致（$b_{i,l} \neq b_{i+1,k}$），则将其物理运动学代价 $\mathcal{C}_{i \to i+1}$ 乘以惩罚倍率因子 $p_0$（实验中设为 $p_0 = 25$）： $$\tilde{\mathcal{C}}_{i \to i+1}(l, k) = \mathcal{C}_{i \to i+1}(l, k) \cdot \big(1 + (p_0 - 1) \cdot \mathbb{I}[b_{i,l} \neq b_{i+1,k}]\big)$$ 该软约束机制既保留了细胞由于未突变或技术遗失导致的条形码未定义状态的自然传输，又极大地抑制了跨谱系虚假分化跃迁，使流匹配学习到的二阶力场严格内生化谱系分化拓扑结构。

损失函数 / 训练策略¶

完成速度补全与 SOAT 传输计划后，TracingFlow 无需进行任何数值 ODE 积分，直接通过闭式条件加速度场训练神经网络： 1. 条件加速度流匹配损失（CAFM）：以相邻快照耦合点构成的三次轨迹多项式二阶导数 $\bm{a}(\bm{x}, \bm{v}, t | \bm{z}) = 6\bm{c}_3 (t - t_i) + 2\bm{c}_2$ 为真实回归目标，对增广状态加微小高斯扰动采样，最小化均方误差： $$\mathcal{L}_{\text{CAFM}} = \mathbb{E}_{(\bm{x},\bm{v})\sim \rho_t(\cdot|\bm{z}), \bm{z}\sim q(\bm{z}), t\sim\mathcal{U}[t_0, t_K]} \Big[ \|\bm{a}_{\bm{\theta}}(\bm{x}, \bm{v}, t) - \bm{a}(\bm{x}, \bm{v}, t | \bm{z})\|^2 \Big]$$ 2. 初始速度回归损失：由于推断阶段只给定初始细胞位置 $\bm{x}_0$，需额外训练轻量残差网络 $\bm{v}_{0, \bm{\xi}}(\bm{x})$ 拟合第 0 步的补全速度 $\bm{v}_0$： $$\mathcal{L}_{\bm{v}_0} = \mathbb{E}_{(\bm{x}, \bm{v}) \sim \rho_0} \Big[ \|\bm{v}_{0, \bm{\xi}}(\bm{x}) - \bm{v}\|^2 \Big]$$ 理论证明（Theorem A.1），生成分布与真实分布在任意时刻的 Wasserstein-2 距离上界直接被网络损失项线性控制：$\mathcal{W}_2^2(\hat{\rho}_{t_i}^{(\text{pos})}, \mu_i^{(\text{pos})}) \le 2e^{\kappa(t_i - t_0)} \big( (t_i - t_0) \mathcal{L}_{\text{AFM}} + \mathcal{L}_{\bm{v}_0} \big)$。

实验关键数据¶

主实验¶

论文在 2D 模拟轨迹数据集、Cite-seq 多时间点高维单细胞数据集（降至 5 维与 100 维）上进行了跨时间点分布传输评估，对比基线包括一阶流匹配（OT-CFM、SF2M）、二阶动量桥（3MSBM）、多边际流匹配（MMFM）、分层流（HRF）及恒加速度流（CAF）。评价指标为生成分布与真实分布间的平均 1-Wasserstein（$\mathcal{W}_1$）与 2-Wasserstein（$\mathcal{W}_2$）距离。

数据集	指标	OT-CFM (一阶)	SF2M (一阶)	3MSBM (二阶)	MMFM	HRF	TracingFlow (本文)
2D Simulation	$\mathcal{W}_1 \downarrow$	0.5270 ± 0.0612	0.9926 ± 0.1857	1.1855 ± 0.2636	2.3689 ± 0.1830	1.5390 ± 0.1130	0.3748 ± 0.0505
	$\mathcal{W}_2 \downarrow$	0.5922 ± 0.0813	1.1318 ± 0.1734	1.2767 ± 0.2215	2.6466 ± 0.1769	1.6631 ± 0.1214	0.5180 ± 0.0360
Cite 5D	$\mathcal{W}_1 \downarrow$	0.8298 ± 0.0162	0.7650 ± 0.0375	3.6369 ± 0.3508	2.5346 ± 0.1520	1.4016 ± 0.0437	0.5861 ± 0.0688
	$\mathcal{W}_2 \downarrow$	0.9155 ± 0.0160	0.8956 ± 0.0408	3.7703 ± 0.3178	2.7576 ± 0.1877	1.5313 ± 0.0430	0.6386 ± 0.0775
Cite 100D	$\mathcal{W}_1 \downarrow$	10.5498 ± 0.0291	12.3426 ± 0.0735	15.7708 ± 0.1473	12.5353 ± 0.6268	10.1532 ± 0.0643	8.0700 ± 0.1537
	$\mathcal{W}_2 \downarrow$	10.6220 ± 0.0310	12.4757 ± 0.0783	15.8833 ± 0.1832	12.7868 ± 0.7061	10.2336 ± 0.0707	8.2530 ± 0.1889

消融实验与谱系先验分析¶

在含有真实条形码标记的 3D 模拟谱系数据集与 50 维造血发育单细胞数据集（Hematopoiesis）上，评估各方法在保持谱系一致性方面的表现（采用加权谱系 Wasserstein 距离 $\mathcal{W}_1^{(L)}, \mathcal{W}_2^{(L)}$），同时验证谱系先验模块对 TracingFlow 的关键贡献：

配置 / 方法	3D Sim-Lineage ($\mathcal{W}_1^{(L)}$)	3D Sim-Lineage ($\mathcal{W}_2^{(L)}$)	Hematopoiesis ($\mathcal{W}_1^{(L)}$)	Hematopoiesis ($\mathcal{W}_2^{(L)}$)	说明
OT-CFM (1st Order)	2.2383 ± 0.0205	2.3162 ± 0.0216	14.9945 ± 0.3442	15.3971 ± 0.3180	一阶基线，无法解耦交叉分化
SF2M	1.5255 ± 0.0513	1.5435 ± 0.0515	18.7170 ± 0.1124	18.9756 ± 0.1168	一阶薛定谔桥流匹配
3MSBM (2nd Order)	2.3126 ± 0.4756	2.4905 ± 0.4582	18.3313 ± 0.9834	18.6427 ± 0.9413	启发式初速度与多次迭代耦合
TracingFlow (w/o Bio-Prior)	1.5292 ± 0.1714	1.5532 ± 0.2303	16.6166 ± 0.1229	17.0110 ± 0.1456	去掉谱系惩罚因子，仅靠二阶动力学
TracingFlow (Full Model)	0.4538 ± 0.2039	0.5328 ± 0.2066	14.0470 ± 0.1530	14.4589 ± 0.1815	完整模型：二阶流匹配 + 谱系惩罚 ($p_0=25$)

在胚胎小体（EB 5D）时序留一插值（Hold-One-Out）实验中，TracingFlow 取得平均 $\mathcal{W}_1 = 4.4974 \pm 0.0950$ 与 $\mathcal{W}_2 = 4.7859 \pm 0.1857$，显著优于 OT-CFM（$5.1205 / 5.6910$）和 3MSBM（$7.1901 / 7.4697$），证明了二阶动力学在预测未观测时间节点分布时的外推与内插优越性。

关键发现¶

二阶表征彻底解决相交轨迹瓶颈：可视化证明在 2D 模拟交叉环流中，OT-CFM 强行产生互相排斥的平行流场导致几何失真，而 TracingFlow 成功重构了具有不同速度方向的重叠相交路径，$\mathcal{W}_1$ 相对一阶模型降低近 30%。
生物先验惩罚带来质的飞跃：在 3D 模拟谱系任务中，引入谱系惩罚使 $\mathcal{W}_1^{(L)}$ 从 1.5292 骤降至 0.4538（降低 70.3%），说明基于条形码的代价修正能强力纠正形态相近但谱系不同的细胞间假配对。
计算扩展性与批次敏感度：在 Cite 5D 数据集上的 Minibatch-OT 消融显示，当 Batch Size 从全数据集降至 2000 和 1000 时，平均 $\mathcal{W}_1$ 仅由 0.5348 微升至 0.5936 和 0.6175，依然优于所有一阶全量基线，证实了分批最优传输在大规模细胞集上的落地可行性。

亮点与洞察¶

增广相空间构建与三对角速度解析求解：将不可观测速度的连续变分问题巧妙化为分段三次样条节点速度的紧支撑三对角线性方程组，避免了传统高阶动力学中对潜在变量的黑盒生成建模。
无仿真二阶流匹配理论完备性：证明了增广空间非交叉性、Marginal 与 Conditional 加速度场梯度的严格等价性，并给出了推断误差关于网络回归损失的严格理论上界，为流匹配走向高阶系统提供了扎实的数学支撑。
跨领域迁移价值：该“状态增广 + 物理导数约束 + 离散先验软惩罚”范式可直接迁移至机器人运动规划、自动驾驶轨迹预测以及物理守恒流体仿真等存在二阶牛顿惯性与先验拓扑约束的动态系统。

局限与展望¶

初始速度近似与随机性缺失：作者指出当前实现将条件速度分布 $\omega_i(\bm{v}|\bm{x})$ 简化为单值 Dirac 期望分布，忽略了同类细胞速度的异质随机波动，未来可引入变分高斯混合或条件扩散模型建模多峰速度分布。
SOAT 预计算的大规模算力消耗：虽然 Minibatch-OT 能缓解压力，但对于百万级单细胞图谱，反复求解高维 Sinkhorn 耦合矩阵仍是时间瓶颈。
物理作用量意义的进一步探讨：目标函数 $\int \frac{1}{2}\|\bm{a}\|^2 \mathrm{d}t$ 对应于最小化加速度平方积分（Jerk 动能松弛），在传统经典力学中并非标准作用量（需引入 4d 维哈密顿体系解释），与真实细胞能量代谢成本的生物物理映射有待深入阐明。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将二阶动力学与无仿真流匹配无缝统一并给出严密数学证明，提出三对角缺失速度补全。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 2D/3D 仿真、高维 Cite-seq、EB 留一法插值、造血干细胞真实谱系示踪及批大小敏感性分析。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密清晰，从物理变分法到生物学先验的过渡自然流畅，附录证明详尽。
价值: ⭐⭐⭐⭐⭐ 为单细胞时空动力学推断与复杂物理系统轨迹生成确立了全新的二阶流匹配基准。

数据集	指标	OT-CFM (一阶)	SF2M (一阶)	3MSBM (二阶)	MMFM	HRF	TracingFlow (本文)
2D Simulation	\(\mathcal{W}_1 \downarrow\)	0.5270 ± 0.0612	0.9926 ± 0.1857	1.1855 ± 0.2636	2.3689 ± 0.1830	1.5390 ± 0.1130	0.3748 ± 0.0505
	\(\mathcal{W}_2 \downarrow\)	0.5922 ± 0.0813	1.1318 ± 0.1734	1.2767 ± 0.2215	2.6466 ± 0.1769	1.6631 ± 0.1214	0.5180 ± 0.0360
Cite 5D	\(\mathcal{W}_1 \downarrow\)	0.8298 ± 0.0162	0.7650 ± 0.0375	3.6369 ± 0.3508	2.5346 ± 0.1520	1.4016 ± 0.0437	0.5861 ± 0.0688
	\(\mathcal{W}_2 \downarrow\)	0.9155 ± 0.0160	0.8956 ± 0.0408	3.7703 ± 0.3178	2.7576 ± 0.1877	1.5313 ± 0.0430	0.6386 ± 0.0775
Cite 100D	\(\mathcal{W}_1 \downarrow\)	10.5498 ± 0.0291	12.3426 ± 0.0735	15.7708 ± 0.1473	12.5353 ± 0.6268	10.1532 ± 0.0643	8.0700 ± 0.1537
	\(\mathcal{W}_2 \downarrow\)	10.6220 ± 0.0310	12.4757 ± 0.0783	15.8833 ± 0.1832	12.7868 ± 0.7061	10.2336 ± 0.0707	8.2530 ± 0.1889