PACE: Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference¶

会议: NeurIPS 2026
arXiv: 2605.18587
OpenReview: https://openreview.net/forum?id=2PQtt7QLHJ
代码: https://github.com/AI4Science-WestlakeU/PACE
领域: 计算生物
关键词: 单细胞轨迹推断, 最优输运, 神经薛定谔桥, 各向异性黎曼度量, 动力学重构

一句话总结¶

针对破坏性单细胞时间序列缺乏细胞对应关系与连续路径的本质病态问题，PACE 构建状态与时间依赖的各向异性流形度量惩罚法向运动，在度量路径作用量下交替迭代神经桥与跨时刻耦合，并将动力学蒸馏为全局连续速度场以高精度恢复发育轨迹。

研究背景与动机¶

理解细胞如何在表型状态之间连续迁移是单细胞生物学与发育生物学中的核心课题。诸如胚胎发育、干细胞分化、重编程以及免疫活化等生物学过程，本质上都是高维分子状态空间中连续演化的细胞群体动态迁移，而非孤立离散的细胞类型转换。然而，现有的单细胞测序实验通常具有破坏性，实验人员无法对同一个细胞进行跨时刻的连续追踪，只能在不同离散观测时间点捕获细胞群体的分布快照。这种破坏性采样使得单细胞轨迹推断面临严重的数学不适定性：仅凭边缘经验分布无法唯一确定跨时刻的细胞对应关系或中间连续运动路径。

为了从离散快照中推断连续动态，最优输运（OT）与连续流匹配（Flow Matching）等方法被广泛应用。但绝大多数现有方法都建立在观测空间或低维潜在空间的欧氏几何假设之上，通过欧氏距离计算端点代价，并使用欧氏动能或平滑度正则化连续速度场。然而，生物发育通常是异步的，相同实验时间采样的细胞往往处于完全不同的潜在伪时间阶段；同时，欧氏空间中几何接近的两个细胞可能属于完全互斥的命运分支或受阻状态。简单的欧氏近邻极易将不兼容的发育程序或甚至将凋亡细胞与快速增殖的干细胞错误配对。虽然引入 RNA 速率或谱系追踪等辅助信息能够缓解该问题，但在很多通用单细胞数据集中这类额外实验信号难以获取。

本文的切入角度是：虽然真实的跨时刻对应与中间轨迹未被观测，但局部时空邻域内蕴的流形切空间仍为细胞的合理迁移方向提供了强有力的几何先验——沿着局部流形切向的运动在生物学上远比横穿流形法向的跳变更为合理。核心 idea：构建状态与时间依赖的各向异性时空黎曼度量，将法向速度分量施以高代价惩罚，在诱导的几何路径作用量下交替优化端点保持神经桥与跨时刻最优输运耦合，并将学习到的桥动力学蒸馏为无须端点对条件的全局连续时间速度场。

方法详解¶

整体框架¶

PACE 面向离散破坏性单细胞快照点云序列，旨在重构跨时刻细胞输运耦合并推断驱动细胞群体演化的连续时间速度场。方法的整体流水线包含四个核心步骤：首先，从离散时空邻域中估计局部法向投影子，插值构建状态与时间依赖的各向异性黎曼度量；其次，参数化端点保持的神经桥模型；随后，在路径作用量代价下交替迭代更新神经桥内部形变与跨时刻 OT 耦合；最后，将学习到的条件桥动力学回归蒸馏为全局速度场，并通过数值积分实现任意细胞状态的前向轨迹推断。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离散时间单细胞快照点云<br/>{X^(k)} at t_k"] --> B["时空切空间与各向异性度量构建<br/>局部PCA估计法向投影子 P_N 与核插值 G_k(x,t)"]
    B --> C["端点保持神经桥参数化<br/>γ_θ(x,y,τ) = (1-τ)x + τy + τ(1-τ)ψ_θ"]
    C --> D["交替优化：黎曼作用量桥训练<br/>固定耦合 π_k 最小化度量作用量 L_bridge(θ)"]
    D --> E["交替优化：路径作用量耦合更新<br/>计算 c_ij^path(θ) 并求解最优输运更新 π_k"]
    E -->|周期性交替迭代 R_0 轮| D
    D --> F["全局连续时间速度场蒸馏<br/>最小化 L_distill(ϕ) 将条件桥映射为 v_ϕ(x,t)"]
    F --> G["ODE数值积分推断连续细胞轨迹<br/>dx_t/dt = v_ϕ(x_t,t), x_0 ~ ρ^_0"]

关键设计¶

1. 时空切空间与各向异性黎曼度量构建：将生物流形先验转化为方向敏感的动力学代价

现有方法多使用各向同性的欧氏度量或单纯基于数据密度的流形支持度量，无法区分沿发育流形的切向演化与切穿流形的非法法向移动。PACE 在每个离散时间点的锚点细胞 $x_r$ 处，通过空间 $k$ 近邻的高斯核加权协方差矩阵提取最小方差主成分，构建局部法向投影算子 $P_N^{(r)} = n_r n_r^\top$（高维下取正交补空间）。随后，为了在任意时空查询点 $(x, t)$ 上获得光滑的度量张量，PACE 采用自适应带宽的时空高斯核对所有锚点的法向投影子进行加权平均，构造平均法向投影场 $C_N^{(k)}(x, t) = \sum_r \omega_r^{(k)}(x, t) P_N^{(r)}$，并据此定义各向异性黎曼度量张量： $$G_k(x, t) = I + \alpha C_N^{(k)}(x, t), \quad \alpha > 0$$ 在理想正交投影下，该度量赋予速度向量 $v = v_T + v_N$ 的切向分量标准欧氏代价 $\|v_T\|^2$，而将横穿流形的法向分量惩罚加权为 $(1 + \alpha)\|v_N\|^2$。由此，两点间的输运代价不再取决于简单的端点欧氏几何距离，而是由沿路径速度与发育切向对齐程度的黎曼路径作用量所决定。

2. 端点保持神经桥与路径作用量自洽交替优化：打破固定欧氏配对与孤立路径拟合的割裂

在连续轨迹推断中，跨时刻细胞配对与点对之间的中间演化路径存在双向依赖：配对不准导致插值路径失真，而缺乏连续路径几何又无法正确评估跨时段细胞的演化难度。PACE 提出显式满足端点约束的神经桥参数化形式 $\gamma_\theta(x, y, \tau) = (1 - \tau)x + \tau y + \tau(1 - \tau)\psi_\theta(x, y, \tau)$，利用 $\tau(1 - \tau)$ 衰减因子严格锁定两端边界条件 $\gamma_\theta(0)=x, \gamma_\theta(1)=y$，使得神经网络 $\psi_\theta$ 仅需自适应学习路径内部的弯曲形变。PACE 将求解过程构造成块坐标下降迭代：在固定当前耦合 $\pi_k$ 时，通过梯度下降最小化由度量诱导的路径作用量损失 $\mathcal{L}_{\mathrm{metric}}(\theta)$（配合跨段连续性正则项与法向抑制项）；在固定神经桥模型后，在离散时间网格上积分计算候选细胞对的精确路径作用量： $$c_{ij}^{\mathrm{path}}(\theta) = \frac{1}{M}\sum_{m=1}^M u_\theta(x_i^{(k)}, x_j^{(k+1)}, \tau_m)^\top G_k\left(\gamma_\theta(x_i^{(k)}, x_j^{(k+1)}, \tau_m), t_{k, \tau_m}\right) u_\theta(x_i^{(k)}, x_j^{(k+1)}, \tau_m)$$ 并重新求解离散最优输运问题更新耦合矩阵 $\pi_k$。这种闭环引导使得模型自发倾向于选择稍微长但沿着发育流形切向平滑迁移的路径，彻底剔除在欧氏距离上虽短但切穿分化屏障的虚假配对。

3. 连续时间全局速度场蒸馏：消除对未观测端点对的依赖以支持任意状态前向演化

交替优化训练得到的神经桥模型 $\gamma_\theta(x, y, \tau)$ 本质上是针对特定已知端点对 $(x, y)$ 的条件插值器；但在实际轨迹推断与下游生物学模拟中，推断任意空间位置的流向不能假设已知其未来目标点 $y$。为了得到自洽且独立的动力学系统，PACE 将神经桥产生的局部速度动力学蒸馏到统一的全局时变速度场网络 $v_\phi(x, t)$ 中。在每个时间段 $k$ 内，采样当前耦合下的端点对 $(x, y) \sim \pi_k$ 及插值时间 $\tau \sim \mathrm{Uniform}[0, 1]$，将中间状态 $x_{t_{k,\tau}} = \gamma_\theta(x, y, \tau)$ 作为输入，并通过均方误差将 $v_\phi$ 对齐到物理时间速度： $$\mathcal{L}_{\mathrm{distill}}(\phi) = \mathbb{E}_{k, (x, y) \sim \pi_k, \tau} \left[ \left\| v_\phi(x_{t_{k,\tau}}, t_{k,\tau}) - \frac{u_\theta(x, y, \tau)}{\Delta t_k} \right\|^2 \right]$$ 其中 $\Delta t_k = t_{k+1} - t_k$ 完成归一化内部时间到物理时间的缩放转换。蒸馏完成后，全局速度场脱离配对限制，可直接借助任意常微分方程（ODE）求解器从初始单细胞分布 $\hat{\rho}_0$ 出发数值积分出全时间跨度的连续演化轨迹。

损失函数 / 训练策略¶

PACE 的两阶段训练策略具体如下： 1. 第一阶段：神经桥交替训练。目标函数为 $\mathcal{L}_{\mathrm{bridge}}(\theta) = \lambda_{\mathrm{metric}}\mathcal{L}_{\mathrm{metric}} + \lambda_{\mathrm{coh}}\mathcal{L}_{\mathrm{coh}} + \lambda_{\mathrm{orth}}\mathcal{L}_{\mathrm{orth}}$。其中度量作用量项 $\mathcal{L}_{\mathrm{metric}}$ 计算时间网格采样的二次型速度能耗；跨段一致性损失 $\mathcal{L}_{\mathrm{coh}}$ 在相邻段边界处对空间接近的生成速度进行余弦对齐；法向抑制损失 $\mathcal{L}_{\mathrm{orth}}$ 显式惩罚生成路径在局部法向投影上的能量泄露。每隔 $R_0$ 个 epoch，利用当前桥网络计算全矩阵作用量 $c_{ij}^{\mathrm{path}}(\theta)$ 并通过 Sinkhorn 或标准 OT 算法求解更新 $\pi_k$。 2. 第二阶段：全局流蒸馏与推断。冻结优化好的神经桥 $\theta$ 与耦合 $\pi_k$，以 $\mathcal{L}_{\mathrm{distill}}(\phi)$ 训练全局速度场网络 $v_\phi$。推断时采用标准数值积分器（如 100 步 Euler 更新），从初始时刻状态群体模拟前向推断未观测留出时刻的细胞分布。

实验关键数据¶

主实验¶

论文在受控二维数据集（Ocean 涡旋、Two-Branch 分叉）以及多个经典生物单细胞分化与重编程基准（EB 拟时序胚胎体分化、iPSC-Liu 人诱导多能干细胞重编程、Schiebinger2019 小鼠胚胎成纤维细胞重编程、高维 OP-Cite/OP-Multi 多组学时间序列）上进行了严格的留出时间点（held-out timepoints）分布重构评测。在所有评测中，仅使用训练时间点数据，留出时间点仅作为真实标签计算分布距离（MMD、$\mathcal{W}_1$、$\mathcal{W}_2$）。

下表汇报了低维单细胞基准数据集上的时间平均分布重构性能对比（EB PHATE 留出 $t=3$；iPSC-Liu 留出 $t\in\{4, 16\}$；Schiebinger2019 留出 $t\in\{6, 11, 16\}$）：

数据集	指标	PACE (本文)	最佳基线 (Prev. SOTA)	提升幅度
EB PHATE	MMD ↓	0.0674	0.1056 (Aligned CFM)	36.2%
	$\mathcal{W}_1$ ↓	0.2087	0.2657 (Aligned CFM)	21.5%
	$\mathcal{W}_2$ ↓	0.2803	0.3243 (Aligned CFM)	13.6%
iPSC-Liu	MMD ↓	0.4385	0.5228 (OT-CFM)	16.1%
	$\mathcal{W}_1$ ↓	0.4853	0.6594 (Aligned CFM)	26.4%
	$\mathcal{W}_2$ ↓	0.6238	0.8320 (Aligned CFM)	25.0%
Schiebinger2019	MMD ↓	0.2084	0.3030 (DMSB)	31.2%
	$\mathcal{W}_1$ ↓	0.2067	0.2774 (DMSB)	25.5%
	$\mathcal{W}_2$ ↓	0.2840	0.3662 (DMSB)	22.4%

在中高维单细胞基准（iPSC-Liu 10D 与 50D PCA 空间）中，PACE 同样在所有指标上保持显著优势：

维度	方法	MMD ↓	$\mathcal{W}_1$ ↓	$\mathcal{W}_2$ ↓
10D	Action Matching	0.5555	3.7801	4.3743
	Aligned CFM	0.4644	2.9405	3.3825
	CURLY	0.5369	3.6128	4.2102
	DMSB	0.7713	6.3088	6.4593
	MFM	0.5273	3.4066	3.8382
	OT-CFM	0.5145	3.3850	3.9447
	PACE (本文)	0.4095	2.9226	3.1879
50D	Action Matching	0.3125	7.6126	8.4868
	Aligned CFM	0.2624	7.3133	8.0552
	CURLY	0.3144	8.1144	8.9096
	DMSB	0.5638	10.9368	11.2423
	MFM	0.2853	7.7924	8.5454
	OT-CFM	0.2898	7.8392	8.5954
	PACE (本文)	0.2550	6.7537	7.2836

消融实验¶

在 Schiebinger2019 小鼠重编程留出数据集上，作者对各个关键模块进行了系统消融（汇总于正文第 4.5 节与图 5 分析）：

配置	核心改动	表现变化与说明
Full PACE	完整模型（各向异性度量 + 局部邻域 + 时空平滑核 + 耦合交替再匹配）	取得最稳定且最低的 MMD、$\mathcal{W}_1$ 与 $\mathcal{W}_2$ 误差
w/o rematching	冻结初始配对，不进行路径作用量更新耦合	误差显著上升，表明桥路径与耦合联合优化的关键引导作用
Euclidean action ($\alpha = 0$)	退化为标准欧氏动能作用量，无切向/法向区分	路径丢失对发育流形几何的约束，流形外漂移增加
w/o local neighborhood	投影子构建采用全量样本而非局部邻域	局部几何切空间估计失真，法向约束出现错误引导
w/o spatio-temporal kernel	移除空间或时间高斯核自适应平滑插值	各向异性度量在时空查询点缺乏平滑性，训练稳定性受损

关键发现¶

各向异性流形先验消除欧氏近邻伪配对：消融表明，若退化为欧氏作用量（$\alpha=0$）或移除重新配对（w/o rematching），模型重构误差明显激增。这证实了在发育异步与分叉状态下，利用切空间敏感度量引导耦合是保证动态合理性的核心机制。
高维空间的几何各向异性抵抗测度集中：在高维单细胞 PCA 空间（10D 至 100D）中，虽然欧氏范数集中现象（如范数变异系数 $CV_{\mathrm{norm}} < 0.3$）使得全局欧氏距离对比度急剧降低，但每个单细胞快照内部的局部各向异性结构仍然高度稳定。PACE 提取的局部切空间能有效提供定向几何约束，打破维度灾难。
外部动力学验证无需速度监督：在 EB PHATE 数据集上，即便未在训练中引入 RNA 速率（RNA velocity）监督，PACE 蒸馏所得速度场与真实 RNA 速率参考方向的余弦距离达到 0.3287（Normalized $L_2$ 为 0.6630），相比其他流匹配基线（例如 MFM 的 0.6930 / 1.0261）降低了 15.4% 以上的误差，证明其推断的流向切实符合生物物理转录方向。

亮点与洞察¶

各向异性度量与自洽神经桥的闭环设计：提出将单细胞流形先验通过局部 PCA 法向投影算子转化为可微黎曼度量，并使用端点保持参数化将连续路径作用量回传给 OT 耦合重加权。这一设计打通了从几何表征到输运对应再到连续动力学的端到端数学通路。
“路径作用量驱动耦合”超越“端点欧氏距离驱动耦合”：传统的动态 OT 往往使用直连欧氏距离初始化或固定配对；PACE 证明哪怕欧氏距离很小，若路径必须跨越流形法向障碍，其作用量代价依然极高，从而天然排除了分化死胡同与错误跨支交联。
高维生物表征的局部切空间有效性：深刻洞察了高维嵌入中全局距离对比度失效但局部协方差主成分依然具有强方向指导性的数学特性，为高维多组学数据动力学建模提供了普适思路。

局限与展望¶

度量构建高度依赖局部几何采样密度：局部主成分分析与法向投影子的质量依赖于单细胞快照的局部覆盖度；若在细胞极为稀疏的瞬间过渡态区域，协方差矩阵可能无法稳定估计真实的切空间方向。
交替优化计算开销较大：在每个交替迭代周期中，评估所有潜在细胞对在离散时间网格上的路径作用量代价矩阵复杂度较高，扩展至超大规模单细胞数据集（数百万级细胞）时需要更高效的批次采样或局部邻域近似。
未来方向：作者指出未来工作可引入更鲁棒的局部度量估计器、结合贝叶斯不确定性量化机制，并进一步拓展至部分监督、谱系条形码整合及复杂空间转录组等多模态时空动态分析中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙结合局部切空间黎曼度量与端点保持神经桥，首次实现无配对离散快照下的几何自洽跨时段动力学推断。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖控制 2D、分化胚胎体、高维人/小鼠重编程及多组学九个留出评测，同时包含外部 RNA 速率客观对齐与高维几何诊断。
写作质量: ⭐⭐⭐⭐⭐ 概念表述清晰透彻，数学推导完整严密，问题背景与解决方案层层递进。
价值: ⭐⭐⭐⭐⭐ 为计算生物学中高维离散单细胞轨迹推断提供了极具影响力的全新几何范式，代码开源且易于复现拓展。

配置	核心改动	表现变化与说明
Full PACE	完整模型（各向异性度量 + 局部邻域 + 时空平滑核 + 耦合交替再匹配）	取得最稳定且最低的 MMD、\(\mathcal{W}_1\) 与 \(\mathcal{W}_2\) 误差
w/o rematching	冻结初始配对，不进行路径作用量更新耦合	误差显著上升，表明桥路径与耦合联合优化的关键引导作用
Euclidean action (\(\alpha = 0\))	退化为标准欧氏动能作用量，无切向/法向区分	路径丢失对发育流形几何的约束，流形外漂移增加
w/o local neighborhood	投影子构建采用全量样本而非局部邻域	局部几何切空间估计失真，法向约束出现错误引导
w/o spatio-temporal kernel	移除空间或时间高斯核自适应平滑插值	各向异性度量在时空查询点缺乏平滑性，训练稳定性受损