WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport
🤖 AI Summary
Existing Wasserstein–Fisher–Rao (WFR) solvers face significant limitations in handling dynamic unbalanced optimal transport problems, including high computational cost, numerical instability, and poor scalability. This work proposes WFR Flow Matching (WFR-FM), which for the first time unifies the flow matching framework with WFR geometry by jointly regressing a displacement vector field and a scalar growth rate function to directly learn continuous-time WFR geodesics. Theoretically, we prove that minimizing the WFR-FM loss exactly recovers the true WFR dynamics, enabling joint modeling of state evolution and mass variation. Empirically, WFR-FM substantially outperforms existing methods in single-cell trajectory inference, achieving notable improvements in computational efficiency, numerical stability, and reconstruction accuracy, and demonstrates strong performance in generative modeling of unbalanced data distributions.
Application Category
📝 Abstract
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.
Problem
Research questions and friction points this paper is trying to address.
Wasserstein-Fisher-Rao
unbalanced optimal transport
dynamic OT
trajectory inference
mass evolution
Innovation
Methods, ideas, or system contributions that make the work stand out.
Wasserstein-Fisher-Rao
flow matching
unbalanced optimal transport
dynamic modeling
mass evolution
🔎 Similar Papers
No similar papers found.
