WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Population Growth Modeling | 50D CITE t=2 | W1 Score27.478 | 12 | |
| Population Growth Modeling | 50D CITE t=1 | W127.831 | 12 | |
| Population Growth Modeling | 50D CITE t=3 | W134.784 | 12 | |
| Cell population dynamics prediction | 50D Mouse (t=1) | W1 Score5.486 | 11 | |
| Trajectory Inference | 100D EB dataset t=1 | W19.941 | 11 | |
| Trajectory Inference | 100D EB dataset t=2 | W111.04 | 11 | |
| Trajectory Inference | 100D EB dataset (t=3) | W1 Error11.516 | 11 | |
| Trajectory Inference | Dyngen t=1 | W1 Score0.11 | 11 | |
| Trajectory Inference | Dyngen (t=2) | W10.098 | 11 | |
| Trajectory Inference | Dyngen t=4 | W10.121 | 11 |