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

2026-01-28

AI Summary

This work addresses the computational bottleneck in reconstructing dynamic evolutionary processes from sparse observations in single-cell biology, where existing methods rely on trajectory simulation and suffer from slow inference. The authors propose Wasserstein–Fisher–Rao Mean Field Matching (WFR-MFM), a novel framework that—by leveraging the geometry of the Wasserstein–Fisher–Rao metric—introduces mean field matching to dynamic, non-equilibrium optimal transport for the first time. By jointly modeling the mean velocity field and the mass growth field, WFR-MFM enables one-step generation over arbitrary time intervals without trajectory simulation. Evaluated on both synthetic and real single-cell RNA sequencing data, the method achieves high reconstruction accuracy while accelerating inference by several orders of magnitude over current baselines, and efficiently scales to large-scale perturbation response prediction involving thousands of conditions.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Matrix & Tensor Methods
• Intelligent Robots: State Estimation

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Reconstructing dynamical evolution from limited observations is a fundamental challenge in single-cell biology, where dynamic unbalanced optimal transport 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 optimal transport 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.

Problem

• dynamic unbalanced optimal transport
• single-cell biology
• trajectory simulation
• scalable inference
• Wasserstein-Fisher-Rao

Innovation

• unbalanced optimal transport
• one-step inference
• mean flow matching
• Wasserstein-Fisher-Rao
• single-cell dynamics