WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport

Reconstructing dynamical evolution from limited observations is a fundamental challenge in single-cell biology, where dynamic unbalanced optimal transport (OT) provides a principled framework for modeling coupled transport and mass variation. However, existing approaches rely on trajectory simulation at inference time, making inference a key bottleneck for scalable applications. In this work, we propose a mean-flow framework for unbalanced flow matching that summarizes both transport and mass-growth dynamics over arbitrary time intervals using mean velocity and mass-growth fields, enabling fast one-step generation without trajectory simulation. To solve dynamic unbalanced OT under the Wasserstein-Fisher-Rao geometry, we further build on this framework to develop Wasserstein-Fisher-Rao Mean Flow Matching (WFR-MFM). Across synthetic and real single-cell RNA sequencing datasets, WFR-MFM achieves orders-of-magnitude faster inference than a range of existing baselines while maintaining high predictive accuracy, and enables efficient perturbation response prediction on large synthetic datasets with thousands of conditions.

Reconstructing how cells change over time is an important problem in biology and medicine, especially because experiments usually capture only a few snapshots of cells rather than continuous observations. Existing computational methods can model these changes, but they often require slow step-by-step simulations, making them difficult to scale to large datasets. In this work, we introduce a new framework that predicts how cell populations evolve much more efficiently. Instead of simulating every intermediate step, our method directly summarizes the overall movement and growth of cells between observations, allowing fast one-step prediction. Based on this idea, we develop WFR-MFM, a method for modeling both changes in cell states and changes in population size. Across synthetic and real single-cell RNA sequencing datasets, WFR-MFM produces accurate predictions while greatly reducing computation time compared with existing approaches. The method also scales to large perturbation studies involving thousands of experimental conditions, enabling faster analysis of cellular responses in complex biological systems.