The Wasserstein-Fisher-Rao (WFR) metric extends dynamic optimal transport (OT) by coupling displacement with change of mass, providing a principled geometry for modeling unbalanced snapshot dynamics. Existing WFR solvers, however, are often unstable, computationally expensive, and difficult to scale. Here we introduce WFR Flow Matching (WFR-FM), a simulation-free training algorithm that unifies flow matching with dynamic unbalanced OT. Unlike classical flow matching which regresses only a transport vector field, WFR-FM simultaneously regresses a vector field for displacement and a scalar growth rate function for birth-death dynamics, yielding continuous flows under the WFR geometry. Theoretically, we show that minimizing the WFR-FM loss exactly recovers WFR geodesics. Empirically, WFR-FM yields more accurate and robust trajectory inference in single-cell biology, reconstructing consistent dynamics with proliferation and apoptosis, estimating time-varying growth fields, and applying to generative dynamics under imbalanced data. It outperforms state-of-the-art baselines in efficiency, stability, and reconstruction accuracy. Overall, WFR-FM establishes a unified and efficient paradigm for learning dynamical systems from unbalanced snapshots, where not only states but also mass evolve over time.
翻译:Wasserstein-Fisher-Rao(WFR)度量通过将位移与质量变化相耦合,扩展了动态最优传输(OT),为建模非平衡快照动力学提供了原理性的几何框架。然而,现有的WFR求解器往往不稳定、计算成本高且难以扩展。本文提出WFR流匹配(WFR-FM),一种免仿真的训练算法,将流匹配与动态非平衡OT相统一。与仅回归传输向量场的经典流匹配不同,WFR-FM同时回归用于位移的向量场和用于生灭动力学的标量增长率函数,从而在WFR几何下产生连续流。理论上,我们证明最小化WFR-FM损失可精确恢复WFR测地线。实证中,WFR-FM在单细胞生物学中实现了更准确、更稳健的轨迹推断,重建了包含增殖与凋亡的一致性动力学,估计了时变生长场,并适用于不平衡数据下的生成动力学。其在效率、稳定性和重建精度上均优于当前最先进的基线方法。总体而言,WFR-FM为从非平衡快照中学习动力系统建立了一个统一且高效的范式,其中不仅状态随时间演化,质量亦随之变化。