Inferring stochastic dynamics with growth from cross-sectional data
Authors: Stephen Y. Zhang
Suryanarayana Maddu
Xiaojie Qiu
Victor Chardès
Venue: NeurIPS 2025 poster ()
Keywords: systems biology
single cell
stochastic dynamics
population dynamics
Primary area: machine_learning_for_sciences
TL;DR: We consider the problem of learning drift, diffusion and growth rate from snapshot observations of a population of individuals that may divide or die, and propose a novel inference scheme to learn these terms from data.

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.

Inferring stochastic dynamics with growth from cross-sectional data
Stephen Zhang!
Suryanarayana Maddu
School of Mathematics and Statistics, University of Melbourne
Center for Computational Biology, Flatiron Institute
Xiaojie Qiu
Victor Chardès
Department of Genetics, Stanford University School Medicine
Center for Computational Biology, Flatiron Institute
Abstract
Time-resolved single-cell omics data offers high-throughput, genome-wide mea- surements of cellular states, which are instrumental to reverse-engineer the pro- cesses underpinning cell fate. Such technologies are inherently destructive, al- lowing only cross-sectional measurements of the underlying stochastic dynami- cal system. Furthermore, cells may divide die in addition changing their molecular state. Collectively these present a major challenge to inferring real- istic biophysical models. We present novel approach, unbalanced probability flow inference, that addresses this challenge for biological processes modelled as stochastic dynamics with growth. By leveraging Lagrangian formulation of the Pokker-Planck equation, our method accurately disentangles drift from intrinsic noise and growth. We showcase the applicability of our approach through evalu- ation on a range of simulated and real single-cell RNA-seq datasets. Comparing several existing methods, find our method achieves higher accuracy while enjoying a simple two-step training scheme.
Introduction
Single-cell measurement technologies offer an unbiased, single-cell resolution view of the molecular processes dictating cell fate. These methods predominantly rely on microfluidic devices to capture and read molecules in individual cells, destroying the cell in the measurement process. Therefore the state of any cell can be measured only at a single time. Experimental study of temporal biological processes must therefore proceed by time-course studies, in which time-series of population snap- shots is obtained. At each time-point, a measurement is obtained from a distinct population of cells, assumed be identically prepared at the initial time. In practice these data can be obtained by serial sampling from a larger population, or from biological replicates measured at different times [II. 2] The challenge with time-resolved population snapshot data is that direct observation of dynamics in terms of longitudinal information is lost, and must be reconstructed from static population profiles. This motivates the inverse problem of reconstructing an underlying dynamical system which not only interpolates the data, but also captures the key biophysical phenomenon at play in the regulation of gene expression, notably intrinsic noise and cellular proliferation. On the other hand, ignoring these aspects can result in incorrect identification of regulatory mechanism [E. Al].
While the task of learning stochastic dynamics interpolating high-dimensional distributions is now routinely solved in generative modelling tasks [5), Ki], doing so while ensuring faithfulness to bio- physically relevant features like intrinsic molecular noise well division and death remains
CCoresponding authors: syz@syz. id. vehardeseflatiiionnnstttttte
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Dritt-diffsion dynamics
matching 8,(0))
characteristic OGE (M) Learned characterisiks SDE trajectories
lug palt)
g1 Z
fd
Ban
s(se
S a Plii) (ang, Myy 1-> (st, m7))
(nda V (AA
(p D drit-dlffusion
growth
Figure I: Overview. (i) Stochastic population dynamics with growth, governed by a Fokker-Planck equation with source. (ii) Score matching trains a neural network to model the contribution of noise. (iii) Neural-ODE based learning of the Polkker-Planck characteristics. (iv) Learned characteristics and corresponding SDE trajectories.
challenge. Similarly, while the dynamical inference problem has seen significant interest from the single cell analysis community, the issues of the noise and growth remain open for the most part [R, 9. 71 8]. For instance, most existing methods assume either noise [[]] or constant isotropic diffusivity [2] I0I, and either no growth [[T]] or the availability of prior information such as growth
rate estimates [2] or lineage tracing
[ 12
With this mind, a flexible framework called probubility flow inference (PFT) was recently intro- duced [13, [4]] tool to infer stochastic dynamical models accounting for intrinsic noise. Lever- aging the Lagrangian formulation of the Fokker-Planck equation, also known as the probability flow formulation [TS], PFI infers the drift of an Itô diffusion interpolating between the distributions successive times. limitation of PFI is that it does not account for growth (i.e. proliferation and death), a key feature of realistic biological systems that must be appropriately modelled for the in- ference of accurate trajectories. This was notably observed in reprogramming experiments, where neglecting cellular death led to the inference of incorrect state transitions between apoptotic and pluripotent cells [Z]. In this paper, we show how the PFI approach naturally extends to account for cellular proliferation, and we propose new algorithm called unbalanced probubility flow inference (UPFI), which allows to disentangle drift from intrinsic noise and growth in diffusion processes.
Stochastic population dynamics with growth following an Itô process of the form
We consider population of individuals evolving
dX2, v. (X,) dt + or (X,)d,,,
(I)
with d.B, the increments of dddimensional Brownian motion, (x())) the drift encoding determin- istic aspects of the system and o,(X,, the strength of the noise, which can depend on the state X,, well as the drift.
To include growth in the model (II), we model individuals as undergoing division and death at rates (X,) and dx(X,) respectively. More precisely, during short time interval < 1, an individual with state X, divides with probability br(X,)) + o(T) and dies with probability dx(X,)TT + o(T). Upon a division event, both descendants inherit their parent's 's state X,, With these prescriptions, the Itō process (M)) along with the division and death mechanism, defines a branching diffusion process [16]] The density of individuals with state at time t, denoted (<()), follows Fokker-Plande equation (FPE) with a source term [[2]]
Ōepe(x)) = - V [pe(x))u(z)) - (p))) + (x)))
(2)
the net fitness associated with the gene expression æ. Regions where gx(x) < 0 correspond to regions of net death, while regions (xx) > correspond net growth. The total number of cells where D((x) (x) sd is the diffusivity matrix and gx( (x(x) (x) is
|Ptl (p(())d is thus not conserved over time, and Pt is not probability density. In Fig. Mii we illustrate the evolution of Pt for a bistable process in two dimensions, where growth has the effect of accumulating the mass on one of the two branches, even though both branches are equally favourable in terms of the energetic landscape. Section EL we will analyse in more detail high-dimensional version of this bifurcating process.

---

To motivate the relevance of (II) for cell dynamics, we point out that a SDE of this form with coupled drift and noise terms naturally arises from the Chemical Langevin Equation, obtained by coarse- graining of the biophysically principled Chemical Master Equation [[]]] The importance of multi- plicative noise in biological systems also pointed out by [3] which highlights the nontrivial impact that complex noise models can have on dynamics. Finally, multiplicative noise models also arise in non-biological settings such as the Gox-Ingersoll-Ros model in finance [[S]]
On the other hand, the importance of modelling division and death is readily apparent for applica- tions such as cell dynamics, ecology, and disease modelling [19 20] In biological systems, organ- ised cell division and death are key control mechanism that work alongside regulatory dynamics, and are thus an essential component that must be modelled to ensure accurate inference results.
Problem statement Population snapshot measurements arise in settings where tracking of individ- uals is either impossible or expensive, such as in single-cell RNA sequencing (scRNA-seg) studies. Longitudinal information on the stochastic trajectories of individuals and their progenitors lost. Instead, one has access only to statistically independent, cross-sectional measurements of the pop- ulation density (e(ar). In other words, we observe the temporal marginals of the process (II) with birth and death. Given the knowledge of pe(az), our goal is to infer the drift vt, the fitness Jr and the diffusivity Dx(x) of the underlying branching diffusion process, or in other words, to invert (2). While our approach applies to the simultaneous inference of these three quantities, Di(x), or the form of its functional dependency on the drift in the case of a Chemical Langevin Equation [I is known. In what follows, we restrict ourselves inferring the drift and the fitness.
assume that
Consider K + 1 ceosonnectional measurements taken from (α) at successive times to <...<< tx = T. Each measurement i consists of N, sampled states (IK.,, < k & N, We assume that N₁/N, estimates the ratio of individuals |pt.1/lpol at all times 4, & & K. With this assumption, we can estimate Pt; up to multiplicative constant Ipol/No using the samples available: pol EN 8(æ xx,()). Practically, the knowledge of this constant term unnecessary because the rescaled density Nopt//Pol satisfies the same equation as Pt so we set it to one.
Pt. No
Related work Approaches based on optimal transport for reconstructing dynamics typically either ignore cellular birth and death entirely [[ [. [4. [i]] handle growth using an unbalanced
relaxation of entropic optimal transport (EOT) and prior knowledge on growth [2. S IO, 12. 22. In particular, while unbalanced EOT can easily be solved computationally using modified Sinkhorn algorithm [23)], it proceeds via a relaxation of the static transport problem and does not readily offer natural dynamical interpretation in the sense of the continuous- dynamics of (2). Recently a significant amount of theoretical progress was made on the closely related Schrödinger bridge problem for branching Brownian motion [I6],, however designing computational approaches for its solution in high dimensional settings remains an open problem. Several recent works aim to learn neural approximations to dynamics with growth, but either assume deterministic dynamics [24]], global fitness linking drift to growth [25] (see also the discussion in Appendix C,), or rely in the end on unbalanced EOT for conditioning flow matching [26]. These methods are therefore subject to limitations rendering them inapplicable for the problem we describe. One work that addresses our problem setting is [27]] which proposes deep learning based method, DeepRUOT. While the algorithm is designed with the same aims as ours, it requires an multi-stage training procedure prone to instability, first relying on flow matching and then on simulation saaued training.
2 Methodology
Unbalanced Probability Flow Inference (UPFI) To solve this inverse problem we build upon the Probability Flow Inference (PFT) approach of [133, 14), in which the Pokker-Planck equation is fitted in the Lagrangian frame of reference. This formulation leverages the fact that the drift-diffusion term (L) can be re-written as a transport term. our setting, such rewriting reads
(x) = [xx (o,(x:)
D.(x)= - Da( V log (e())) + ge (x))
(3)
The phase-space velocity in the transport term now depends on the score V llggtt(:). which is independent of the total mass of the measure since V Vog P: V log (Pr/(pe)). Interestingly, (G) can be solved in the same Lagrangian frame of reference as without growth 9: O. Indeed, the
3

---

solution (ze, m.) of the following system of ODEs
dæa ve (x())) - V D.(x,) (z( V pr)))())) =: ue(xe),
dr
2o ~ Po.
(4)
dm ge (x))
ma (o(
produces a weak solution to (2) [28. Proposition 2.1]. where we have written u, to denote the phase- space velocity field. In other terms, by going to the Lagrangian frame of reference, can trade solving PDE in d dimensions with solving an ODE system in + 1 dimensions. This gain comes cost, since it requires learning the score function V log Pt (x)) beforehand. However, this can efficiently be done offline for high-dimensional datasets because the score is independent of the parameters of the dynamical model in (H)
We remark on the importance that (e) yields a weak solution to (2) solving this system for finitely
many particles yields weighted sum of point masses Pr. approximating the measure Pr. dr but does
not provide the density Pt. The density evolution is governed by d(log Pr (((z)))/dl = 9t (zx) ) - (V
ut) u,)(x;) i.c additional term arises from the divergence of the phase-space velocity. While this is the setting arising in the continuous normalising flows literature [29] and is also used by [22!] in the context of growth, in practice computing the divergence term is expensive as relies on computing the trace of Jacobian.
to be computationally
Our approach, UPFI, consists of two steps: (i) estimating a time-dependent score function (xx)) - V log (z():) from available snapshots, and (ii) learning the coefficients of the ODE system (e) that fit the observed snapshots {Pt,,] 1 < i < Kj. For the first step of the algorithm, we estimate the score with denoising score matching (5%, 30), as it has shown best performance terms of accuracy and scalability on similar problems [[4]]. For the second step, we push the observed samples and their associated mass from time t; to time f following (El). The explicit update equations read
He = Tkit + St, (Va (vu()k,u) - V (.. Du( su (x(k,))) du ma,1 exp (Ia, gu gu(#x,
(5)
(6)
with &e [ta, f. This allows us to construct an inferred marginal distribution Pruceo which reads
(x)) = Σ ma,1.42 8( σ).
We can then optimise jointly over the drift and the growth by minimising discrepancy D(p Pt.11) between the inferred and the true measures. In practice, we minimise the total discrepancy across all snapshots, (u) (x)) (x)) = arg (a)) ΣK= D(pr,, Pt.). While in principle various choices of D could be appropriate, we opt to use the unbalanced Sinkhorn diver- gence [3I]]] which possesses desirable geometric and computational properties [23, B2]: D = Se.y- where E, > 0 specify the entropic regularisation level and soft mass constraint respectively. Most importantly, this choice works directly with discrete measures, allowing us to bypass entirely com- putation of densities.
We summarise the UPFI algorithm in Alg. MI These successive steps are illustrated for a two- dimensional bifurcating process in Fig. (ii)- In Fig. m[vv) we show the difference between the resulting probability flow lines inferred with PFI and with UPFI. Because PFI doesn't it explicitly account for proliferation, infers flow lines which correct for the mass imbalance between the two bifurcating branches by incorrectly connecting them. On the other hand, UPFI recovers the appro- priate flow lines which do not connect the two bifurcating branches. The inference of drift biased by proliferation, as exemplified with PFI in IID, IID, is an instance of wider issue regarding the ability to infer jointly drift of proliferation.
The problem of identifying simultaneously drift and fitness Even in the absence of cellular proliferation, the drift term (x()) is in general not uniquely identifiable [L B3, [4]]. Naturally, these same arguments extend to the simultaneous inference of the drift and growth, and furthermore we can expect to mistake drift for growth and vice versa. While the approach outlined in Algorithm applies to any nonlinear force field, here we aim to gain theoretical insight into the identifiability problem by analytically studying an Crnstein-Uhlenbee which the population density remains Gaussian measure. Although OU processes can not capture bifurcations, they have been successfully used to infer non-bifurcating processes in single-cell RNA- seq data [35]. For these processes, the drift is linear function vx = A,æ + e, and the matrix
(OU) process with quadratic fitness, in

---

Algorithm Unbalanced Probability Flow Inference (UPFI)
Input: + statistically independent snapshots J=- sampled from {pt,) i.. the
solution to (2) at successive times to.t1,...tx .tt.
Estimate score: 8&(t, x) V log pr(a)) using score matching
Initialise: Force (t, x) growth rate æ)), diffuusivity 2t),
while not converged do
for k {0, ................................ K 1} do
(ggg Me )
odesolve (Eq. (e), x0 (Fe me,t. )
[L 1++1])
Ple+ D1 me.44 8(x - x(,(x-1)
Predicted marginal at tx+1 Sinkhorn loss on marginals
4 (3,7(Pt441 P()-
C K-1 Σx --1 bx
Compute total loss Update
Update 0, VoL)
return Ug
Autonomous
Non-autonomous
-25
0.0
2.
2.5
00
2.5
22 68
25
25
50
25
23 00
25
04
25
-25
20
25
-25
0.0
2.5
25
63
25
-25
25
Figure 2: Non-identifiability of growth and drift in the Gaussian case. 8 d-dimensional OU process with quadratic growth. Left: autonomous drift and growth produces the same marginals as non- autonomous system.
A, directly interpretable the gene regulatory network driving the process. We prove in the appendix the following result for OU processes with quadratic fitness. Proposition 2.1 (OU process with quadratic fitness). Consider an OÙ process with fime whose density satisfies
1 2 (,
@)p((x) = - ((Axt e()(()) D.
+ (b. + c. x
(8)
where A, E RRdxd, et, Ce E Rd,, and by E R are generic, D₁ Rdxd is symmetric positive definite, and It E Raxd is symmetric negative semi-definite. If po maN(Ho, En) with mo 0, then P: = m,N(H, Ex) for all t > 0, where
Σ, A,E, + E,A, +2D, E,T,E,
(9)
#. (Ai E,T)) + E.c; + Ct,
(10)
ni m, J4: T.p. + c, + + b, tr(T,
(11)
Using this result, uniquely, and that many solutions fit equally well the data. This made concrete in the follow- ing result.
observe that it is impossible to identify the gene regulatory network
Corollary 2.2. Consider OU process with drift ve(x) Aax and time dependenn fitness as in (8). Let K. E IRnx be an arbitrary matrix. Then there exists E > such that the system is indistinguishable from another OU process with drift vz(x)) = A, + I + eK₁ and time-dependent quadratic fitness.
Importantly, this result holds even if enforce the drift to be autonomous: both symmetric and asymmetric parts of the drift matrix can be confounded with cellular growth, and vice-versa. We show in Fig. 0 simple illustration of this fact the same sequence of distributions can arise from autonomous linear drift with quadratic growth, but also non-autonomous linear drift alone.
Loss function and uniqueness of the minimum The identifiability issue discussed previously calls for the use of a regularisation when optimising for the drift (x)) This regularisation ensures that, even the method is not consistent in general, leads at least to a unique solution. While different regularisers could be used, here opt for the Wasserstein- PeasserttinnFisheer-re Rao energy function
5

---

H2
as a natural choice. The loss function reads, with > 00,1 100000 0:
Ch -= - Se, Spn + A(t, - t f" (110)())) + alg((1)) dpr(a)) dt. (12) 24-1
K
Once again, can gain some insight into the role of the regularisation by studying the linear- quadratic case from Prop. 2 Assuming t4 - 1 = At, we show in the appendix that in the non-entropic limit
(in which case the Sinkhorn divergence becomes the Gaussian Hellinger- Kantorovich distance (GGI) and when At - O with T = K At fixed, the loss tends to a functional function of the parameters defining the fitness and the drift.
which has unique minimum
Theorem 2.3 (Loss function for OU processes with quadratic fitness). Consider the true process as well as the inferred processed to both be OU processes with quadratic filness i.e. (8). Denoting q 2/7, with K At fixed, when At 0 we have L L. where L is the continuous time loss function:
4 = 9-1 2 dt mi ((1)
- -6llx;
- + 2 (1+ )2,)²
(13)
20,00 jan
I
+ 9 (o)
Σ
W., -B.))B)))))))) + a
W
Ra dt
of o.
jo
i.
(14)
where E: = Σ; o., Wi is the eigendecomposition of the covariance Ex, of = 1 + go, for
all i, X, 25, + q We also have
B. B, = =E, Á,) + (A, - - + D.),
-,TT.-f,)E.
V.- ú, ((A: Ad) + ET; -o.-rr))ee [.)) : E.(e: - ē.) (e: (e:-és ē(,,,
(15)
h4 - h, = z4: (T,-f.)u + (a: - - + - b,) (b.
1
2 tr( -f)),,,
and RG is a strongly convex function of the parameters defining the fiteess and the drift. For X > 0, C has unique minimiser t -
-A; (Ai, ēt Bi 2 Fi).
Scalability The unbalanced Sinkhorn divergence Se,y can be computed with a per-iteration plexity of O(B²) and a dimension-indee -innppeonnnndeeedeent sample complexity of O(B-1/2) [[]]] provided the batch size B is sufficiently large. In contrast, denoising score matching has per-iteration complex- ity of O( ((B,), where d denotes the dimensionality. The computational cost of ODE integration can be reduced, at the expense of some accuracy, by taking fewer Euler steps between time points. In practice, we find that two to three steps are typically sufficient. This imposes a mild limitation on the scalability to high dimensions, and we find the UPFI algorithm can handle up moderately high dimensions as illustrated in the results section below. This moderately high-dimensional regime particularly relevant for modelling cellular processes such haematopoietic stem cell differentia- tion, where only a few dozen key transcription factors drive cell fate decisions [M]] Additionally, PCA dimensionality reduction is a routine preprocessing step in single cell analysis workflows [[7]]
3 Results
High-dimensional bifurcating system We first consider a bistable system in Rd where the dy- namics are driven by potential landscape
- VV, V(z) = 0.9/ al x bl3 10 E x?, where attractors are located at a, 6 = [ + e2). We impose a birth rate 6(x) = } (1+ tanh (2))) that individuals closer the positive attractor divide at high rate (see schematic in Fig. (()). For d = {2. 5. 10, 25, 50) we simulate the system and sample population snapshots at T = 5 timepoints (see Fig. (ia) a) for d 10). To reconstruct the dynamics, we apply UPFI (Alg. W. as well as PFI J. E.. For additional comparisons, we consider principled deterministic fitmess- based method (fitness-On ODE, see Section (.3). a method inspired by TIGON [24]], DeepRUOT [27]], optimal transport flow matching (OTFM) [3%] and its unbalanced variant, UOTFM [26]] The imple- mentation of TIGON as described and provided by [[4]] is difficult to use in practice we therefore re-implemented the method with several simplifications which we believe improves its performance,
and d(x))

---

(6)
DeepRUOT
Force and fate probabilras
0-
True
Growthh
Sampled
433
20
01
Figure 3: 10-dimensional bistable system. (a) Population snapshots over time, shown in the first coordinate zn- (b) True and inferred force fields shown in (1), coloured by fate probabili- ties. (c) True and inferred growth rates shown in (rgg x1). (d) Sampled trajectories from true and learned dynamics with growth suppressed (pure drift-diffusion process). The fraction of trajectories terminating in the upper (resp. lower) regions are indicated. Without growth both branches are cquiprobable.
UPFI
PFI
PIINESS-OEE
TIGON+
DEEPRUOT
OTFM
UOTFM
Path energy distance (.)
0.14 + 0.09
1.41 0.16
0.30 = 0.18
0.46 0.12
2,15 0.01
1,16 0.13
0,42 0,13
0.04
1 0.03
1.34 4: 0.06
0.30 : 0.14
0.63 0.16
0.47 = 0.04
1.07 0.11
0.36 0.10
10
0.05
0.04
1.03 0.18
0.29
0.15
0.61 0.06
1.32
1.09 0.19
0.38 0.08
25
0.22 + 0.20
1.51 = 0.10
0.20
= 0.04
0.74 - 0.04
0.14 0.00
1.39 0.09
0.57 = 0.22
0.62
Force errer (casine ()
2
0.07 1: 0.02
0.08 + 0.01
0.34 : 0.04
0.26 0.04
0.44 0.00
0.41 i 0.03
0.44 0.06
10
0.10
+ 0.00
0.12 + 0.00
0.37 + 0.05
0.35
0.02
0.45 0.00
0.45
+ 0.01
0.45 0.0
25
0.06
0.00
0.09
0.00
0.19 1 0.04
0.27 0.01
0.37 0.00
0.60 0.00
0.63 0.022
50
0.06
0.00
0.07
0.00
0.26 0.03
0,25
0.01
0.44 0.00
0.55 0.01
0.52 0.01
Force error (L", 1)
1.17
1. 0.03
1.81 1: 0.09
1.70 :: 0.04
1.62 0.03
2.84 0.00
1.77 0.06
1.68 0.04
0.11
0.18
3.83
0.04
0.00
0.01
10
4.31
+ 0.08
4.76
0.14
5.87
= 0.13
5.89
0.04
6.17 =.000.00
6.09 0.04
5,96 000
25
6.74
1 0.05
6.95
1 0.05
8.94
: 0.20
9.58 0.03
9.88 0.00
10.32 0.02
10.28 0.04
50
9.45
0.06
9.41
0.03
13.37
0.22
13,43 0.04
14.45 0.00
14.60
0.04
14.49
0.01
correlation, f)
0.98
0.01
0.57 0.04
0.94 1 0.01
0.91 0.03
0.51
0.00
0.79 0.01
0.96 1 0.01
0.99
=
0.76
10
0.99
1 0.00
0.65
0.02
0.93
1 0.01
0.84 0.03
0.77 0.00
0.76 0.02
0.95
0.00
0,95
0.70
0.01
50
0.98
0.00
0.63 0.03
0.91 = 0.01
0.80 0.02
0.000000
0.73
0.93
0:9000011
Table I: Numerical evaluation results for d-dimensional bistable system, de (2, 5, 10, 25, 50).
and we thus refer to TIGON+ (see Section C for in-depth discussion). We provide details on flow matching in Section MGG
Fig. Bib) shows in the (zo, plane: (i) the vector field bilities by colour, obtained from the true and reconstructed dynamics. We define the fate probability of state (t,, x) towards an attractor A the frequency of sampled trajectories starting from (t,, x) that terminate closest to that attractor. (see Section for full experiment details)
streamplot, and (ii) the fate proba-
We observe that UPFI recovers vector fields and fate probabilities that almost perfectly resemble the ground truth. On the other hand PFI, unable to account for the presence growth, misattributes the shift in distribution to a spurious drift towards the faster dividing branch. On the other hand, fitmess- ODE and TIGON+ are able to pinpoint the bifurcation but produce binary fate probabilities since they are deterministic. Furthermore, the inferred vector fields are qualitatively different from the ground truth for the same reason. DeepRUOT, which in theory models both growth and stochasticity, unable to recover good dynamics. This is potentially owing instabilities in their training scheme,

---

(add)
(muit.)
(add.)
(maat)
A
LPTS SOAL AAM
IMO -
16100
1000
R401
Figure 4: Simulated regulatory networks. (a) (i) True and inferred vector fields for 7-dimensional bifurcating system, coloured by fate probabilities. (ii) Learned causal graphs using neural graphical model within UPFI (PFT) frameworks with true interactions shown in red. (iii) Precision-recal curve quantification prediction accuracy. (b)(i-iii) Same as (a) for I--dimensional HSC system.
(add.)
(mult.
(add.)
(mult.)
Bifurcating
Path energy distance (L)
1.22 0.38
1.69
0.14
1.91 1.04
4.62 + 1.14
Force error (L J)
12.59 0.39
11.78 0.30
18.96 3.25
16.49 1.62
Pearson fale correlation (1)
0,97 0.00
0.98 0.00
0.66 0.17
0.62 0.06
Haematopoletic
Path energy distance (L)
0.99 0.26
0.98
0.15
1.42 + 0.42
1.86 + 0,30
Force error (L?,, L)
19.58 0.33
18.13 0.25
26.80 3.71
26.84 1.48
Fate prediction Table 2: Numerical evaluation results for CLE systems.
0.07 0.00
0.23 0.03
0.27 0.04
which involves multiple distinct stages and losses [27,, Algorithm 1J compared to the two-step UPFI approach (Alg. M.. Visualising the true and inferred growth rates (Fig. ((c)) we find good agreement. Finally Fig. Bid) we show sampled trajectories over time for the true system and each method, and these further illustrate conclusions from before. Because of space constraints, show sampled trajectories for OTFM and UOTEM in Fig. f in the Appendix.
In Table MI we show numerical evaluation metrics for each method as d varies; here we also com- pare to results for OTFM and UOTFM. At the level of trajectories, we quantify the energy distance [39] between sampled paths from the ground truth and inferred processes treated empirical dis- tributions in L (To, 1], R.). At the level of the force, we measure the reconstruction error on v; in terms of the L and cosine distances. We also report the Pearson correlation between the true and estimated fate probabilities. In all cases, we find that UPFI recovers dynamics with high accuracy, while remaining competitive in terms of runtime, show in Table Ø Finally, in Table N we also present regularisation ablation experiments which show limited effect the growth regularisation term We interpret this as the evidence of implicit regularisation stemming from relatively small neural network sizes.
Simulated gene regulatory networks Next we study more complex systems where dynam- ics arise from chemical reaction networks via the Chemical Langevin Equation (CLE) [[4] 20]: we consider 7-gene bifurcating system (Fig. Baa) and a 11-gene haematopoietic stem cell (HSC) system (Fig. (b) in which the simulated dynamics reflect known biology [[I]] Cells in the bifurcating system branch and randomly proceed to one of two stable states. The HSC system is multifurcating and cell states evolve towards one of four stable states, [Monocyte (Mo), Granulocyte (G), Megakaryocyte (Mk), Erythrocyte (E)). Dynamics both sys-
(Jac)
Bifurcating 0.07
AUPR (1)
0.52
0.09
AUROC (1)
0.79
0.02
0.74
0.03
0.65 0.07
0.63 + 0.08
Haematopoietic
AUPR (1)
0.59 0.04
0.58 1 0.02
0.53 : 0.05
0.51 1 0.04
AUROC (1)
Table Accuracy of directed graph prediction using neural graphical model for CLE systems.

---

velechy
TGON
UPF
ODE
TIGON
Lineage
Figure 5: Monocyte-neutrophi development. (a) Temporal snapshots of monoeyte-neutrope fate determination, shown using SPRING coordinates and celltype annotations from the original publication. (b) Learned vector fields: RNA velocity vector field learned from spliced-unspliced data, all others from temporal snapshots. Cosine distance (relative to RNA velocity field) for Mon, Neu cells shown. (c) Learned growth rates. (d) Fate probabilities empirically estimated from lineage tracing data and predicted from learned dynamics. Pearson correlation (relative to lineage tracing data) shown.
tems are of the form
dx, = (u(X,) v(X,)) dt
1 V + vxX,) dB, VV
(16)
where specify production and degradation rates and V > Os the system volume. In particular, the CLE involves a multiplicative noise model where the diffusivity varies with æ. This is in contrast to the additive noise setting of Fig. 3 where the diffusivity is fixed. For the bifurcating system, specify a birth rate (zx) causing cells in one branch divide faster than others. In the HSC system, we specify b, similarly that cells in the Erythrocyte branch divide faster. We then apply UPFI and PFI: for each inference method, we consider both additive and multiplicative noise models
dx, vo (x:) dt + dB,,
(add.)
dx, (fo(æ) (o(x)) dt + folx) go(x:) d,,..
(mult.)
We show in Fig. Ma b)(i the learned vector field for the bifurcating and HSC systems for each method and model, and highlight fate probabilities towards the fast growing branch. We find PFI in general infers a spurious force towards the fast-growing branch as in the previous example and con- sequently incorrect fate probabilities. On the other hand, UPFI with both additive and multiplicative noise models produce qualitatively similar results that resemble the ground truth. Quantitatively we confirm this in Table and also find that UPFI with the multiplicative noise model mostly outper- forms the additive noise model by small but consistent margin.
Since UPFI and PFI work with neural parameterisation of the force v, this allows for plug-in of interpretable architectures. To illustrate this, we make use of the Neural Graphical Model (NGM) architecture [ from which sparse and unsigned "causal graph" can be extracted and interpreted as the learned regulatory network. For simplicity we use the additive noise model, although use of NGM in the multiplicative noise case is also straightforward. We show in Fig. H (ii) the learned adjacency matrices along with the true interactions highlighted in red, and (iii) the precision-recall curves quantifying the network inference accuracy. both cases UPFI recovers more accurate network than PFI - see Table ( for full evaluation results. This highlights that modelling of growth essential making accurate inferences on dynamics and thus gene interactions [1), as well as the possibility for incorporating custom architectures within the UPFI approach.
Lineage-tracing single-cell RNA-seq dataset We now apply UPFI an experimental time- course of semocyte-nneuroplee development in vitro [M]] across 3 timepoints over 4 days (Fig. B Additional information on the dynamies are available in two forms: (i) lineage tracing, where de- scendants of a common ancestor cell are distinguished by unique molecular labels ("barcodes") and (ii) "RNA velocity" information [[3], which provides partial information on gene expression dynamics from relative abundances of spliced and unspliced RNA transcripts. This information, although limited, independent from the population snapshots and so provides a point of compari- son for the dynamics inferred from snapshots. We preprocess the data following standard pipelines

---

UPF1
PFI
FITNESS-OD -DDE
TIGON+
Celltype
Cosine distance (J)
Monocyte
0.26 0.01
0.33 + 0.02
0.2 + 0.04
0.30 + 0.02
Neutrophil
0.27 0.02
0.31 0.01
0.44 + 0.02
0.26 0.01
Monocyte
0.26 1: 0.01
0.32 1 0.01
0.46 1 0.04
0.28 1 0.02
Neutrophil
0.24 0.01
0.01
Monocyte
0,27 0.01
0.36 1 0.01
0.41 - 0.03
0.33 0.02
Neutrophil
0.31 0.01
0.33 1 0.01
0.38 1 0.01
0.27
0.01
Metric
Fate correlation (1)
Kendall's
0.22 0.03
0.07 0.03
-0.2 1 0.01
0.19 119-1000 0.04
Pearson's is
0.26
0.04
0.09
+ 0.04
-0.02 + 0.01
0.19 0.05
Spearman's
0.27 0.04
0.08 0.04
-0.03 0.01
0.20 0.05
Table 4: Evaluation results lineage tracing dataset [.]].
P | and use the leading 10 principal components represent cell state. We apply UPFI with an additive noise model as well as several other methods to this data and show results for the learned drift alongside the RNA velocity field Fig. BX)).. Since the magnitude of RNA velocity estimates are unreliable [[5] we measure the cosine similarity between the RNA velocity field and the inferred force from snapshots (Table E). We restrict this analysis the monocyte and neutrophil cell clusters, corresponding to cells that are already committed to their respective lineages. We find that UPFI and TIGON++ perform comparably.
Inferred growth rates (Fig. [(c)) show that UPFI predicts that cells in the earlier progenitor states have higher division rate, which is consistent with the expected biology [M] 244] On the other hand, the fitness-ODE yields the opposite while TIGON++ predicts a very low growth rate for most cells. The availability of lineage tracing data across timepoints provide us with a ground truth for measuring accuracy of fate predictions (Fig. H): fate probabilities predicted by UPFI agree with the lineage tracing data more closely than other methods. In particular, PFI and fitness-ODE exhibit bias towards monocyte lineage. While TIGON+ is able to predict a mixture of both fates, as in the bistable system produces only deterministic outcomes. Similarly, quantitative results (Table show that UPFI performs best in several different fate correlation metrics.
Finally, we note that this dataset was also studied in [ 27] which introduced TIGON and Deep- RUOT respectively. However, in these papers all analyses of this dataset were carried out using woo-dimensional embeddings computed by force-directed layout algorithm [T]] The setting we
consider (1 (10-diz (O-dimensional PCA embeddings) is more challenging, and also more realistic scenario since nonlinear dimensionality reduction methods are known introduce distortions and are thus unsuitable for downstream quantitative analysis [[6]].
4 Discussion
We study the problem of inferring stochastic dynamics from cross-sectional snapshots popula- tion subject to drift, diffusion well birth and death. To this end propose UPFI, a method based on the probability flow characterisation of the Fokker-Planck equation. Our approach is flexi- ble enough to handle generic noise models as well as growth, and crucially applies to the practically important scenario where growth rates are unknown and must be learned together with the drift. We provide theoretical treatment in the qunear-quadratic case, demonstrating the nature of the non- identifiability that arises from growth. We conclude by showcasing our method's 's efficacy using both simulated and experimental single cell data.
There are overall several lessons to summarised from our work, which may be of relevance more generally to the problems of reconstructing dynamics from snapshots. As shown by our theoretical analysis and supported by results, in the fully general setting where the drift and growth components non-autonomous, there is little reason to expect to accurately separate drift from growth effects. In the absence of additional information to aid with inference, we advocate to impose an autonomous drift and possibly also autonomous growth field. This is consistent with cell cell-autonomous models of biological dynamics, i.e. ignoring cell-cell interactions or temporally varying environments. Subject to the modelling limitations incurred by these assumptions, we argue that this serves strong inductive bias that can help inference.
10

---

References
[1] Weinneh, A Kodriiuuzz-Fraticlli, FD Camargo, AM Klein, Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science 367, eaaw3381 (2020).
[2] Schiebinger, et al., Optimal-transport analysis of single-cell gene expression identifies devel- opmental trajectories in reprogramming. Cell 176, 928 928-943 (2019).
[3] MA Coomer, I. Ham, MP Stumpf, Noise distorts the epigenetic landscape and shapes cell-fate decisions. Cell Systems 13, 83-102 (2022).
[4] B Bonham-Carter G Schiebinger, Cellular proliferation biases clonal lineage tracing and tra- jectory inference. Bioinformatics 40, btae483 (2024).
[5] Y Song, S Ermon, Generative modeling by estimating gradients of the data distribution. Ad- vances in neural information processing systems 32 (2019).
[6] Song, et al., Score-based generative modeling through stochastic differential equations. arXiv preprint 2rXiv:2011.13456 (2020).
[7] EB Gunnarsson, J Foo, K Leder, Statistical inference of the rates of cell proliferation and phenotypic switching in cancer. Journal of Theoretical Biology 568, 111497 (2023).
[8] DR Sisan, M Halter, JB Hubbard, AL Plant, Predicting rates of cell state change caused by stochastic fluctuations using a data-driven landscape model. Proceedings of the National Academy of Sciences 109, 19262-19267 (2012).
[9] A Tong, Huang, G Wolf, D Van Dijk, Krishnaswamy, Trajectorynet: A dynamic opti- mal transport network for modeling cellular dynamics in International conference on machine learning. (PMLR), pp. 9526-9536 (2020).
[10] GHT Yeo, SD Saksena, DK Gifford, Generative modeling of single-cell time series with pre- scient enables prediction of cell trajectories with interventions. Nature communications 12, 3222 (2021).
[11] Hashimoto, D Gifford, T Jaakkola, Learning population-Ievel diffusions with generative mns in International Conference on Machine Learning. (PMLR), pp. 2417-2426 (2016).
[12] E Ventre, et al., Trajectory inference for branching sde model of cell differentiation. arXiv preprint arXiv.2307.07687 (2023).
[13] V Chardès, S Maddu, MJ Shelley, Stochastic force inference via density estimation. arXiv preprint arXiv:2310.02366 (2023).
[14] S Maddu, V Chardès, MJ Shelley, Inferring biological processes with intrinsic noise from cross-sectional data. arXiv preprint arXiv:2410.0750/ (2025).
[15] MS Albergo, NM Boffi, E Vanden-Eijnden, Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint 2003v02303.08877 (2023).
[16] A Baradat, H Lavenant, Regularized unbalanced optimal transport entropy minimization with respect to branching brownian motion. arXiv preprint arXiv:2111.01666 (2021).
[17] R Grima, P Thomas, AV Straube, How accurate are the nonlinear chemical fokker-planck and chemical langevin equations? The Journal of chemical physics 135 (2011).
[18] JC Cox, JE Ingersoll, SA Ross, et al., A theory of the term structure of interest rates. Econo- metrica 53,
(1985).
[19] DS Fischer, et al., Inferring population dynamics from single-cell ma--eeuencingg time series data. Nature biotechnology 37, 461-468 (2019).
[20] M Pariset, YP Hsieh, C Bunne, A Krause, De Bortoli, Unbalanced diffusion schr\" odinger bridge. arXiv preprint arXiv:2306.09000 09099 (2023).

---

[21] C Bunne, L Papaxanthos, A Krause, M Cuturi, Proximal optimal transport modeling of popu- lation dynamics in International Conference on Artificial Intelligence and Statistics. (PMLR), pp. 651 1-6528 (2022).
[22] L Chizat, S Zhang, M Heitz, G Schiebinger. Trajectory inference via mean-field langevin in path space. Advances in Neural Information Prucessing Systems 35, 16731-16742 (2022).
[23] L. Chizat, G Peyré, B Schmitzer, FX Vialard, Scaling algorithms for unbalanced optimal trans- port problems. Mathematics of computation 87, 2563-26099 (2018).
[24] Y Sha, Y Qiu, PZouu, Q Nie, Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence 6, 25-39 (2024).
[25] K Neklyudov, R Brekelmans, D Severo, A Makhzani, Action matching: Learning stochastic dynamics from samples in International conference on machine learning. (PMLR), pp. 25858- 25889 (2023).
[26] L Eyring, et al., Unbalancedness in neural monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100 (2023).
[27] Z Zhang, T Li, P Zhou, Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. arXiv preprint arXiv:2410.00844 (2024).
[28] A Chertock, A practical guide to deterministic particle methods in Handbook of numerical analysis. (Elsevier) Vol. 18, pp. 177-202 (2017).
[29] W Grathwohl, RT Chen, Bettencourt, I Sutskever, D Duvenaud, Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367 (2018).
[30] Vincent, A connection between score matching and denoising autoencoders. Neural compu- tation 23, 1661-1674 (2011).
[31] T Séjourné, Feydy. FX Vialard, A Trouvé, G Peyré, Sinkhorn divergences for unbalanced optimal transport. arXiv preprint 1910.1299.2 (2019).
[32] L Chizat, G Peyré, Schmitzer, FX Vialard, An interpolating distance between optimal trans- port and fisher-rao metrics. Foundations of Computational Mathematics 18, 1 1-44 (2018).
[33] G Schiebinger, Reconstructing developmental landscapes and trajectories from single-cell data. Current Opinion in Systems Biology 27, 100351 (2021).
[34] C Weinreb, S Wolock, BK Tusi, M Socolovsky, AM Klein, Fundamental limits on dynamic inference from single-cell snapshots. Proceedings of the National Academy of Sciences 115, E2467-E2476 (2018).
[35] L Wang, et al., Dictys: dynamic gene regulatory network dissects developmental continuum with single-cell multiomics. Nature Methods 20, 1368 1378 (2023).
[36] Janati, Ph.D. thesis (Institut Polytechnique de Paris) (2021).
[37] MD Luecken, FJ Theis, Current best practices in single-cell raa-se analysis: a tutorial. Molec- ular systems biology 15, c8746 (2019).
[38] A Tong, et al., Improving and generalizing flow- generative models with minibatch opti- mal transport. arXiv preprint arXiv:2302.00482 (2023).
[39] ML Rizzo, GJ Székely, Energy distance. Wiley interdisciplinary reviews: Computational statis- ties 8, 27-38 (2016).
[40] A Pratapa, AP Jalihal, JN Law, A Bharadwaj, T Murali, Benchmarking algorithms for gene regulatory network inference from single-cell transcriptomic data. Nature methods 17, 147- 154 (2020).
[41] J Krumsiek, C Marr, T Schroeder, FJ Theis, Hierarchical differentiation of myeloid progenitors is encoded in the transcription factor network. PloS one 6, c22649 (2011).
12

---

[42] A Bellot, K Branson, M van der Schaar, Neural graphical modelling in continuous time: con- sistency guarantees and algorithms. arXiv preprint arXiv:2105.02522 (2021).
[43] Qiu, et al., Mapping transcriptomic vector fields of single cells. Cell 185, 690-711 (2022).
[44] FA Wolf, P Angerer, FJ Theis, Scanpy: large-scale single-cell gene expression data analysis. Genome biology 19, 1-5 (2018).
[45] SC Zheng, G Stein-OBrien, L. Boukas, LA Goff, KD Hansen, Pumping the brakes velocity by understanding and interpreting ra velocity estimates. Genome biology 24, 246 (2023).
[46] Chari, Pachter, The specious of single-cell genomics. PLOS Computational Biology 19, e1011288 (2023).
[47] DL Lukes, DL Lukes, Differential equations: classical to controlled, Mathematics in science and engineering. (Academic Press, New York), (1982).
[48] S Särkkä, A Solin, Applied Stochastic Differential Equations, Institute of Mathematical Statis- ties Textbooks. (Cambridge University Press, Cambridge), (2019).
[49] JE Potter, Massachusetts Institute of Technology, United States, A matrix equation arising in statistical filter theory, NASA CR- :270. (National Acronautics and Space Administration, Washington, D.C.), (1965).
[50] GA Seber, AJ Lee, Linear Regression Analysis. (John Wiley Sons), (2003).
[51] Dacorogna, Direct Methods in the Calculus of Variations, Applied Mathematical Sciences eds. F John, JE Marsden, L Sirovich. (Springer, Berlin, Heidelberg) Vol. 78, (1989).
[52] J Feydy, al., Interpolating between optimal transport and mmd using sinkhorn divergences The 22nd International Conference on Artificial Intelligence and Statistics. pp. 2681-2690 (2019).
[53] TO Gallouët, Monsaingeon, A jko splitting scheme for kantorovich-fisher--e gradient flows. SIAM Journal on Mathematical Analysis 49, 1100-1130 (2017).
13

---

A General results for Ornstein-Uhlenbeck processes
A.1 OU process with afline fitness
Proposition A.I. Consider an OU process with affine fitness, whose density satisfies
P -VA--- + et)pt - D.Vpt) + (be + c/ x)pt.
(17)
where A, E Raxdd Ct, Ct E Rd, and by R are general, D, Redxd is symmetric positive definite. If Po moNoHo, En), then Pt m Ex) for all t 3 0, where
Σ, A,E, + E,A] 2Di,
(18)
i Alt + E,G + et:
(19)
m cj pu bb.
(20)
my
Proof.
We denote (k) the Fourier transform of pr(zz). We have the following identities
F(V (A,x + e()p.) = =k' A.V + ie, kpr,
(21)
F(V (D.Vp() - -k.DD
(22)
Such that the Fourier transform of the PDE reads
eipe = (k A, + ic. VV - k Dekpt + (b. ie, k)
(23)
We define the characteristics as the solution of
dk
A! k. ics.
(24)
ds
such that
(ks)) ds
dk.
T (k.) + (ko)
(25)
ds
- D.kaD. + (b, 1 ie_ kg)
(26)
Denote ((t, to) the state transition matrix solution to the homogeneous ODE (2E)) with initial condi- tion fo. It exists because the coefficient are locally integrable on I. | The state transition matrix has the following properties [[8]:
( u) = A D(s u) d u) ds
((ss A!. 1 (t, 0 = Φ(t, s). s1 = (s, t), (s. I. Denote (ss,0) = 4, for simplicity, we have, using the initial condition ky:
(27)
du
(28)
k. 2.k0 - i [ (s u)c,du.
(29)
We introduce T. 00 (sss, u)cydu. such that we have the Fourier transform along the character- istics reads, with an initial Gaussian distribution
$ Du - 2ikg S" $ Duuu,du
(k.) exp
iku æn
ki
Enko
.. Da Bu buddu - i ((e e T Sf² e @u
(30)
We introduce the following quantities
V, [ + D., du,
S( 4 D.,ddd
q.- e eT u. = jo
(31)
7o [
a. Dr,dDidd, ws S e Vw.
(32)
14

---

Inverting the equation for characteristics we have ky = - k, + i So ((s, 4) = Φ K2 i We can directly see that the Fourier transform is quadratic in k, such that it is the Fourier transform Gaussian measure. By considering one after the other the terms O(k²), O and O(1), we can find analytical expressions for the covariance, the mean and the mass of the Gaussian measure as a function of the state transition matrix. They read:
Σ. = 4 T E + 2Φ, 1. 22
(33)
M4 1 T (nO + 2u, +ga) E.
(34)
log m, = 8 ,, 2 8 J. + wis + + S dy
(35)
Taking derivatives of these quantities and using the properties of the state transition matrix we re- cover the expected results.
L
A.2 OU process with quadratic fitness
Proposition 2.1 (OU process with quadratic fitness). Consider an OU process with quadratic fiteess, whose density satisfies
Cepe(x)= = -
D.V + (b. + c, px))
(8)
2
where A, E Rdxd es, Cr E Rd, and R are generic, D, Iddd is symmetric positive definite, and T3 E Raxd is symmetric negative semi-definite. If Po = moN(Po, En) with mo > 0, then Pt mWN(ps Ex) for all t 0, where
Σ. A,E, + E,A, + 2D, E,,
(9)
Pr (A, E,T.), Σ,c. + et,
(10)
m,
m 211 T Tip. + c. Pe + b, tr(T,E,).
(11)
Proof. Let's use the ansatz pr(x) uz(x) exp [x / where Ge is symmetrie. The equation for 2 now reads
6i + uux dG.
dt
+ 1 - C D.Vu)
- ((A
(r.D)r(G G. + x I + G.D.G. - AT))
u.
(36)
The quadratic term vanishes when G, verifies for all the equation x' TTT,/2 + G,D.G. A, G,) = f G,a. This verified for Ge satisfying the matrix Riccati equation
dG. dl
+ 2G.DGG, A, G.A,
(37)
Provided the coefficients locally integrable Rt, that I, & 0 and D > 0. if Go & 0, then there exists a unique solution G2 on I for (G77), and G; & 0 for all > [_.]. We consider this unique solution with Go Then, the equation for u4 becomes affine in growth
0.u- = -V.
((Ar te,
-DD, - - D.Vu,)+ (b. L.) + c, x - e G. ue- (38)
We redefine À, A,- 3y = by - tr(G. and ट5. C: - G.e, The equation now reads
Oc4 - ((A, + e:)u: Dj D.vu)) + (b. + टे x uc.
(39)
We can apply the results of Prop. A. Since uo( = po (w) is a Gaussian measure, it remains Gaussian, and we denote its covariance Σ, its mean H; and its mass my. It follows that pe(x) is
15

---

Gaussian measure for all t > 0 since -
G. is positive definite as the sum of positive definite and
positive semi-defiitte terms. The covariance of this Gaussian measure is then Σ, (É,-) --,))+
We have the following fact for any differentiable matrix valued function B, invertible for all t:
dB,
dB dl
B,
B
(40)
di
Applying this to the covariance, we find:
C Et
dEi
E.
=EE,, (A,E, Σ,A + +2D.) Σ E.
dG.
E.
(41)
dl
dl
dl
Substituting A. -2D,G, and G. T; 2G,DGG A. - G.A, we have dΣ,
A,E, E,A, + 2D, + Eg
(42)
dt
Completing the square in high high-dimensions allows us to write the mean of the overall process as
Pt
Σ G.)
E E Tr +I+ D.G.)
(43)
Using Prop. A. we have
dx
D:
G. E, dG.
dt pr (I+ E,Gi + Σ, + 2.)
(44)
dt
dt
We replace A, A, 2D, and use the Riccati equation in the terms multiplying jú, In the other terms we replace č = Ci - G.e, and e, = é. We find:
diu (A, ETD E,c, ec.
(45)
We can derive the fitness following the same approach and using lengthy simplifications. We only need to gather together the terms remaining after completing the square, as well as adjust for the change in covariance in the normalisation factor. However, this result is more easily obtained by using a result for the mean over a Gaussian probability distribution of quadratic form [50, Theorem 1.5]. We find that
1 XT b, + cJ + tr(T,E,) + 214 Tip vi = Bb. + c X +
(46)
2
where the expectation is understood to be taken with respect to Nxxx,, E.).
Corollary 2.2. Consider OU process with drift ve(x) A,æ and time dependent fitness as in (8). Let K. E R2x" be an arbitrary matrix. Then there exists e > such that the system is indistinguishable from another OU process with drift vi(x)) = A, I cK₁ and time deee-depnnennt quadratic fitness.
Proof. Let's define the following growth parameters: b. b. tr((I + K.)) " 14, and r. r, (II - Σx E, (I+ K.)). With these and the drift (ær), the solution to the system of ODE above is unchanged. We take as the largest value such that T, is negative definite. This value is strictly larger than zero thanks to the identity. This derivation still holds the drift autonomous: if vx( Az, then there for any K, there exists e > such that the system is indistinguishable from another OU process with 6, (x) A + I + eK and dimedeepnndett quadratie fitness.
-1
B Continuous-time loss for OU processes with quadratic fitness
Theorem 2.3 (Loss function for OU processes with quadratic fitness). Consider the true process as well as the inferred processed to both be OU processes with quadratic fitness, i.e. (.). Denoting
16

---

q/2//, with K At fixed, when At 0 we have At L. where L is the continuous time loss function:
C q-' [/² dt mi
2
8,1xx,1
2 (he -Pr,²2
(13)
Wi, - (B))a))) + x {
+
Σ tj
Boath + of Rng
RG di
(14)
where E; Li i,t W.F is the eigendecomposition of the covariance E, oz = + go? for
all i, X, = -2E; + -9- We also have
B. B, E,(A, AD) + (Ai - A,)E, E, (TT,-f,)E, + Di),
Vi 0, ((A A.) + +T; - -T:)) Ei(e: - ē.) (e: - -ei),
(15)
h4 Tu 2/14 (T,-f.) + (c+ - 2 + (be - b.) + マ Lr ir((T: -f.)..))
and R₃ is strongly convex function of the parameters defining the fitness and the drift. For X > has unique minimiser t (Ái, et 82 G, ri).
Proof. When the entropie regularisation Gaussian measures reduces to the Gaussian- Meeasiinn-ellinger--a Kantorovich Kantonng distance Su,7 GHKy. Be- tween the inferred and the true process time & reads:
0, the unbalanced Sinkhorn divergence between two
GHK (m,,N(HH,, E.,), Noe E..)
mg del J
1 TT, ( (e:))
=2q mr + me: -2
exp
(e:
(47)
2
where we have
X = Er, + Ei, + 1 J= - Ei. (I Ex, Ex. 2 2,,))
1-11 )2/2),
(48)
with
Σ E L + I, Et,g = Ee,, +
(49)
Without loss of generality consider the case At, and we perform Taylor expansion in At of the GHK loss. Because the final expansion will be of order At2, only the terms of order At of the covariances play role.
The hard part in this expansion the expansion of detJ. We denote Est = Eg + ALA and Ea, 5-1
E.t En + Ā and Eo.g I + q Eg. For this, we need to compute ΣxI Ar.g second order in At. We have that
Ar. up to
EA.D Ac. = (Eo + AtA) /
- AΣ
+ +2
At
AΣ A
(50)
EoEryy A1(AE 1.00 1A=11
(51)
+ A6 En. AΣT! AE
AΣT1
0.9 0.9.
AΣ
(52)
We also have
Σ A
Ex. (E 100---11 At 12g
0.9 + 9² At 20 AE
AΣ A.
0.
(Eo + Att)
(53)
= -ETTE + de 0.4 DD.
0..
AΣ
0.9 En) 0.9
(54)
+ At AA ÅΣT! AE Eo
Ep.
(55)
We can now gather the 0(1), O ( separately. We denote them respectively M. Hi, H2355 We define the operation 'tt.' as the 'tilde transpose' operation, which is applied
17

---

to the term directly to its left. We therefore have
M-
En
0,0
(56)
H Eg Σ 2 A
0.0
9
+ tt.
(57)
H2315 = -oo 2
D0
Σ 0,y A A A) + tt.
(58)
0.4
+ 0,1
2 Ā + Σn
2
A
2 AEO Eo
(59)
0,7
0,
- qAE
0.9
10.
Σn + tt.
(60)
Using the Woodbury formula we also have Eg 2
1 7
2ogg). We have that Eo and 2
commute, and that M Eg 20.g to the following Lyapunov equation
q (1 2. We introduce S(Y) the unique solution X
2 1/2 X + -M
(61)
This solution is expressed in terms of an integral (and is linear in Y), ie.
S(Y)
[- e M - Ye 2 dt
e-
(62)
As a result, we have
1/2 Aig Ar. Est)
ESID S-1
2 + AIS(H)) + At (S(22ss)))) - S(S(H))
(63)
We have I - 9M -ogg such that
det 1- EEs 2r.9 Ar. )11)
1/2 )(²) = det De. (I AtL + AFG),
(64)
where = - o Hi) ((H)) and G -q-o. (S( (5(H2 - S(S(H))))) We then have the fol-
lowing expansion at second order in At
det (I + AIL + APG) = + AtrLL + At2
2 ( 2lr
(65)
So we need to compute LrL, trG, L? We have, using the fact that Eo.g commutes with 1/2
2 E M2 HI Hee-M ( - 2 1/2 EM,
trL qtr((2 -qt((Eogg
(66)
! 1ir(E 20 H1) --2 (Ir(A + A) qlr((A+ + Ā) E -0)
(67)
2
2 tr (tr (A+ A)).
(68)
Computation of trG
Using the same trick as for trL and standard properties of the trace we have
tr(-9Zo.g
(2345)) ((
2 (f=tr(AE AΣ o. + tt.) qtr(En AE 4)
1
(69)
To compute tr( S(SHH, ()) we denote 1/2 = E, uw the cigedeccooponinio 4/2 We also denote W. = w.w. Therefore we have
S(Hi) Σ i.j
W.H.W.
(70)
Us Uj
such that
(1- 2 uz W. (u.ā + unA) We (ueA+ + wsA) W,.
((SH))))))) J.
Σ
(m: un (u; + u+) uu (
uy)
(71)
18

---

Taking the trace we have
(1 - u1 - un1) ir
tr ((
A+
(.aa). (72)
20;(u; (ui + uz)²
2
Additionally, we have that
tr (Wi(uAA u, u, (u, (w -uuyir (Wi(A- (oygrr(W(44
(73)
(A -A) + (ui + uz) (W.AW,A)
(74)
As result we have
trG 2" ( AΣ 20. tt.)
9 Σ
(1 u,g) uj tr (W,(A A) A))
2
(u, + u;)²
(75)
Computation of Lr Similarly we have
S0, Σ
(1 (1-ujj) u;9) u1 + uz)
W; (u:A+ + ( Wj.
(76)
Using the same approach for trG find
u.g) -ug)),uu ir (Wi(A- ( - A))
(((Eo,S(HH1)) =
(77)
(w: + uy)²
1
+ tr 0.0 A A).
(78)
As result have
(1 uzq) - u.q))uu tr (Wi(A- ÂAW5(A- -A))
t tr AE A) +0²) Σ
(u: u2)²
(79)
Computation of det( (ESt.GEEEs.)) 1/2
We have
Eac. D.L. Ei,g
+ E) o2 9 212
(80)
Taking the determinant and Taylor expanding we have
det Esi. E.i.g = det (2u.g)
(81)
x 1 + [+ A + A)) + 9²
At"
((LP(2 (A A)) inEo.AAEo A + uu.))
2
(82)
Expanding the square root we find
dct (Esi, E 1/2 = det (E.gg)
(83)
x
( 1 ? Ir( 1 + A)) + q²
Al
At2
8
(r2
( (A+ (A+ A) - 2ir(E AIrEEAAA551
AΣ
0.4
A + u1.))
(84)
19

---

Computation of det J
Going back det J, are left with
del J
1 +
Af Ir D 0.0
(A+
A)) + 9² At2
(17(20)) + A)) A u..))
(85)
At 2 (t'L-trL'' tr + + 2tr
x (1+Attr +
(86)
Using the result for trL we see that the terms in At cancel, and the final expansion is of order At?. As result, after simplifications we are left with
det J (1 11- 4
Al2 (-trI² + 211G))
E
AE 0,
A+ tt.)
(87)
Simplifying further we find that
det J =1 AP2q 9 (E ( -À))
(88)
2
0.9
+ uj (1 u: u1 g(uu (w))++ ir (Wa(A- A -A){
(89)
(us + uz)²
This can be simplified
y tr
E 0.9
(A
A
A) 0,91 (A- -\))
(90)
2
u5(1 (1 - u1 (u, uz)2
q(uj + wi)) tr - (W:(A-A)W,(A- -1))
+
(91)
4
Σ
(1 q14)(1 - qua))
9
q 22 2 (u) tr (Wi(A A)W,(A- -A))
(92)
(n + u2)²
2
which leaves us with the following simplification
At
u: (1 - u.g)(1 ugg)
(r (W:(A- (W((A-A)),(AA- (O,((AAi)) )
det J = 1+
(93)
2
(u; + uz)
Using Eo Σ, o W. we can compute the eigenvalue decomposition function of o; Using
the notation o. + go? we have
u,(1 u1 u,g)
of ja (o o + to? 2 o
(94)
(w: + 12
Finally, at first non-zero order in At the determinant reads
Ata
det J - -
Σ
o o lr (W.(A- (A A))
(95)
2
( o
Taylor expansion of the GHK term
We now expand the masses and mean at first order in At, MA = mn + Atmoh, MAi = mg + Atmph MAt
P4 + Atv. Pss no + Atě. We have at first non zero order in At
At
exp - - (ust
T T -ass)
2 (v 0) T X (v 8).
(96)
20

---

We denote X. = + Expanding the remaining terms, the zeroth and first order terms cancel,
leading to the expansion
GHK (MaN MAN.
MaN Eail)
(97)
-mog--aa (0- X (v- 8) + 1 +z(h =2(h-hi) hi)²)
(98)
o0 tr WW W -Aa)
q 1.
(0 + o
(99)
Integrating over all snapshots, the continuous time loss reads
1 2 hi)2) ((1-- X (v- 8) +
==
mgg
oj t (Wi(A-A)W, (A A)) )at
(gg + o
S f (p(x) (2) + alg((=))²)) dxdt.
(100)
The first order expansions A, h, and e are obtained directly using the ODEs in Prop. 2. giving the final result.
Let's denote 0g R-N the vector of the N parameters defining the growth and the drift at time t. We study the functional [[O.] which reads
[ L (F)(O.) + (R.(O.))dL
(101)
where F' is the part of the integrand coming from the expansion of the Coussian-Hellingerr Kantorovich, and Rt is the regularisation.
Let's take [0, TJ. Because 0 + (0) is the composition of convex functions and of affine maps, 0 - (0) is also convex. Let's show that 9 - (e) a strongly convex function.
Lemma B.1. Let w : (0,x) RN x IR - Hd be function continuous in æ and linear in 0. Let f Rd + I be a continuous, strictly convex function and p a continuous function from R to 10, +oo We then have that
8:0ERN - p(a)s(((O x))de
(102)
strictly convex.
Proof. Let's take 0, 02 and ej0, 1[. By linearity, we have Vz: w(a01 + - a)02, æ) aww(01 + (1 - a) (62.
(103)
Because 8 - f(=((( x)) is a the composition convex function and of a linear map. it is convex for all Then, Var,
f(w(af1 + (1-a)02, m)) (1
& (1
x)).
(104)
Because w is linear in 0, each its component in Rd is multivariate polynomial function of 0 of degree one. By uniqueness of the coefficient of polynomial functions, because 0 * O2 there exists To such that w(O,, my) * w (022, zg). By continuity, this is also true on an open set U IR" centred in Zo. Therefore, since is strictly convex, the inequality is strict for all U. i.e.:
f( (w(af); + a1-2)62,220 x)) (aw x)) (1-a) a) w (2. x)) < af( x)) + (1 a) x))
(105)
By multiplying by p(x), which is strictly positive, and by integrating, we keep the inequality strict and we find
f(a0) + (1 - a)02)= < al(O)) + (1 a))(()))
(106)
proving that l is strictly convex.
21

---

Hidden layers
Batch size
Learning rate
Iterations
Bistable (d & 10)
64, 64, 64
256
0.003
25,000
Bistable (d > 10)
256, 256, 256
256
0.001
100,000
Bifurcating CLE
64, 64, 64
256
0.01
10,000
Haematopoietic CLE
128, 128, 128
256
0.01
10,000
Lineage tracing
128, 128, 128
256
0.01
10,000
Table 5: Hyperparameter settings: score networks
Hidden (force)
Hidden (growth)
Batch
LR
Iterations
Bistable (d & 10)
64, 64,
64, 64
256
0.003
5,000
seo details
5
UPFT
Ristable (d 10)
256, 256, 256
256, 256, 256
256
0.001
25,000
details
5
Bif CLE
64, 64, 64
64
256
0.003
10,000
0.001
5
CLE
128, 128
128
0.003
0.001
Lineage
128, 128, 128
128, 128, 128
256
0.001
10,000
0.001
25
Bistable (d 10)
64, 64, 64
256
0.003
5,000
details 8ec details
Bistable (d > 10)
256, 256, 256
256
0.001
25,000
PFI
Bif CLE
64. 64,
256
0.003
0.001
Huem CLE
128, 128, 128
256
0.003
10,000
0.001
Lineage
128,
0.001
10,000
0
ODE Bistable (d < 10)
Coupled to growth
64, 64, 64 256, 256, 256
256
0.003
5,000
see details details
5
Bistable (d 10)
Coupled growth
256
0.001
25,000
5
Lineage
Coupled growth
128, 128, 128
256
0.001
10,000
0.001
0.1
25
TIGON+
Bistable (d 10)
64, 64, 64 256, 256, 256
64, 64, 64
256
0.003
5,000
details
5
Bistable (d 10)
256, 256, 256
256
0.001
25,000
see details
5
Lineage tracingg
128, 128
128, 128, 128 256
0.001
10,000
0.001
0.1
Table 6: Hyperparameter settings: dynamics. For bistable system certain hyperparameter'' chosen
as function of d, see Section
for details.
We have pr(x:) > 0 for all because it is Gaussian density. The drift and the fitness being linear in 0 this shows that
0 R(()) = pc(x)) (Ive(=))1 + (l)( do
(107)
is a strictly convex function 0. Additionally, it is multivariate polynomial function of 0 of degree two, so its Hessian is a definite positive and constant, proving that 0 - (2( is strongly convex, and is Ft + Rr. Along with the fact that both (t. 0) (0) and (t, 0) Rs (0) are continuous in [0, 2' x I this is enough to ensure the existence and uniqueness of a minimum - 07 for the functional [[O,] [А]]]
C Implementation and experiment details
We implement Alg. M using PyTorch and employ the GeomLoss package [52] for computation of the unbalanced Sinkhorn divergence. All model training was carried out using NVIDIA L40S GPU. Code is available at http:///gtt.tb..oooo////ee////1
C.1 Score matching
While principle any score matching approach to learn 8, (a: V (z) can be used within Alg. III in practice we employ denoising score matching [30] within the noise-conditional score network framework introduced [[]] For K + snapshots taken at times (t,) we parameterise the time-dependent score using multilayer perceptron (MLP) sy(t, x) = NN&(d 2, d) where is the dimension and 17 is the noise level, and train using the algorithm described in [[]. Section 4.2]. While a range of noise levels To
< πι used for training the score, subsequently for training the probability flow we use the smallest noise scale r) = 7oo representing our final estimate of the score.
In what follows, we use a default of L = 5 noise levels logarithmically spaced between (exp(-2) 1). Score networks (x) are parameterised using MLPs with ReLU activations. In all case, score training was carried out using noise-conditional denoising score matching [5, 60] using the AdamW optimiser with the hyperparameter choices listed in Table B.
22

---

C.2 Training: UPFI and PFI
Additive noise models For UPFI, in all cases we parameterise an autonomous force (y(x) using a MLP. This is because, in all simulated systems, the true force is also autonomous. In the exper- imental lineage tracing dataset we reason that an autonomous force would be consistent with the biologically motivated model of Waddington's landscape [[=]]. We parameterise separately the force no (1)))) and growth go go(x:) a:) using MLPs with ReLU activations. We train UPFI models using the AdamW W optimiser, our architecture and hyperparameter choices are listed in Table f
For PFI, motivated by the observations in the Gaussian case we reason that an autonomous force is insufficient to fit the data if growth is not accounted for. We therefore employ a non-artonomons force, parameterising vo(t, x) NNo + 1, (t, æ) with ReLU activations. We train PFI models using the Adam W optimiser, our architecture and hyperparameter choices are listed in Table B
Multiplicative noise model For the multiplicative noise models considered in Fig. 5. we pa- rameterise an autonomous force in two components, corresponding to production and degradation terms in the model (IG)). Specifically, we let
fo(æ) NNo(d, (x)), go(x) (N)(ddd)(az))
and the resulting force is (fo - ( (z)) To constrain the output of both networks to be non-negative, opt for Softplus activation the final layer of outputs. For all other layers the ReLU activation as a default choice. Architectures for (f.g) g) and all other hyperparameter choices are as given in Table B
Scaling regularisation with dimension for bifureating system For the bifurcating system of Figure B and Table M1 it is necessary scale the regularisation parameters (A, a) as d ranges in {2, 5, 10, 25, 50}. We reason that ||v|2 Dd of grows roughly linearly with increasing dimension d, while the growth rate lgl does not depend on d since it is a scalar-valued field. This motivates the rescaled regularisation term
M
algl²)
In practice we use 0.001 and a for all the values of d considered.
Additionally, we alter aspects of training such as network size, learning rate, and number of training iterations depending on d. For d 10 used wider network. paired with smaller learning rate and more training iterations. These choices were consistently applied to all methods under consideration except for DeepRUOT due to difficulties with modifying the implementation.
C.3 Training: fitness-ODE
Motivated by issues pertaining to the identifiability of dynamics involving both drift and growth, we propose well-known dynamical model baseline model for inference. Let {pe) he be continuous distributional path satisfying mild conditions in the space measures describing some population evolution. Then there exists unique scalar field U2(2.) such that Pt satisfies
(x))
VU,))(((( + (z(x))(x).
(108)
That this the case can be read from [25] Section A.3] or [5... Proposition 2.2]. This can be interpreted continuous dynamics where U+(x) is the fitness of state The rate which agents reproduce is prescribed by Uz) and agents migrate to regions of higher fitness following (x)) Theoretically, for regular enough sequence of population snapshots {prhe. single time-dependent fitness function Ug is sufficient to generate the path t + Ps via these dynamics. While perhaps not obvious, single quantity, the fitness U,, is enough to generate the full path Pt in the space of measures. As a baseline, we therefore propose a neural parameterisation of U, and to learn U, following the TIGON+ setup, but with vy(a:) VU+(x) and g: U+) We parameterise U,(x) NN₂(d + 1, (t, x) using ReLU activations. All hyperparameter choices are listed in Table g
C.4 Training: TIGON+
The problem of dynamical transport for systems with mass imbalance was previously studied in the work [[4]] In this work, the authors consider deterministic systems only and allow both the force
23

---

(x(2)) and growth (x( to be time dependent. However, the TIGON algorithm relies on kernel density estimation (KDE) from the input data [24. Methods] which sensitive the choice of kernel bandwidth and suffers from the curse of dimensionality as the dimension increases. The choice of data-fitting loss, in the form of minimising squared discrepancies between the predicted and KDE densities, adds to these difficulties: this loss relies point-wise on estimated densities and is thus not "gcometry-aware" in the sense of optimal transport based losses [52]]. Finally, propagating densities under flow models e.g. (H) are well known to be computationally costly. For these reasons, training the TIGON algorithm was infeasible for most of our numerical experiments
Because we needed a robust comparison baseline, we decided to implement the same model used by TIGON (deterministic transport with growth), but train it with the UPFI training procedure. With UPFI, we circumvent the need for density estimation by using the probability flow formulation of the Fokker-Planck equation. Together with the use of the unbalanced Sinkhorn divergence as the data fitzing loss, we believe that this substantially improves the TIGON method and also makes for a more rigorous baseline to test against UPFI. We call TIGON+ this re-implementation of TIGON.
Specifically, we parameterise a drift Up (t, x) = (NN(dd1) dt, æ) and growth ga x) NN+(d+
its state (t) with ReLU activations. For sampled data point (æt, mi 1) at (we m.)) at each timepoint 43 is simulated by forward integration of the system
. = (t 2x), mh = go(t, 3))me,
(109)
and we form the empirical distribution Pr. (æ) = E m,,,, 0( x) for each timepoint 4, For the rest of the training procedure we use the same data data-fitting loss as UPFI, i.e. (I2).
C.5 DeepRUOT
We use the existing DeepRUOT implementation provided by [27]], which parameterises the force, growth rate and score function. Different to our method, however, the individual components of the dynamics are coupled via a physics-informed neural network (PINN)-type loss (127, Section 5.3)) that aims to incorporate information from the governing Fokker-Planck equation. The training procedure DeepRUOT consists of multiple stages and involves first neural ODE training, followed by flow matching, and again neural ODE training. For the bistable system example (Fig. 3 their Py PyToch implementation of [27 Algorithm 1]. For each of the drift, growth and score networks, three hidden layers of size 128 were used. For further details on DeepRUOT implementation and training we refer the reader to [27] and accompanying code.
C.6 Flow matching
We consider optimal transport-condittonedd flow matching (OTFM), a class of simulation-free meth- ods for training flows that approximate dynamical OT [38]. Compared to the other methods, OTFM training does not require numerical integration during training. This comes at the cost of requiring the transport plan Ti or coupling. between successive snapshots (pt. Pt. ) to be computed in advance of training the flow model. Given a pair of distributions (p.p') with a OT coupling т on t E [0, 1], OTFM trains a flow network vo (1, æ)) by minimising a nonlinear least squares objective:
min E ((r,ar)~~ lI x) (t (1 t)x + tx²)
(110)
where æ is the optimal coupling obtained by solving the entropic OT problem: 2 B(e,w")==ll
min
2'112 + <T))))p))
(111)
We refer the reader to [38] for an in-depth discussion. We implement OTFM for dynamics inference by applying (III) across each pair of snapshots (t,, f441) to learn single network un uttt,z)). To generate results shown in Table I we compute entropic OT couplings between pairs of snapshots using a squared Euclidean cost and pick (2: 1:).
By default, OTFM treats the balanced case of transport and is not appropriate for modelling dy- namies with growth. As additional bascline, we try an unbalanced variant, UOTFM, where the coupling x is obtained by solving the unbalanced relaxation [23]] of (III), where the hard marginal constraint is replaced with soft penalty with weight À > 0:
min (E 1| 2'112 + (L p¹) (KL(T1|p) (LL(T11p'')
(112)
*
24

---

OTFM
UOTFM
0:11
0.25
Figure 6: Sampled trajectories for trained flow matching models for bifureating example, 10 (see Figure 2
In practice we take = 2- [ -=11 x2112, i.. the mean of the cost matrix. This is essentially the approach advocated for in [26]] In Figure we show sampled trajectories from trained OTFM and UOTFM models for the bifurcating system of Section B. As expected, OTFM results in trajectories biased towards the upper branch, while this is partially mitigated using UOTFM.
C.7 Forward simulation
Bistable system
We consider a potential-driver dynamics in dimension d (2, 5, 10) specified by
VV,
(()) 0.9 æ al3 x bl3 10 10 Σ 27
b(x:) $ (1 + tanh(22a)).
d(x:) 0.
For fe [0, 1] we simulate particles following d.,, v(X,) dt + a dB, using the Maler-Maruyama method. We set the noise level to 1/2, and use the initial condition Xo No 0.017). At each Euler step, simulated particles divide with probability b( (t) We simulate starting from No 500 particles, and the total population size grows over time following the prescribed dynamics. Population snapshots are taken from independent realisations the process at K + timepoints uniformly spaced between [0, 1].
The bifurcating and HSC reaction networks were taken from previous literature [HI [II], corresponding the networks BF and HSC in the collection of BoolODE bench- marking problems [[[]] The original implementation, however, modelled both gene and protein expression levels and as a result does not strictly fall in the modelling framework we consider. This is because protein levels are not observed and thus are hidden variables. We re-implemented each of these systems to involve only gene expression dynamics, and also change the noise model: in the original implementation an ad-hoc square-root noise model was used, i.e. axx) oye. We choose to use a more biophysically motivated noise model (I6), and take
Reaction network systems
oi Vf(x)) λDZ
where f(x) is a vector-valued function state-dependent production rates for each gene x;, and each X; is the corresponding degradation rate. For the growth rates, we consider scenario where cells one branch of the system trajectory divide a faster rate than the others:
In the bifurcating network,
(tanh( (aanh(5(27 0.7)) +1) (1- tanh(5(21- 0.7))
(xx)
2
HSC network, we set
ttan([rg 1.5) +
8(x) 5.5
2
25

---

We simulate both systems using the same code as for the bistable system, with the volume parameter
1/vV = 0.5, starting with a population of 500 cells and capturing バ + 1 = 10 timepoints. For bifureating and HSC network we use simulation time intervals of < 1.25 and & 1 respectively.
Fate probability computation We provide a straightforward definition of fate probability in what follows. Let zy be the state of an observed cell or individual at time t. Let 4اي2 be a partitioning of the state space (e.g. some subset of Rd) where each Ω2 is understood to correspond to well-defined, stable states of the system at some final time, say = 1. In the biological setting, this is typically thought of as mature cell "type" the conditional probability:
[I] 34). Then the fate probability of æ; towards Ω2 is defined as
P P(X, E S X =x,).
form partitioning of the state space into two regions by running k-means with k = 2 on the final snapshot from the system. Given any query state 2; at time we empirically estimated true and inferred fate probabílities by forward simulation of either the ground truth SDE or inferred dynamics:
In practice, such as for the bistable system of Fig.
PP x
Ω2 X = xx) M-1
1a ( xz):
where M the number of trials to sample.
C.8 Neural graphical model
For the Neural Graphical Model (NGM) example of Fig. HL we use the architecture introduced in [2] as drop-in parameterisation of the autonomous force vz(x). For each output variable, use two hidden layers with sizes [64, 64J. We use a group lasso regularisation strength ÄGL = 0.03 and employ the proximal update scheme outlined in [4 Section C.4.1] with a learning rate of 0.003 and train for 5, 000 iterations. We use the score networks that were already pre-trained for the additive and multiplicative UPFI models. All other training details are taken to be the same as for the earlier UPFI training.
C.9 Single cell lineage tracing data
Preprocessing Data for the study of [IT] are available from the original publication using the GEO database with accession number GSE140802, Starting from raw counts, expression data is nor- malised using byn.pp.ud wwwppppppppp....pppppppppp.p..pp recipe_monocle function from the Dynamo package [[3]. In brief, raw monocle
gene expression values are per-cell normalised and then log(1 (1 x) -r--tffsformedd For all our exper- iments we use the 10-dimensional PCA embedding of cell gene expression profiles. Spliced and unspliced transcript counts were obtained from reanalysis of the raw sequencing data of [M]] and RNA velocity estimates were subsequently obtained using the Dynamo package [E3]. Scripts and datasets for this re-analysis are available upon request.
From the full dataset, 86,416 cells deemed to be contributing to the "Neutrophil-Monoeyye trajec- tory (as determined by the original publication [Π]) were selected. Using the 10-dimensional PCA embedding for these data, we apply UPFI, PFI, fitness-ODE and TIGON. We do not include Deep- RUOT in this analysis since resulted in out-of-memory error the initial stages of training. Noting also that the original publication [27]] considered only the 2D SPRING layout, be believe that further modification of its training pipeline may be necessary.
Training For UPFI, we train a time-dependent score model with hidden dimensions [128, 128, 128] for 10, 000 iterations with batch size of 256 and learning rate of 10- We adopt an additive noise model and parameterise an autonomous force Up (x) and growth ge(x) each with a MLP with hidden dimensions [128, 128, 128]. We set 25.0, À 0.001, a 1.0 and we set o = 0.5. Note that the choice of 7 is not scale-invariant, and we found that the typical length scale in the lineage tracing data is larger than in the simulation data. We train for 10, 000 iterations with batch size of 256 and learning rate 10 10-3
We train PFI with the same hyperparameter choices UPFI, except we use a non-autonomous force as done in earlier examples. Finally, we train TIGON and fitness-ODE following the training procedure outlined in Sections Me and e and the same hyperparameters as for UPFI and PFI.
26

---

Score matching
UPFI
PF1
ODE
TIGON
25
249.73
0.50
439
1.51
333.51 4: 2.02
861.41 1: 11.43
313.79 +: 3.71
50
250.71
1.52
443.52 + 2.23
335.94 + 3.65
868.93 + 11.80
316. + 4.17
d
OTFM
UOTEM
25
100.68 0.84
52.64 0.91
50
99.61 1 0.68
52.95 1. 0.75
Table 7: Runtimes (s) for the different methods on the bistable system, for d = 25, 50.
Cosine similarity force field
Growth rate correlation
= 0.0
0.01
0.1
0.0
0.01
0.1
a 0
0.105
0.059
0.078
0
0,923
0.917
0.917
0.01
0.105
0.059
0.060
0.01
0.923
0.917
0.899
0.1
0.105
0.086
0.203
0.923
0.909
0.349
Table 8: Results for regularisation ablation experiments for the bistable system with 10.
C.10
Remark on UPFI in the case when only frequencies are available
When No/No is not good estimator for the ratio l/pol/ we can only build an estimator for the normalised density pt.//Pt,
Pt. Pr. N.
Σx= 8(x axst.)
(113)
This poses limitations on the fitness which can be inferred. Indeed, writing 7e Pe/lpel, we have
O.Pt = de
1 - (de log (p())..
[pe]
Substituting back the PDE governing we find that Pe is also governed by a drift-diffusion PDE, but with time-dependent bias in the source term:
e.p. - [p: (v. - V D: - D. log P.)] (g. è; log (et)
Therefore, when we don't it have access to the absolute number of individuals present in a population given time, we can only hope infer the fitness Je up to a time-dependent bias. this case, the UPFI approach can also be applied using a large enough mass conservation strength q in the unbalanced Sinkhorn distance.
C.11 Ablation experiments
We performed ablation experiments for the regularisation on the bifurcating system with d = 10. In Table we show for varying a and À the cosine error for force recovery as well Pearson correlation for growth rate recovery. When À = 0, there is no regularisation in the loss function, and the resulting error is larger than for unregularised model performs better, albeit at the cost of higher force field error.
0. From the growth rate perspective, however, the
These results suggest that the growth rate is less sensitive to regularisation, which probably stems from the implicit regularisation arising from the relatively small neural network sizes. Since the additional loss terms included only to ensure uniqueness, motivated by the theorem in the main text, this ablation study suggests simple rule of thumb for UPFI: use little or no regularisation when employing moderately sized neural networks.
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect our paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification:
The paper discusses the theoretical limitations of the method, in particular issues of identifiability.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-tex system might not be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach ad- dress problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
28

---

Answer: [Yes]
Justification: We provide in the supplementary materials a complete proof for each of the claims made in the main text.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Yes]
Justification: We provide full details on architectures and hyperparameters necessary to reproduce the results presented in experiments. These details concern not only the method novel this paper, but also other methods to which it is compared. The paper also fully discloses the algorithm used for training, and provides guidelines in appendix on how to train it. The code used for all experiments will be made public on a github repository after the revision process.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
29

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The lineage tracing data used in the experiments is publicly accessible on the GEO database with the query code GSE140802. The code used for all experiments will be made public on github repository after the revision process.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https:
SooeeSubme/ ssiiiioliiiiicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper specifies in appendix all the architectures and hyperparameters used in the experiments. This information is sufficient to interpret and reproduce the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The paper provides error bars for each of the comparison carried out. The error bars are also consistently defined as the standard deviation of the observable, mea- sured repeated experiments using different random seeds.
30

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The neural networks used in this paper are small and can run on a single GPU at time. This information is disclosed in the appendix related to the experimental settings.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with NeurlPS Code of Ethics
Answer: [Yes]
Justification: We have reviewed the code of ethics, and the research conducted in this paper conforms to its requirements.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
31

---

Answer: [NA]
Justification: Our work addresses some fundamental inference problems for population dynamics models, and we potential for direct societal impacts of our works, either positive or negative.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (c.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
allow for controlled use of the model, for example by re-
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All assets used in our work are properly attributed in the main text supple- mental material.
Guidelines:
The answer NA means that the paper does not existing assets.
32

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, Saperswithcode com coo/datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification: A new algorithm is detailed, along with the necessary training details (in the appendix).
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then much detail possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection. cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Not applicable.
33

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research
with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva-
lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy htct:s///eeerrreeeec////eeeereee for what should or should not be described.
34


Reviews:
[{"review_id": "OH6XBy2H9N", "reviewer": "Reviewer_wcim", "summary": "The 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.", "questions": "I 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.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The 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.", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n### Scalability\n> 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.\n> I would like to hear more about alorithm scaling. My sense is currently this approach might be hard to apply to some data sets.\n\n**Theoretical complexity**\n\nThe 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`.\n\n### Bifurcating ex. (Table 1) to $d = 25, 50$\n\n**Cosine error**\n|$d$|UPFI|PFI|ODE|TIGON|DeepRUOT|OTFM|UOTFM|\n|-------:|:--------------------|:--------------------|:----------------|:----------------|:------------------|:----------------|:----------------|\n|25|**0.06±0.00**|0.09±0.00|0.19±0.04|0.27±0.01|0.37±0.00|0.60±0.01|0.63±0.02|\n|50|**0.06±0.00**|0.07±0.00|0.26±0.03|0.25±0.01|0.44±0.00|0.55±0.01|0.52±0.01|\n\n**Fate prob.**\n|$d$|UPFI|PFI|ODE|TIGON|DeepRUOT|OTFM|UOTFM|\n|---:|:--------------------|:----------------|:----------------|:----------------|:----------------|:----------------|:----------------|\n|25|**0.97±0.01**|0.65±0.02|0.90±0.02|0.77±0.01|0.95±0.00|0.74±0.02|0.93±0.01|\n|50|**0.98±0.00**|0.63±0.03|0.91±0.01|0.80±0.02|0.92±0.00|0.73±0.01|0.93±0.01|\n\n**Use of PCA**\nThis 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.**\n\n**Practical runtime**\n\nOn 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. \n\n### Runtime (in seconds, for Table 1)\n\n|dim|Score|UPFI|DeepRUOT|\n|------:|:------------------|:------------------|-------------:|\n|2|61.32±0.50|86.83±0.62|1156.21|\n|5|61.21±0.43|86.99±0.38|1036.71|\n|10|61.26±0.42|87.22±0.51|1002.07|\n|25`new!`|249.54±1.00|442.58±2.01|1194.7|\n|50`new!`|249.25±0.34|440.84±1.87|1079.56|\n\nNote 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.\n\n### Runtime (in seconds, for Fig.5)\n\n|Score|UPFI|\n|----------:|----------:|\n|24.19±0.13|207.97±1.63|\n\nWe emphasize that the lineage example has 86,416 cells on which UPFI runs in a matter of a few minutes. \n\n### Motivation for the growth model\n> 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.\n\nWe 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.\n\nWe 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.**\n\n\n### Additionnal comments\n> 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.\n\nThe 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.**\n\n\n```\n[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.\n[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.\n[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.\n[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.\n```"}, {"review_id": "JIvWeM1F21", "reviewer": "Reviewer_wbX3", "summary": "The 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.", "questions": "- 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?", "limitations": "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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Overall, 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.", "was_revised": true, "final_justification": "Thank you for engaging in the discussion, the rebuttal answers my questions and addresses the concerns.", "rebuttal": "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.\n\n### Scalability and comparison to Flow Matching\n> 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.\n\nWe 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. \n\nWe also implement OT-conditional flow matching (OTFM) [Tong2023] and an unbalanced variant (UOTFM) [Eyring2023].\n\n### Bifurcating ex. (Table 1) to $d = 25, 50$\n\n**Cosine error**\n|$d$|UPFI|PFI|ODE|TIGON|DeepRUOT|OTFM|UOTFM|\n|-------:|:--------------------|:--------------------|:----------------|:----------------|:------------------|:----------------|:----------------|\n|25|**0.06±0.00**|0.09±0.00|0.19±0.04|0.27±0.01|0.37±0.00|0.60±0.01|0.63±0.02|\n|50|**0.06±0.00**|0.07±0.00|0.26±0.03|0.25±0.01|0.44±0.00|0.55±0.01|0.52±0.01|\n\n**Fate prob.**\n|$d$|UPFI|PFI|ODE|TIGON|DeepRUOT|OTFM|UOTFM|\n|---:|:--------------------|:----------------|:----------------|:----------------|:----------------|:----------------|:----------------|\n|25|**0.97±0.01**|0.65±0.02|0.90±0.02|0.77±0.01|0.95±0.00|0.74±0.02|0.93±0.01|\n|50|**0.98±0.00**|0.63±0.03|0.91±0.01|0.80±0.02|0.92±0.00|0.73±0.01|0.93±0.01|\n\n\nWe 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:\n\n### Runtime (in seconds, for Table 1)\n\n|dim|Score|UPFI|DeepRUOT|\n|------:|:------------------|:------------------|-------------:|\n|2|61.32±0.50|86.83±0.62|1156.21|\n|5|61.21±0.43|86.99±0.38|1036.71|\n|10|61.26±0.42|87.22±0.51|1002.07|\n|25`new!`|249.54±1.00|442.58±2.01|1194.7|\n|50`new!`|249.25±0.34|440.84±1.87|1079.56|\n\nNote 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.\n\n### Runtime (in seconds, for Fig.5)\n\n|Score|UPFI|\n|----------:|----------:|\n|24.19±0.13|207.97±1.63|\n\nCompared to UPFI, DeepRUOT is much slower to train and is less accurate.\n\n**We will discuss more thoroughly these scaling issues and include these additional experiments in the paper at the camera ready stage.**\n\n> 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]).\n\nWe 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.**\n\n### Model and regularization choices\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\nThe 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.\n\n> Is the model always an OU process with quadratic fitness? This is not entirely clear from Algorithm 1.\n\nThe 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.**\n\n\n### Additional comments\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\nYes 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.\n\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\nWe 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.\n\n```\n[Mroueh2020] Mroueh Y, Rigotti M. Unbalanced sobolev descent. Advances in Neural Information Processing Systems. 2020;33:17034-43.\n[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\n[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.\n[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."}, {"review_id": "DQsp0ng0ZD", "reviewer": "Reviewer_rQy3", "summary": "The 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.", "questions": "1. **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.", "limitations": "The 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 1, "originality": 3, "strengths_and_weaknesses": "#  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.", "was_revised": true, "final_justification": "The rebuttal addressed scalability concerns as well as provided additional context in terms of other methods.", "rebuttal": "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.\n\n### Contribution and positioning\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.\n> Positioning - How does the framework position itself relative to the recent literature on flow matching [1-6]? \n\nWe 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.\n\nTo 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**. \n\nAs 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].\n\nTo 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.\n\nOn 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].\n\nOn 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.\n\nFor 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.\nThus, 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.\n\nMoving towards more scalable solutions is indeed an important direction and we believe our paper provides a useful solid theoretical foundation for later works.\n\n**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.**\n\n### Scalability\n> Limited scalability - relies on an ODE solver within the main training loop (Algorithm 1) raises questions about the method's computational scalability. \n> 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?\n\nTo 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**.\nAn improvement of our method is that, as opposed to TIGON and DeepRUOT, UPFI does not perform density estimation, although \nwe agree with the reviewer that UPFI cannot scale to thousands of dimensions. \n\nRegarding 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`.\n\nRegarding 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.\n\nWe 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.\n\n**See table in our response to `wcim`: _Bifurcating ex. (Table 1) to $d = 25, 50$_**\n\nTo 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:\n\n### Runtime (in seconds, for Table 1)\n\n|dim|Score|UPFI|DeepRUOT|\n|-:|:-|:-|-:|\n|2|61.32±0.50|86.83±0.62|1156.21|\n|5|61.21±0.43|86.99±0.38|1036.71|\n|10|61.26±0.42|87.22±0.51|1002.07|\n|25`new!`|249.54±1.00|442.58±2.01|1194.7|\n|50`new!`|249.25±0.34|440.84±1.87|1079.56|\n\nNote 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.\n\n### Runtime (in seconds, for Fig.5)\n\n|Score|UPFI|\n|-:|-:|\n|24.19±0.13|207.97±1.63|\n\nWe emphasize that the lineage example has **86,416 cells** on which UPFI runs in a **matter of minutes**. \n\nThese 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.**\n\n> Lack of ablations\n\nWe 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.\n\n### Regularization ablation\n**Force recovery, cosine**\n||$\\lambda$=0.0|0.01|0.1|\n|-:|-:|-:|-:|\n|$\\alpha$=0|0.104997|0.059309|0.0782117|\n|0.01|0.104997|0.0593571|0.0599831|\n|0.1|0.104997|0.0856431|0.203179|\n\n**Growth rate recovery, Pearson**\n\n||0.0|0.01|0.1|\n|-:|-:|-:|-:|\n|0|0.922827|0.91711|0.916789|\n|0.01|0.922827|0.916712|0.898708|\n|0.1|0.922827|0.909402|0.349477|\n\nFor $\\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 .** \n\n### Additional comments\n> Velocity prior: Could kernel-smoothed RNA-velocity guide or evaluate the drift, as in [4]?\n\nYes 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]\n\n```\n[Sejourne2019] Thibault Séjourné, Jean Feydy, Francois-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport.\n[Coomer2022] Coomer MA, Ham L, Stumpf MP. Noise distorts the epigenetic landscape and shapes cell-fate decisions. Cell Systems\n[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.\n[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.\n[Gorin2022] Gorin G, Fang M, Chari T, Pachter L. RNA velocity unraveled.\n```"}]

Author rebuttal:

Meta-review: This paper studies a multimarginal stochastic transport problem with applications to various simulated systems and to single-cell RNA sequencing time series data. The transport problem studied allows for division and death coefficient, a model that is more suitable for some applications than the measure-preserving models studied in most past work. It is solved using a simulation-based algorithm that minimises Sinkhorn divergences between target marginals provided by samples and simulated marginals pushed forward from a previous time step.

All reviewers are positive about the paper. They particularly praise the solid technical contribution (including theoretical guarantees of uniqueness of the solution), the clarity of exposition, and the quality of the results.

Main weaknesses:
- Concerns about scaling to higher dimennsion and computation cost (wcim, wbX3, rQy3). This was answered in the rebuttal by new experiments that show the results hold up when the dimensionality of synthetic bifurcating data is increased from 10 to 50 and timing results. The authors are encouraged to include these in the final version and to study the scaling of the biological data experiments with more PCA components as well, which would answer wcim's concerns more fully.
- Review of related work on inference of dynamics from snapshot data (wcim, rQy3). The authors promised extended discussion in the revision.
- Lack of comparison with simulation-free dynamic OT methods such as flow matching (wbX3, rQy3). New experiments added in the rebuttal answered this concern.

I follow the reviewers' recommendation and recommend acceptance.
Decision: Accept (poster)