Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport
Authors: Zhenyi Zhang
Tiejun Li
Peijie Zhou
Venue: ICLR 2025 Oral (2025)
Keywords: optimal transport
Schrödinger bridge
trajectory inference
single-cell
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Reconstructing dynamics using samples from sparsely time-resolved snapshots is an important problem in both natural sciences and machine learning. Here, we introduce a new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. Based on the RUOT form, our method models these dynamics without requiring prior knowledge of growth and death processes or additional information, allowing them to be learned directly from data.  Theoretically, we explore the connections between the RUOT and Schrödinger bridge problem and discuss the key challenges and potential solutions. The effectiveness of our method is demonstrated with a synthetic gene regulatory network, high-dimensional Gaussian Mixture Model, and single-cell RNA-seq data from blood development. Compared with other methods, our approach accurately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape. Our code is available at: [https://github.com/zhenyiizhang/DeepRUOT](https://github.com/zhenyiizhang/DeepRUOT).

Published as a conference paper at ICLR 2025
LEARNING STOCHASTIC DYNAMICS FROM SNAP- SHOTS THROUGH REGULARIZED UNBALANCED OPTI- MAL TRANSPORT
Zhenyi Zhang', Tiejun Lit.2.,, and Peijie Zhou 2,3,4,3,*
LLAM and School of Mathematical Sciences, Peking University.
Center for Machine Learning Research, Peking University. NELBDA, Peking University.
'Center for Quantitative Biology, Peking University. *AI for Science Institute, Beijing.
zhenyi zhang pku. edu. cn, {tieli, pjzhou)@pku.eede
ABSTRACT
Reconstructing dynamics using samples from sparsely time-resolved snapshots is an important problem in both natural sciences and machine learning. Here, we introduce a new deep learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dy- namics from observed snapshots. Based on the RUOT form, our method mod- els these dynamics without requiring prior knowledge of growth and death pro- cesses or additional information, allowing them to be learned directly from data. Theoretically, we explore the connections between the RUOT and Schrödinger bridge problem and discuss the key challenges and potential solutions. The ef- fectiveness of our method is demonstrated with synthetic gene regulatory net- work, high-dimensional Gaussian Mixture Model, and single-cell RNA-seq data from blood development. Compared with other methods, approach accu- rately identifies growth and transition patterns, eliminates false transitions, and constructs the Waddington developmental landscape. Our code is available at: https://githu..com/ http:h/wwww....ecaeiiaea//eeeeee
INTRODUCTION
In machine learning and natural sciences, a key challenge is coupling high-dimensional distribu- tions from observed samples, exemplified by Variational Autoencoders (VAEs) (Kingma & Welling, 2013) which map complex data to simpler latent spaces. It is also important in multi-modal anal- ysis for integrating diverse data types (Lahat al., 2015), particularly in biology through aligning multi-omics data into unified cellular state representations (Cang & Zhao, 2024; Cao et al., 2022; Demetci al., 2022: Gao al., 2024b). Recently, there has been growing interest in understand- ing the dynamics of how distributions are coupled over time, such as diffusion models (Ho et al., 2020: Sibl-Dickstein et al., 2015) and stochastic differential equations (SDEs) (Song et al., 2021). The task is useful for interpolating paths between arbitrary distributions and learning the underlying dynamics (De Bortoli et al., 2021; Neklyudov al., 2023; Tong et al., 2024b; Wang et al., 2021).
Due to the destructive nature of technology, the analysis of time series data in single-cell RNA sequencing (scRNA-seq) provides an important application scenario for the high dimensional prob- ability distribution coupling and dynamical inference problem (Bunne et al., 2023b; 2024; Jiang & Wan, 2024; Jiang et al., 2022; Lavenant al., 2024; Peng et al., 2024; Schiebinger et al., 2019; Sha et al., 2024; Tong al., 2023; Zhang et al., 2021). Trajectory inference in scRNA-seg data have been extensively studied (Saclens et al., 2019), and optimal transport (OT) methods has emerged a central tool for datasets with temporal resolution (Bunne al., 2024; Klein et al., 2023a;b; Schiebinger et al., 2019). Often, there is a need to learn the continuous dynamics of cells over time and fit the mechanistic model that transforms the initial cell distributions into the distributions later temporal points. This could be solved through the dynamical formulation of OT (Benamou & Brenier, 2000), also known as the B-B form. However, the formulation has fully taken the
"Corresponding author.

---

Published as a conference paper at ICLR 2025
stochastic dynamical effects into account, especially the intrinsic noise in gene expression and cell differentiation (Zhou et al., 2021a;b), which is prevalent in biological processes on single-cell level (Elowitz et al., 2002).
By incorporating stochastic dynamics, the Schrödinger bridge (SB) problem aims to identify the most likely stochastic transition path between two arbitrary distributions relative to reference stochastic process (Léonard, 2014), and has been applied in a wide range of contexts, including scRNA-seq analysis and generative modeling (Liu al., 2022a; Pariset et al., 2023). Meanwhile, recent regularized unbalanced optimal transport (RUOT) offers a promising approach for mod- eling both stochastic unbalanced continuous dynamies (Baradat Lavenant, 2021: Buze & Duong, 2023: Chen et al., 2022b; Janati et al., 2020), which can be viewed as an unbalanced relaxation of the dynamic formulation of Schrödinger bridge problem. However, computational methods for learn- ing RUOT or such nigh-dimensional unbalanced stochastic dynamics from snapshots are relatively lacking, especially when there is no prior knowledge of unbalanced effect.
Here we develop new deep learning method (DeepRUOT) for learning general RUOT and inferring continuous unbalanced stochastic dynamics from samples based on the derived Fisher regularization form without requiring prior knowledge. We demonstrate the effectiveness of DeepRUOT on both synthetic and real-world datasets. Compared to the common SB method, our approach accurately identifies growth and transition pattems, eliminates false transitions, and constructs the Waddington developmental landscape of scRNA-seq data. Overall, our main contributions can be summarized as follows:
We reformulate RUOT with a Fisher regularization form and explore the connections be- tween RUOT and unbalanced SB. The formulation transforms the SDE into the ordinary differential equation (ODE), which is computationally more tractable.
We propose DeepRUOT, the neural network algorithm for learning high dimensional un- balanced stochastic dynamics from snapshots. Through the neural network modeling for growth and death, our framework does not require prior knowledge of these processes.
We validate the effectiveness of DeepRUOT on both synthetic data and real scRNA-seq datasets, showing its promising performance compared with existing approaches.
2 RELATED WORKS
Deep Learning Solver for Dynamical OT To tackle the dynamical OT (i.e. B-B form) in high dimensions, many methods (Albergo & Vanden-
2023; Cheng et al., 2024a;b; Chow al.,
2020; Gao et al., 2024a; Huguet et al., 2022; Jiao et al., 2024b; Jin et al., 2024; Lipman et al., 2023; Liu et al., 2021; 2022b; 2023b: Pooladian et al., 2024; Ruthotto al., 2020; Tong et al., 2020; 2024a; Wan et al., 2023; Wu et al., 2023; Zhang et al., 2024a) have been developed based on continuous normalizing flow and neural ODE formulation either in original latent space. To account for sink and source terms in unnormalized distributions, (Peng et al., 2024; Sha et al., 2024; Tong et al., 2023) formulated the neural-netwont based solver the unbalanced dynamical OT setup. However, the appropriate formulation along with an effective deep learning solver to simultaneously account for unbalanced ferm and stochastic effects in dynamical OT remains largely lacking.
Computational Methods for Schrödinger Bridge Problem Many methods have recently been developed to solve the static SB problem (Chizat et al., 2022; De Bortoli et al., 2021; Gu et al., 2024; Lavenant et al., 2024: Liu et al., 2022a; Pariset et al., 2023; Pooladian & Niles-Wees 2024: Ventre et al., 2023). To tackle the dynamical Schrödinger Bridge, methods based on neural SDE solver, neural ODE solver with fisher information regularization or flow matching (Albergo et al., 2023; Bunne et al., 2023a; Chen et al., 2022a; Jiao et al., 2024a; Koshizuka & Sato, 2023; Liu et al., 2023a; Maddu et al., 2024; Neklyudov et al., 2023; 2024; Tong et al., 2024b; Wang et al., 2021; Zhang et al., 2024b; Zhou et al., 2024a;b) have been proposed. However, these methods either fail to account for unnormalized distributions resulting from cell growth and death, or require prior knowledge of the processes (growth/death rate) (Chizat et al., 2022; Lavenant et al., 2024; Pariset et al., 2023; Schiebinger et al., 2019) or additional information (c. (c.g.cell lineage) (Ventre al., 2023).
2024; Shi et al.,
2

---

Published as a conference paper at ICLR 2025
Study of RUOT The RUOT has recently formulated (Chen et al., 2022b), also known unbal- anced Schrödinger Bridge. The existing studies are mostly focused on the theoretical side. For instance, (Janati et al., 2020) derived closed-form formula for entropic OT between unbalanced Gaussian measures. (Baradat & Lavenant, 2021) investigates relations between RUOT with branch- ing Schrödinger bridge. (Buze & Duong, 2023) investigates different formulations of the problem.
PRELIMINARIES AND BACKGROUNDS
In this section, we provide an overview of stochastic effects and unbalanced forms within the dy- namical framework. Specifically, considering only stochastic effects leads to the Schrödinger bridge problem (Section 3.1). whereas addressing solely the unbalanced aspect results in unbalanced dy- namic optimal transport (Section 3.2). By integrating these two perspectives, we motivate the for- mulation of the Regularized Unbalanced Optimal Transport (RUOT) framework.
3.1 STOCHASTIC EFFECT: SCHRÖDINGER BRIDGE PROBLEM
The Schrödinger bridge problem aims to identify the most likely evolution between a given initial distribution 1/ and terminal distribution V1 (assumed to have density in this paper), relative given reference stochastic process. Formally, this problem can be formulated as the minimization of the Kullback-Leibler (KL) divergence in the optimal control perspective (Dai Pra, 1991) as below:
min
DEL (H00,)llu(
(1)
H =vb, pf =M
X where 110 denotes the probability measure induced by stochastic process X, (0 < t S 1) defined
on the space of all continuous paths C((O, 1]. Rd), with the distribution of X; at given time charac- terized by the measure u with density function ((, t). In this work, we consider X, the stochastic process characterized by the following stochastic differential equation (SDE):
dx, (X,,,t dt + a (X,
and the reference measure M/o.1 is chosen as the probability measure induced by the process dY, o oY,,, t,W e. where W, E Rd is the standard multidimensional Brownian motion defined on a probability space (2,F,P) with P the Wiener measure for the coordinate process W((w) = w(),, and or (xx t) E gdddd is the diffusion matrix which is typically assumed bounded, coercive and invertible. Under this formulation, the solution to Eq. (1) is also referred to as the diffusion Schrödinger bridge. The diffusion Schrödinger bridge also possesses dynamic formulation, which can be formally stated as the following theorem.
Theorem 3.1. Consider the diffussio Schrödinger bridge problem (1) where Mo.1 the reference measure induced by dY, oY :Yd, Then (1) is equivalent
fee [ (p.b)
bT (2. t)a tatttaaaaaa (xx t)b(x, t) p(zx, t)dadt,
inf
(3)
where the infimum is taken over all function pairs (p, b) such that p(., = Vor p(,,1) = 21 p(x, absolutely continuous, and
0.plæ, 1) -v. (p(xr, t)b(ar, ()) V2 (a(æ, ())((
(4)
where V2 (a(æ, 1)((a, ()) = Ens Ozg (aiji)) and a(x, () (æ ()oo (a, 1), coupled with van- lim p(æ, O.
ishing boundary condition:
lac +00
This theorem or its variants have been stated and proven in various forms, such in (Chen et al., 2016; Dai Pra, 1991; Gentil et al., 2017). Here, we provide a simple direct proof for illustration in Appendix D. We note that in this dynamic formulation, it accounts for the stochastic aspects but does not incorporate unbalanced effects. As we will discuss below, RUOT can be motivated by incorporating unbalanced effects in the dynamic formulation of SB.
3

---

Published as a conference paper at ICLR 2025
3.2 UNBALANCED EFFECT: UNBALANCED DYNAMIC OPTIMAL TRANSPORT
The optimal transport problem has been extensively studied in various fields. Given two prob- ability distributions E R and [3 E [44 its primary goal is to find the optimal cou- pling π E R transport a given distribution of mass (or resources) from one location to another while minimizing the cost associated with the transportation. The static optimal transport problem can be mathematically formulated as min ( c), where II = ( E I I Q, n l B, क 2 0}. The cost matrix c E R defines the transporta- tion cost between each pair of points, where Cjj == c(a,, v2) represents the cost of moving a unit mass from point æ; to point y. We refer to (Peyré et al., 2019) for more details. Next we briefly state some well-known results on the dynamical formulation of the OT.
Dynamical Optimal Transport The formulation is also known the Benamou-Brent formula- tion (Benamou & Brenier, 2000), which can be stated as follows:
/ / s.t. 0.p + V (b(x, t)p) = 0, p = vo, p|+ = V1-
inf W (vo, 11) (pla..t).((..0))
2 b(x, 1)) t)dædt,
Compared to SB, in OT the distributions are connected by deterministic transport equation instead of diffusion. It can be shown that this dynamical formulation corresponds to a static Kantorovich's optimal transport problem with cost function c(x, y) || 2112.
Regularized Optimal Transport The regularized optimal transport is defined by a dynamical form of the general Schrödinger bridge by taking diffusion rate as constant and scaled by o² in (3):
/ / _1,(2,, 1)) t)dxdtt
inf W (vo, v1) (o(a..)).b(ooe)))
2 b(xx
0² Ap= = 0, pl--0 = un pla- V1, 2
s.t. 0.p + V (b(x,
It can be demonstrated that a² approaches zero, the solution to this problem converges to that of the Benamou-Brenies problem (Mikami & Thieullen, 2008). And it is equivalent to the SB problem (Baradat & Lavenant, 2021; Gentil al., 2017: Léonard, 2014).
Unbalanced Dynamic Optimal Transport To account for unnormalized marginal distributions and effects such as growth and death, an unbalanced optimal transport problem with Wasser- -ishn-Fisher-Rao (WFR) metric has been proposed (Chizat et al., 201 or its extensions (Gangbo et al., 2019). Here we adopt the WFR unbalanced optimal transport:
L.' J (21bbe t)ls + alg(a, 1) p(æ, f)dxdt,
inf W (vo, v1) = (((m),).b(m.)))...)
s.l. 0.p + (b(x, t)p) = g(æ, t)p, pl:=0 = vo, Po. 11.
Here g(x, t) is a scalar function that denotes the growth or death rate of particles at the state æ and time i, and is also optimized in the total energy term. is the hyperparameter of weight. We should also note that this case, Mo and 1/1 are not necessarily the normalized probability densities, but are generally densities of masses.
4 REGULARIZED UNBALANCED OPTIMAL TRANSPORT
To simplify the notation and illustrate a commonly used setting, we take special case considering a(x, t) = o²( I,, and the general case is left to Appendix D.4 for further discussions. Inspired by unbalanced dynamic optimal transport, the dynamical formulation Theorem 3.1 suggests a natural approach to relaxing the mass conservation constraint by introducing a growth/death term gp in Eq. (t) gp. Meanwhile, also define a loss functional in Eq. (5) which incorporates the growth penalization and Wasserstein metric considered in Eq. (3) (with rescaling by o² (t)). We refer to this formulation as the regularized unbalanced optimal transport.
(4) as a,p V_ (pb) } V2
4

---

Published as a conference paper at ICLR 2025
Definition 4.1 (Regularized unbalanced optimal transport). Consider
inf L.¹ / [1b( 0)122 ()dxdt + S' J atl av (1(æ, ()) plæ, t)dadd, (p.b.g)
(5)
where V/ : R + 10, too corresponds to the growth penalty function, the infirum is taken over all
pairs (p, b) such that p(,,0) 0) = Mo.pl p(- 1) = V1, "(,p( t) absolutely continuous, and
O.p V. (pb) v2 (o²())Tp) + 9p
(6)
with vanishing boundary condition:
lim plar, 1)
|a 00
Note that in the definition if ((g)= +00 unless g 0 and (0 = then it implies g(x, t) and the RUOT degenerates to the regularized optimal transport problem. If (t) 7 and Ve (g(x, t)) (g( this degenerates to the unbalanced dynamic optimal transport with WFR metrics. Mcanwhile when (t) is constant, it coincides with the definition of RUOT provided in (Baradat & Lavenant, 2021). Then, we can reformulate Definition 4.1 with the following Fisher information regularization.
Theorem 4.1. The regularized unbalanced optimal transport problem (5) is equivalent to
inf 1 /v/ (p.v.))
o²(t)
Plg o²(
1 do)
1
2 (1 p) g
logp + at (g)
pdædt,
2 dt
(7)
where the infinum is taken over all triplets (p, v. g) such that p(-,0) 0) vo, p.,p(:;)) v1, absolutely continuous, and
0,p - (po( f]) + g( t)p
(8)
with vanishing boundary condition:
lim p(æ, t) 0.
Here v(z, t) represents new vector field. The proof is left to Appendix D.2. In (Baradat & Lavenant, 2021), they proposed another Fisher regularization form of RUOT, expressed the for- mula Eq. (15), In their formulation, computing the cross term logoo (t)()) required calcu- lating the derivative of the function log p with respect to æ and performing vector multiplications with v. The formulation here is equivalent to theirs but more computationally tractable, as we avoid differentiation and vector multiplication in our cross-term.
Remark 4.1. When g = O, (0) 0 and (t) is constant, then Eq. (7) is the same as the dynamic entropy- reguaa-eeg optimal transport form as discussed in (Bunne et al., 2023a; Gentil et al., 2017; Léger& & Li, 2021; Li et al., 2020; Neklyudov et al., 2024; Pooladian Niles- Weed, 2024).
the Fisher information. Notably, when considering growth/death factors, Eq. (7) includes not only the Fisher- Rao metric but also an additional cross-term Jg-- - = (t) (1 + log p(æ, t)) g(xr, t)p(æ, t)dæ.
Remark 4.2. The term I(p) fR== V log p(æx, 1)]1p((x, t)dx in Eq. (7) referred
Remark 4.3. From the Takker-Plank equation and the proof of Theorem 4.1, the original SDE
dx2, = (b (X,, ()) dt + (f) dW, can be transformed into the probability flow ODE
dx, = b (xx,,)) 0 (
(1)V poogg((Xp pogr(X.,))) dt.
x(x..t)
Conversely, the probability flow ODE's drift v(ær, t), the diffusion rate (t) and the score function Ve log p(x, t) are known, then the the drift term b(x,t) of the SDE can be determined by b(x, t) v(x, t) + ( log plgp t). Thus, to specify an SDE is equivalent to specifying the probability flow ODE and the corresponding score function V p(g,((tt, t) (Tong et al., 2024b).
In (Buze & Duong, 2023), the authors defined RUOT problem with nonlinear Fokker-Planck equation constraints. In Appendix D.3, we show that the proposed form is indeed consistent with the RUOT defined here when at)) is constant and V (g(x,t)) takes the quadratic form (i.e., W (g(x, t)) = 19(x, (3)). We then explore the connections between RUOT and SB problem in Appendix E.
5

---

Published as a conference paper at ICLR 2025
5 LEARNING RUOT THROUGH NEURAL NETWORKS
Given unnormalized distributions at discrete time points, X; b4 for fixed timepoints i E (O....., 1}. aim learn continuous stochastic dynamics satisfying the RUOT from data. Similarly, to simplify the exposition and illustrate a commonly used case, we consider Defini- tion 4.1 with the stochastic dynamics dX,, (Xx,, ()) dt (t) dW .. As previously discussed, we approach this by transforming the problem into learning the drift v(z, t) of the probability ODE and its score function 2/2 (t) log p(æ, in Theorem 4.1. We parameterize v(z,t), ((x t), and (t) log p(æ, t) using neural networks va, go and sg respectively (Fig. 1). To solve Theorem 4.1, the overall loss is composed of energy loss, reconstruction loss, and the Tokker-Planee constraint:
L Fiergy + X LRecons + A LFp.
The Energy loss promotes the least action of kinetic energy Eq. (7). The reconstruction loss CRecons promotes the dynamics to match data distribution at later time point (i.e. p(- 1) 11), and the LFP promotes the three parameterized neural network to satisfy Bikker-Plancee constraints Eq. (8).
5.1 ENERGY LOss
To compute the integral in Eq. (7), the direct caleulation is infeasible due to the high dimensionality. Thus, we need to transform it into an equivalent form that can be evaluated using Monte Carlo methods. Adopting the approach in (Sha et al., 2024) (see Appendix A.4), Eq. (7) is equivalent to the following form
(-"
(02(1)
(o)
so + a (oo) wt)dtt (2
Eieergy
-Broo -po
1ool2
2
| /
+ 80 90
(10)
where (t) eñó go (ant),,)) and (t) satisty dæ/dt v(x, t)dt. We compute through Monte Carlo sampling and Neural ODE solver.
5.2 RECONSTRUCTION Loss
The reconstruction loss aims to match the final dis- tribution in Theorem 4.1 (i.c., p(-, 11)). Many works use the balanced optimal transport to evalu- ate the distance between two distributions. However due to the unnormalized effect, here we aim to use the unbalanced optimal transport instead. realize this we need to tackle two parts:
A
growh
death
dp
Vo (00) 3 (62(6189))
My
L= Lowey Lisces + Lep
nolse
velocity
LRsccms Cmsss + AaCor
(11)
growth
where Am and Xa hyperparameters. The mass matching loss Lasss promotes to align the number of cells. We then normalize the distributions according to the matched masses. The Lor uses these weights to perform optimal transport matching.
ODE solver
Figure 1: Overview of DeepRUOT.
For LMas. we propose local mass matching strategy. We denote A (i = 0. 1) as the dataset observed at different time points, and define of Rd x T + RdÍli as the trajec- tory mapping function of the Neural ODE dæ//// va. which inputs a given starting point as the initial condition and outputs the observed particle coordinates at T following ODE dynam- ies, with T being set of time indices. Starting with an initial set Ag. the Neural ODE function o% predicts the subsequent sets of data points over the time indices in T. Specifically, the pre- dicted sets AI,..., ÄT-I range of time indices from to T 1, i.e. A Ar-1 o. (Ag. (1. (1...., -1)- Simi- larly, we define og IR x T + Rl the particle weight mapping function of the Neural ODE d w;(t), (t)/dt outputs the sampled weights at T following ODE dynamics. For simplicity of notation in algo- rithm description Algorithm 1, we also w(A) denote the of weights of particles in set
obtained by applying the Neural ODE function Ao and the full
go (z, t), which inputs given weight of particle as the initial condition and

---

Published as a conference paper at ICLR 2025
A. Assume the sample size No at the initial time point, with the relative total masses at subse- quent time points denoted as nt, = NN/No for l = 0,1, . - 1. sampled particle i has weight wx(t) with the initial condition w;(0) Along the trajectory, each 1/No- We then estab- lish a mapping hy from points at time la to the sampled particles (total number No) predicted at time tgs which links the real data points in A& to its closest point in sampled particles at As, i.e., hy Ag + Ak. (a))) argmin yeA, |Fte v11.. The mass matching error at time l is then defined as Mr - Em [wx(la)) card (ae= (Ix))) No | Here card(A) ()) denotes the cardinal- ity (i.e., number of elements for the finite set) of set A. This error metric ensures the matching of local masses. So the LMass E Mj Furthermore, when M, 0, we have E wa((s) EN card (h, (xx())) /No = N = ntg- So the local matching loss also encourages the matching of total mass, since the left-hand side is the unbiased estimator of fga p(y. 4 Jdy. Once we have determined the weights w; (t), we utilize these weights to perform optimal transport match- ing of the distributions Lor := k= (2 w(tx)), where v² (1/Ne, 1/Nk, .... 1/Nx)) is the uniform distribution of A& time t, w(tx) = (wi (tx), wz uz((x)..... WN. (tx))/ EN w,sts) is the predicted weight distribution of sample particles at time ly and W2 represents the Wasserstein distance between Ag and Ag with normalized distribution defined by ê² and w(x)).
5.3 FOKKER-PLANCK CONSTRAINT AND TWo STAGE TRAINING
In addition to the energy loss defined in Eq. (10) and the reconstruction loss defined in Eq. (11), it necessary to incorporate a physics-informed loss (PINN-loss) (Raissi et al., 2019) to constrain the relationships among the three neural networks, i.c., the Fokker-Planck constraint Eq. (8). Here we utilize a Gaussian mixture model to estimate the initial distribution (Appendix A.3), ensuring that satisfies the initial conditions Pu. and the PINN-loss is defined as
LFP |1:.p + Va (DWDo) gePell + Am |po 0
(12)
where pa exp _ so. Therefore, the total loss function (9) is defined the weighted sum of the energy, the reconstruction eтror and the Fokker-Planch loss with the loss weights as hyper- parameters. This allows to develop a neural network algorithm for solving the RUOT problem (Algorithm and Appendix 1).
We adopt two-stage training approach to deal with multiple loss terms and stabilize the training process. For the pre-training stage, initially, we use reconstruction loss only to train va and ga, ensuring a required matching as the initial value. Subsequently, we fix va and go and employ flow- matching to learn an initial and well-optimized log density function (sa t)). Specifically, it is conducted by conditional flow matching (Lipman et al., 2023; Tong et al., 2024a;b) (Appendix A.2). As part of this pre-training stage, we incorporate a hyperparameter scheduling strategy to further enhance stability (Appendix C.2). Based on these warmup steps the pre-training stage, in the training stage, we use the obtained vg. Ia. and the optimized log density function as the starting point, then get the final result by minimizing the total loss (9). In summary, by integrating all the methodologies, we derive the Algorithm I for training the regularized unbalanced optimal transport (Appendices A.2 and C.3). We discuss the loss weighting strategy and settings in Appendix C.2.
NUMERICAL RESULTS
This section evaluates the Deep solver's ability to recover growth and transitions accurately and produce realistic stochastic dynamics for constructing a Waddington's 's developmental landscape. We take V (g(x, (x, 1)) = gg( (x, 1 and (l) is constant in the following computations.
Synthetic Gene Regulatory Network Inspired by (Sha et al., 2024), we adopt the same three-gene simulation model (Appendix B) to explore the stochastic dynamics of gene regulation, as illus- trated in Fig. 2(a). The resulting gene regulatory and cellular dynamics are illustrated in Fig. 2(b), (c), where a quiescent region and an area exhibiting both transition and growth can be observed. We focus the projection of these dynamics onto the two-dimensional space of (X,, since on X2 it remains quiescent. In Fig. 2(d-e), we compare DeepRUOT with the balanced diffusion Schrödinger bridge method as described in (Tong et al., 2024b). We find that neglecting the growth
7

---

Published as a conference paper at ICLR 2025
Algorithm Training Regularized Unbalanced Optimal Transport
Require: Datasets Ag..... ... Ar-1. 1, batch size N, maximum ode iteration Rode. maximum log density iteration denssiy initialized ODE vo. growth go and log density So
Ensure: Trained neural ODE ve, growth function go and log density function sp-
1: Pre-Training Stage:
2: for = 1 to Redde do
Distribution Reconstruction training
3:
for 4= to T 2 do
Ac+1 t 06 (A, w/ ((A+1)) + 08 (((((A)))) (w(Aa))+=).
4:
5:
LRecons f- LRecons + Am Me + >aW2 (i', w(l)) (1I), update ve and Go W.r.t. CRscons with hyperparameter scheduling (Appendix C.2).
6:
for = to 2 do
7:
A++1 <- 03 (A,
b Generating samples from learned vo-
8: for i = 1 to dolggdonsity do
D CFM Score matching (Tong et al., 2024b) p(æ, (ær, a1)) at generated datasets Ag.... Ar Lscore + V + 211 (13), update so w.r.t. the loss Lecore
9:
(zo, x1) (( æ1); ~~ 4(0, 1);
10: Training Stage:
11: for i to Rede do
12:
Estimate the initial distribution through Gaussian Mixture Model (Appendix A.3).
13:
for t = 0 to T - do
Ac+1 + 0% (e,t + 1), 20( ((((11)) + 0g (w(A)), t
14:
LEnorgy +- EmErpE Set1 2 11 1 + | 2012 - (
(a"(t)) so + av (ow) w(2)dz
15:
g0
16:
CRocOnS t PRecOms + Am Mr + >aM22 10 vn(() (11)
17:
LFP t 110ipe + Væ (pove( x t)) - go (x, pell + Au (lpo 0) po lI (12)
18:
L + Linergy + Ar RRecoms + L (9), update uo. 9o and so w.r.t. L
in the balanced diffusion Schrödinger bridge leads to a false transition, incorrectly attracting cells in the quiescent state to the transition and growth region. The underlying reason is that the growth factor causes an increase in the number of cells. If the growth factor is ignored, the regions with increasing cell numbers will attract cells transition into them to maintain balance. Our approach, which explicitly incorporates growth dynamics, indeed eliminates false transitions and yields results consistent with the ground truth dynamics Fig. 2(b). Furthermore, the growth rates inferred by Deep- RUOT (Fig. 2(f)) closely match the ground truth after normalization. demonstrating the robustness and accuracy of the DeepRUOT solver (Fig. 2(c)). In Table we present quantitative metrics (W1 and W/2 Appendix C.1) to evaluate the performance of our proposed algorithm in comparison with several representative baseline methods (Balanced OT, Balanced SB, Unbalanced OT, Unbalanced Action Matching (AM), Unbalanced SB). The results demonstrate that DeepRUOT is quantitatively more accurate than the competing methods. Furthermore, we observe that our algorithm can bene- fit from the incorporation of stochasticity, outperforming the solver without diffusion (Unbalanced OT). We include the detail ablation studies in Appendix B.5.
Table I: Wasserstein distance (W) and W2) of predictions at different points across five runs on gene regulatory data (Appendix C.I). We show the mean value with one standard deviation.
Model
Wi
Wg
W
W2
31//
W2
MIOHow (Huguet et al., 2022)
0.313ancce
p:260l.0mm
1.2854n.000
0.421.0.000
0.G244
0.802.000
SP (Tong at al., 2024bb
0.174aarom
0.500
0:450goo07
0.719
0 0666623001
1000
0.871
L.222
Unbalanced SB (Pariset al., 2023)
0.729
u.546
0.747a0m
0.612_2m1
0.725
0.572
0.594420
nAM (Nakyyudov et al., 2023)
0.448
0.4962n.000
0.6502oma
1.00120000
0.661
0.740
0.672
0.86420000
Unbalanced OT (Sha al., 2024)
0.047
0.050aness
0.050.0mm
1.08S
0.073
0.054,0000
0.307
0.124
DeeRUUOT (ours)
0.044
0.058
0.050
0.064grom
0.071
0.065
0.104
1.121
Learning Waddington Developmental Landscape In biophysics, Waddington's landscape metaphor a well-known model for representing the cell fate decision process. Constructing such potential landscape has been widely studied (Bian al., 2023; 2024: Li & Wang, 2013; 2014; Shi et al., 2022; Wang et al., 2010; Zhao et al., 2024: Zhou & Li, 2016; Zhou et al., 2024c;d), however, it still remains a challenging problem in the single-cell omics data. The energy landscape
8

---

Published as a conference paper at ICLR 2025
(4)
Transition Grovth
True dynamics
(c)
True growth rate
1a (9)
X1
X2
Gene Xs
Genee Kg
Quiescent
Gene X1
(d)
SF2M (balanced)
(e)
DeepRUOT (ours)
(F)
Learnt growth rate
Gene Xg
Gane 3S
Genne
00
X;
X;
Figure 2: (a) Illustration of the synthetic gene regulatory dynamics. (b) The ground truth cellular
dynamics project (X1, X2). (c) The ground truth growth rates. (d) The dynamics learned by
balanced Schrödinger bridge (SF2M (Tong et al., 2024b), a = 0.25). (e) The dynamies learned by
DeepRUOT solver (o 0.25). (f) The growth rates inferred by our DeepRUOT solver. (g) The
Waddington developmental landscape learned t = 1. (h) The constructed landscape at =
defined by U -o2 log log Pss/ where Pss is the stcady-state PDF satisfying the Bokker-Plance
equation -V (p..b) - Aps gPss. In this work, leveraging our model, we can naturally in- fer the time-varying potential energy landscape through the learned log density function, and use -o²(() log p(xr, 1)/2 to represent the landscape at time l. The lower energy function indicates more stable cell fates in the landscape. Fig. 2(g-h), we can observe the temporal evolution the poten- tial landscape in the synthetie gene regulation system. The quiescent cells always occupy potential well, indicating its stability. Interestingly, the location of potential wells belonging to the transitional cell population is moving in consistent direction with ground truth cell-state transition dynamics. Overall, the results suggest the accuracy and usefulness of DeepRUOT algorithm.
Synthetic Gaussian Mixtures Inspired by (Ruthotto al., 2020). we employed a high- dimensional Gaussian mixture model to evaluate the scalability of DeepRUOT. Initially, we gen- erated a 10-dimensional Gaussian mixture distribution. The initial density represented as a Gaus- sian mixture derived from the average of two Gaussians, with their projections onto the (æ1, x22) plane depicted in Appendix B.2. The final density is also modeled as mixture of three Gaussians, with its projection onto (æ1, (2) planes. Fig. 5 (Appendix B.2) demonstrates that our model effec- tively incorporates growth while learning the stochastic dynamics that transition from the initial to the target distribution. We observe that the cells in the upper region exhibit proliferation without transport. the true cell dynamics indeed follow this pattern, models that neglect the growth fac- tor, such conventional balanced methods, would fail to capture these dynamies. Conversely, if the true dynamics involve transport from the lower to the upper cells, our model can still recover such dynamics, provided more detailed temporal resolution data is available. This highlights the robustness of our approach in capturing complex dynamical behaviors involving both growth and transport. Similarly, we visualized the energy landscapes at the initial and final time points Fig. (Appendix B.2). The results reveal that the cells ultimately differentiate into three distinct fate states.
Real Single-Cell Population Dynamics Next, we evaluate our algorithm on a real ScRNA-sed dataset. We use the same dataset as in (Sha et al., 2024; Weinreb et al., 2020), which involves mouse hematopoiesis analyzed by using a lineage tracing technique. After batch correction across different experiments, the data was projected onto 2D reduceed-dimension force-directed layouts (SPRING plots). A clear bifurcation is observed where early-stage progenitor cells differentiate into two distinct fates (Fig. 3(a)). Using the RUOT- based approach, we leam the underlying stochastic dynamics of the data, along with the growth rates and developmental landscape at different times (Fig. (dd). We similarly evaluated quantitative metrics to compare method with several

---

Published as a conference paper at ICLR 2025
baseline approaches (Table 2). We find that method outperforms others (Appendix B.3). We then test our method on a time-series single cell dataset from an A549 cancer cell line, where cells exposed to TGFB1 induce EM (Sha et al., 2024). We project the data with PCA ten- dimensional latent space as the input of our algorithm. The results are shown in Appendix B.4, which indicates that DeepRUOT remains effective and applicable in higher-dimensional settings, demonstrating its versatility and robustness in modeling complex single-cell dynamics.
Table 2: Wasserstein distance (W1 and W2) of predictions at different points across five runs on scRNA-seq data (Appendix C.I, 0.25). We show the mean value with standard deviation.
1
1=
Model
W/
W
MIOFlow (Huguct et al., 2022)
0.276725 c0. m0at
0.312102 2c.
0.307610, a nusco
0:102190..n.0m
SF2M (Tong et al., 2024b)
0. 0.166447
0.2 0.213489
0.190020
0.2 0.241516 000
Unhulanced SB (Poriset al. 2023)
0.3876538
0.4006603=20.oo
0.128254200000
0.188639-60.200
uAM (Neklyudov al., 2023)
0.72/1918 18..09 mm.00
0.777237 auuco
0 0.889627 =0.000)
Unbulanced OT (Sha et al., 2024)
n.3135522
0.3966 37
0.3442330
0.4099422 .0.000
DeepRUOT (ours)
0.145026gg0.00
0.172878goo0z
0.1.341 1:0000
0.167328
Learnt growth rate
1.0
(a)
(b)
Ta.o
11.0
hecicted
300
5755
5215
748
330
30
Figure 3: Application in hematopoiesis scRNA-seq data (a 0.25). (a) The stochastic dynamics learned by RUOT. (b) The growth rates learned by DeepRUOT. (c) The constructed Waddington developmental landscape O. (d) The landscape 2.
7 CONCLUSION
We have introduced DeepRUOT for learning regularized unbalanced optimal transport (RUOT) and continuous unbalanced stochastic dynamics from time-series snapshot data. By leveraging Fisher regularization, our method transforms an SDE problem into an ODE constraint. Through the use of neural network modeling for growth and death, our framework models dynamics without requiring prior knowledge of these processes. We have demonstrated the effectiveness of our method on syn- thetic gene regulatory network, high-dimensional Gaussian Mixture Model, and single-cell RNA-seq data, showing its ability to eliminate false transitions caused by neglecting growth processes. Future directions involve extending the methodology to learn the cell latent embedding space jointly with the dynamics and developing to address these aspects and improve trajectory inference problems in various scenarios, including multi-omics data and other scenarios beyond biology in machine learning.
computationally efficient algorithms. We expect future work
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Dr. Hugo Lavenant (Bocconi) for helpful discussions on Section E.1 and Qiangwei Peng (PKU) for helpful suggestions on the algorithm. We thank Prof. Qing Nie (UCT) Prof. Chunhe Li (Fudan), and Juntan Liu (Fudan) for their insightful discussions. This work was supported by the National Key R&D Program of China (No. 2021 YFA 1003301 to T.L.) and National Natural Science Foundation of China (NSFC No. 12288100 101 to T.L.& & P.Z., and $206100646, T2321001 to P.Z.). We acknowledge the support from the High-pe begh-performance Computing Platform of Peking University for computation. We thank the anonymous referees for their valuable feedback and constructive suggestions.
REFERENCES
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint 1.Xiv:2303.08797 2023.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Aymeric Baradat and Hugo Lavenant. Regularized unbalanced optimal transport as entropy mini- mization with respect branching brownian motion. arXiv preprint arXiv:2111.01666, 2021.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Shirui Bian, Yunxin Zhang, and Chunhe Li. An improved approach for calculating energy landscape of gene networks from moment equations. Chaos: An Interdisciplinary Journal of Nonlinear Science, 33(2), 2023.
Shirui Bian, Ruisong Zhou, Wei Lin, and Chunhe Li. Quantifying energy landscape of oscillatory systems: Explosion, pre-solution, and diffusion decomposition. arXiv preprint arXiv:2401.06955 2024.
Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The schrödinger bridge be- tween gaussian measures has closed form. In International Conference on Artificial Intelligence and Statistics, pp. 5802-5833 PMLR, 2023a.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo. Mitch Levesque, Kjong- Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature methods, 20(11):159-1176, 2023b.
Charlotte Bunne, Geoffrey Schiebinger, Andreas Krause, Aviy Regev, and Marco Cuturi. Optimal transport for single-cell and spatial omies. Nature Reviews Methods Primers, 4(1):58, 2024.
Maciej Buze and Manh Hong Duong. Entropic regularisation of unbalanced optimal transportation problems. arXiv preprint arXiv:2305.022. 02410, 2023.
Zixuan Cang and Yanxiang Zhao. Synchronized optimal transport for joint modeling of dynamics across multiple spaces. arXiv preprint 0:X113 2024.
Kai Cao, Qiyu Gong, Yiguang Hong, and Lin Wan. A unified computational framework for single- cell data integration with optimal transport. Nature Communications, 13(1):7419, 2022.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary dif- ferential equations. Advances in Neural Information Processing Systems, 2018.
Tianrong Chen, Guan-Horg Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Represen- fations, 2022a.

---

Published as a conference paper at ICLR 2025
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:671 099:671-611,000 2016.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced marginals. SIAM Journal on Control and Optimization, 60(4):2016-2039) 2022b.
Qing Cheng, Qianqian Liu, Wenbin Chen, and Jie Shen. A new flow dynamic approach for wasser- stein gradient flows. arXiv preprint arXiv:2406. 14870, 2024a.
Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, and Yao Xie. Convergence flow-based generative models via proximal gradient descent in wasserstein space. IEEE Trunsactions Information Theory, 2024b.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. An interpolating distance between optimal transport and fisher-rao metrics. Foundations of Computational Math- ematics, 18:1-44, 2018a.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced opti- mal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 27)(11); 3090-3123 2018b.
Lénaïc Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. Trajectory inference via mean-field langevin in path space. Advances in Neural Information Processing Systems, 35; 16731-16742, 2022.
Shui-Nee Chow, Wuchen Li, and Haomin Zhou Wasserstein hamiltonian flows. Journal of Differ- ential Equations, 268(3): 1219, 2020.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-329, 1991.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695 2021.
Pinar Demetci, Rehecca Santorella, Björn Sandstede, William Stafford Noble, and Ritambhara Singh. Scot: single-cell multi-omics alignment with optimal transport. Journal of computational biology, 29(1):3-18, 2022.
Michael B Elowitz, Arnold J Levine, Eric D Siggia, and Peter S Swain. Stochastic gene expression in a single cell. Science, 18835111) 2002.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021.
Wilfrid Gangbo, Wuchen Li, Stanley Osher, and Michael Puthawala. Unnormalized optimal trans- port. Journal of Computational Physics, 399:108940, 2019.
Yuan Gao, Jian Huang, Yuling Jiao, and Shurong Zheng. Convergence continuous normalizing flows for learning probability distributions. arXiv preprint arXiv:2404.0055), 2024a.
Zizhan Gao, Kai Cao, and Lin Wan. Graspot: A graph attention network for spatial transeriptomics data integration with optimal transport. bioRxiv, pp. 2024-02, 2024b.
Ivan Gentil, Christian Léonard, and Luigia Ripani. About the analogy between optimal transport and minimal entropy. In Annales de la Faculté des sciences de Toulouse: Mathématiques, volume 26, pp. 569-600, 2017.
12

---

Published as a conference paper at ICLR 2025
Anming Gu, Edward Chien, and Kristjan Greenewald. Partially observed trajectory inference using optimal transport and a dynamics prior. arXiv preprint arXin:2406.07475, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transoort flows for trajectory inference. Advances in neural information processing systems, 35:29705-29718, 2022.
Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport be- tween unbalanced gaussian measures has a closed form. Advances in neural information process- ing systems, 33: 33:10468-10479, 2020.
Qi Jiang and Lin Wan. A physics-informed neural SDE network for learning cellular dynamics from time-series scRNA-seg data. Bioinformatics, 20-i1120-i127,, 09 2024. ISSN 1367-4811.
Qi Jiang, Shuo Zhang, and Lin Wan. Dynamic inference of cell developmental complex energy landscape from time series single-cell transcriptomic data. PLoS Computational Biology, 18(1): e1009821, 2022.
Yuling Jiao, Lican Kang, Huazhen Lin, Jin Liu, and Heng Zuo. Latent schrödinger bridge diffusion model for generative learning. arXiv preprint arX1v:2404.13309, 2024a.
Yuling Jiao, Yanming Lai, Yang Wang, and Bokai Yan. Convergence analysis of flow matching in latent space with transformers. arXiv preprint (xXix:2404002536 2024b.
Yijie Jin, Shu Liu, Hao Wu, Xiaojing Ye, and Haomin Zhou. Parameterized wasserstein gradient flow. arXiv preprint S4ti:22404.11333 2024.
Diederik Kingma and Max Welling. arXiv: 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Lactitia Mleng-Papaxanting Michael Sterr, Aimée Bastidas-Ponce Marta Tarquis-Meding et al. Mapping cells through time and space with moscot. bioRxiv, pp. 2023-02220000 05, 2023a.
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Generative entropic neural optimal transport to map within and across spaces. arXiv preprint arXiv:23110092254 2023b.
Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023.
Dana Lahat, Tülay Adali, and Christian Jutten. Multimodal data fusion: overview of methods, challenges, and prospects. Proceedings of the IEEE, 14999);449-1147, 2015.
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, Geoffrey Schiebinger, et al. Toward mathe- matical theory of trajectory inference. The Annals of Applied Probability, 34(1A):428-55) 500, 2024.
Flavien Léger and Wuchen Li. Hopf-cole transformation via generalized schrödinger bridge prob- lem. Journal of Differential Equations, 274:788-98 827, 2021.
Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems-Series A, 34(4):1533-1574, 2014.
Chunhe Li and Jin Wang. Quantifying cell fate decisions for differentiation and reprogramming of human stem cell network: landscape and biological paths. PLOS computational biology, 98): el e1003165, 2013.
Chunhe Li and Jin Wang. Landscape and flux reveal new global view and physical quantification of mammalian cell cycle. Proceedings of the National Academy of Sciences. :11(39):14130-U1 14130-14135, 2014.
13

---

Published as a conference paper at ICLR 2025
Wuchen Li, Jianfeng Lu, and Li Wang. Fisher information regularization schemes for wasserstein gradient flows. Journal of Computational Physics, 416:109449, 2020.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schrödinger bridge. In Advances in Neural Information Processing Systems, volume 35, pp- 9374-9388, 2022a.
Guan-Horn Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. Ib Image-to-image schrödinger bridge. arXiv preprint arXiv:2302.05882, 2023a.
Shu Liu, Shaojun Ma, Yongxin Chen, Hongyuan Zha, and Haomin Zhou. Learning high dimensional wasserstein geodesics, arXiv preprint arXiv:2102.0299) 2021.
Shu Liu, Wuchen Li, Hongyuan Zha, and Haomin Zhou. Neural parametric cokker-plancke equation. SIAM Journal on Numerical Analysis, 60(3): 1385-133814449 2022b.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Repre- sentations, 2023b.
Suryanarayana Maddu, Victor Chardès, Michael Shelley, et al. Inferring biological processes with intrinsic noise from cross-sectional data. arXiv preprint arXiv:2410.07500 2024.
Tingwei Meng, Zongren Zou, Jérôme Darbon, and George Em Karniadakis. Hj-sampler: A bayesian sampler for inverse problems of a stochastic process by leveraging hamilton-jacobi pdes and score-based generative models. arXiv preprint 2xXi::2409.09114 2024.
Toshio Mikami and Michele Thicullen. Optimal transportation problem by stochastic optimal con- trol. SIAM Journal Control and Optimization, 47(3):1127-11333 2008
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learn- ing stochastic dynamics from samples. In International conference machine learning, pp- 25858-25889, PMLR, 2023.
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. A computational framework for solving wasserstein lagrangian flows. In Forty-first International Conference Machine Learning, 2024.
Matteo Pariset, Ya-Ping Hsieh, Charlotte Bunne, Andreas Krause, and Valentin De Bortoli. Unbal- anced diffusion schrödinger bridge. ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
Qiangwei Peng. Peijie Zhou, and Tiejun Li. stver: Reconstructing spatio-tempon dynamics of cell development using optimal transport. bioRxiv, pp. 2024-06, 2024.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trendy in Muchine Learning, 11(5-6)355-6077 2019.
Aram-Alexandre Pooladian and Jonathan Niles-Weed. Plug-in estimation schrödinger bridges.
arXiv preprint urXiv:2408. 1686, 2024.
Aram-Alexandre Pooladian, Carles Domingo-Enrich Ricky T. Q. Chen, and Brandon Amos. Neural optimal transport with lagrangian costs. In The 40th Conference Uncertainty in Artificial Intelligence, 2024.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
14

---

Published as a conference paper at ICLR 2025
Lars Ruthotto, Stanley J Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving nigh-dimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, 117(17):1133-1195 2020.
Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature biotechnology. 37(5):547-554 2019.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Arych Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal- analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 943, 2019.
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.
Jifan Shi, Kazuyuki Aihara, Tiejun Li, and Luonan Chen. Energy landscape decomposition for cell differentiation with proliferation effect. National Science Review, 9(8): 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger
bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2265.2265 PMLR, 2015.
Yang Song. Jascha
Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526-9536. PMLR, 2020.
Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P San Juan, Laura Rangel, Brandon Zhu, John G Lock, Christine I. Chaffer, and Smita Krishnaswamy. Learning transcrip- tional and regulatory dynamics driving cancer cell plasticity using neural ode-based optimal trans- port. bioRxiv, pp. 2023-03, 2023.
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024a. ISSN 2835-8856. Expert Certification.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-frre schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pp. 1279-1287. PMLR, 2024b.
Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. Trajectory inference for branching sde model of cell differentiation. arXiv preprint arXiv:2307.0768/ 2023.
Wei Wan, Yuejin Zhang, Chenglong Bao, Bin Dong, and Zuoqiang Shi. A scalable deep learning approach for solving high- nigh-dimensional dynamic optimal transport. SIAM Journal on Scientific Computing, 45(4): B5633, 2023.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International conference on machine learning, pp. 10894, PMLR, 2021.
Jin Wang, Chunhe Li. and Erkang Wang. Potential and flux landscapes quantify the stability and robustness of budding yeast cell cycle network. Proceedings of the National Academy of Sciences, 107(18):8 109881980000 2010.
15

---

Published as a conference paper at ICLR 2025
Caleb Weinreb, Alejo Rodriiuez-FFaaticlll, Fraticelli, Fernando D Camargo. and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367(6479); caaw3381, 2020.
Hao Wu, Shu Liu, Xiaojing Ye, and Haomin Zhou. Parameterized wasserstein hamiltonian flow.
arXiv preprint wwXiv:2306.00191 2023.
Jiaqi Zhang, Erica I Larschan, Jeremy Bigness, and Ritambhara Singh. scNODE: generative model for temporal single cell transcriptomic data prediction. Bioinformatics, 40(Supplement.2:ii ii15, 09 2024a. ISSN 1367-4811.
Peng Zhang, Ting Gao, Jin Guo, and Jinqiao Duan. Action functional as early warning indicator in the space of probability measures. arXiv preprint arXiv:2403./04000 2024b.
Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. Optimal transport analysis reveals trajectories in steady-state systems. PLoS com- putational biology, 17112)ee0094466 2021.
Yue Zhao, Wei Zhang, and Tiejun Li. EPR-Net: constructing non-equilibrium potential landscape
via variational force projection formulation. National Science Review, 11(7), 2024.
Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models. In The Twelfth International Conference on Learning Representations, 2024a.
Mo Zhou, Stanley Osher, and Wuchen Li. Score-based neural ordinary differential equations for computing mean field control problems. arXiv preprint arXiv:2409.164777 2024b.
Peijie Zhou and Tiejun Li. Construction of the landscape for multi-stable systems: Potential land- scape, quasi--poenttial, a-type integral and beyond. The Journal of chemical physics, 144(9), 2016.
Peijie Zhou, Xin Gao, Xiaoli Li, Linxi Li, Caoyuan Niu, Qi Ouyang, Huiqiang Lou, Tiejun Li, and Fangting Li. Stochasticity triggers activation of the s-phase checkpoint pathway in budding yeast. Physical Review X, 1101),011044, 2021a.
Peijie Zhou, Shuxiong Wang, Tiejun Li, and Qing Nie. Dissecting transition cells from single-cell transcriptome data through multiscale stochastic dynamics. Nature communications, 12(1) ):5609, 2021b.
Peijic Zhou, Federico Bocci, Tiejun Li, and Qing Nic. Spatial transition tensor of single cells. Nature Methods, pp. 1-10, 2024c.
Ruisong Zhou, Yuguo Yu, and Chunhe Li. Revealing neural dynamical structure of c. elegans with deep learning. Iscience, 27(5), 2024d.
A TRAINING REGULARIZED UNBALANCED OPTIMAL TRANSPORT
A. TRAINING OF REGULARIZED UNBALANCED OPTIMAL TRANSPORT
After the pre-training stage, once we have obtained initial væ ge and sø, then we can continue training by minimizing the total loss:
L -Ennery + Xy CRecons Xy Crp.
To compute the each component of the loss, the temporal integral and ODEs were numerically solved using the Neural ODE solver (Chen al., 2018). The gradients of the loss function with respect to the parameters in the neural networks for v(x, t), g(æ, f) and s(x, were computed us- ing neural ODEs with a memory memory-efficienn between discrete distributions, efficient implementation. For computing the Wasserstein distance utilize the implementation provided by the Python Optimal Trans- port library (POT) (Flamary et al., 2021). The algorithm is detailed in Algorithm (Training Stage).
16

---

Published as a conference paper at ICLR 2025
A.2 TRAINING LOG DENSITY FUNCTION
To utilize conditional flow matching (CFM) learn initial log density, first choose pair samples (æo, 21) from the optimal transport plan g( v1) and then construct Brown-
is constant. The log density function is then matched with these Brownian bridges, i... p(æ, (zo,
bridges between these pair samples. We first consider at)
N (a;; la0 (1 D) a²1(1 0) and V log, (aro, m11) (en +11 +(1-)
)- E [0, 1].
We then utilize the fully connected neural networks so so(x, t) to approximate the log density function o² log p(z, with weighting function Ag. Then we have
Lue 2 V
2 V vo
And the correspongding CFM loss is
Lwcrr
Eqr V
v800
= vn log dogp(a,, (ro, o. 1)12
2
where Q' = (1 1)) & ( æi) @ p(xr
(m 11). We take the weighting function
2 2
(t) σ²
0 /
-
Then we have
0² Luorre =\,(1)² Voso(ar11) -
Va logp(x 2o.
2
V -\>((1V050000011) -a(()
(13)
logp (xc
2o, x2
2
where e (0, I). This also numerically stable. For the case where (t) is not constant, similar approach can be applied. One may refer to (Tong et al., 2024b) for further details.
A.3 ESTIMATING INITIAL DISTRIBUTION FOR FOKKER- PLANCK EQUATION
The initial distribution is estimated by a Gaussian mixture model to generate density
No exp (- (a (lo))) Σ (x ar((1)))
p(x, 0) =po(x)
E = ol E Radd
E
A.4 ENERGY LOSS
In (Sha et al., 2024), it introduces the following theorem to transform the high-dimensional integral to the Monte Carlo integral.
Theorem A.I. If smooth density p(æ, () Rd x |0, 1] + velocity fildd v(x, () Rd (0, 1] + Bd
and growth rute g(æ, 1) Rdd x [0. 1] satisfy
0.plx, t) + V (v(a, ()p 1)) = g(x, t)p(x,t), p(x, 0) po(x),
for all 0 S t < 1 with
dmi(t) di v(x, t) and (0) = wo. then for any measurable function f(x, t) Rd [0, + Rdd have So Sae ((x (),( t)dard Eo So f(x t)esd s(s,s)ds
17

---

Published as a conference paper at ICLR 2025
B ADDITIONAL RESULTS
B.1 SYNTHETIC GENE REGULATORY NETWORK
The system dynamics are governed by the following set of stochastic ordinary differential equations (ODEs):
dX,
01 x
81X1 + 7
dt
1+ aa + X2 + 73 +
dX2
Xx + 3
82X₂ + T15 1+ + 2 + 73 +
d/
dX3
03X3
1+a3X3 6gX3 + mahe.
dt
Genes X, and X2 mutually inhibit each other and self-activate, forming a toggle switch, An external signal 3 uniformly activates both X1 and X2 independently of gene expression levels. Gene X3 in- hibits the expression of both 11 and X2. Here, (t) represents the concentration of gene i at time t, with a; and "/ denoting the strengths of self-activation and inhibition, respectively. The parameters 5, represent gene degradation rates, while 77.i accounts for the stochastic effects with additive white noise. The probability cell division correlates positively with the expression of X2, calculated
x %. Upon division, cells inherit the gene expression states (t) Xt(,,, xt)))) 9 agg+tt
of the parent cell, subject to independent perturbations ndN((, 1) for each gene, and transition independently thereafter. The hyper-parameters are listed at Table 3. The initial cells cho- sen independently and identically distributed from two normal distributions ([2, 0.2, 0J. 0.01) and N([0, 0, 2), 0.01). At each step, we corrected the negative expression value to
Table 3: Simulation parameters on gene regulatory network.
Parameter
Value
Description
Q1
0.5
Strength of self-activation for X1
71
0.5
Strength of inhibition by Xg on X1 Strength of self-activation for X2 Strength of inhibition by X3 on X₂ Strength of self- seff-activation for X3
1
2/2
1
73
10
Half-saturation constant for inhibition terms
81
0.4
Degradation rate for Xi Degradation rate for X₂ Degradation rate for X3 Noise intensity for X; Noise intensity for X₂ Noise intensity for X3
82
0.4
83
0.4
17
0.05
1/2
0.05
133
0.01
Tld 0.014
Noise intensity for cell perturbations
8
1
External signal activating X1 and X₂
dt
1
Time step size
Time Points
[0. 8, 16. 22]
Time points at which data is recorded
We conducted a comprehensive evaluation of the methodologies proposed by (Neklyudov et al., 2023; Pariset al., 2023) utilizing our simulated dataset. As detailed in Table 1, our approach con- sistently outperforms the referenced methods across multiple quantitative metrics, thereby demon- strating its superior ability to capture the underlying dynamics of the system. Furthermore, Fig. provides the outcomes produced by the UDSB (Pariset et al., 2023), While their method predicts an overall increase in cell population, a closer examination reveals the presence of false dynamics within the predicted transitions.
B.2 SYNTHETIC GAUSSIAN MIXTURES
For the initial distribution, we generated 400 samples from the Gaussian located lower in the (æ) x2) plane, and 100 samples from the Gaussian positioned higher. For the final distribution, generated
18

---

Published as a conference paper at ICLR 2025
(a)
(b)
Alive cells (f)
Gene X2
Initial
mass
lime
Gene X1
Figure 4: Results obtained by UDSB (Pariset et al., 2023) on gene regulatory network. (a) The trajectory learned UDSB, where black dots indicate particle death, red dots signify particle growth, orange dots represent the target distribution, dark blue dots denote the source distribution, and gradient blue dots illustrate particle trajectories. (b) Predicted changes in cell population intermediate time points, with dots representing the actual mass.
1,000 samples from the upper Gaussian, and 200 samples each from the two lower Gaussians. We then tested the RUOT-based model's ability to learn the stochastic dynamics using the samples generated in this manner.
(a)
(b)
Learnt growth rate
Ground Fredicted Trajectory
TDo0
11.0
(c)
(d)
25
2445
3305
5.2
16
-E5
66
1.6 0.0
-11
Figure 5: Results of DeepRUOT on Gaussian mixtures (o 0.1, 10D). (a) The learned trajectory by DeepRUOT (a 0.1). (b) The growth rate inferred by our model. (c) The Waddington devel-
opmental landscape t 0 (o 0.1). (d) The Waddington developmental landscape at t
(o = 0.1).
19

---

Published as a conference paper at ICLR 2025
To further evaluate the scalability of our algorithm, simulated Gaussian mixture models with 20, 50, and 100 dimensions, as illustrated in Fig. We similarly visualize the results on the first and second dimensions. The results demonstrate that our model remains effective and applicable across these higher-dimensional settings.
Similar to the frameworks such TrajectoryNet (Tong al., 2020), MIOflow (Huguet al., 2022) and (Koshizuka Sato, 2023), which involve simulating an ODE/SDE and performing optimal transport for distribution matching, the computational cost and scalability of our method are com- parable to these existing approaches. Notably, the score matching in our method is conducted via conditional flow matching (Lipman et al., 2023; Tong al., 2024a;b), which is simulation- free and thus highly efficient. Consequently, our approach does not introduce significant additional com- putational overhead, since the cost of each component similar to previous works. Meanwhile, when it extends to thousands of dimensions, our approach similarly encounters challenges, necessi- tating the development of simulation-fre dimensional settings effectively (Lipman et al., 2023; Tong et al., 2024a;b). However, in the field of single-cell biology, dimensionality reduction of gene expression data routinely employed, making hundreds of dimensional space sufficient for practical applications.
training methods akin to flow matching handle higher-
Trajectory
Learnt growth rate
t=0 landiscape
landscape
20D
(b)
50D
(c)
100D
foo
Figure 6: Results of DeepRUOT on higher dimensional Gaussian mixtures (a = 0.1). The first column presents the trajectories learned by DeepRUOT (a 0.1), the second column displays the growth rates inferred by our model, the third column shows the Waddington developmental land- scape at = (a 0.1), and the fourth column depicts the Waddington developmental landscape at f 1 (a 0.1). (a) 20 dimensions, (b) 50 dimensions, (c) 100 dimensions.
B.3 SINGLE CELL DYNAMICS IN MOUSE HEMATOPOIESIS
We conducted a comparative evaluation of the methodologics proposed by (Neklyudov et al., 2023; Pariset al., 2023) using our mouse hematopoiesis dataset. As detailed in Table 1, our approach consistently outperforms the referenced methods across range of quantitative metrics, underscor- ing its superior efficacy in capturing the underlying biological dynamics. In Fig. 7, we present the results generated by the UDSB (Pariset et al., 2023). Fig. 7(a) displays the learned trajectory, where black dots indicate particle death, red dots signify particle growth, orange dots represent the target distribution, dark blue dots denote the source distribution, and gradient blue lines illustrate the parti- cle trajectories. This visualization highlights the dynamic trajectory cell population changes over time. Fig. 7(b) illustrates the predicted changes in cell population intermediate time points, with
20

---

Published as a conference paper at ICLR 2025
each dot representing the actual mass observed. Finally, Fig. 7(c) shows the predicted cell distribu- tions at various time points in red, while the blue distribution at the final time point represents the initial distribution. This panel underscores the model's 's ability to track and predict the evolution of cell distributions from the initial state to subsequent stages.
Notably, although this method achieves marginally higher Wi metric at the final time point, closer examination reveals that the predicted distributions at intermediate time points, as illustrated in Fig. 7(c), deviate from the data distributions. This indicates that the (Pariset al., 2023) method may have challenges to accurately model the transitional dynamics that occur between the initial and final stages of hematopoiesis.
Furthermore, Fig. 7(a) and Fig. 7(b) demonstrate that the growth locations predicted by the UDSB algorithm are consistent with those identified by our model. This consistency provides form of cross- vooss-valiittion, affirming that our algorithm effectively captures both the growth and migration dynamics inherent in the cell population. So the ability of our model to maintain accurate predictions across all time points, including the intermediate stages, demonstrates its enhanced capacity for modeling unbalanced dynamics within complex biological systems.
(a)
(b)
Alive cells
Umap2
Initici
Lime
Umap1
(c)
Time
Figure 7: Results obtained by UDSB (Pariset et al., 2023) mouse hematopoiesis. (a) The tra- jectory learned by UDSB, where black dots indicate particle death, red dots signify particle growth, orange dots represent the target distribution, dark blue dots denote the source distribution, and gradi- ent blue dots illustrate particle trajectories. (b) Predicted changes in cell population at intermediate time points, with dots representing the actual mass. (c) Red denotes the predicted cell distributions at various time points, while the blue distribution at the final time point represents the initial distri- bution.
B.4 SINGLE CELL DYNAMICS IN EMT
Subsequently, applied DeepRUOT to time-series single-cell RNA sequencing dataset derived from the A549 cancer cell line. In this study, cells were treated with TGFB1 to induce epithelial- mesenchymal transition (EMT) during the initial four-time points (Sha et al., 2024). Cells harvested
21

---

Published as a conference paper at ICLR 2025
at each time point were cultured in vitro with identical initial cell numbers, ensuring that the cell counts directly reflect the dynamics of the cell population over time.
We employed Principal Component Analysis (PCA) to reduce the dimensionality of our single-cell RNA sequencing (scRNA-seq) data to a ten-dimensional latent space. This latent representation was subsequently utilized the input for DeepRUOT. To aid in the interpretation of the results, we projected the algorithm's 's outputs onto the first and second principal components, which capture the majority of the variance in the data and thus provide meaningful insights into the underlying biological phenomena.
The trajectories inferred by DeepRUOT exhibited consistent and biologically plausible transition dynamics, effectively mapping the progression of cells through different developmental states. Fur- thermore, we visualized the growth rates estimated by the algorithm alongside the developmental landscapes at various time points, providing comprehensive view of both the quantitative and qualitative aspects of cellular dynamics.
Notably, the growth patterns inferred by DeepRUOT displayed elevated values during the initial and intermediate stage epithelial- meseehhhmmlll transition (EMT) compared the epithelial (E) and mesenchymal (M) stages, as illustrated in Fig. 8. This observation aligns with previous studies that have reported enhanced stemness and proliferative capacity in cells at the intermediate stage.
Moreover, our landscape analysis revealed that the inferred developmental landscapes are consistent with the observed data distributions at corresponding time points. The developmental landscapes provide representation of the cell states, illustrating how cells traverse through different regions of the latent space as they progress through various stages of differentiation and migration.
(a)
(b)
Learnt growth rate
Coear
(c)
(d)
-0
33204
0 0.1090
01805
33304
3.8
32214
0.A
-0.1292
0.1755
0.174
0:2
32234
64
0.2
0.0
1.0
0.0800
0.2
37354
Figure 8: Results of DeepRUOT on EMT SeRNA-seg data 0.05). (a) The stochastic dynam- ies learned by RUOT (a = 0.05). (b) The growth rates learned by DeepRUOT (a = 0.05). (c) The constructed Waddington developmental landscape at (a 0.05). (d) The landscape at
(o = 0.05).
22

---

Published as a conference paper at ICLR 2025
B.5 ABLATION STUDIES
We conducted ablation experiments the gene regulatory dataset to evaluate the contributions of key components in our algorithm. Specifically, we examined the impact of the growth term g(r, t)- Then we assessed the necessity of the mass loss term LMas within the reconstruction loss function RRoon and CFp in the total loss (9). These analyses demonstrate the critical roles that the growth term, the mass loss component, and the Fokker-Plann constraint play in enhancing the performance and accuracy of our model.
Table 4: Ablation studies across five runs on gene regulatory data (o 0.25). We show the mean value with one standard deviation.
Model
W
11/2
11/2
W2
W;
W2
SF2M (Tong et al., 2024b)
0.17
0.303
0.490
0.719
0.686
1.050
a
1.242
DeepRUOT wio growth
a.177
0.198
0.32440
0.371
0.48110ma
0.376.20m:
0.772.0000
DeppRUOT
0.150
0.186
0.308
0.349
0.465
0.553
0.600
0.806=0.00
DeegRUIC for
gov00
0.107
0.108200m
0.132
0.135
0.156
0.160
0.197-mom
DeepRUOT pre training
0.737 =aoce
1.113g0017
1.910
1.529
1.542
1.005_000
DeepRUCT wie score matching
1.291 aoucer
1. 400
2,051
2.18740004
2.675
2 2.758:uumm
3.301-0000
DeepRliCT wio mainngg
0.089
0.108
0.102
0.1324urnz
0.134
0. 0.158a00ms
D. 0.1632m0mm
0.2000
DeepRLiCT
(10195
0.115.g0.00
0.104g0
0.130
0.130gg
0.1440.00:
0.168
We observe that when our algorithm does not account for growth, i.e., by setting g(x,t) 0, it reduces to regularized optimal transport, which is equivalent to the Schrödinger bridge problem. Consequently, SF2M (Tong et al., 2024b) corresponds to the scenario where g(æ, f) 0 within our approach. Our experiments have demonstrated that omitting the growth factor leads to inaccurate trajectory reconstructions, underscoring the necessity of modeling unbalanced dynamics. To fur- ther substantiate this finding, we conducted an explicit ablation study with g(x, t) (Table 4). The results from this experiment align with those of the SF2M, confirming that disabling the growth rate adversely affects performance by neglecting essential growth and death processes. This com- parison reinforces the importance of the growth rate component in accurately capturing unbalanced stochastic dynamics, thereby validating the enhanced performance of DecpRUOT.
Next, we assess the impact of the mass loss term, LMsss, within the reconstruction loss function CRecons = Am Mass + Cor during all training stage. By setting Am = 0, we observed that the ex- clusion of this component results in inaccurate trajectory reconstructions (see Table 4). This finding underscores the crucial role of incorporating local mass matching in unbalanced settings, demon- strating that the omission of Lass significantly impairs the model's ability to accurately capture the system's is dynamics.
Subsequently, we investigate the training stage using the loss funetion L Energy + X, Cecoons + Xy Crp. By setting 'y 0, thereby excluding the Fokker-Planck constraint, the results presented in Table reveal that incorporating this loss term enhances our model's performance, particularly at the final time point. These results collectively highlight the importance of both mass matching and the Fokker-Plann constraint in improving the accuracy and reliability of trajectory reconstructions within our framework.
Furthermore, we investigated the impact of different stages in our training process on the final re- sults. Our training algorithm (Algorithm 1) comprises two stages: the pre-training stage and the training stage. To assess their respective contributions, we conducted experiments under three sce- narios: (1) omitting the pre-training stage, (2) excluding score matching during the pre-training stage, and (3) removing the training stage. As presented in Table 4, we observed that the algorithm failed to converge in the first two scenarios. Specifically,, scenario (2) yielded poorer performance due to inaccurate score matching, which caused particle trajectories to deviate from the correct dy- namical paths. In scenario (1), the necessity to balance multiple loss functions resulted in the vector field and score not being adequately learned, leading particles are inclined to remain near their initial positions. scenario (3), the absence of the training stage slightly degraded the performance met- rics, especially at later time points. These ablation studies demonstrate that the pre-training stage is critical within our framework, while the training stage contributes to the refinement of final results, ensuring we are addressing the regularized optimal transport problem, which is theoretically neces- sary. Empirical results indicate that the pre-training stage significantly enhances the performance of our algorithm.
23

---

Published as a conference paper at ICLR 2025
C EXPERIENTIAL DETAILS
C.I EVALUATION METRICS
We evaluate the empirical 1- I-Wasserstein distance (Wi) and 2- Wassersstein distance (W2) on the gene data and scRNA-seq data. It is defined as follows:
/ velddip vu)))) (nEI((p-))
1/2
W2(p.9)
min
and
/
(p min =Ell(ppg)
||x y122d
where p and q represent empirical distributions.
We evaluate our method by first learning the dynamics using data from all time points. We then apply these inferred dynamics to the initial data point generate data for subsequent time points. Next, we compute the W1 and W2 distances between the generated data and the real data. For the gene regulatory network data, we ensure that the number of generated data points matches that of the real data since we have the true dynamics. In the scRNA-seq dataset, since the true dynamics are unknown, we apply the inferred dynamics to the initial data and compute the W1 and W22 distances between the generated data and all the data at that time point. Because our method could assign weights to the sample points, we compute the weighted Wi and W/ distances accordingly in the evaluation of our methods. We take the weight computed by DeepRUOT without diffusion to use for simplicity. For other methods we used the default uniform distribution evaluate.
To evaluate the performance of SF2M and DeepRUOT on scRNA-seq datasets we take the same diffusion coefficient as a = 0.25. To evaluate the performance on synthetic gene regulatory data, we choose 0.05 for our DeepRUOT method and a = 0.25 for SF2M. This is due to we observe some numerical instability when using SF2M with o = 0.05. So to ensure a fair comparison, we choose the parameter as default in SF2M.
To evaluate the performance of (Sha et al., 2024) on gene regulatory network data and single-cell RNA sequencing (scRNA-seq) mouse hematopoiesis data, observed certain instabilities on the datasets, to ensure a fair comparison, we independently reimplemented the method and conducted additional evaluations (by setting = 0 in DeepRUOT).
To evaluate the performance of (Neklyudov et al., 2023) and (Pariset et al., 2023) on gene regulatory network data and single-cell RNA sequencing (scRNA-seq) mouse hematopoiesis data, we utilized their default parameter settings. Specifically, for (Pariset et al., 2023), whose default configuration involves three time points, we selected samples at time points 2, and from the gene regulatory network dataset as inputs for the algorithms.
C.2 HYPERPARAMETERS SELECTION AND LOss WEIGHTING
We use one A100 Nvidia GPU along with 16 CPU cores for computation at a shared high- performance computing cluster. For the two-stage training of DeepRUOT, each stage necessitates the selection of hyperparameters tailored to its specific requirements, and here we will provide a guideline based on empirical studies.
In the pre-training phase, we need to select the parameters in Croons i.e., Am and A in Lrocnms Am MMsss + AaCoT- Parameter selection important during this stage, as the ablation studies we have shown that both components are important. To effectively manage these parameters, we implement a hyperparameter scheduling strategy during the pre-training stage. Initially, we assign predefined values to An and Aa and train the model for specified number of epochs. Subsequently, we adjust Am to zero (Am, = 0) and conduct additional epochs of training. This strategy is grounded in the rationale of first aligning growth patterns and subsequently refining the transition process based on the established growth rates. The parameter scheduling procedure is elucidated in Table 5, where the arrow notation (-) signifies the rescheduling of hyperparameter values during the pre- training stage. Specifically, each arrow indicates the change from an initial value to a subsequent value, reflecting the strategic adjustment of hyperparameters during the pre-training stage.
24

---

Published as a conference paper at ICLR 2025
During the training stage, the pre-training phase provides robust initial values, thereby reducing the sensitivity to parameter variations in L = Eieerg + X, Recons + Ly As a result, we utilize the set of parameters and train the model for 10 epochs consistently across different datasets.
As shown in Table 5, the primary hyperparameters requiring adjustment during the pre-training stage are mainly the number of training epochs selected during the Am resetting process, specifically the number of epochs trained before Am resetting and the number of epochs trained after Am resetting. As previously discussed, the pre-training stage crucial for stabilizing the training process. Deter- mining the number of epochs relies on the complexity and scale of the problem, including factors such as the number of time points and the size of the dataset. For the first epoch number, a default of 30 epochs is recommended. However, in the case of the Gene Regulatory Network dataset, we opted for 20 epochs due to larger number of time points and greater data volume. Similarly, for the EMT dataset, 20 epochs were chosen as it had already achieved convergence.
Table 5: Parameter Settings for Different Datasets Across Two Training Stages (Regulatory Network (a = 0.25). Mouse Hematopoiesis (a 0.25), EMT (a 0.05). Gaussian Mixtures (a = 0.1)). The arrow (-) indicates the scheduling of parameter values within the Pre-Trainin Stage, where the numbers on the left indicate the hyperparameter values before resetting, and the right ones indicate the values after resetting. For example, the (1.0, 0.1, 20) - (0.0, 0.1, 10) entry in (Am, λd, Epochs) row denotes 20 epochs training when Am was 1.0, Aa was 0.1, and another 10 epochs training after Am was reset to 0.0.
Regalabery Netwok
EXIT
SED
PoeTinading Stge Diperpiranen Schilligg
18
11
63
R1 81
In Table 6, we evaluated the GRN dataset by maintaining the second epoch number (trained after Am resetting) as 10 while varying the first epoch number (trained before resetting) among 10, 20, and 30, with the other hyperparameters chosen as Table 5. The results indicated that the first epoch number value of 10, which was insufficient, deteriorated performance, whereas the first epoch num- ber values of 20 and 30 showed negligible differences. This suggests that increasing the number of pre-training epochs does not adversely affect the training process; hence, we recommend training for a higher number of epochs when feasible. The selection of the second epoch number varies across datasets based on convergence criteria. For instance, in the GRN dataset, training with the to terminate further train-
second epoch number set to 10 was sufficient for convergence, leading ing. Conversely, for the EMT dataset, convergence was achieved during the first stage, negating the need for additional second epoch number training. Generally, a higher number of epochs in this stage enhances training efficacy. In the case of Gaussian Mixtures, the number of the second epoch number increases with dimensionality. However, for 50D and 100D, the second epoch number de- creases because, for 100D, training can extend beyond 120 epochs for improved results, while 50D converges around 100 epochs, making further training unnecessary. Therefore, for the adjustment of the second epoch number, we recommend that users train until convergence based on their specific data requirements.
Table 6: Results of different epoch scheduling during the pre-training stage on gene regulatory dataset. Here we use the arrow to represent the number of epochs trained before(left) and (fter(right) the Am resetting, for example, 10 + 20 denoting 10 epochs of training when Am and 20 epochs of training when Am 0.
t
3
4
Epoch Scheduling
W
W2
W1
W9
Wi
W2
W/
W2
10 - 10
0.138 goom
0.173 jounee
0,243 Lo.00s
0.278 Lacos
0.353 Lom
0.106
0. 1555_ooo
10
0. 2.000
0.115
0.104gn.ozz
0. 43000
2.
0.1400noorr
0.108 goomr
30 + 10
0. 107 #0.006
0.
0.128
0. 0.120
0. 1.2
0. 130
0. 1..0000000
0.154g0012
0.188 1818
In the training stage, we empirically identified A as a hyperparameter requiring tuning. We found that Xn may relate the noise amplitude a. Specifically, for the Gene Regulatory Network (GRN)
25

---

Published as a conference paper at ICLR 2025
and Mouse Hematopoiesis datasets, where o is set to 0.25, smaller Aw is necessary to maintain computational stability; setting Ag to in these cases resulted in instability. For the remaining datasets, we have found that selecting Aw within the range of [0.1, 1] is feasible, variations within this range do not significantly impact the training outcomes. In Table 7, we present the results of training with Xw values of 0.1, 0.4, 0.7, and 1.0 EMT dataset. We computed the values of the trained results (v(æ,t) t), s) at each data point, using the results obtained with Au = as the baseline. The distances between the outcomes of other Aw settings and the baseline were measured using Euclidean distance for inferred v(z,t) t) and mean squared error (MSE) for g and As illustrated in Table 7, we observed that as long as Au ensures computational stability, the outcomes for different A. values are consistent. Therefore, we recommend selecting smaller A with a default value of 0.01 (for smaller σ)) or 0.001 (for larger σ), to achieve stable and consistent training results.
Table 7: Results with different A settings on the EMT dataset (a = 0.25). Distances for v are measured using Euclidean distance, while g and evaluated using mean squared error (MSE).
Am
v Dist
MSE
8 MSE
0.1
2.64 x 10-8
6.61 x 10-16
1.63 x 10-9
0.4
2.20 x 10-8
3.80 x 10-16
7.23 x 10-10
0.7
2.11 x 10-8
3.46 10-16 16
1.80 10-10
1.0
0.00
0.00
0.00
C.3 TRAINING INITIAL LoG DENSITY FUNCTION
In our gene regulatory network example training, we utilize the following steps to stabilize train- ing. Here since we parameterize e" log p(æ, t) using sx(x,t) when we compute Fokker-Planck constrained loss, we need to do the exponential operations, which may cause numerical instability during computations. To mitigate this issue, we augment the CFM loss with an additional penalty
term defined as L Openally max(s5((x t), 0) prior to training so sa(x, t). This penalty encourages sp sg(2, t) to adopt negative values, thereby preventing instability caused by the subsequent exponen- tial function. Our training procedure is conducted in two stages:
1. Initial Training Phase:
We optimize the combined CFM loss and penalty term for 3,000 epochs. This stage ensures that remains negative, stabilizing the training process by avoiding potential numerical issues associated with positive values of 8.
2. Secondary Training Phase: After the initial phase, remove the penalty term and con- tinue training solely with the CFM loss for an additional 6,000 epochs. This allows the model to fine-tune the parameter without the constraint imposed by the penalty, facilitat- ing more accurate learning.
For the last two examples, such stabilization measures unnecessary, as the dynamics of these cases do not induce the same level of numerical instability. Consequently, we proceed with training using only the CFM loss without incorporating the penalty term.
This issue may be avoided we use the Vamiton--Jccobi-ellln (HJB) equations to set the con- straints rather than the direct Fokker-Planck equations (Jiang & Wan, 2024; Meng et al., 2024; Zhao et al., 2024; Zhou et al., 2024b).
D TECHNICAL DETAILS FOR SCHRÖDINGER BRIDGE AND RUOT
D.I PROOF OF DYNAMICAL FORMULATION OF SCHRÖDINGER BRIDGE
Proof. For Schrödinger bridge problems, the optimal solution can be found within the class of SDEs:
X: 1 dX, = b (Xe,, dt oo (X,,t) dW, The Pokker-Planck equation for the SDE is
0,p(x,t) = Va (p(a, (2x t)) + 2
v2 (a(2,t)p(a,t)
26

---

Published as a conference paper at ICLR 2025
where (x oax,t)o (x, Then we need to compute the KL divergence of these two
stochastic processes. Define the auxiliary variable
(xx,)) = (x. 1) (b(x,,
and apply the Girsanov's theorem, we get
DKL. x (H0.I)llH0.11) = Er E ((r,
-Er
In t)12at)
(14)
where (& hejua an exponential martingale with respect to the Brownian filtrations (Fi} defined
& = exp (L 7( D.W, 1
2 J. In a1
is well-known that the exponential martingale & satisfies the SDE
de, = (,n(Y,,))W))
Eo =1
So we obtain
& 6=
+
Since S/ yY,, ()dW, = O, we get
= L/ (r(W), L rrxx..oaw...
E Jn 7(Y ((x,,t)d,,
E [ & (1l7(Y,,))1)22
E (I ()) F.) dt.
E
E EE (1 011) dr,
&, I In((Y,,))
Combine with Eq. (14), we obtain that
DeL. (, = EE E1 (f/))
2 ( .001 at ]
/ 11
EE
(1 1)11 dt.
(x(x,t))l2d
= S/'' J b t)(2 t) b(0,
1) ((x),
p( t)ddd..
The proof is done.
D.2 PROOF OF FISHER REGULARIZATION OF RUOT
Proof. From (6) obtain
O.p = -V - (pb) + - V2 (o²(1)Tp) + +g(Ø, t)p,
1
2
= - ((0=
0²
(t) - logp p + g(x, t)p.
2
Using the change of variable v(x t) b(x, t) 2 logp, we see that is equivalent
0.p V. (pr(ar, ()) g l)p.
27

---

Published as a conference paper at ICLR 2025
Correspondingly, the integrand in the objective functional becomes 1 (v(2,, logp) + av (0) p(a, ()ddd..
! ///B Rd
1 1 || 1)22+ V log +
(t)
(15)
Letting H (p(x, 1)) = Sas p(a, 1) log l) be the entropy, we have
(1)H (p(a, 1)) (0) (p(x,, 0) S'' àe (t) (p(2, t)) dt,
I J (1
ad dt p( log 1) (1) (1 + log 1)) 0.p(ar, l)dxdd,
I' Jzd (1) (~ (p(x, t) t)) + gp(2, t)) +
do² dt p(a, t) log t)dadt,
L. J.c. (V)pla, o(x v)x (0) log pp
do p(a, log p(,, t)drd..
dt
Therefore,
J0 [(oP(t)V> vog, f))] t)dxdt (o (1)H (p(x: 1)) (p(x))(p((x
S.' / o2) (1 + log 1)) gpdard
L / do²(t)
di p(æ 1) log 1 dxdd
D.3 RUOT WITH NONLINEAR FOKKER- POKKER-PLANCK EQUATION CONSTRAINTS
Note that when & (g(æ,t) t)) take the quadratic form i.c. v (g(x, t)) (9(x, t)} then by canceling the cross term Theorem 4.1 we will get another Fisher information regularisation form as below. Theorem D.1. When v (g(x, t)) Ig( x1 t)l2 and ott) is constant, the regularized unbalanced optimal transport problem (7) is equivalent to
ot
inf (p.0-.)
log + al)( 0)
16a (1 logr) plwdd, (16)
where the infimum is taken all pairs (p, u. Φ) such that p(-,0) uu,p((; 1) u1,p(0, absolutely continuous, and
0.p V (po(x, f)) + J(x t)p d= (1+ =1+logg)
a²
(17)
4a
coupled with vanishing boundary condition:
lim p(æ, ()
0.
|ac 00
Proof. Based on the formulation of Theorem 4.1, we proceed by introducing the following change of variables:
== 4a
logp)
(18)
Substituting this variable into Eq. (7), then it follows immediately.
|1
From Theorem D.1, can find that if we introduce term in the continuity equation Eq. (6)
as:
0,p = V. (pb) + 1 v2 (d²Ip) + + (1 + logp)p.
d²
4a
Then Eq. (17) now yields the standard continuity equation:
app (pv(x, t)) ((x,t)p
So introduce new definition.
28

---

Published as a conference paper at ICLR 2025
Definition D.1. We introduce
inf (vo. vn) L l 2 [bb 1)112 pdladde S.' J 10( 1)22 plaadd,
1
UOTD (vo, u1) (pC
where, for (p, b, belonging to appropriate functions spaces,
(Vo, = b,g Opp = ~ (pbb) + 2 v2 (o² + +9P
(1 log
4a
where vo and P1 = V1-
Remark D.1. Note that Definition D.1 is consistent with the form provided in (Buze & Duong, 2023), and it has been conjectured that it is the higher order approximation scheme to the unbal- anced optimal transport problem as O.
D.4 RUOT IN THE GENERAL CASE
For the general case, we consider the following equation
1 v2a (a(æ, t)p) + JP. (pb) +
0,p - - Ve
(19)
2
with the corresponding metric
[ Jea b(x, t)ll p(x, t)dxdt + L.' J are (g(x, t)) p(æ, t)dædt,
(20)
where and R [0, +xo is the growth penalization. Then we can define the regularized unbalanced optimal transporl problem in the same way.
Definition D.2 (Regularized unbalanced optimal transport II). Consider
inf [(' [ee [1b(z, 1)ll p(ar, t)dxdt + L/ Jee av (g(2, ()) p(æ,, t)dardt,
(21)
(p.b.g)
where the infimum is taken all pairs (p, b) such that p(-:0) 0) = vo, po.p(:1)) = v1, v1, t) absolutely continuous, and
1 2 (a(z, t)p) + gp.
- (pb) +
(22)
2
coupled with vanishing boundary condition:
lim p(æ, t)
|20 -
We similarly can reformulate Definition D.2 with the following Fisher information regularization.
Theorem D.2. Consider the regularized unbalanced optimal transport problem Definition D.2, is equivalent to
( [
1 1 |[(x 1)l p(x, t)dædt + 3 ||o V log (a(x, t))) p(æ, t)dædt
inf ()
(23)
+ 3 (0(x, t), alxxtt(V V = log (a(x, t)p)) p(æ, t)d + at (g(x, t)) p( x, f doddt
where the infimum is taken all pairs (p. u. such that p(-,0) uo, vo. 1) v, p(x, t) absolutely continuous, and
0.p V (po(xx, f)) g(x t)p.
(24)
coupled with vanishing boundary condition:
lim p(æ,1 t) 0. b=
Proof. From Definition D.2 we know that the Bokker-Plancll equation for the SDE is
0,p - (pb) + 2 (a(æ, t)p) + g(x, t)p,
2 a
= = ((D- a(x t) V log (a(x, t)p) p) + g(x, 1)p.
29

---

Published as a conference paper at ICLR 2025
with minimization problem
inf / f [bb p(z, ()dadd + S/¹ Lee at (s(æ, ()) p(a, l)dadd
(p.b)
Using the change of variable v(x,t) t) b(x,t) fa(2 V log (a(æ, t)p), we see that it is equivalent to
Oep V (po(a, t)) g t)p.
log (a(x; t)p) + 2 (v( t), fa(x, log (a(x, t))), the integrand in the objective then becomes
On the other hand, since [16(x, t)l12 ||, t)(1 +
f gd
1 || t)ll p(a, t)dædt+ 1 t) log (a(x, t)p) p(z, t)dxdt
+ 2 (u(z, f), t) log (a(2, t)p)) p(=; t)dxdt + at (g(x, t)) p(z, d)ddd
Similarly, from the proof of Theorem D.2, the original SDE (b(X,, di+ can be transformed into the probability flow ODE
dx,- b (X.. -a a( V log (o)x. dt.
x(,).f)
Conversely, if the probability flow ODE's drift v( t) the diffusion rate a(x, and the function Ve log (a(æ, (æ) t)) are known, then the the SDE's drift term b(x, t) can be determined by
1 b(x, t) = v(x,t) t) + aaz V log (a(x, t)p(x,
2
So to specify an SDE is equivalent to specifying the probability flow ODE and the corresponding log density function.
CONNECTIONS WITH SCHRÖDINGER BRIDGE
Although the RUOT problem can be derived from the dynamical formulation of Schrödinger bridge relaxation, the relationship between the RUOT problem and the original static Schrödinger bridge problem, the choice the growth penalty form in RUOT as well as its microscopic interpretation, remains unsolved. In this section, will discuss two potential microscopic interpretations with the choice of the growth penalty form and identify the persisting problems and challenges associated with each.
E. CONNECTIONS WITH BRANCHING SCHRÖDINGER BRIDGE
In (Baradat & Lavenant, 2021), the authors explore the relationship between the RUOT and the branching Schrödinger bridge, with the corresponding reference process identified as branching Brownian motion. We will begin by reviewing their principal findings and subsequently highlight the associated issues and challenges.
A Branching Brownian Motion (BBM) can be formally described by the diffusivity s, branching rate À (Each particle is associated with an independent exponential distribution with parameter A), and the offspring distribution denoted by p = (PK)xEN P(N) (The probability measures on N), where is the probability of producing k offspring at branching event. We define 9k Ap& and introduce the generating function of as 2/gg
V/g =
Σ KEN
Pk *
KeN
30

---

Published as a conference paper at ICLR 2025
The branching mechanism is characterized by q, since ΣKEN Ps: = and X, p be recovered from q by EKEn 9 and pk: 9 for l N. Therefore, is sufficient fully determine the BBM with v, Ap, and initial distribution vo. So R ~ (BM(s, q, Mo) represent such BBM. Their core result indicates that consider Definition 4.1 where (t)) = a is constant and ε = o² then the growth penalty term is defined as follows, it can correspond to branching Schrödinger process.
+:0
V V.gg( (s) = =E(U' (es/s e ( -
Σ gh exp ((k-1) -1].
(25)
km
and we define V Ve,((9) = SUPERR g8 V (s)) So by following their result, can present some examples.
The case of only loss
This example has been included in their work. We consider = 1, Po = 1, and Ex22 Pk: = O, which means the particle can only die at the branching event.
Substituting it into Eq. (25), we have
(i.g(s) exp (-s)
Then by computing its Legendre transform obtain
(1) = 1 + log(-g),
where g < 0. This form is consistent with the formula discussed in (Chen et al., 2022b).
The case of only gain mass Similarly, if we consider À = 1. P2 1, which means the particle can only divide into two parts at the branching event. Then by a similar calculation, we obtain
V (ig((g) g + g log(g),
which g 0.
The case of both gain mass and loss mass This is also presented in their work. We consider x 1, Po = P2 3. which means the particle can divide into two parts or die at the same probability at the branching event, which we have
V (s) cosh(s)
However in this case the Legendre transform would be more complicated but it has some asymptotic properties discussed their work.
General case Although we have illustrated several examples, challenges persist, particularly in the context of more general scenarios. Specifically, when the growth penalty term cannot be expressed V V..g. the connection to the corresponding stochastic process remains unclear. For instance, this ambiguity arises in cases where the RUOT the T (g) takes the quadratic form i.e., the WFR metric, when (g) is in linear form, such as lg). These issues still require specific mathematical treatment and to our knowledge, remain an open problem.
E.2 CONNECTIONS WITH SCHRÖDINGER BRIDGE ON WEIGHTED PATH
From an alternative perspective on the microscopic interpretation of RUOT, we aim to consider particles on the weighted path space, namely:
(' exp ( 9(x s)ds d'
where " represents the probability measure induced by the SDE (Eq. (2)). Here for example we consider a simple case when g( 1) = C. Then following the procedure of Theorem 3.1.
( (.n)) ( Eng.n (L (('xXX..))aa
(I((,,,))l +tc).t+....)
(So II(C.. Cexp(C),
= E (Ei exp(C)) (' 2 (). at
(26)
= Eexp(C)
31

---

Published as a conference paper at ICLR 2025
Here Pfo.1j is the distribution of stochastic process induced by dY, ,Y c. Our derivation reveals that, from path weighting perspective, even when the growth term is constant, it still does not correspond to the form RUOT. This suggests that the microscopic interpretation from this angle may need further investigation determine whether RUOT problem could be recovered from this starting point.
32


Reviews:
[{"review_id": "cOSffi1RDa", "reviewer": "Reviewer_FkRe", "summary": "The paper introduces DeepRUOT, a new deep-learning approach for solving regularized unbalanced optimal transport (RUOT) and inferring continuous unbalanced stochastic dynamics from observed snapshots. The method models dynamics without requiring prior knowledge of growth and death processes, allowing these to be learned directly from data through a Fisher regularization formulation. The authors explore theoretical connections between RUOT and the Schrödinger bridge problem while reformulating RUOT to transform stochastic differential equations into more computationally tractable ordinary differential equations (a constraint). When tested on both synthetic gene regulatory networks and real single-cell RNA sequencing data, DeepRUOT outperformed existing methods in constructing meaningful developmental landscapes.", "questions": "- is not knowing the death/growth rates (what you call prior knowledge) such a big advantage?\n- Why did you project the data onto a 2d manifold for the \"Real Single-Cell Population Dynamics\" experiment? Isn't it possible to be done on a higher dimensional space, e.g. gene expression of PCA space?\n- The method involves multiple neural networks (v, g, s) and complex loss functions, which could require substantial computational resources as well as necessary tuning efforts, potentially limiting its practical application in real-time analysis. Would be great to get a comment on this.", "limitations": "", "rating": 10, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The experimental results seem very convincing\n- The paper provides thorough theoretical foundations with formal theorems and proofs connecting RUOT to the Schrödinger bridge problem.", "weaknesses": "- it remains unclear how the loss function weights (cost, reconstruction error and the PINN loss) are set\n- Figure 1 and its caption are not that helpful\n\nTypos:\n- Theorem 3.1: \"probelm\" -> \"problem\"\n- page 8: \"a increase\" -> \"an increase\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vJgnbivy7Z", "reviewer": "Reviewer_RQfu", "summary": "This paper proposes a new approach for learning dynamics from the observation of distributions at different time steps. The paper proposes to use an unbalanced optimal transport approach that allows modeling particle growth or death over time, which is an important aspect of single cell dynamics modeling. The authors compare their approach on datasets from syntethic gene regulatory networks and single cell trajectories dynamics. The paper shows favorable performance results compared to previous methods.", "questions": "- Regarding the connection with SB, could you either make the impact of the connection more explicit, or considering removing that part from the paper for clarity ?\n\n- Could you please restructure Section 5 to help the readers follow the reasoning until the final loss and algorithm ? Having a more detailed pseudo-code description would already help significantly.\n\n- Can you clarify Figure 2 - the predictions dots in the ground truth panel ? \n\n- Could you please add the results of an ablation experiment where you disable the learnable growth rate in your model ?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Population dynamics from snapshot data is an important problem, especially in biological application such as single cell trajectories inference.\n- The proposed method appears to surpass the performance of previous methods.\n- The related works section is extensive.", "weaknesses": "- One of the most important weaknesses of this paper at this stage is its clarity. In its current form, the paper is poorly written, and the ideas are poorly motivated. The main limitation from prior works identified is the unabalanced aspect e.g. \"However, computational methods for learning high-dimensional unbalanced stochastic dynamics from snapshots without prior knowledge of growth and death or additional information are still lacking.\". Yet, the paper only indirectly addresses that issue, which confuses the reader. Despite the favorable experimental results, the current presentation of this paper prevents it from being appreciated by the community.\n\n- Section 5, which is the core of your method, is extremly confusing. For instance, the reconstruction loss is not introduced previously such that it's extremely challenging to follow the development until the final loss in Equation 10. The algorithm 1 does not provide any help either in much needed clarification. The algorithm is too laconic and does not refer to corresponding parts of the text.\n\n- It's not clear to me why the connection with the Schrodinger bridge is necessary. As far as my reading goes, it seems that you only need to leverage the dynamic version of regularized unbalanced OT. As such, it does not seem that the connection with Schrodinger bridges helps the reader understand the method, and you don't leverage that connection later on in your method.\n\n- Theorem 4.1 uses $v$ which is not introduced previously. It's later undertstood that it corrresponds to the vector field of the corresponding ODE but this hinders the clarity of the paper.\n\n- The presentation of the results in Figure 2 raises questions. Why does panel (b), with the true dynamics has ground truth, predicted and trajectories, if it is the ground truth dynamics ? What is the \"predicted\" in this case ? \n\n- Because the most important contribution of this paper relies on modeling the growth rate, an ablation study is required where the growth rate part of the method is disabled (e.g. g = 0). This would allow to understand the impact of it on the final performance. \n\n- Despite the extensive related works section, the experiments only compare with two baselines. Baselines such as Bunne et al. or Yutong et al. could also be considered.\n\nSha, Yutong, et al. \"Reconstructing growth and dynamic trajectories from single-cell transcriptomics data.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HoAmrNX6Jg", "reviewer": "Reviewer_FFk6", "summary": "This paper introduces DeepRUOT, a method for reconstructing dynamics from snapshots using regularized unbalanced optimal transport. This work tackles both the unbalanced and stochastic settings simultaneously, which few works have attempted thus far. An algorithm is proposed to learn drift, score, and growth functions from snapshot data including both synthetic and single-cell transcriptomics data.", "questions": "It is unclear to me exactly what the computational / numerical cost for allowing varying growth rates is in this framework. From my understanding both the deterministic and stochastic setting are extremely efficient, but it has thus far been numerically challenging to incorporate varying growth rates. How scalable is DeepRUOT? Can it be applied to the higher dimensional settings (20,50,100,1000D) in previous works?\n\nI don’t understand the motivation behind the reconstruction loss R_d. What is its purpose and why does it correspond to a reconstruction loss vs. say an OT loss?\n\nComments:\n\n\\Psi(g) is not explicitly defined (although is fairly clear from context).\n\n> So to specify an SDE is equivalent to specify the probability flow\nODE and its score function.\n\nProbably missing an article here. \n\n### Overall\n\nOverall this work has great potential, however I believe further empirical study would greatly improve its value to the community. I would be eager to reconsider my evaluation given further empirical study relative to existing unbalanced dynamic OT methods and a better empirical understanding of the various components of training and loss weighting.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* Provided code provides additional clarity on experimental setting.\n* Tackles an important problem in trajectory inference from population level data.", "weaknesses": "* I found it extremely difficult to understand the multistage training procedure for DeepRUOT. There are some stages described in the main text, but then initial log density training also in the appendix? While there are a number of losses and initializations proposed to stabilize training, it is unclear to me what effects these have on final performance and how difficult RUOT is to train in practice. The paper would be much stronger if it was made clear which parts of the training procedure are important practically potentially through an ablation study on different components of the training procedure. \n* For a largely empirical work proposing a new algorithm to tackle this problem, the empirical studies are limited. The DeepRUOT is only benchmarked against balanced transport methods. It would greatly strengthen this paper if it was also compared to other unbalanced transport methods such as Action matching [Neklyudov et al. 2023] and unbalanced diffusion Schrödinger bridge [Pariset et al. 2023]. It would also be very useful to know how the many different weighting parameters affect performance, and how sensitive training is to them.\n* It is difficult to tell what parts of the theory are novel and which are minor adaptations from prior work. Theorem 4.1 seems like a subclass of Baradat & Lavenant for the Fischer information case. I don’t believe varying sigma(t) over time changes anything theoretically and as far as I understand it is not used in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)