We thank the reviewer for their thorough review and appreciation of our efforts to provide a clear and theoretically justified method. Below we address comments about the i) relevance of our contribution and its positioning in the literature, ii) scalability of our algorithm, and iii) ablations.

Contribution and positioning

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. Positioning - How does the framework position itself relative to the recent literature on flow matching [1-6]?

We would like to point out that [6] became available online May 19 2025 and represents a concurrent work of which we could not have been aware at the time of our submission.

To better position UPFI and its limitations in the literature, we would like to stress one important point: UPFI is a principled approach to learn, from cross-sectional samples, Fokker-Planck equations with growth and arbitrary diffusion tensor.

As such, UPFI can fit not only processes with additive noise $\sigma=\mathrm{const}$ (as is typical of OT-based methods), but also models with state and drift-dependent multiplicative diffusion $\sigma(t, x)$. The relevance of such models in biological systems is well known [Coomer2022].

To our knowledge, such models are are yet out of reach of any Flow Matching (FM) approach currently available, and UPFI is the first approach able to simultaneously account for growth and arbitrary, possibly multiplicative, intrinsic noise. This aspect of our work is illustrated on the simulated HSC example in Fig.4 and Table 2, for which fitting with additive noise ensures good interpolation, but does not extract correct regulatory interactions. We believe that this versatility of UPFI explains apparent limitations with respect to more recent FM methods.

On the other hand, while FM enables simulation-free interpolants between distributions, these approaches sacrifice flexibility for scalability. A crucial ingredient in any FM approach is the choice of coupling between distributions, since this prescribes the dynamics. For balanced transport, Sinkhorn couplings are a reasonable choice thanks to Benamou-Brénier and reciprocal property of the Schrodinger bridge [Tong2023].

On the other hand, we are not aware of theoretical arguments for a good choice of couplings in the unbalanced case, that have been applied to FM. [Eyring2023] consider using unbalanced OT with soft marginal constraints, under which the dynamics occur in two distinct steps: (i) a "reweighting" procedure involving inflation/deflation of mass density, and (ii) transport under the FM framework. Effectively, this isolates the effect of transport and growth. In the biological setting, and for general Fokker-Planck equations with source terms, however, transport and growth occur simultaneously and can interact. UPFI allows for these to be jointly modeled.

For a rigorous extension of FM to the inference of diffusions with arbitrary intrinsic noise and growth, rigorous mathematical insight on the construction of coupling and choice of interpolant is still lacking. Thus, while we agree that UPFI lacks the scalability of FM, it can model systems of biological and physical relevance for which principled, theoretically justified FM approaches are unavailable at present.

Moving towards more scalable solutions is indeed an important direction and we believe our paper provides a useful solid theoretical foundation for later works.

Thanks to this feedback, we now add a comparison to OT flow matching and an unbalanced variant [Eyring2023], see below and our response to other reviewers. We will expand discussion of related work, and emphasize what sets UPFI apart from FM methods.

Scalability

Limited scalability - relies on an ODE solver within the main training loop (Algorithm 1) raises questions about the method's computational scalability. 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?

To contextualize the scalability issue, we first want to point out that our method already improves over existing methods like DeepRUOT and TIGON, which applied their approach on the Lineage Tracing dataset in a two-dimensional force layout embedding. An improvement of our method is that, as opposed to TIGON and DeepRUOT, UPFI does not perform density estimation, although we agree with the reviewer that UPFI cannot scale to thousands of dimensions.

Regarding the scalability and computational complexity of UPFI, we consider each step: i) score estimation, ii) sinkhorn loss and iii) ODE fitting. See also our response to wcim.

Regarding the denoising score matching, its complexity is $\mathcal{O}(B\times d)$ where $d$ is the dimension, B is the batch size. For large entropic regularization, the unbalanced Sinkhorn divergence can be computed efficiently with a $\mathcal{O}(B^2)$ cost and with a dimension-independent sample complexity $\mathcal{O}(B^{-1/2})$ [Sejourne2019], provided $B$ is large enough for these scalings to old. We agree with the reviewer's concerns regarding the scalability of the ODE integration step, which in theory can limit the inference of dynamical models in high-dimensions. In practice, we mitigate this issue by taking ~2 to 3 Euler steps in between each time points. Overall, while our approach can't scale to hundreds or thousands of dimensions, we can reach moderately high dimensions provided we increase the size of the neural networks. We performed additional experiments with ~50 dimensions of the bifurcating example to showcase the ability of UPFI to scale to these moderately high-dimensions.

We also compared UPFI with (unbalanced) flow matching and unbalanced flow matching (U)OTFM. While unbalanced OTFM offers indeed a fair improvement upon OTFM, we find that UPFI is consistently more accurate for vector field estimation and fate probabilities.

See table in our response to wcim: Bifurcating ex. (Table 1) to $d = 25, 50$

To give an idea of the runtimes of UPFI, we time for our method associated to the experiments of Table 1 and Fig 5, which we report below:

Runtime (in seconds, for Table 1)

Note we add results for $d=25,50$, for which we increase score matching iters from 25k to 100k, and UPFI iters from 5k to 25k.

Runtime (in seconds, for Fig.5)

We emphasize that the lineage example has 86,416 cells on which UPFI runs in a matter of minutes.

These timings show that UPFI runs for realistic $N$ and $d$ in short time. In comparison, DeepRUOT is an order of magnitude slower whilst less accurate. We will add a discussion regarding the scalability and computational cost, as well as timing results, at the camera ready stage.

Lack of ablations

We add regularization ablations for bifurcating system in $d = 10$. For varying $\alpha$ and $\lambda$ we show the cosine error for force recovery as well as Pearson correlation for growth rate recovery.

Regularization ablation

Force recovery, cosine

Growth rate recovery, Pearson

For $\lambda = 0$ there is no regularization, and the error is larger than when $\lambda > 0$. From the growth rate perspective, no regularization performs better, while this incurs a larger error on the force field. From these measurements we find that the growth rate is less sensitive to the regularization. We interpret this as resulting from an implicit regularization arising from parameterisation with relatively small neural nets. The motivation for regularization arises from theoretical considerations to ensure uniqueness, as motivated by our Theorem. We can draw from these ablation study a simple rule of thumb for UPFI: use small or no regularization with moderately sized neural networks. Based on this comment, we propose to add these ablations experiments and discuss their consequences on the architectures choices in the manuscript at the camera ready stage .

Additional comments

Velocity prior: Could kernel-smoothed RNA-velocity guide or evaluate the drift, as in [4]?

Yes this could be done, however we find that RNA velocity is not reliable enough to be used directly since it requires strong biophysical assumptions to be inferred in the first place [Gorin2022]

[Sejourne2019] Thibault Séjourné, Jean Feydy, Francois-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport.
[Coomer2022] Coomer MA, Ham L, Stumpf MP. Noise distorts the epigenetic landscape and shapes cell-fate decisions. Cell Systems
[Eyring2023] 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.
[Tong2023] 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.
[Gorin2022] Gorin G, Fang M, Chari T, Pachter L. RNA velocity unraveled.

We thank the reviewer for their thorough review. Below we adress the reviewer's comments regarding i) the scalability of our approach and its comparison to Flow Matching, ii) the choice of the regularization, and we clarify a number of points.

Scalability and comparison to Flow Matching

A 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.

We agree with the reviewer that in its current form the paper misses a discussion about the scalability and computational cost of the approach. Concerning the scalability, we mitigate the neural ODE cost by only taking 2 to 3 Euler steps in between snapshots. This trades off accuracy for scalability. The unbalanced Sinkhorn loss benefits from a good sample complexity (see [Sejourne2019] and our response to Reviewer wcim). For the scales of data available in typical time-resolved single-cell RNA-seq datasets ($\leq 10^5$ cells with $\sim 5$ time-points), we can scale our approach to moderately high dimensions of ~50. To demonstrate this, we run the bifurcating example (Table 1 of our paper) with 25 and 50 dimensions.

We also implement OT-conditional flow matching (OTFM) [Tong2023] and an unbalanced variant (UOTFM) [Eyring2023].

Bifurcating ex. (Table 1) to $d = 25, 50$

Cosine error

Fate prob.

We can see that while UPFI still outperforms the other methods, UOTFM does a decent job at recovering the fate probabilities, certainly because it handles appropriately the unbalancedness of the two branches. We also computed the timings for our method on the different experiments we ran, and we see that the force training and score training take a similar amount of time, which are very reasonable:

Runtime (in seconds, for Table 1)

Note we add results for $d=25,50$, for which we increase score matching iters from 25k to 100k, and UPFI iters from 5k to 25k.

Runtime (in seconds, for Fig.5)

Compared to UPFI, DeepRUOT is much slower to train and is less accurate.

We will discuss more thoroughly these scaling issues and include these additional experiments in the paper at the camera ready stage.

However, 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]).

We agree with the reviewer that we lack a proper comparison with Flow Matching approaches. We performed additional numerical experiments to fill in this gap (see response to previous question). For the bifurcating example they are reported as the last two last columns of the table, and they confirm that although performing comparably to other methods, balanced and unbalanced flow matching (resp. OTFM and UOTFM) still underperform with respect to UPFI. We would like to stress that UPFI, although not as scalable as Flow Matching, allows to jointly infer drift and growth, while simultaneously accounting for multiplicative (state and drift dependent) noise. Such models, which are biologically meaningful are not currently supported in existing Flow Matching frameworks. We will include these new results at the camera ready stage, as well as to improve the literature review of flow matching approaches. This will also be an opportunity to clarify the positioning of our approach, and what sets it apart from these more scalable methods.

Model and regularization choices

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.

The regularizer is the Wasserstein-Fisher-Rao energy, as introduced in [Chizat2018]. This distance generalizes the Benamou-Brénier dynamic formulation of optimal transport to unbalanced optimal transport. Overall, our loss function includes a data fitting constraint via the Sinkhorn divergences - which can be seen as a soft penalty on the marginals at all times points - and a WFR loss in between time points. With this choice of regularizer, our loss resembles an approximation of a multi-marginal unbalanced optimal transport. However, it is important to note that our method diverges from unbalanced optimal transport since we do not penalize the phase-space velocity (which for Fokker-Planck flows has a contribution stemming from the score, see Eq. 3 or Eq. 4), but only the force field. In practice, because of concerns regarding the identifiability problem, we also enforce the force field and growth rate to be autonomous. This additionally differentiates our approach from dynamic unbalanced optimal transport.

Is the model always an OU process with quadratic fitness? This is not entirely clear from Algorithm 1.

The OU/quadratic process is a toy system for which we can provide theoretical analysis. Alg. 1 applies in the general (i.e. non-linear) setting. We will clarify the distinction at the camera ready state.

Additional comments

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?

Yes the loss functional on the marginals needs to be able to handle positive (empirical) measures with differing mass. While different losses can be used (and not necessarily unbalanced OT, see e.g. [Mroueh2020]), we opt for UOT as a natural choice. In practice, we use the GeomLoss implementation which uses Sinkhorn-type iterations.

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.

We thank the reviewer for this spot on suggestion. This is what the precision-recall curves shown in Fig.4(iii) measures as a summary over all possible thresholds. Given a threshold value, we identify the entries in the matrix that are larger than this threshold. Among these thresholded entries, the precision computes the proportion of true interactions. The recall on the other hand computes the proportion of true interactions recovered among all true interactions in the unthresholded matrix. A large recall and large precision are synonym of good prediction. As the threshold level is not prescribed, we vary it to build the precision-recall curve, and compute the area under the curve. A larger area under the curve indicates better recovery of the underlying true interactions.

[Mroueh2020] Mroueh Y, Rigotti M. Unbalanced sobolev descent. Advances in Neural Information Processing Systems. 2020;33:17034-43.
[Chizat2018] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In Advances in neural information processing systems, pages 3036–3046, 2018
[Eyring2023] 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.
[Tong2023] 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.

We thank the reviewer for their appreciation of our efforts to provide a clear and mathematically rigorous approach. Below we address the different reviewer's comments regarding i) the scalability of the approach, ii) the motivation for the inference of growth.

Scalability

I suspect the use of 10 principal components is due to challenges in scaling the algorithm, I would have liked to hear more about this. I would like to hear more about alorithm scaling. My sense is currently this approach might be hard to apply to some data sets.

Theoretical complexity

The unbalanced Sinkhorn divergence can be computed efficiently with per-iteration complexity $\mathcal{O}(B^2)$, with a dimension-independent sample complexity $\mathcal{O}(B^{-1/2})$ [Sejourne2019], provided $B$ is large enough, where $B$ is the batch size. The denoising score matching on the other hand has a per-iteration complexity $\mathcal{O}(B\times d)$ where $d$ is the dimension. For the ODE integration, we can mitigate the complexity at the price of accuracy by simpling taking less Euler steps in between time points. In practice we find that taking 2 to 3 steps is satisfying. While this setting still prevents us from working with thousands of dimensions, it isn't preventing us from scaling to ~50 dimensions, as we illustrate with additional results on the bifurcating example. Note the addition of (U)OTFM (flow matching), as per request of wbX3 and rQy3.

Bifurcating ex. (Table 1) to $d = 25, 50$

Cosine error

Fate prob.

Use of PCA This has consequences for the analysis of real biological datasets like Hematopoietic Stem Cells differentiation, for which only a few tens of critical transcription factor are responsible for the differentiation of cells [Weinreb2020]. On the Lineage Tracing dataset, using additional lineage barcoding information, it was shown that there exists at least ten genes whose expression levels are better predictors of fate probability than any state-of-the-art algorithms working in reduced dimensional PCA space [Weinreb2020]. For this reason we believe that the scaling limitations of UPFI do not completely hinder its predictive potential for the prediction of cell fate decision making, since the difficulty seemingly resides more in building a relevant list of genes or a better dimensional reduction method than PCA. Motivated by this comment, we will to include a discussion about the scalability and computational cost at the camera stage.

Practical runtime

On the practical front, we find that UPFI is as fast or faster to run compared to competing methods. Below we provide runtimes (in seconds, over 5 repeats) for the bifurcating and lineage (Figs 3,5 of our paper) and compare to DeepRUOT.

Runtime (in seconds, for Table 1)

Note we add results for $d=25,50$, for which we increase score matching iters from 25k to 100k, and UPFI iters from 5k to 25k.

Runtime (in seconds, for Fig.5)

We emphasize that the lineage example has 86,416 cells on which UPFI runs in a matter of a few minutes.

Motivation for the growth model

Is 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.

We thank the reviewer for raising the issue of the necessity of a growth model, this is an important point which we didn't elaborate enough upon. In our cellular differentiation example we infer that growth is decreasing with time, with progenitor cells and stem cells dividing more than terminal cell types (Monocytes and Neutrophiles). While this is expected [Sha2024], this indicates that growth has a rather limited impact on the inferred drift since both terminal states which emerge at a similar time in the differentiation process have comparable growth rates. We observe however that taking growth into account improves by a factor ~3 the cell fate prediction accuracy with respect to PFI. A more robust and precise quantification of the role of the growth rate would be necessary to understand why taking the growth rate into account improves the prediction. This is work we plan to do, but which goes beyond this method-oriented paper.

We can however provide examples where growth plays a clear role. Such a situation, which happens when growth is heterogeneous over cell types, is illustrated on our toy bifurcating example. In real systems, it was observed when attempting to learn Optimal Transport trajectories summarizing a cell reprogramming experiment in [Schiebinger2019]. In this experiment, differentiated cells where reprogrammed towards induced pluripotent stem cells (iPSCs) via an experimental protocol. During the course of the reprogramming, intermediate stromal cell expressing apoptotic genes (inducing cell death) emerged. The misinterpretation of the death of these stromal cells as a transition to another cell state (transport) lead to incorrect reprogramming trajectories connecting them to the iPSCs emerging at the end of the reprogramming experiment. This mistake could only be corrected by accounting for cell death in the dynamical model. To motivate our method, we propose to include this example in the introduction of the paper.

Additionnal comments

I'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.

The reviewer is right that a more in-depth discussion of related work is warranted. In this instance we were not able to provide this due to the page limit, and we provided an in-depth discussion of TIGON and DeepRUOT in the appendix. We propose to rework the introduction to include an improved related work section.

[Sejourne2019] Thibault Séjourné, Jean Feydy, Francois-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. arXiv preprint arXiv:1910.12958, 2019.
[Weinreb2020] Caleb Weinreb, Alejo Rodriguez-Fraticelli, Fernando D Camargo, and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367(6479):eaaw3381, 2020.
[Schiebinger2019] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928–943, 2019.
[Sha2024] Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25–39, 2024.