WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport
This work addresses the challenge of modeling dynamical systems with evolving mass, such as in single-cell biology, by providing a more efficient and stable method for learning from unbalanced snapshots, though it appears incremental as it builds on existing flow matching and optimal transport concepts.
The paper tackles the problem of unstable and computationally expensive solvers for the Wasserstein-Fisher-Rao metric in dynamic unbalanced optimal transport by introducing WFR-FM, a simulation-free training algorithm that unifies flow matching with this geometry, resulting in more accurate and robust trajectory inference in single-cell biology with improved efficiency, stability, and reconstruction accuracy over state-of-the-art baselines.
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. The Python code is available at https://github.com/QiangweiPeng/WFR-FM.