Title: Inferring stochastic dynamics with growth from cross-sectional data
Venue: NeurIPS 2025
Accepted: yes
Keywords: systems biology

single cell

stochastic dynamics

population dynamics
Abstract: Time-resolved single-cell omics data offers high-throughput, genome-wide measurements of cellular states, which are instrumental to reverse-engineer the processes underpinning cell fate. Such technologies are inherently destructive, allowing only cross-sectional measurements of the underlying stochastic dynamical system. Furthermore, cells may divide or die in addition to changing their molecular state. Collectively these present a major challenge to inferring realistic biophysical models. We present a novel approach, unbalanced probability flow inference, that addresses this challenge for biological processes modelled as stochastic dynamics with growth. By leveraging a Lagrangian formulation of the Fokker-Planck equation, our method accurately disentangles drift from intrinsic noise and growth. We showcase the applicability of our approach through evaluation on a range of simulated and real single-cell RNA-seq datasets. Comparing to several existing methods, we find our method achieves higher accuracy while enjoying a simple two-step training scheme.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes **Unbalanced Probability Flow Inference (UPFI)**, an extension of Probability-Flow Inference that incorporates a **growth (birth-death) term** into the Lagrangian formulation of the Fokker–Planck equation. UPFI proceeds in two steps:  \n1. **Score estimation** of ∇log ρ_t via denoising score matching -- using denoising score matching to estimate the time-dependent score of the data distribution \n2. **Neural-ODE fitting** of drift *v*(x,t) and growth *g*(x,t) by pushing weighted particles through the characteristic ODE and minimizing an **unbalanced Sinkhorn divergence** plus a Wasserstein–Fisher–Rao (WFR) regularizer -- learning the drift and growth terms by solving a characteristic ODE system.  \n\nThe authors provide a theoretical analysis, demonstrating the non-identifiability between drift and growth. The method is evaluated on a series of simulated systems, including a high-dimensional bistable system and chemical Langevin equation models, as well as on a real-world single-cell RNA-seq dataset of monocyte-neutrophil development.\n\nSTRENGTHS AND WEAKNESSES:\n#  Strengths\n* **Problem Relevance** - tackles a crucial and challenging problem in computational biology: learning cell dynamics from static snapshots, accounting for population growth and death.\n* **Identifiability analysis** - provides an analytic example (linear OU + quadratic fitness) and proves uniqueness under WFR. This provides a solid theoretical foundation for the necessity of WFR regularization.\n* **Unbalanced OT loss** - handles unequal cell counts directly, an important practical feature for real lineage data. \n* **Clarity** -  well-written, clearly structured, and easy to follow\n\n\n# Weaknesses\n \n* **Limited scalability** -  relies on an ODE solver within the main training loop (Algorithm 1) raises questions about the method's computational scalability. While the method is presented as a general solution, the experiments are restricted to low-dimensional PCA embeddings (d=2, 5, 10 for simulations; d=10 for real data).\n* **Contribution feels narrow** - While the problem is important, the proposed solution feels like a very incremental improvement. Given the severe scalability concerns and the lack of engagement with state-of-the-art generative modeling techniques (like flow matching), the paper's impact seems limited. The recent work [1-6] suggests that the field is moving towards more scalable solutions, questioning the long-term relevance of the proposed method.\n* **Related-work gap** - paper does not discuss the advancing field of flow matching [1-5], while it is far more relevant to the paper's goal of learning vector fields between distributions than some of the older cited works\n* **Limited Ablations** - paper does not provide ablation along regularisation, which makes it difficult to understand how much hyper tuning is required given method is already expensive\n\n\n[1] Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G. and Bengio, Y., 2023. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00482.\n\n[2] Tong A, Malkin N, Fatras K, Atanackovic L, Zhang Y, Huguet G, Wolf G, Bengio Y. Simulation-free schr\\\" odinger bridges via score and flow matching. arXiv preprint arXiv:2307.03672. 2023 Jul 7.\n\n[3] Kapusniak, K., Potaptchik, P., Reu, T., Zhang, L., Tong, A., Bronstein, M., Bose, J. and Di Giovanni, F., 2024. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 37, pp.135011-135042.\n\n[4] Petrović, K., Atanackovic, L., Kapuśniak, K., Bronstein, M.M., Bose, J. and Tong, A., 2025, March. Curly flow matching for learning non-gradient field dynamics. In Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2025.\n\n[5] Eyring L, Klein D, Uscidda T, Palla G, Kilbertus N, Akata Z, Theis F. Unbalancedness in neural monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100. 2023 Nov 25.\n\n[6] Wang, D., Jiang, Y., Zhang, Z., Gu, X., Zhou, P. and Sun, J., 2025. Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling. arXiv preprint arXiv:2505.13413.\n\nQUESTIONS:\n1. **Positioning** - How does the framework position itself relative to the recent literature on flow matching [1-6]? These methods have been extended to stochastic and unbalanced settings. A discussion of the relative merits would strengthen the paper. \n2. **Compute budget** - Could the authors provide a more detailed analysis of the computational complexity and scalability of UPFI? Specifically, how does the training time scale with the number of samples, dimensions, and the number of time points?\n3. **Velocity prior:** Could kernel-smoothed RNA-velocity guide or evaluate the drift, as in [4]?  \n4. **Regulariser sensitivity:** How do drift and growth fields vary with different λ and α values? Please provide an ablation.\n\nLIMITATIONS:\nThe authors have adequately addressed the primary theoretical limitation of their work: the non-identifiability of drift and growth. The significant computational cost and scalability issues inherent in their simulation-based training scheme is not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 1, "ORIGINALITY": 3, "DATE": "2025-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscript focuses on the task of predicting single-cell trajectories and inferring their underlying dynamics. The authors adapt the PFI approach to the unbalanced case, allowing the model to capture both cell growth and cell death. They present results on both toy and single-cell datasets.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, the paper and method are clear and easy to follow. The approach is well motivated by biological considerations—such as cellular growth and death—that are sometimes omitted in other methodologies. The authors propose an unbalanced version of PFI, which demonstrates improvements over the original PFI on both synthetic and real datasets. They also present a theoretical contribution by showing the existence of a unique minimizer for the continuous version of their loss function.\n\nHowever, I believe comparisons against flow matching-based methods are missing. For instance, it would be informative to include a vanilla flow matching model or even an unbalanced one (akin to [1]).\n\n[1] Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., & Theis, F. (2023). Unbalancedness in neural Monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100.\n\nQUESTIONS:\n- Lines 60–61: Should one of the conditions be g(x) < 0.\n- Could you provide more intuition for Equation (12)? For example, it seems that as alpha tends to zero, it converges to the dynamic OT formulation.\n- Is the model always an OU process with quadratic fitness? This is not entirely clear from Algorithm 1.\n- In Figure 4 (ii), if you rank the top-k elements of matrix A, how many correspond to true interactions (the ones shown in red)? Reporting this would help compare the two learned matrices more easily.\n- Is the algorithm robust to the “Sinkhorn loss on marginals” step in Algorithm 1? Have you experimented with different solvers for this step? Specifically, does this step require an unbalanced Sinkhorn solver?\n\nLIMITATIONS:\nA more detailed discussion of limitations should be included, or they should be highlighted through experiments—for example, the computational cost associated with using a neural ODE.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper looks at the challenge of reconstructing dynamical systems from snapshots. In particular, for single cell data, the sample is destroyed by the measurement. The authors present \"unbalanced\" probability flow inference. Which has somewhat of the feel of a method for dealing with censored data that takes account of the underlying dynamics. The paper modifies the probability flow inference technique to deal with a limitation where proliferation and death isn't handled.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper does a nice job of explaining the problem, being clear about what previous work does and where the contribution has come in.\n\nThe algorithm is illustrated with a range of experiments that help both in understanding of the set up and the quality of results.\n\nFormatting and description of algorithm is a good balance between technical detail and clarity of explanation\n\nThe problem in single cell data is well motivated and the method is compared with other methods from the literature and performs in the manner that would be expected given enhancements.\n\nExperimental set up is clearly explained and easy to follow. Sensible simulations are given before the full experiments.\n\nIn an ideal world we would have seen more from real RNA-seq data or other data sets with similar problems. \n\nI suspect the use of 10 principal components is due to challenges in scaling the algorithm, I would have liked to hear more about this.\n\nQUESTIONS:\nI would like to hear more about alorithm scaling. My sense is currently this approach might be hard to apply to some data sets.\n\nIs there an example where growth is critical and current methods cannot be deployed as a result? Even if it's a motivational example it would be good to hear it.\n\nI've been out of the loop for comp bio for a bit, and we were discussing snapshot data some years ago. I was a little surprised there wasn't a review of progress in that direction that goes beyond listing references.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-07"}