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. The Python code is available at https://github.com/QiangweiPeng/WFR-FM.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Reconstruction of discrete branching dynamics | 5D EB t=1 | W10.324 | 24 | |
| Reconstruction of discrete branching dynamics | CITE 50D (t=3) | W19.182 | 24 | |
| Trajectory Interpolation | EB 50D (held-out time points) | Mean W16.586 | 24 | |
| Trajectory Inference | EB dataset 5D (test) | W1 (t=1)0.324 | 23 | |
| Trajectory reconstruction | Gaussian Mixtures 1000D | W1 Distance2.233 | 18 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=4 | W1 Distance0.121 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=1 | W1 Distance0.11 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen (t=2) | W1 Distance0.098 | 15 | |
| Discrete Branching Dynamics Reconstruction | Dyngen t=3 | W1 Distance0.211 | 15 | |
| Discrete Branching Dynamics Reconstruction | 50D EB t=2 | W18.659 | 12 |