Reconstructing dynamical evolution from limited observations is a fundamental challenge in single-cell biology, where dynamic unbalanced optimal transport provides a principled framework for modeling coupled transport and mass variation. However, existing approaches rely on trajectory simulation at inference time, making inference a key bottleneck for scalable applications. In this work, we propose a mean-flow framework for unbalanced flow matching that summarizes both transport and mass-growth dynamics over arbitrary time intervals using mean velocity and mass-growth fields, enabling fast one-step generation without trajectory simulation. To solve dynamic unbalanced optimal transport under the Wasserstein-Fisher-Rao geometry, we further build on this framework to develop Wasserstein-Fisher-Rao Mean Flow Matching (WFR-MFM). Across synthetic and real single-cell RNA sequencing datasets, WFR-MFM achieves orders-of-magnitude faster inference than a range of existing baselines while maintaining high predictive accuracy, and enables efficient perturbation response prediction on large synthetic datasets with thousands of conditions.
翻译:从有限观测中重构动态演化是单细胞生物学中的一个基本挑战,其中动态非平衡最优传输为建模耦合传输与质量变化提供了一个原则性框架。然而,现有方法在推断时依赖于轨迹模拟,使得推断成为可扩展应用的关键瓶颈。本文中,我们提出了一种用于非平衡流匹配的平均流框架,该框架利用平均速度场与质量增长场来概括任意时间间隔内的传输与质量增长动态,从而无需轨迹模拟即可实现快速一步生成。为了在Wasserstein-Fisher-Rao几何下求解动态非平衡最优传输问题,我们基于此框架进一步构建了Wasserstein-Fisher-Rao平均流匹配方法。在合成与真实的单细胞RNA测序数据集上的实验表明,WFR-MFM在保持高预测精度的同时,其推断速度比一系列现有基线方法快数个数量级,并且能够在包含数千种条件的大型合成数据集上实现高效的扰动响应预测。