Multiscale Supervised Unbalanced Optimal Transport Flow Matching¶

会议: NeurIPS 2026
arXiv: 2605.16529
OpenReview: https://openreview.net/forum?id=15nSsWpSmo
领域: 图像生成
关键词: 非平衡最优传输, 流匹配, 多尺度监督, 单细胞轨迹推断, Wasserstein-Fisher-Rao

一句话总结¶

提出了多尺度监督非平衡最优传输流匹配框架（MUST-FM），利用分层数据结构在无模拟条件下求解连续非平衡输运与质量变化，通过块内独立提升策略消除了细粒度点对点耦合的存储与求解开销，实现了百万级单细胞图谱的快速轨迹推断。

研究背景与动机¶

从非成对的时间截面单细胞转录组数据中重构连续细胞发育轨迹是计算生物学与动态生成建模的前沿课题。基于最优传输（Optimal Transport, OT）与流匹配（Flow Matching, FM）的方法凭借无模拟（simulation-free）训练的优势，已成为学习连续动力学的重要范式。然而，真实生物发育系统本质上是非质量守恒的，细胞的分化伴随着广泛的增殖（proliferation）与凋亡（apoptosis），导致分布间总质量随时间剧烈波动。基于 Wasserstein-Fisher-Rao（WFR）几何的非平衡最优传输（UOT）虽然在理论上能协同建模空间位移与质量源汇变化，但求解大尺度数据集上的连续动态 UOT 面临极高的计算门槛。

现有的大规模 UOT 近似计算策略（如 Sinkhorn 正则化、低秩分解和随机小批量 mini-batch OT）大多将细胞群体抽象为扁平且无结构的离散点云。这种盲目的点云假设直接忽略了现代单细胞图谱中普遍固有的多层级层次化注释（例如粗粒度胚层谱系、中粒度细胞类型、细粒度亚型）以及明确的生物学先验知识。如果直接沿用传统平衡监督 OT 施加刚性转移先验，严格的边缘分布匹配往往会导致约束优化问题在数学上无解；而随机小批量切片则会破坏大尺度的生物拓扑结构，导致跨时间步推断发生谱系错配与假阳性转移。

面对计算效率与生物先验约束难以兼顾的矛盾，本文提出从数据的层级结构与弱约束非平衡特性切入。核心 idea：将非平衡动态最优传输投影到多尺度分层空间，自顶向下利用生物先验和粗粒度耦合构建稀疏掩码，并在倒数第二层进行块内独立提升采样，以零细粒度求解开销引导无模拟流匹配联合学习空间速度场与质量演化。

方法详解¶

整体框架¶

MUST-FM 整体由两个协同阶段构成：第一阶段为全批次多尺度离散先验预计算，第二阶段为基于闭式测地线靶标的无模拟连续流匹配训练。首先将未配对的细胞快照构建为多层级的加权质心测地表示；自粗至细（coarse-to-fine）逐层利用生物学转移先验矩阵与上一层的粗粒度耦合概率构建稀疏转移掩码，并在倒数第二层执行遮蔽最优熵传输（Masked OET）求解；随后通过提出的块内独立提升（Within-Block Independent Coupling Lifting）策略生成隐式细粒度半耦合分布；最后利用行进高斯（Traveling Gaussian）的解析 WFR 测地线闭式解构造连续速度与质量增长率靶标，驱动条件非平衡流匹配目标函数的端到端训练。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["单细胞时序非配对快照<br/>含多层级分层注释与谱系先验"] --> B["阶段1：多尺度层级表示与先验掩码构建<br/>提取各层加权质心并逐层求交掩码"]
    B --> C["阶段2：分层遮蔽最优熵传输求解<br/>自粗至细稀疏 OET 求解至倒数第二层"]
    C --> D["阶段3：块内独立提升隐式半耦合采样<br/>块内经验分布快速分发无需细粒度显式矩阵"]
    D --> E["阶段4：基于行进高斯的条件非平衡流匹配<br/>沿解析 WFR 测地线回归时变速度场与增殖率"]
    E --> F["连续动力学生成推断<br/>时间条件速度场与质量预测"]

关键设计¶

1. 多尺度层级表示与先验掩码构建：逐层收敛可能转移空间的拓扑约束

针对百万级单细胞数据直接计算转移成本矩阵导致的内存与计算爆炸，MUST-FM 依据细胞固有的 $L$ 层分层注释，在每一层 $l$ 将经验测度近似为加权质心测度： $$\hat{\mu}_{0}^{(l)}=\sum_{i=1}^{K_{0}^{(l)}}w_{0,i}^{(l)}\delta_{\bm{c}_{0,i}^{(l)}},\quad\hat{\mu}_{1}^{(l)}=\sum_{j=1}^{K_{1}^{(l)}}w_{1,j}^{(l)}\delta_{\bm{c}_{1,j}^{(l)}}$$ 其中权重 $w$ 严格保留未归一化的原始细胞绝对计数值，以真实反映 WFR 动力学中的总质量变化。为了注入先验，该方法在各层引入双重约束掩码的逐元素逻辑与操作 $M^{(l)}=B^{(l)}\odot H^{(l)}$。其中 $B^{(l)}$ 为专家或图谱已知的生物谱系转移二值先验矩阵，剔除不可能的跨谱系跳跃；$H^{(l)}$ 则是粗到细拓扑传播掩码，定义为 $H_{ij}^{(l)}=\mathbb{I}(\gamma_{\text{pa}(i),\text{pa}(j)}^{(l-1)}/w_{0,\text{pa}(i)}^{(l-1)}\geq\epsilon)$。该机制将粗层的输运质量概率筛选传递至子簇，随着分辨率下移使得有效转移空间呈指数级稀疏化。

2. 分层遮蔽最优熵传输求解：在非平衡松弛下保障解的可行性与极速收敛

传统的监督平衡 OT 在面对刚性转移约束和精确边缘匹配时极易产生空可行域。MUST-FM 将半耦合问题转化为带有非负测度支撑集限制 $S_M$ 的遮蔽最优熵传输（Masked OET）问题： $$\min_{\gamma^{(l)}\ge 0,\;\gamma_{ij}^{(l)}=0\,\text{if}\,M_{ij}^{(l)}=0}\sum_{i,j}C_{ij}^{(l)}\gamma_{ij}^{(l)}+\text{KL}\big(\gamma^{(l)}\mathbf{1}\,\|\,w_{0}^{(l)}\big)+\text{KL}\big((\gamma^{(l)})^{\top}\mathbf{1}\,\|\,w_{1}^{(l)}\big)$$ 其中成本矩阵元素为 WFR 几何核 $C_{ij}^{(l)}=-2\ln\overline{\cos}(\|\bm{c}_{0,i}^{(l)}-\bm{c}_{1,j}^{(l)}\|_2/(2\delta))$。理论证明，经该遮蔽 OET 极小化得到的耦合解 $\gamma^{(l)}$ 通过行/列和归一化投影后，能严格构成静态监督 WFR 问题的最优半耦合对 $(\gamma_0, \gamma_1)$。通过改进底层 POT 求解器跳过被掩蔽掉的零元素运算，将高层宏观输运与中观亚型对齐的收敛速度提升了数倍至数十倍。

3. 块内独立提升隐式半耦合采样：消除细粒度显式矩阵求解的常数级扩展

当单细胞图谱扩展到数百万规模时，即便在精细层求解稀疏 OET 仍面临显著的显存瓶颈。MUST-FM 敏锐地指出，生物转移先验通常定义在可解释的亚型或谱系层级（即 $L-1$ 层），因此无需在单个细胞层面求解最优匹配。算法在第 $L-1$ 层停止 OET 优化，提出块内独立假设（Within-Block Independence），将每个粗传输块直接按子节点的经验计数值概率分布展开： $$\widetilde{\gamma}_{ij}^{(L)}=\gamma_{IJ}^{(L-1)}\frac{w_{0,i}^{(L)}}{w_{0,I}^{(L-1)}}\frac{w_{1,j}^{(L)}}{w_{1,J}^{(L-1)}},\quad I=\text{pa}_0(i),\,J=\text{pa}_1(j)$$ 更重要的是，在随后的流匹配训练中完全无需在显存中实例化该高维矩阵。理论推导表明， finest-level 采样可以直接等价为两步隐式采样：先以概率 $(I,J)\sim\gamma_0^{(L-1)}$ 采样父级转移块，再各自在块内独立按归一化子细胞权重采样具体细胞对 $(i,j)$。采样对之间的端点质量比值满足 $\rho_{IJ}=\gamma_{1,IJ}^{(L-1)}/\gamma_{0,IJ}^{(L-1)}$，子节点权重在除法中严格精确抵消。这使得整个耦合采样过程在训练期间退化为常数时间查找，相对开销随细胞数量增加渐近归零。

4. 基于行进高斯的条件非平衡流匹配：连续空间位移与质量增殖的无模拟协同学习

为了避免神经常微分方程（Neural ODE）反复数值积分的高昂训练代价，模型采用行进高斯（Traveling Gaussian）条件测度路径 $\rho_t(\bm{x}|\bm{x}_0,\bm{x}_1)=m_t\mathcal{N}(\bm{x}|\bm{\eta}_t,\sigma_t^2\mathbf{I})$。对于隐式采样的端点对，设置初始相对质量 $m_0=1$ 和目标相对质量 $m_1=\rho_{IJ}$，均值轨迹 $\bm{\eta}_t$ 与质量演化 $m_t$ 精确遵循 Dirac-to-Dirac WFR 解析测地线： $$m_t=At^2-2Bt+m_0,\quad \bm{\eta}_t=\bm{x}_0+\frac{\bm{\omega}_0}{\sqrt{m_0 A-B^2}}\left(\arctan\frac{At-B}{\sqrt{m_0 A-B^2}}-\arctan\frac{-B}{\sqrt{m_0 A-B^2}}\right)$$ 神经网络 $\bm{v}_{\bm{\theta}}(\bm{x},t)$ 与 $g_{\bm{\phi}}(\bm{x},t)$ 分别对齐诱导出的瞬时传输速度 $\bm{u}_t$ 与增殖率 $g_t=\frac{\mathrm{d}}{\mathrm{d}t}\ln m_t$，并以瞬时演化质量 $m_t$ 进行自适应加权回归： $$\mathcal{L}_{\text{CUFM}}=\mathbb{E}\left[\left(\|\bm{v}_{\bm{\theta}}(\bm{x},t)-\bm{u}_t\|_2^2+\kappa(g_{\bm{\phi}}(\bm{x},t)-g_t)^2\right)m_t(\bm{x}_0,\bm{x}_1)\right]$$ 该设计使深度模型能够在多时间截面之间自洽还原全局连续细胞流动与细胞群落净增殖动态。

损失函数 / 训练策略¶

对于包含 $K+1$ 个离散时间点 $t_0 < \dots < t_K$ 的多时间步数据集，各时段独立执行预先的离散层级 OET 求解。在连续流训练时，采用时间条件多层感知机（MLP，5 层，256 隐藏神经元，LeakyReLU 激活）统筹跨区段的速度场 $\bm{v}_{\bm{\theta}}$ 与增殖率标量场 $g_{\bm{\phi}}$。训练批量大小设为 $b$，归一化时间 $t\sim\mathcal{U}[0,1]$ 映射到实际区间 $t^{(k)}=t_k+(t_{k+1}-t_k)t$，空间摄动高斯带宽 $\sigma$ 设为预定值，超参数 $\kappa$ 用于平衡速度均方误差与增殖率回归损失。优化器选用 Adam，全局端到端平滑拟合所有时序快照。

实验关键数据¶

主实验¶

在评估耦合恢复精度的合成多尺度基准（含 2.7 万细胞、3 个宏观簇、27 个微观簇，已知精确平移真实解）上，MUST-FM 展现出极高精度与极优运行效率；在跨多模态与生物单细胞基准（Gene 2D, Gaussian 1000D, 真实小鼠造血小鼠数据集 Mouse 2D, 人胚体细胞分化 EB 50D）的动态输运评估中，以 1-Wasserstein 距离（$W_1$）和相对质量误差（Relative Mass Error, RME）全面优于现有先进基线。

数据集	指标	MUST-FM (Sparse)	MUST-FM (Independent)	WFR-FM (Prev. SOTA)	VGFM	DeepRUOT	TIGON
Gene 2D	$W_1$ ↓ / RME ↓	0.019 / 0.001	0.019 / 0.001	0.019 / 0.001	0.046 / 0.006	0.043 / 0.017	0.045 / 0.014
Gaussian 1000D	$W_1$ ↓ / RME ↓	2.212 / 0.053	3.035 / 0.194	2.233 / 0.044	3.010 / 0.037	3.785 / 0.303	2.263 / 0.127
Mouse 2D	$W_1$ ↓ / RME ↓	0.045 / 0.001	0.053 / 0.001	0.042 / 0.001	0.105 / 0.031	0.139 / 0.171	0.328 / 0.151
EB 50D	$W_1$ ↓ / RME ↓	6.155 / 0.005	6.155 / 0.010	8.892 / 0.005	8.891 / 0.049	9.092 / 0.312	9.415 / 0.092

在合成多尺度基准上针对真实耦合结构的还原测试（Q1）数据如下： - Exact OT: 目标差 2.35e-5%，单点精度 0.268，微观精度 0.510，宏观精度 1.000，耗时 89.91s - Mini-batch OT: 目标差 0.15%，单点精度 0.002，微观精度 0.328，宏观精度 0.988，耗时 8.86s - Low-rank OT: 目标差 13.5%，单点精度 0.000，微观精度 0.076，宏观精度 0.695，耗时 151.92s - MUST-FM (Sparse): 目标差 1.78e-14%，单点精度 0.999，微观精度 1.000，宏观精度 1.000，耗时 10.50s - MUST-FM (Independent): 目标差 0.0399%，单点精度 0.001，微观精度 1.000，宏观精度 1.000，耗时仅 0.025s

消融实验¶

在涵盖 130 万单细胞的小鼠器官发生细胞图谱（MOCA，包含 E9.5 至 E13.5 共 5 个发育阶段、38 种大类与 655 种细分亚型）上开展留一法（Leave-One-Out）跨时间步插值预测消融，检验生物谱系先验对轨迹恢复的不可替代性：

配置	E10.5 阶段 $W_1$ ↓	E11.5 阶段 $W_1$ ↓	E12.5 阶段 $W_1$ ↓	说明
MUST-FM w/ Priors (完整模型)	1.817	2.405	2.079	结合层级谱系转移先验引导，全时期重建误差最低
MUST-FM w/o Priors (无先验消融)	1.942	2.825	2.614	移除先验后，各留一插值阶段分布对齐误差显著上升

扩展性分析表明，在 130 万细胞全尺度下： - MUST-FM (Sparse) 在 20% 细胞规模以上即因细粒度矩阵显存超限而崩溃（OOM）。 - MUST-FM (Independent) 成功扩展至 100% 全图谱（1.3M 细胞），且由于层级元信息数量远小于细胞总数，峰值 CPU/GPU 内存随采样比例扩增基本保持平稳，其耦合构建耗时在总体训练开销中的占比随着数据规模扩大而显著衰减。

关键发现¶

分层拓扑优先于细粒度几何对齐：在细胞分化等复杂非线性过程中，单细胞层面的点对点配对往往充斥噪声与测序失真。MUST-FM (Independent) 虽在理论上放宽了最精细层内的点匹配精度（0.001），但由于严格锁定了宏观与中观亚型拓扑（两者精度达 1.000），在连续下游流匹配动力学学得的分布指标（EB 50D $W_1$ 达到 6.155，相较基线的 8.89 降低约 30%）反而与稀疏精准求解持平甚至更稳健。
生物先验有效纠正无约束流匹配的跨谱系漂移：在 MOCA 留一法实验中，缺少先验约束的无监督 UOT 会沿着欧氏测地线产生跨胚层/跨谱系的虚假流动（短路），而注入先验后预测误差在各时期平均下降 10%–20%，显著提升了未知时间步推断的生物可信度。

亮点与洞察¶

块内独立提升与隐式采样构思精巧：将最优传输问题截断在生物概念具有明确可解释性的中观亚型层级（$L-1$ 级），利用两步采样分解配合解析权重抵消（$\rho_{IJ}=\gamma_{1,IJ}/\gamma_{0,IJ}$），在数学上实现了完全免存最细粒度传输矩阵的常数时间采样，彻底化解了百万级图谱的显存爆炸死结。
从平衡 OT 走向非平衡流匹配的自洽闭环：不仅利用 WFR 距离闭式测地线实现了无模拟高效生成，还克服了传统监督最优传输在严格质量守恒与强先验下易导致不可行解的固有缺陷，以熵传输软约束完美包容了单细胞的增殖与消亡动态。
范式迁移潜力广阔：该多尺度掩码加隐式提升的思想可无缝迁移至具备分层组织架构的其它高维科学生成任务中，例如基于器官-组织-细胞三级标注的全身空间转录组动力学建模、基于骨架拓扑约束的人体三维动作生成流匹配等。

局限与展望¶

高度依赖预定义先验与注释质量：算法的有效性建立在高质量的多层级细胞聚类注释与准确的谱系先验二值图之上；若实验标注存在谱系偏见或先验遗漏，硬性掩码可能直接扼杀真实存在的新型转分化（transdifferentiation）或罕见发育旁路。
硬约束机制缺乏数据驱动的动态自适应性：当前先验矩阵 $B^{(l)}$ 采用静态的二值化硬过滤机制，未来值得探索引入软先验约束、先验松弛惩罚项或根据流匹配反向传播联合微调谱系连通性的可微先验演化机制。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将多尺度层级结构、生物谱系先验掩码与无模拟非平衡流匹配深度结合，巧妙提出无需显式存贮矩阵的隐式块内提升采样算法。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖高维合成点云、精确平移基准、小鼠与人单细胞真实测序数据，直至 130 万细胞的全尺度 MOCA 图谱留一法与扩展性极限测试，证据链完备。
写作质量: ⭐⭐⭐⭐⭐ 数学形式推导极其严谨规范，理论定理与附录采样算法推导闭环自洽，架构清晰。
价值: ⭐⭐⭐⭐⭐ 为单细胞乃至空间组学中百万级至千万级细胞图谱的连续非平衡发育轨迹推断树立了全新标杆，极具落地与参考价值。

数据集	指标	MUST-FM (Sparse)	MUST-FM (Independent)	WFR-FM (Prev. SOTA)	VGFM	DeepRUOT	TIGON
Gene 2D	\(W_1\) ↓ / RME ↓	0.019 / 0.001	0.019 / 0.001	0.019 / 0.001	0.046 / 0.006	0.043 / 0.017	0.045 / 0.014
Gaussian 1000D	\(W_1\) ↓ / RME ↓	2.212 / 0.053	3.035 / 0.194	2.233 / 0.044	3.010 / 0.037	3.785 / 0.303	2.263 / 0.127
Mouse 2D	\(W_1\) ↓ / RME ↓	0.045 / 0.001	0.053 / 0.001	0.042 / 0.001	0.105 / 0.031	0.139 / 0.171	0.328 / 0.151
EB 50D	\(W_1\) ↓ / RME ↓	6.155 / 0.005	6.155 / 0.010	8.892 / 0.005	8.891 / 0.049	9.092 / 0.312	9.415 / 0.092