Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action
Authors: Yuhao Sun
Zhenyi Zhang
Zihan Wang
Tiejun Li
Peijie Zhou
Venue: NeurIPS 2025 poster ()
Keywords: unbalanced optimal transport
Schrödinger bridge
trajectory inference
single-cell
Primary area: machine_learning_for_sciences
TL;DR: Var-RUOT is a new framework improving RUOT for recovering dynamics from high-dimensional snapshots. It incorporates first-order optimality conditions, achieving better performance on single-cell datasets.

Abstract: Recovering the dynamics from a few snapshots of a high-dimensional system is a challenging task in statistical physics and machine learning, with important applications in computational biology. Many algorithms have been developed to tackle this problem, based on frameworks such as optimal transport and the Schrödinger bridge. A notable recent framework is Regularized Unbalanced Optimal Transport (RUOT), which integrates both stochastic dynamics and unnormalized distributions.  However, since many existing methods do not explicitly enforce optimality conditions, their solutions often struggle to satisfy the principle of least action and meet challenges to converge in a stable and reliable way. To address these issues, we propose Variational RUOT (Var-RUOT), a new framework to solve the RUOT problem. By incorporating the optimal necessary conditions for the RUOT problem into both the parameterization of the search space and the loss function design, Var-RUOT only needs to learn a scalar field to solve the RUOT problem and can search for solutions with lower action. We also examined the challenge of selecting a growth penalty function in the widely used Wasserstein-Fisher-Rao metric and proposed a solution that better aligns with biological priors in Var-RUOT.  We validated the effectiveness of Var-RUOT on both simulated data and real single-cell datasets. Compared with existing algorithms, Var-RUOT can find solutions with lower action while exhibiting faster convergence and improved training stability. Our code is available at https://github.com/ZerooVector/VarRUOT.

Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action
Yuhao Sun
Zhenyi Zhang Zihan Wang3., Tiejun Li and Peijie Zhou 13,4,,.t Center for Machine Learning Research, Peking University.
LMMAM and School of Mathematical Sciences, Peking University.
3 Ceetter for Quantitative Biology, Peking University.
NELBDA, Peking University. SAI for Science Institute, Beijing,
Emails: (2501111524 zhenyizhang, jackwzh @stu. pku. edu. cn, {tieli, ,il, pwwpppohhappppppp @pku.. edu.
Abstract
Recovering the dynamics from few snapshots of high-dimensional system is challenging task in statistical physics and machine learning, with important appli- cations in computational biology. Many algorithms have been developed to tackle this problem, based on frameworks such as optimal transport and the Schrödinger bridge. A notable recent framework Regularized Unbalanced Optimal Transport (RUOT), which integrates both stochastic dynamics and unnormalized distribu- tions. However, since many existing methods do not explicitly enforce optimality conditions, their solutions often struggle to satisfy the principle of least action and meet challenges to converge
stable and reliable way. To address these issues, we propose Variational RUOT (Var-RUOT), new framework to solve the RUOT problem. By incorporating the optimal necessary conditions for the RUOT problem into both the parameterization of the search space and the loss function design, Var-RUOT only needs to learn a scalar field to solve the RUOT problem and can search for solutions with lower action. We also examined the challenge of selecting a growth penalty function the widely used Wasserstein- -aieeein--ihhe--aa metric and proposed solution that better aligns with biological priors in Var-RUOT. We validated the effectiveness of Var-RUOT on both simulated data and real single-cell datasets. Compared with existing algorithms. Var-RUOT find solutions with lower action while exhibiting faster convergence and improved training stability. Our code is available at https: / //g thub.
Introduction
Inferring continuous dynamics from finite observations is crucial when analyzing systems with many particles (Chen et al. 2018). However, many important applications such as single-cell RNA sequencing (scRNA-seq) experiments, only few snapshot measurements are available, which makes recovering the underlying continuous dynamics a challenging task (Ding al. 2022). Such task of reconstructing dynamics from sparse snapshots is commonly referred to as trajectory inference in time-series scRNA seq modeling (Zhang al. 2025b; Ding al. 2022; Heitz al. 2024; Yeo et al. 2021a; la; Schiebinger et al. 2019a 9a; Bunne et al. 2023b; Zhang et al. 2021) or the mathematical problem of ensemble regression (Yang et al. 2022).
A number of frameworks have been proposed to address this problem. For example, in dynamical optimal transport (OT), particles evolve according to the ordinary differential equations (ODEs) with
'These authors contributed equally. 'Coresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

RUOT Configuration Manifold
To Achieve Least Action, You can use Variational Method to derive Optimal Conditions
(g)) Represents Biological Prior
Torminal
For given timepoint 1'(g) > 0 means g increases along x(x,
de(g) dg
u
Initial
RUOT in Probability Manifold
gt 2 |,
v 2g
mls) -
Only Learns One Scalar Field ((x, Faster Training, Faster Convergence!
41 [.L. Liwnseess nhe- L bout
C
ARaoo Cuocnn + λSIR LHIR + CAtion
((()) < means g decreases along =xx(t)
Previous RUOT Solver
?
Use 3 Neural Networks
Hard to Achiove Least Action
vralocity
anlver
Figure 1: Overview of Variational RUOT (Var-RUOT).
the objective of minimizing the total action required to transport the initial distribution to the terminal distribution (Benamou and Brenier 2000). Unbalanced dynamical OT further extends this framework by adding a penalty term vb(g) on the partiele growth or death processes in total transport energy (namely the Wasserstein-Fisher--ee
metric or WFR metric) to handle unnormalized distributions (Chizat et al. 2018a; Chizat et al. 2018b). Moreover, stochastic methods such as the Schrödinger Bridge adopt similar action principles while governing particle evolution via stochastic differential equations (SDEs) (Gentil al. 2017: Léonard 2014). Recently, the Regularized Unbalanced Optimal Transport (RUOT) framework generalizes these ideas by incorporating both stochasticity and particle birth-death processes (Lavenant et al. 2024; Ventre et al. 2023; Chizat et al. 2022; Pariset et al. 2023; Zhang et al. 2025a). In machine learning, generative models such as diffusion models (Ho et al. 2020; Song et al. 2021; Sohl-Dickstein al. 2015; Song al. 2020) and flow matching techniques (Lipman et al. 2023; Tong et al. 2024a; Liu et al. 2022) have also been adapted to solve transport problems. However, these approaches face two major challenges: 1) They usually do not explicitly enforce optimality conditions, leading to solutions that violate the principle of least action, and they meet challenges to converge reliably; 2) Selecting an appropriate penalty function (gg) that aligns with underlying biological priors remains challenging.
To overcome these challenges, we propose Variational-RUOr (Var-RUOT), Our algorithm employs variational methods to derive the necessary conditions for action minimization within the RUOT framework. By parameterizing a single scalar field with neural network and incorporating these optimality conditions directly into our loss design, Var-RUOT learns dynamics with lower action. Experiments on both simulated and real datasets demonstrate that our approach achieves competitive performance with fewer training epochs and improved stability. Furthermore, we show that different choices of the penalty function for the growth rate g yield distinct biologically relevant priors in single-cell dynamics modeling. Our contributions summarized as follows:
We introduce a new method for solving RUOT problems by incorporating the first-order optimality conditions directly into the solution parameterization. This reduces the learning task to a single scalar potential function, which significantly simplifies the model space.
We show how incorporating these necessary conditions into the loss function and architecture enables RUOT to consistently discover transport paths with lower action, providing a more efficient and stable training process for RUOT problem.
We address a key limitation in the classical Wasserstein- Boassrstein-FFshe--RRee metric, which can yield biologically implausible solutions due to its quadratic growth penalty term. We propose the criterion and practical solution to modify such a penalty term, therefore enhancing the more realistic modeling of single-cell dynamics.

---

2
Related Works
Deep Learning Solver for Trajectory Inference Problem There are a large number of deep learning-based trajectory inference problem solvers. For example, there solvers for Optimal Transport using static OT solver, Neural ODE or Flow matching techniques (Tong et al. 2020; Huguet et al. 2022; Wan et al. 2023; Zhang et al. 2024a; Tong et al. 2024a; Albergo et al. 2023; Palma et al. 2025; Rohbeck et al. 2025; Petrović al. 2025; Schiebinger al. 2019b; Klein et al. 2025; Wang et al. 2025). well solvers for the Schrödinger bridge that utilize either its relation to entropy regularized OT or its optimal control formulation(Shi et al. 2024; De Bortoli et al. 2021: 1: Gu et al. 2025; Koshizuka and Sato 2023; Neklyudov et al. 2024; Zhang al. 2024b; Bunne et al. 2023a; Chen et al. 2022a; Zhou et al. 2024; Zhu et al. 2024; Maddu et al. 2024; Yeo et al. 2021b; Jiang and Wan 2024; Lavenant et al. 2024; Ventre et al. 2023; Chizat et al. 2022; Tong et al. 2024b; Atanackovic et al. 2025; Yang 2025; You et al. 2024). However, these methods typically employ separate neural networks parameterize the velocity and growth functiinn, without leveraging their optimality conditions or the inherent relationship between them(Zhang et al. 2025d) with few exceptions such Action Matching (Neklyudov al. 2023; Neklyudov al. 2024)) This poses challenges in achieving optimal solutions that minimize the action energy. Our work generalizes Action Matching (Neklyudov al. 2023; Neklyudov al. 2024) to handle the setting where unbalanced distributions and stochastic effects coexist.
HJB equations in optimal transport Methods that leverage the optimality conditions (e.g., Hamilton-Jace JaconnnJJa-----aa Beelmm- (HJB) equations) of dynamic OT and its variants, have been proposed (Neklyudov et al. 2024; Zhang et al. 2024b; Chen et al. 2016; Benamou and Brenier 2000; Nek- lyudov al. 2023; Wu et al. 2025; Chow et al. 2020). However, these approaches typically do not simultaneously address both unbalanced and stochastic dynamics.
WFR metric in time-series scRNA-seq modeling In computational biology, several existing works model both cell state transitions and growth dynamics simultaneously temporal scRNA-ses datasets by minimizing the action in the WFR metric i.e.solving the dynamical unbalanced optimal transport problem (Sha et al. 2024; Tong et al. 2023; Peng et al. 2024: Eyring et al. 2024b) or its variants (Pariset et al. 2023; Lavenant al. 2024; Zhang al. 2025a; Zhang et al. 2025c). However, these works usually adopt the default growth penalty function (gg)) = 1 in the WFR metric and have not investigated the biological implications of different choices for ((g)).
3 Preliminaries and Backgrounds
Dynamical Optimal Transport The Dynamical Optimal Transport, also known as the Benamou- Brenier formulation, requires minimizing the following action functional (Benamou and Brenier 2000):
inf (-' f p.u
1 [u,xx t)l p(x, t) dx dt,
where P and u are subject to the continuity equation constraint:
ap(x. t)
+ Vx (u(x, t))(x,))) = O, pl:,0) = poo p(:,1) p1.
ØL
Unbalanced Dynamical OT and Passerstein-Fisher--e unnormalized probability densities in practical problems (for example, to account for cell proliferation and death in computational biology), one can modify the form of the continuity equation by adding a birth-death term, and accordingly include corresponding penalty term in the action. This leads to the optimal transport problem under the Wasserstein-Fisher Rao (WFR) metric (Chizat et al. 20; Chizat et al. 2018b).
(WFR) metric In order to handle
inf J. [ (bllutx t)I (Gll((x..)/++ p.n.g
+ p(x, 1) dx
with p. u. and g subject to the unnormalized continuity equation constraint 0p(x,t) xx (u(xx t)px, g(xx t)p(x, =0, p(-.0) = po. p(:11) Pi-
0t
3

---

Schrödinger Bridge Problem and Dynamical Formulation Schrödinger bridges aims to find the most likely way for a stochastic system evolve from an initial distribution μo to terminal distribution /ri Formally, let Mi0.1] denote the probability measure induced by the stochastic process (t). & t $ 1, and let Hj..1] denote the probability measure induced by a given reference process Y (1), Sssss 1. The Schrödinger bridge seeks solve min
((o.1) DKL (((o.) ((.1) In
particular, if X, follows the SDE dX, u(X,, t) dt + a (X,,t) t) dW&, &, where W. E I2 is standard Brownian motion, a (x. t) Rdxd is a given diffusion matrix, and the reference process is defined as dy, = a (Y,, !) dW,, then the Schrödinger bridge problem is equivalent to the following stochastic optimal control problem (Chen al. 2016; Gentil al. 2017):
inf L/ P.u
Jo (2 (2u7 [2"(xxt)
a (x, a '(x, u(x, t)) p(xx t) dx
where P and u are subject to the Fokker-Planck equation constraint
0p(x, + Vx:
V (a(x,)) p(x, 0, pl:0) = po: p(:1) ph
2
Here, t(x,t)= (oxxt)o" (x, and V (a(x,t) ()) = Eis Oig (ass p(x,t))
Regularized Unbalanced Optimal Transport If we consider both unnormalized probability densities and stochasticity simultaneously, we arrive the Regularized Unbalanced Optimal Transport (RUOT) problem (Chen et al. 2022b; Baradat and Lavenant 2021; Zhang et al. 2025a).
Definition 3.1 (Regularized Unbalanced Optimal Transport (RUOT) Problem). Consider minimizing the following action:
(' [ee (x, (xx,,x, () p(x, () dx dl, P.0.9
inf
where db IR 10, +oc) is growth penalty function, a is penalty coefficient, and the quantities p. and g are subject to the following constraint, which is an unnormalized continuity equation:
0p Ot
Vx
9xx, (a(x.)))) (ab))))))--a
(() (a(x,t)
p(:,0) = pos p(:,1) P1.
Remark 3.1. When modeling the particle dynamics in these OT variants, one can use either autonomous systems (where g does not explicitly depend t) non-autonomous systems (where v or explicitly depends on t). Here, we follow methods such as DeepRUOT et al. 2025a) choose non-aulonomous systems. Indeed non-aulonomous systems offer greater expressive power from a learning perspective and are less prone to underfitting the data.
4 Necessary Optimality Conditions for RUOT
To simplify our problem we adopt the assumption isotropic time-invariant diffusion, i.e., a(x, II. We refer to RUOT problem in this scenario as the isotropic Lime-invariant RUOT problem.
Definition 4.1 (Isotropic Time-Invan Invariant (ITI) RUOT Problem). Consider the following minimum- action problem with the action functional given by
1u D)1 p(x, + S. [ a p(x, dx
inf I -= = (p.n.g)
Here, ib : R + [0, +00) is the growth penalty function, and the triplet (p. u, g. subject to the constraint of the Fokker-Planck equation
1 0p(x, f) - (p(x, 1) u(x, ()) - V 1) g(x, p(x,
01
Additionally, satisfies the initial and terminal conditions pl: 0) 14. p(,, p1.

---

In particular, if ()(x t)) = = (x. t), then this problem is referred to as the unbalanced dynamic optimal transport with WFR metric. We can derive the necessary conditions for the action functional to achieve a minimum using variational methods.
Theorem 4.1 (Necessary Conditions for Achieving the Optimal Solution in the ITI-RUOT Problem). In the problem defined in Definition 4.1, the necessary conditions for the action I to allain minimum are
dv((g) X, dg
ax Ot
1 - Ag a - =0 2
==
(3)
Here, A(x, 6) is a scalar field. The proof of this theorem can be found in Appendix A.I.I.
Remark 4.1. Substituting the necessary conditions satisfied by u and g into the Fokker Planer- equation, the evolution of the probability density p(x,t) is determined by Bo(x.))
-V
- (V(xx1)Vxx(xx11) denotes the inverse function of v/!.
1) + (('))"((() (Nox.t) p(x,t), where of delg) and (1-')-
dg
Remark 4.2. If we choose the growth penalty function to take the form used in the WFR metric,
i.e., ((9)= = " and set a = o = O, then the above optimal necessary conditions immediately degenerate
Vx g = X, ax + ; x|1? + 1 0, which is same the form
at
derived in (Neklyudov et al. 2024) under the WFR metric. If we let = O and (0) = 0 it becomes
VXA,
2 -Z 0, which is same the form derived in (Neklyudov et al. 0A + 1 V
2
2024; Zhang el al. 2024b; Chen el al. 2016) under the Schrödinger Bridge problem.
From Theorem 4.1 and Remark 4.1, the vector field u(x, t) and the growth rate g(x. t) can be directly obtained from the scalar field X(x, t). Moreover, since the initial density p(- 0) is known, once the necessary conditions are satisfied, the evolution equation (i.e., the Planck-Planck equation)
completely determined by A(x, t). Thus, the scalar field A(x, t) fully determines the system's evolution, and we only need to solve for one X(x, l), which simplifies the problem.
However, these necessary conditions introduce coupling between u(x, I) and g(x, (). and this coupling could contradict biological prior knowledge. In biological data, it is generally believed that cells located at the upstream of trajectory are stem cells with the highest proliferation and differentiation capabilities, and thus the corresponding values should be maximal. Along the trajectory, as the cells gradually lose their "stemness," the g values would decrease. Under the necessary conditions, however, whether g(x, t) increases or decreases along u(x, t) at given time t depends on the form of the growth penalty function.
Theorem 4.2 (The relationship between u and Biological prior). At a fixed time I, is dw)(() > 0, then g(x, f) ascends in the direction of the velocity field u(x, 1) (i.e., u(x, (V ()) 0); otherwise, it descends.
The proof is given in Appendix A.1.2. According to this theorem, to ensure the solution complies with biological prior, i.e., that at a given time the cells upstream in the trajectory exhibit the higher
value, it is necessary to ensure that that d' O.
dg
5 Solving ITI ROUT Problem Through Neural Network
Given samples from distributions Pt K discrete time points, / {Ti,... TRr. aim to recover the continuous evolution process of the distributions by solving the ITT RUOT problem, that is, by minimizing the action functional while ensuring that p(x, t) matches the distributions Pt the corresponding time points, Since the values of u(x. t) and 9 t), as well as the evolution of p(x, t) over time, are fully determined by the scalar field A(x, t) in variational form (Appendix A.1.1), we approximate this scalar field using a single neural network. Specifically, we parameterize X(x, t) as A t), where f represents the neural network parameters.
5.1 Simulating SDEs Using the Weighted Particle Method
Directly solving the high-dimensional RUOT with PDE constraints is challenging. Therefore, we reformulate the problem by simulating the trajectories of a number of weighted particles.
5

---

Theorem 5.1. Consider weighted particle system consisting of N particles, where the position of particle i time is given by X R. and its weight by wy w;(t) > 0. The dynamics of each particle are described by
dx; = --xx( t) dt + ott dW..
(4)
dw,- ( wy dt,
x R is a growth rate function, : (0, - [0, +x) is time-varying diffusion coefficient, and W. is an N- dimensional standard Brownian motion with independent components in each coordinate. The initial conditions are X p(x, 0) and (0) 1. In the limit as N + oc, the empirical measure (t(() 6(x
where u R x [0, - R is a time-varying vector field,
μ (xx t) = L
X converges to the solution of the following Fokker-Plance
equation:
0p(x,1)
- - (u(x,t) + 2 (0) 6) 9(x, p(x,t),
(5)
Ot
with the initial condition p(x, 0) Pu(x))
The proof is provided in Appendix A. This theorem implies that we can approximate the evolution of p(x, 1) by simulating N particles, where each particle's 's weight w; is governed by an ODE and its position X; is governed by an SDE. The evolution of the empirical measure ve"' (x, () thereby approximates the evolution of p(x,t).
5.2 Reformulating the Loss in Weighted Particle Form
The total loss function consists of three components such that C LRecon +HBB CHIB + + Ccction- Here, LRccnn ensures that the distribution generated by the model closely matches the true data distribution, LHи enforces that the learned Ag (x, t) satisfies the HJB equation in the necessary conditions, and Action minimizes the action as much as possible.
Reconstruction Loss by the model is consistent with the real data distribution. Since in the ITI RUOT problem the proba- bility density p(x, f) is normalized, we need to match both the total mass and the discrepancy between the two distributions. Our reconstruction loss is given by LRconn = "Mass LMsss + LOT, where, time point k the true mass is fgu p(x, Tk) dx (Tw,,, and the weight of particle is w; (Tx). Thus, the total mass of the model- generated distribution is (((Tk) N
Minimizing the reconstruction loss guarantees that the distribution generated
E wy (T+... 2
The mass reconstruction loss is then defined as LMass = ΣK (M(TL) - (((H) Let the true
distribution at time point be p(x, Ts.. Its normalized version is given by p(x, Th) fat p(x.Tx) dx while the normalized model-generated distribution is r(x, Tx)
p(x.T2)
A N.N
(Tv)
The distri-
M(T2)
bution reconstruction loss is then defined as Lor = Σ Wz (p(,Te) ((,TT)) where Mass is hyperparameter that controls the importance of the mass reconstruction loss.
HJB Loss Minimizing the HJB loss ensures that the learned Xe t) obeys the HJB equation constraints specified in the necessary conditions. Since the gradient operator the HJB equation is local operator, we compute the HJB loss by integrating the extent to which do t) violates the HJB equation along the trajectories. When using N particles, the HJB loss given by: LA
Tk
(t)
(8.(((.1) 0t
+
1 + ozV + go (Xt, a ()())) dt
Here
ΣE EE,00 (t)
dev(a)) dgo
Ge is obtained from the necessary condition a
Ag.
Remark
5.1. The TK Jam (8(xx.) ot
expectation
of
Loss is BICNo + g(x,1) a M(s)) dxdl where p(x,
HJB
+V
+
p(x,t)
is the probability density by normalizing p(x, 1). The proof is left in Appendix A. 1.1.4
Jas p(x, f)dx

---

Action Loss Since the variational method provides only the necessary conditions for achiev- ing minimal action rather than sufficient ones, we need incorporate the action into the loss so that the action is minimized as much as possible. The action loss is also com- puted via simulating weighted particles. When using N particles, it is given by C
Action
1 N 1 N So |uo (X; (tt)dt + +So ())
Here u and is obtained from
the necessary condition ug Vx)0(Xx, t),
dvige) Xe.
dge
Remark 5.2. The expectation of action loss is exactly the action defined in the IT1 RUOT problem
(Definition 4.1): B[C T. The proof is left in Appendix A.1.5.
Overall, the training process of Var-RUOT involves minimizing the total of three loss terms described above to fit Ag. The training procedure is provided in Algorithm 1.
5.3 Adjusting the Growth Penalty Function to Match Biological Priors
As discussed in Theorem 4.2, the second-order derivative of (g)) encodes the biological prior: (g) > O, then at any given time f, g increases in the direction of the velocity field, and vice versa. Therefore, we choose two representative forms of (g)g) for our solution. Given that (ggg) penalizes nonzero 9, it should satisfy the following properties: (1) The further 3 deviates from 0, the larger (g) becomes, i.e., de(g) 0. (2) Birth and death penalized equally when prior knowledge is
dlsl
absent, i... (g)) = -b(-g))
Case 1: ("(g) > O In the case where () > O, typical form that meets the requirements (g))) = Cg p E Zt, C> O. We select the form used in the WFR Metric, namely, (0) 29.h The optimality conditions are presented in Appendix A.2.
Case 2: (g) < O For the case where (g) < 0, a typical form that meets the conditions is (g)) = C (22p)/(2r+1) where p. p.qe E Z+ and 2p < 2q + In order to obtain a smoother gA))
relationship from the necessary conditions, and as a illustrative example, we choose (2((g) = g2/15 The optimality conditions are also presented in Appendix A.2.
6 Numerical Results
In the experiments presented below, unless the use of the modified metric is explicitly stated, we utilize the standard WFR metric, namely, (1((g) = 2
6.1 Var-RUOT Minimizes Path Action
To evaluate the ability of Var-RUOT to capture the minimum-action trajectory, we first conducted experiments on a three-gene simulation dataset (Zhang et al. 2025a). The dynamics of the three-gene simulation data are governed by stochastic differential equations that incorporate self-activation mutual inhibition, and external activation. The detailed specifications of the dataset are provided in Appendix B. The trajectories learned by DeepRUOT and Var-RUOT are illustrated in Fig. and the Wi and W2 losses between the generated distributions and the ground truth distribution, as well as the corresponding action values, are reported in Table 1. In the table, we report the action of the method that utilizes the WFR Metric. The experimental results demonstrate that Rur- achieves a lower path action while maintaining distribution matching accuracy. To further assess the performance of RUOT on nigh-dimensional data, we also conducted experiments on an Epithelial Mesenchymal Transition (EMT) Mataset(Sha et al. 2024: Cook and Vanderhyden 2020). This dataset was reduced to a f0-dimensional feature space, and the trajectories obtained after applying PCA for dimensionality reduction are shown in Fig. Both Vur-RUOT and DeepRUOT learn dynamics that can transform the distribution at t = 0 into the distributions at t = 1, 3. Var-RUOT learns the nearly straight-line trajectory corresponding to the minimum action, whereas DeepRUOT learns a curved trajectory. The results of Wi, Wi distance and action, summarized in Table 2, Rur- also learns trajectories with smaller action while achieving matching accuracy comparable to that of other algorithms.
7

---

Table I: On the three-gene simulated dataset, the Wasserstein distances (Wi and W2) between the predicted distributions of each algorithm and the true distribution at various time points. Each experiment was run five times to compute the mean and standard deviation. (Wi and W2 are relative metries and should therefore only be compared on the same dataset.)
Action
Model
74
W
W:
11:
11g
SF2M (Tong ct 20048)
0.3253 wooee
0,470 soven
0,7648 monee
10070
2.1879uax0:
2.8530
PISDR (lisng and Wan 2004)
0:3212
0:2911 goors
0.5
0,548
(1957
MID Flow (flgus al 2002)
0.12002coom
0:20872r00
0:2063
04565
1.018f
1.7827...m 2:4473
Actin Marching (Nekcrodow
2013)
0.013
0,228
1,54M
0.3035
0.3078
2.4023
0.5867
0.435
D'OTCMMEvvigg al. 3024g)
0,3002
0
0,1653
0
0.2630
0.4
0.7583
WL Mostyoidyee atal 2004)
0,5273aroo
0,8346200m
018571
TIGON (Sho 2M44
0,8231 0:08192000
prowe
0,0053aoom
2.15
2.24%
012201
1,3442
DeegRUOT (Zhang 2025g)
0,0059 louar
0,1125 uooo
0.0811 lour
0.1573 couec
1.1266
0.2158.omm
0.1538
0.2588
1 1,4058
VarRUKO (Chns
0.0452
0.1181
grow
0,0085 govrs
0.1270g
0.1445
0.1144
0.1772
(.2144g:3ma
1.1 1.1165
(a) DeepRUOT
(b) Var-RUOT
Trajectory
Growth
Trajectory
Growth
1.00
1.M0
1.00
Figure 2: Comparison between DeepRUOT and Var-RUOT on dynamics reconstruction on three-gene simulation dataset. a) The trajectory and growth rate learned by DeepRUOT. b) The trajectory and growth rate learned by Var-RUOT on the same dataset. In the figure, the circles and "+' signs denote the generated data and the original data points, respectively.
Table 2: On the EMT dataset, the Wasserstein distances (Wi and W2) between the predicted distributions of each algorithm and the true distribution at various time points. Each experiment was run five times to compute the mean and standard deviation.
Path Action
Wi
14
We
W3]
W12
SF2M (Tong ci al. 202-Ib)
0.2566.0avss
0.26-16.autte
0,28 1 1o00s
0 0.2897
0.2900.xu01
0.3005
PISDE (Jiang and Wan 2024)
0.2694.camr
0.2765-ausie
0.2860
0.2944-20013
0.2790
0.2920
MIO (Hugnel
2022)
0.2e29_mwe
2.26652mmm)
0.2770=awwe
0.2984 20000
Action Matching (Neklyudov et al. 2023)
0.4723
0.4794.x0e
0.638240am
0.6654:30m
0.8453
0.8524.n0m
0.8583
OTCFM (Tong et al. 2024a)
0.2557
0.26492amee
0,2701
D. 1
0.2799
042914
UOTCFM (Eyring
2021a)
0.2629
10200)
0.2797 L2
0.29
WLF (Noklyudov et al. 2024)
0.390] 20000
0.3977-xee
0.4381 100000
0.4466
0.2848
0.2976
TIGON (Sha et al. 2024)
0.2433 cuavo
0.2523 Juov
0,266 Lea00)
0,2766 2o0ic
0.2844
Laroe
0.2989 cozao
0.4672
DeepRUOT (Zhang 2025a)
0,2902
0.2987
0.3193
0.3293
0.3291
0.3410..oar
0-1857
(ar-RUOT(Ous)
0.2540
0.2623-4um1
0.2670
0.2756.30014
0.2683
0.2796
0.3544..x0
6.2 Var-RUOT Stabilizes and Accelerates Training Process
To demonstrate that Rar- converges faster and exhibits improved training stability, we further tested it on both the simulated and the EMT dataset. We trained the neural networks for the various algorithms using the same learning rate and optimizer, running each dataset five times. For each training, we recorded the number of epochs and wall-clock time required for the OT loss related to the distribution matching accuracy to decrease below specified threshold (set to 0.30 in this study). Each training session was capped at maximum of 500 epochs. If an algorithm's 's OT loss did not reach the threshold within 500 epochs, the required epoch was recorded as 500, and the wall-clock time was noted as the total duration of the training session. The experimental results are summarized in Table 3, which lists the mean and standard deviation of both the epochs and wall-clock times required for each algorithm on each dataset. The mean values reflect the convergence speed, while the standard deviations indicate the training stability. Our algorithm demonstrated both a faster convergence speed and better stability compared to the other methods. In Appendix C.I, we further demonstrate our training speed and stability by plotting the loss decay curves.
8

---

(a) DeepRUOT
(b) Var-RUOT
Trijectory
Growth
Trajectory
5.00
Figure 3: Comparison between DeepRUOT and Var-RUOT dynamics reconstruction EMT dataset. a) The trajectory and growth rate learned by DeepRUOT. b) The trajectory and growth rate learned by Var-RUOT on the same dataset. In the figure, the circles and '+' signs denote the generated data and the original data points, respectively.
Table The number of epochs and wall time required for the OT Loss to drop below the threshold for each algorithm. We trained each algorithm five times to compute the mean and standard deviation.
Simulation Gene
EMT
Model
Epoch
Wall Time(See.)
Epoch
Wall Time(Sec.)
TIGON (Sha et al. 2024)
228.40 222.40
11 1142.79
1 110.40+
365.54 555
RUOT w/o Pretraining (Zhang et al. 2025a)
172.00=22..
578.67=
228. 20 as
819.31 +
RUOT with Epoch Pretraining (Zhang et al. 2025a)
653.33-m00.
221.60+22ss
801.18+10.4
Var-RUOT (Ours)
27.60
33.98
5.20
7.37
6.3 Different (g)) Represents Different Biological Prior
To illustrate that the choice of (g)g represents different biological priors, we present the learned dynamies under two selections of (g). We apply our algorithm the Mouse Blood Hematopoiesis dataset (Weinreb et al. 2020; Sha et al. 2024). In Fig. 4(a), the standard WFR metric is applied, i.e., 4 (g) 3.9², from which can be observed that at time points drift vector field u(x, (), 9(x, () gradually increases. In Fig. 4(b), on the other hand, the alternative selection (g) = g mentioned in Section 5.3 is used, and it is evident that at each time point, g(x, t) gradually decreases along the direction of u(x, t). The distribution matching accuracy and the action are reported in Table 4. When employing the modified metric, the corresponding action quantity is not directly comparable those obtained using the WFR metric. Therefore we do not report its action here.
0, 1, along the direction of the
Table Wasserstein distances (W, and W2) between the predicted distributions of each algorithm and the true distribution at various time points mouse blood hematopoiesis. Each experiment was run five times to compute the mean and standard deviation.
2
Path Action
Model
Wi
W2
Wi
W22
Action Matching (Neklyudov et 2023)
0.4719
0.5673=00000
0.8350=0000
0.8936=0.000
4.3517
MIOFLOW (Huguet al. 2022)
0.3546=00000
0.4083=0.00z
0.2772-0.00
0.3459=0.000
SF2M (Tong et al. 2024h)
5.2
0.1202
PISDE (Jiang und Wan 2024)
0.3124 1::124-000ss
0.3499 11.999
0.2531 =a
0.2983-auurg
OTCFM (Tong al. 2024a)
0.3674-00000
0.48 4=
0.3222 nood
=..737=0.000
UOTCHM (Lyyring et al. 2024a)
0,3301 =000000
0,3950 =-00000
0,205 000000
0.2606=0.000
WLF (Neklyudov et al. 2024)
0.3302 220 0oo0
0.3950 2 DO0)
0,205
1 D000
0.2606gnmmm
TIGON (Sha et al. 2024)
0.4499 1010000
0.51
(1000)))
0.4368 ::1 DOO0
0,4852 200000
3.7438
De (Zhang al. 2025a)
0.1456=0000s
0.1807
0019
0.1469
= DOGG
0.1791 =0,000
5,5887
Var-RUOT (Standard WI-R
0.1200 -o.m0
0.1459 000000
0.1431
0002
0.1764
3.1491=00857
Var-RUOT (Modified Metric)
0.2953
0.
0.1917
0.2226
In addition to the three experiments presented the main text, we conducted several ablation studies and supplementary experiments to further validate our method. First, we performed ablation studies on the weights of the HJB and action losses to verify their effectiveness in learning dynamics with minimal action. Additionally, we explored alternative modeling choices, including comparison

---

Standard WFR Metric (()))
Trajectory
[- Growth
Growch
t= Growth
(b) Modified Metric (g)) g2
Trajectory
[= Growth
t= Srowth
Leaar Lo0
Figure 4: Comparison Var-RUOT using different growth metric on mouse blood hematopoiesis dataset. The trajectory and growth at time points t 0, 1, learned using the standard WFR metric. b) The trajectory and growth at time points f = 0, 1, 2 learned using the modified metric. In the figure, the circles and "+' signs denote the generated data and the original data points, respectively.
between using L1 and L2 norms for the loss terms, and an analysis of the impact of different growth penalty functions beyond the ()) = g example where ()) < The detailed results of these studies are provided in Appendix C.2. To assess the model's generalization capabilities, we performed hold-one-out and long-term extrapolation experiments. The results indicate that the Var-RUOT algorithm effectively perform both interpolation and extrapolation, with the learned minimum-actio dynamics leading to more accurate extrapolation outcomes (Appendix C.3). Furthermore, we carried out experiments several high-dimensional datasets, which further validate the effectiveness of Var-RUOT in high-dimensional settings (Appendix C.4), Finally, we compared the path action computed by RUOT against that from traditional OT Dolver(Appendd) C.5). The experimental results show that the path actions obtained by both methods are similar, which helps to validate approach.
Conclusion
In this paper, we propose new algorithm for solving the RUOT problem called Variational RUOT. By employing variational methods to derive the necessary conditions for the minimum action solution of RUOT, we solve the problem by learning a single scalar field. Compared to other algorithms, Var-RUOT can find solutions with lower action values while achieving the same level of fitting performance, and offers faster training and convergence speeds. Finally, we emphasize that the selection of (g)) in the action is crucial and directly linked to biological priors.
Although the Var-RUOT algorithm presented in this paper offers insights for solving the RUOT problem, it is subject to several key limitations. Firstly, because the algorithm based on neural network optimization, can only converge local minimum rather than guaranteeing the attainment of the global minimum of the action. In practice, the HJB loss functions more as regularization term than hard constraint that converges zero. In addition, when using the modified metric, the goodness-on of for the distribution deteriorates, which warrants further investigation. Finally, the choice of the function (g) within the action is dependent on biological priors and is not automated. Future work could address these limitations by systematically investigating more canonical problems, which would provide clearer insight into the algorithm's 's fundamental behavior and limitations, and by automating the selection of (g)) either by approximating it with a neural network by deriving it from caat-rrrciple-bssee microscopic dynamics, such as branching Wiener (Baradat and Lavenant 2021). We discussed more limitations of our work and potential broader impacts in Appendix E.I.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the National Key R&D Program of China (No. 2021 to T.L.), National Natural Science Foundation of China (NSFC No. 12288101 to T.L. & P.Z., and 8206100646, T2321001 P.Z.) and The Fundamental Research Funds for the Central Universities. Peking University. We acknowledge the support from the High- high-performance Computing Platform of Peking University for computation.
References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. "Stochastic interpolants: A unifying framework for flows and diffusions". In: arXiv preprint arXiv:2303.00773 (2023).
[2] Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. "Meta Flow Matching: Integrating Vector Fields the Wasserstein Manifold", In: The Thirteenth International Conference Learning Representations. 2025.
[3] Aymeric Baradat and Hugo Lavenant. "Regularized unbalanced optimal transport as entropy minimization with respect branching brownian motion". In: arXiv preprint arX1v:2111.01666
(2021).
[4] Tean-David Benamou and Yann Brenier. "A computational fluid mechanics solution to the
mass transfer problem".. In: Numerische Mathematik 84.3 (2000), pp. 375-
393.
[5] Jules Berman, Tobias Blickhan, and Benjamin Peherstorfer. Parametric model reduction of mean-field and stochastic systems via higher-order action matching. 2024. arXiv: 2410. 12000
[stat.ML] ML].
[6] Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. "The schrödinger bridge between gaussian measures has closed form". In: International Conference on Artificial Intelligence and Statistics. PMLR. 2023a, pp. 5802-5833.
[7] 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". In: Nature methods 20.11 11 (2023b), pp. 1759-1768.
[8] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. "Neural ordinary differential equations". In: Advances in neural information processing systems 31 (2018).
[9] Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. "Likelihood Training of Schrödinger Bridge using Forward- -aaraadddda SDEs Theory". In: International Conference on Learning Representations. 2022a.
[10] Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. "On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint". In: Journal of Optimization Theory and Applications 169 (2016). pp. 671-691.
[1]] Yongxin Chen, Tryphon Georgiou, and Michele Pavon. "The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced marginals". " In: SIAM Journal on Control and Optimization 60.4 (2022b), pp. 2016-2039.
[12] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. "An inter- polating distance between optimal transport and Fisher-Rao metrics". In: Foundations of Computational Mathematics 18 (2018a), pp. 1-44.
[13] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xaviee Vialard, "Unbalanced optimal transport: Dynamic and Kantorovich formulations". In: Journal of Functional Analysis 274.11 (2018b), pp. -020-3123..
[14] Lénaic Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. "Trajectory infer- ence via mean-field langevin in path space". In: Advances in Neural Information Processing Systems 35 (2022), pp. 16731-16742.
[15] Shui-Nee Chow, Wuchen Li, and Haomin Zhou. "Wasserstein hamiltonian flows". In: Journal of Differential Equations 268.3 (2020), pp. 1205-1219.
[16] David P Cook and Barbara C Vanderhyden. "Context specificity of the EMT transcriptional response", In: Nature communications 11.1 (2020), p. 2142.

---

[17] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. "Diffusion
schrödinger bridge with applications to score-based generative modeling". In: Advances in Neural Information Processing Systems 34 (2021), pp- 17695-17709.
[18] Jun Ding, Naday Sharon, and Ziv Bar-Joseph. "Temporal modelling using single-cell transcrip- tomics". In: Nature Reviews Generics 23.6 (2022), pp. 355-368.
[19] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian Theis. Unbalancedness Neural Monge Maps Improves Unpaired Domain Translation. 2024a. arXiv: 2311 15100 cvs.
[20] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation". In: The Twelfth International Conference
Learning Representations. 2024b. [21] Nicolas Fournier and Arnaud Guillin. On the rate of convergence in Wasserstein distance of the empirical measure. 2013, arXiv: 1312. 2128 [math.PP PR].
[22] Ivan Gentil, Christian Léonard, and Luigia Ripani. "About the analogy between optimal trans- port and minimal entropy". In: Annales de la Faculté des sciences de Toulouse: Mathématiques. Vol. 26. 3. 2017, pp. 569-600.
[23] Anming Gu, Edward Chien, and Kristjan Greenewald. "Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior". In: The Thirteenth International Conference Learning Representations. 2025.
[24] Matthieu Heitz, Yujia Ma, Sharvaj Kubal, and Geoffrey Schiebinger. "Spatial Transcriptomics Brings New Challenges and Opportunities for Trajectory Inference". In: Annual Review of Biomedical Data Science 8 (2024).
[25] Jonathan Ho, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models". In: Advances in neural information processing systems 33 (2020), pp. 6840-6851.
[26] Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. "Manifold interpolating flows for trajectory inference". In: Advances neural information processing systems 35 (2022), pp. 29705-29718.
[27] Qi Jiang and Lin Wan. "A physics-informed neural SDE network for learning cellular dynamics from time-series ScRNA-seq data". In: Bioinformatics 40 (2024), pp. ii120-ii127. ISSN: 1367- 4811.
[28)
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Lama Saber, Changying Jing, et al. "Mapping cells through time and space with moscot". In: Nature (2025), pp. 1-11.
[29] Takeshi Koshizuka and Issei Sato. "Neural Lagrangian Schrödinger Bridge: Diffusion Mod- eling for Population Dynamics".. In: The Eleventh International Conference on Learning Representations. 2023.
[30] Hugo Lavenant, Stephen Zhang, Young Heon Kim, Geoffrey Schiebinger, et al. "Toward a mathematical theory of trajectory inference". In: The Annals of Applied Probability 34.1A (2024), pp. 428-500.
[31] Christian Léonard. "A survey of the Schrödinger problem and some of its connections with optimal transport". In: Discrete and Continuous Dynamical Systems-Series 34.4 (2014), pp. 1533-15721
[32] Xuechen Li, Ting- Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. "Scalable gradients for stochastic differential equations". In: International Conference on Artificial Intelligence and Statistics (2020).
[33] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hame Maximilian Nickel, and Matthew Le. "Flow Matching for Generative Modeling". In: The Eleventh International Conference on Learning Representations. 2023.
[34] Xingchao Liu, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow". In: arXiv preprint 2rXiv:2209.03m8 (2022).
12

---

[35]
Malte Luecken, Daniel Burkhardt, Robrecht Cannoodt, Christopher Lance, Aditi Agrawal, Hananeh Alice, Ann Chen, Louise Deconinck, Angela Detweiler, Alejandro Granados, Shelly Huynh, Laura Isacco, Yang Kim, Dominik Klein, BONY DE KUMAR, Sunil Kuppasani, Heiko Lickert, Aaron McGiever, Joaquin Melgarejo, Honey Mekonen, Maurizio Morri, Michaela Müller, Norma Neff, Sheryl Paul, Bastian Rieck, Kaylie Schneider, Scott Steelman, Michael Sterr, Daniel Treacy, Alexander Tong, Alexandra-Chloe Villani, Guilin Wang. Jia Yan, Ce Zhang, Angela Pisco, Smita Krishnaswamy, Fabian Theis, and Jonathan M Bloom. "A sandbox for prediction and integration of DNA, RNA, and proteins in single cells". : In: Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks. Ed. by J. Vanschoren and S. Yeung. Vol. 1. 2021.
[36]
Suryanarayana Maddu, Victor Chardès, Michael Shelley, et al. "Inferring biological processes with intrinsic noise from cross-sectional data". In: arXiv preprint arXiv:2410.07501 (2024).
[37] Kevin Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. "Visu- alizing structure and transitions high-dimensional biological data". In: Nature biotechnology 37.12 (2019), pp. 1482-1492.
[38] Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. "Action matching: Learning stochastic dynamics from samples". In: International conference on machine learning.
PMLR. 2023, pp.
[39] 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 on Machine Learning. 2024.
[40] Alessandro Palma, Till Richter, Hanyi Zhang, Manuel Lubetzki, Alexander Tong, Andrea Dittadi, and Fabian J Theis. "Multi-Modal and Multi-Attribute Generation of Single Cells with CFGen".. In: The Thirteenth International Conference on Learning Representations. 2025.
[41] Matteo Pariset, Ya-Ping Hsieh, Charlotte Bunne. Andreas Krause, and Valentin De Bortoli. "Unbalanced Diffusion Schrödinger Bridge". In: ICML Workshop on New Frontiers in Learning. Control, and Dynamical Systems. 2023.
[42] Qiangwei Peng. Peijie Zhou, and Tiejun Li. "stVCR: Reconstructing spatio-temporal dynamics of cell development using optimal transport". In: bioRxiv (2024). pp. 2024-06.
[43] Katarina Petrović, Lazar Atanackovic, Kacper Kapusniak, Michael M. Bronstein, Joey Bose, and Alexander Tong. "Curly Flow Matching for Learning Non-gradient Field Dynamics". In: Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2025. 2025.
[44] Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Y. Chen, Aviv Regev, and Romain Lopez. "Modeling Complex System Dynamies with Flow Matching Across Time and Conditions". In: The Thirteenth International Conference Learning Representations. 2025.
[45] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. "Optimal-transpon analysis of single-cell gene expression identifies developmental trajectories in reprogramming" In: Cell 176.4 (2019a), pp. 928-943.
[46] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. "Optimal- analysis single-cell gene expression identifies developmental trajectories in reprogramming". In: Cell 176.4 (2019b), pp. 928-943.
[47] Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. "Reconstructing growth and dynamic trajectories from single-cell transcriptomics data". In: Nature Machine Intelligence 6.1 (2024), pp. 25-39.
[48] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnadd Doucet. "Diffusion Schrödinger bridge matching". In: Advances in Neural Information P Processing Systems 36 (2024).
[49] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. "Deep unsu- pervised learning using nonequilibrium thermodynamics". In: International conference machine learning. PMLR. 2015, pp. 2256-2265.
[50] Jiaming Song, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models". In: arxiv preprint arXiv:2010.02502 (2020).
13

---

[51] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. "Score- Beore Generative Modeling through Stochastic Differential Equations". In: International Conference on Learning Representations. 2021.
[52] 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". In: Transactions on Machine Learning Research (2024a). Expert Certification. ISSN: 2835-8856.
[53] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. "Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics". In: International conference on machine learning. PMLR. 2020, pp. 9526-9533
[54] Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P San Juan, Laura Rangel, Brandon Zhu, John G Lock, Christine L Chaffer, and Smita Krishnaswamy. "Learning transcriptional and regulatory dynamics driving cancer cell plasticity using neural ODE-based optimal transport". In: bioRxiv (2023), pp. 2023-03.
[55] Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet. Guy Wolf, and Yoshua Bengio. "Simulation- Free Schrödinger Bridges via Score and Flow Matching".. In: International Conference on Artificial Intelligence and Statistics. PMLR. 2024b, pp. 1279-1287.
[56] Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. "Trajectory inference for branching SDE model of cell differentiation". In: arXiv preprint arXiv:2307, 220.0777777777 (2023).
[57] Adrian Veres, Aubrey L Faust, Henry L Bushnell, Elise N Engquist, Jennifer Hyoje-Ryu Kenty, George Harb, Yeh-Chuin Poh, Elad Sintov, Mads Gürtler, Felicia W Pagliuca, et al. "Charting cellular identity during human in vitro B-cell differentiation". In: Nature 569.7756 (2019),
pp. 363-373.
[58] Wei Wan, Yuejin Zhang, Chenglong Bao, Bin Dong, and Zuoqiang Shi. "A scalable deep learning approach for solving high-dimensional dynamic optimal transport", In: SIAM Journal on Scientific Computing 45.4 (2023), B544-B563.
[59] Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. "Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling". In: Advances in Neural Information Processing Systems (2025).
[60] Zihao Wang, Datong Zhou, Yong Zhang. Hao Wu, and Chenglong Bao. Wasserstein-Fiste Fisher- Rao Document Distance. 2019. arViv: 1904. 10294 [Cs.LG).
[61] Caleb Weinreb, Alejo Rodriguez-Fraticelll Fernando D Camargo, and Allon M Klein. "Lineage tracing on transcriptional landscapes links state to fate during differentiation". In: Science 367.6479 (2020), eaaw3381.
[62] Hao Wu, Shu Liu, Xiaojing Yc, and Haomin Zhou. "Parameterized wasserstein hamiltonian flow". In: SIAM Journal on Numerical Analysis 63.1 (2025), pp. 360-395.
[63] Liu Yang, Constantinos Daskalakis, and George E Karniadakis, "Generative ensemble regres- sion: Learning particle dynamics from observations of ensembles with physics-informed deep generative models". In: SIAM Journal on Scientific Computing 44.1 (2022). 880-B99
[64] Maosheng Yang. "Topological Schrödinger Bridge Matching". In: The Thirteenth International Conference on Learning Representations. 2025.
[65] Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. "Generative modeling of single- cell time series with PRESCIENT enables prediction of cell trajectories with interventions". In: Nature communications 12.1 (2021a), p. 3222.
[66] Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. "Generative modeling of single-
cell time series with PRESCIENT enables prediction of cell trajectories with interventions".
In: Nature communications 12.1 (2021b). p. 3222.
[67] Yuning You, Ruida Zhou, and Yang Shen. "Correlational Lagrangian Schr\," odinger Bridge: Learning Dynamics with Population-Level Regularization". In: arXiv preprint arXiv:2402. 10227 (2024).
[68] Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. "scNODE: generative model for temporal single cell transcriptomic data prediction".. In: Bioinformatics 40.pupplc-
ment_2 2 (2024a), pp. ii 1115-, ISSN: 1367-4811.
14

---

[69] Peng Zhang, Ting Gao, Jin Guo, and Jinqiao Duan. "Action Functional as Early Warning Indicator the Space of Probability Measures". In: arXiv preprint arXiv:2403. 10405 (2024b). [70] Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. "Optimal transport analysis reveals trajectories in steady- systems". In: PLoS computational biology 17.12 (2021), e1009466.
[71] Zhenyi Zhang, Tiejun Li, and Peijie Zhou. "Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport". In: The Thirteenth International Conference on Learning Representations. 2025a.
[72] Zhenyi Zhang, Yuhao Sun, Qiangwei Peng, Tiejun Li, and Peijie Zhou. "Integrating Dynamical Systems Modeling with Spatiotemporal ScRNA-Seq Data Analysis". In: Entropy 27.5 (2025b). ISSN: 1099-4300.
[73] Zhenyi Zhang, Zihan Wang, Yuhao Sun, Tiejun Li, and Peijie Zhou. "Modeling Cell Dynamies and Interactions with Unbalanced Mean Field Schrödinger Bridge". In: Advances in Neural Information Processing Systems (2025c).
[74] Zhenyi Zhang, Zihan Wang, Yuhao Sun, Jiantao Shen, Qiangwei Peng. Tiejun Li, and Peijie Zhou. "Deciphering cell-fate trajectories using spatiotemporal single-cell transcriptomic data". In: Preprint (2025d).
[75] Linqi Zhou, Aaron Lou, Samar Khanma, and Stefano Ermon. "Denoising Diffusion Bridge Models". In: The Twelfth International Conference on Learning Representations. 2024.
[76] Qunxi Zhu, Bolin Zhao, Jingdong Zhang, Peiyang Li, and Wei Lin. "Governing equation discovery of complex system from snapshots". " In: arXiv preprint 4rX (2024))
15

---

A Technical Details
A.1 Proof of Theorems
A.1.1 Proof for Theorem 4.1
Theorem A.I. The RUOT problem with isotropic and time-invariant diffusion intensity is formulated
as:
min A -( / + p(x,t) tx dx
(6)
P.g
1 u(x, ap(x.t) t) -V (((,,))u(xtt) 2 p
s.t.
+ g(x, t))
(7)
dt
2
In this problem, the necessary conditions for the action A to achieve a minimum are given by:
dub(g)
u= Vx),
X,
a = V
2 2 a =0,
(8)
dg
where ((x, 1) is scalar field.
Proof. In order to incorporate the constraints of the Fokker-Planck equation, we construct an augmented action functional:
Al - 1/// - p(x, t) dx
/ / t(x,t)
(Op(x,t) Ot
1 + v (p(x,t) t(x, 2 - 9(x, t)xxx)) dxdt
We take variations with respect to 9. and p. At the stationary point of the functional, the variation of the augmented action functional must vanish.
Step 1: Variation with respect to u.
Let + őu. The variation of the augmented action is
SA f('
22 (u Su)p + XVx (pou)dxdt
(u=ru)p (:' I (Vx (Apōu) (ondu)-((Vaul) P(Vx)) on)dxd/
3
RRd
[
(u='uu)p+ [" f Sm
1
ds 1//e J P( Sudxdt
(u (n- poudxdd
2d
Here, 200 denotes the boundary at infinity in Rd and dS is the surface element. Based on the assumption that
Ap (ōu)" as = 0,
and using the arbitrariness of Su, we obtain the optimality condition
VxA
Step 2: Variation with respect to g.
Let g + g dg. then the variation of the augmented action becomes
SA - / (a dv((g)
->)
8g dx dt.
dg
16

---

Since og is arbitrary, we immediately obtain the optimality condition
dlb(g) X.
dg
Step 3: Variation with respect to p.
Let p + óp. Then the variation of the augmented action is given by
SAt= [/' / (Glupp (
(
06p 01 + Vx (uôp) V2(00) 220(0( (00)
dx
1/' / (Glul -all++a(0)))) \o) op [ [(
(D(A8p)
0d
Igc
-8p a1
dl dl
01
1/' / (Vx (u\ōp) dxdt u (Vx)) dx
2 1/ / (Vx (AV - (VxA) dx 2
gdd
J." / (Glul22 [|ll|2+00(00)0 - Ag
0s
of Aapti
u"
01
J
V (uXop)
0²
(XVxx)))) dt
2 Ss
1
1/' / (Vx)) (V
+
0²
2
L.' Rd
/
( [1ll| + - Ag Øt
ax
- uu (V op
1
1/' /
+
Vxx ((Vx))6p) - (V 6p) dx dt
2
Rd
/'' / ( [11l112 +a Ag
0x
1 -2²A )2X))oxxxxt op
u (Vxxx)
Rd
0t
2
f (1V22)0p) ds
1/' / (1lll12 +a - Ag at
1 00 V 2
nF (V
dp dt.
Since 8p arbitrary, the corresponding optimality condition is
ax u (Vx))
1 ? 0.
aall + a (g)) Ag
2
Substituting the previously obtained condition
Vxd, we arrive the final optimality condition:
0x
at + | 2 2002 - Ag (b) 0.
17

---

A.1.2 Proof for Theorem 4.2
Theorem A.2. The choice of ()) affects whether g ascends or descends along the direction of the velocity field u given time. Specifically, at a fixed time t, if
d (g)
0.
(9)
dg²
then g(x. t) ascends in the direction of the velocity field u(x, t) (ie. (Vxxx)) 1)) > 0); otherwise, descends.
Proof. Let
dg((g) (g) and d2v((g) (g)
dg
dg²
Using the optimality condition for g from Appendix A.1.1,
dxb(g) X,
a dg
taking the gradient with respect to a on both sides yields
Vxg(x, t) ()) Vx)(x,t).
The condition for to increase along the velocity field is that the inner product between Vx9(x, t) and u(x, positive everywhere. Using the optimality condition for the velocity.
u(x, t) = Vx)(x,t), t),
we have
u(x, t)² Vx0(x,)) 1) =
Vx>((x t)12?
a (g)
Since >> 0, the condition
u(x, 1)2 Vxg(x, t) >
is equivalent to requiring that
(g) > Vg.
A.1.3 Proof for Theorem 5.1
Theorem A.3. Consider a weighted particle system consisting of N particles, where the position of particle is given by X E Rdd and its weight by w:(()) 0. The dynamics of each particle described by
dx
dt u(X, f) dt (t)) dw,
(10)
dw; g(X w. dt,
dt
where 1 x [0, 7] Rd is time-varying vector field, g Rdd x [0, + TR. is growth rate function, o [0, [0, +oo) is a tame-varying diffusion coefficient, and W is an N-dimensional standard Brownian motion with independent components in each coordinate. The initial conditions are X ~ p(x, 0) and w,(0)= = In the limit as N - oc, the empirical measure
ne (x,
Σ w, (1)))((x-x)
(11)
converges to the solution of the following BaakerrPllance equation:
0p(x,t) - (u(x,, 1) p(x, o o (1) 1(xx p(x, 1),
(12)
Öt
with the initial condition p(x,0) 0) px(x).
18

---

Proof. Consider a smooth test function ø o: R. + R. We study the evolution of the expectation
N
Jec ((xx n (x,t)dx =
Σ ((()))
N
By applying Itö's formula, we have
(t)))) = wi(() do(x;) + (xx) dt,(t) + (t) do(x)).
Using Itô's formula to compute do(();), we obtain
d = (V xxx))) ax:
2 (2)( o dt.
Since dw, = g(X, t) w; (t) df contains no stochastic term (i.c., there is no dW), the term dw;(t) dø X!) ) is of higher order and can be neglected. Therefore, we have:
d (ox(1) ((x))(()))))))) - o(x) ol X ()) di
+ (V (u(X),
(2X())) (t)
+ 2
Next, we compute
N
N N
mi (t) oxx:)
E
Σ
=
(((!) g(X: ( t) (t)
dt
N
+ (t xxx t(
x2 0² (1)
+ 2
Thus, the limit as N > oc, and let p(x, 1) = -x0 t), we have
di / o(x) t(x, = / (g(x,)) E) p(x, t)((x) + (Vxox))) uu(x t)
d
+ 0 (2)xx
By integrating by parts, we obtain
J/L p(x, (V (Vx)(x))T u(xxx))) ( dx = / s(xx (x Vx (uxx)))(xx1)) 1))x
and
[ p(x,t) v20(x)) / o(x,t) (2p(x,t) dx.
Hence, we deduce that
di / o(x, p(x, t) dr
d
o(x [ t)p(x, t)) (/
Since Φ(x) is arbitrary, we obtain the Fokker-Plance equation:
tp(x,t) -x (u(xx())(x,))) + 2 () g(x, p(x, 1).
H1.
19

---

A.1.4 Proposition: : the Expectation of HJB Loss
Proposition A.I. Consider the following HJB loss:
BIB
Tw wi(()
(DA(X;,1) ()
|x>||
1 + 22 oXx a (
2
di
D 1i=1 wi(l)
Øt
(13)
where
dx; = u(X;, dt o(t dW,, du, = g(X;,t) w, dt,
The expectation of HJB loss is
TK [ p(x,t)
(QA(x,t) t)
+ V + 1 - + amt))) 2
EE SLCHN
dxde
(14)
p(x,,t JK p(x, t)dx
where p(x,t) t)
is the normalized probability density.
Proof. Taking the expectation of LivB is equivalent to drawing N particles each time to obtain LiBB repeating this process infinitely many times, and computing the average of the L obtained in each instance. Since the particles are independent, this operation directly equivalent taking the number of particles N + oo, thus:
N
TK
wi(()
01 0t
oo -x 2 - de
E[C
lim
+ + |
N N
2j=1 w)(()
i=1
N
JK
wx(()
01
+ ž/ |xA| 1 - - a(()))
2
= lim
dl
N-000
N
Ee= wx(t)
01.
N
TK
lim N
1 (x,
+
|V>1||
|
フタ 2 + - - a
dxdt
JEs N dx
01
2
TTk
lim
(a)
+
1
V
2 , + a <5(g) dxdt
JR.
lim Jza N N-000
N-000
Øt
2
Jo (T 1(x,1) 1)
()(xxt) t)
1
+
,(x,1) 1) vb(0)
dxdl
a1
2
In the final equality of the proof, we employed the previously proven Appendix A.1.3.
A.1.5 Proposition: : the Expectation of Action Loss
Proposition A.2. Consider the following action loss:
N
(
2 | (n (' ())
NN Action
(15)
N 1=1
where
= u(X{, dt + o(t) dW., dw, = g((,, (X,(t) t) w, dt.
The expectation of action loss is equal to the action defined in the RUOT formulation, namely,
E = TTx [ (2 (1xx, t)] + (le(xx))) 01) p(x, t) dx dt.
(16)
20

---

Proof. Taking the expectation of Action is equivalent to drawing N particles each time to obtain Action repeating this process infinitely many times, and computing the average of the Action obtained in each instance. Since the particles are independent, this operation is directly equivalent to taking the number of particles N + 0o, thus:
1 N
Tx
||( 1|2 wi( I."
TK a (o((),)))) wi(1)
B|c
= lim N
2
1=1
N
TK
/ 1
= lim
||((xx t)]2 nN (x) dt
No000
N
1=1
+ Tx
/ ( t)) / (x) dt
Tx /
1 lu t)] a p(x, t) dx dt.
In the final equality of the proof, we employed the previously proven Appendix A.1.3.
A.2 Optimality Conditions Under Different v)(g)
1 9. the optimality
In our Experiment, we use two different (b(g) as examples. When (g) = conditions are
2
A
0% 1
1
V X
V
+
-
-
+
= 0.
2
2 a
When (2(g) 2/15 the optimality conditions are:
13
2a
0d
1 + IV + 00 =2
1
13
2a
Vx,,
0.
15
15
15
Note that in this case the function g(A) exhibits singularity À 0. In fact, given the two properties we imposed on (g))
dxb(g)
d|g > 0 and (g))) = vi(-g) along with the constraint ()())) < 0, follows
that v// (g) must be discontinuous at 0, and hence g(A) = (1/) - Mx.()) necessarily has singularity
at = 0. For the sake training stability, we slightly modify g(A) to remove this singularity, We redefine (() as:
a(A).
(A) =
28 (o(0) - -(--))(A++ +g(((())) -8 O g(-A),
1
1
X << x
where d is small positive constant. In our computations, we set 5 = 0.1. Another possible treatment
for this singularity involves fitting the reciprocal of nu, thereby canceling the singularity
ib In this case, g (2 and thus the singularity near zero points no longer exists. We leave how
to eliminate this singularity more systematically in numerical computation the future work.
A.3 Training Algorithm and Computation of Losses
For the algorithms mentioned above, we use the particle method compute the losses. In Ap- pendix D.2, we discuss the convergence rates of these losses: both CHJB and CAction have
convergence rate of O while for data with high dimensionality (d > 1), the convergence
21

---

Algorithm Training Var-RUOT
... DR. batch size N, training epochs Epoch initialized network \x(x, t). Ensure: Trained scalar field (x(x, t).
Require: Datasets Di.....
for 1 to NEpoch do
2: From the data at the first time point, D1, sample N particles and set all their weights to w, w,( 1, for NN..
3: for 1 to K do
dvb(ge) (x, t) calculate
4:
use optimality conditions u 1) Vx)o(xxt), t),
>x
dgo
(x 1) and go (x, t), where E [T; N Cxction Action +
5:
1
T(+1
Ug (X) w))) fr oxk ()oz(()
N
T:11
wt(t)
CN - CiuB
+
TT
Σ wt(()
6:
1 -2 + go (X - a ((nn)
x
+ 2
V \x
dl
δ1.
7:
CRecon f CRecon + (M(T++1) ( W Pl.,
8:
Lioual = "Actioo Action + "HuB CIUB + LRcco
9:
update Xx( t) w.r.t. Lraal
rate of Lor is O(. To balance training speed, memory consumption, and the algorithm's convergence rate, we recommend using approximately N = 10 to 10 particles to compute CHIB. Caction and CRecOn Furthermore, ensure high accuracy in the integration for Chив and Laction- we advise using a high-order integrator (Berman et al. 2024) (e.g.. Stochastic Runge-method (Li al. 2020) and setting the integration step size to approximately At 0.1. L measures the discrepancy between the mass learned by the model and the true mass. The true mass at time point i (for > 1) is defined as the ratio of the number of data points at that time, N, the number of data points at the initial time, N1.
In particular, computing Limn requires second-order derivatives calculate the Laplacian V2X(ar l). This step can be computationally expensive. However, we only need the trace the Hessian (i.e., the Laplacian), not the full matrix. This can be efficiently estimated using Hutchinson's randomized trace estimator. For a Hessian matrix H, the trace is given by trace Hf) = E [v" v.. where the vector drawn from a distribution with zero mean and unit covariance, such (O,
B Experiential Details
B.1 Additional Information for Datasets
Simulation Dataset In the main text, utilize simulated dataset derived from three-gent -ene regulatory network (Zhang et al. 2025a). The dynamics of this system are governed by stochastic differential equations that incorporate self-activation mutual inhibition, and external activation. The dynamics of the three genes are described by the following equations:
dxi
a1 (Xi)² +
5, x + m &
df
1+ m (Xi)²
02 (X:)² na (X-)² +
dX
02 (x + 3
82 X + 122 Et,
dt
1 (X!)²
02 (X)) 73 (X5)² +
dxj
03 (X
83 X33 n73 Et.
dt
1+ ag (x
where (L) represents the gene expression levels of the ith cell time The coefficients Ors Ti and B control the strengths of self-activation, inhibition, and the external stimulus, respectively.
22

---

The parameters ó; indicate the rates of gene degradation, and the terms & account for stochastic influences using additive white noise.
The probability cell division is linked to the expression level of X2 and is given by
X3
g ag 0 1+ X3
When a cell divides, the resulting daughter cells are created with each gene perturbed by an independent random noise term, 7a-V (0, 1), around the parent cell's gene expression profile (X1 (t), X2 (t), Xa(t)). Detailed hyper-parameters are provided in Table 5. The initial popula- tion of cells independently sampled from two normal distributions, W[[2, 0.2, 0), 0.1) and ((0,, 0, 2]. 0.1). At every time step, any negative expression values are set to zero.
Table 5: Simulation parameters for the gene regulatory network.
Parameter
Value
Description
0.5
Self-activation strength for X1- Inhibition strength exerted by X3 on Xi. Self-activation strength for Xz.
711
0.5
022
1
72
I
Inhibition strength exerted by X3 on X2. Self-activation strength for X3.
03
I
73
10
Half-saturation constant in the inhibition term. Degradation rate for X1.
81
0.4
82
0.4
Degradation rate for X2. Degradation rate for X3. Noise intensity for X Noise intensity for X2.
83
0.4
171
0.05
1/2
0.05
133
0.01
Noise intensity for X3-
0.014
Noise intensity for perturbations during cell division. External signal activating X1 and X2.
B
dt
1
Time step size.
Time Points
[0, 8, 16, 24, 32]
Discrete time points when data recorded.
Other Datasets Used in Main Text In addition to the three-gens simulated dataset, our main text also utilizes the EMT dataset and the Mouse Blood Hematopoiesis dataset. The EMT dataset is sourced from (Sha et al. 2024; Cook and Vanderhyden 2020) and is derived from A549 cancer cells undergoing -inducedd spithelial-mesenchynnee distinct time points, containing total of 3133 cells, with each cell represented by 10 features obtained through PCA dimensionality reduction. Meanwhile, the Mouse Blood Hematopoiesis dataset covers 3 time points and includes 10,998 cells in total (Weinreb al. 2020; Sha et al. 2024) and was reduced to 2-dimensional space using nonlinear dimensionality reduction.
transition (EMT). It comprises data from four
Gaussian Datasets To examine the differences between the optimal transport paths solved by our algorithm and those obtained from traditional OT solvers, and to validate the scalability of our model (i.e., its learning capability on ligh-diimnsional data), we constructed two 2D Gaussian datasets (a balanced version and an unbalanced version) and three high- sigh-dimessional Gaussian datasets. The 2D Gaussian dataset consists of two time points, where the distribution at each time point is mixture of multiple Gaussian distributions, as illustrated Figure 5.The three high-dimensional Gaussian datasets have dimensions of 50, 100, and 150, respectively. The 100-dimensional dataset was adapted from the experiments in RUe (Zhang et al. 2025a) The results of these high-dimensional datasets, after being reduced to two dimensions using PCA for visualization, are shown in Figure 5,
Other High Dimensional Datasets In addition, we employed four real-world datasets for our evaluation. The first is the Mouse Blood Hematopoiesis dataset from (Weinreb et al. 2020), which comprises 49. 19,302 cells collected three time points. We reduced its dimensionality to 50 using PCA, and subset of this processed data was used for the experiments in our main text. The second the Pancreatic 3-cell Differentiation dataset from (Veres al. 2019), consisting of 51. 1,2744 cells sampled across eight time points. We subsequently reduced its dimensionality to 30 via PCA. The third dataset, sourced from (Moon et al. 2019), profiles 16,819 cells sampled at five time points from human embryoid bodies (EBs), which developed via the spontaneous aggregation of stem cells in
23

---

(a) Zdim Gaussian Mixture
(b) 50dim, 100dim Gausslan Mixture
(c) 150dim Gaussian Mixture
Figure 5: Diagram of the Gaussian mixture datasets. PCA is used to reduce the high dimensional data to dimensions.
3D culture throughout 27-day differentiation process. We reduced its dimensionality to 100 using PCA and then standardized each feature. The fourth dataset is from the NeurIPS Challenge (Luecken al. 2021). From this, utilized the gene expression data for the donor with DonorlD 28483. We extracted the 1,079 cells from this donor that were annotated with pseudotime values. To create discrete time points, performed K-Means clustering these cells based their pseudotime. The resulting clusters were then ordered by their centroid pseudotime values and assigned discrete labels t 0, 1,2,3,4 4. This data also standardized per-feature basis.
B.2
Evaluation Metrics
To assess the fitting accuracy of the learned dynamics to the data distribution, we compute the Wi and W2 distances between the data points generated by the model and the real data points. They defined as
Wi(p.g) = =et(ip.)) / y
and
/ 11x-y11 y dx(xx,))
1/2
(2(p,g)
min (Tll(p..)
We compute these two metries using the emd function from the pot library.
To evaluate the action the dynamics learned by the model, we directly compute the action loss. Appendix A.1.5 guarantees that the expectation of the loss is equal to the action defined in the RUOT problem. The action loss is:
N
(
| (() L' (x))
C
N Action
Σ
N 1=1
2
where
dx = u(X t) dt + ott) dW..
dw; o(X, w: dt,
with initial conditions X ~ p(x, 0) and w; (0) = 1. We run our model times on each dataset, to calculate the mean and standard deviation of Wi, W2 and action.
To evaluate the training speed of the model, we use the SamplesLoss class from the geomloss library to compute the OT loss at each epoch during the training process for each method, with the blur parameter set to 00 We sum the OT losses at all time points to obtain the total OT loss. For each model, we perform 5 training runs, recording the number of epochs and the time required for the OT loss to drop below specified threshold. We then compute the mean and standard deviation of these values, with the mean reflecting the training/convergence speed and the standard deviation reflecting the training stability.
For models whose dynamics are governed by stochastic differential equations, the choice of a directly affects the results (both the OT loss and the path action). Therefore, when running the RUOT and our Var-RUOT models each dataset, is set to 0.10.

---

C Additional Experiment Results
C.1 Additional Results on Training Speed and Stability
We plotted the average loss per epoch across five training runs in Fig. 6. Experimental results show that on the Simulation Gene dataset, our algorithm converges approximately 10 times faster than the fastest among the other algorithms (RUOT with 3-epoch pretraining), and the EMT dataset, algorithm converges roughly 20 times faster than the fastest alternative (TIGON).
(a) Training Process on Dataset Simulation Gene
(b) Training Process on Dataset EMT
RUOT w/o Pretraining
TIGON
10 LoSs
RUOT with Pratraining
RUOT w/o Pretraining
RUOT
RUOT with Epoch Pretraining
TIGON
About huster
VanRUOT
Faiter
266
100
Stop
Figure 6: Loss curves of each algorithm over five training runs. The curves are obtained by averaging the losses from the five runs.
C.2 Hyperparameter Selection and Ablation Study
Hyperparameter Selection We used NVIDIA A100 GPUs (with 40G memory) and 128-core CPUs to conduct the experiments described in this paper. The neural network used to fit A(x, t) is fully connected network augmented with layer normalization and residual connections. It consists of our algorithm, the main hyperparameters that need
hidden layers, each with 512 dimensions. tuning include the penalty coefficient a for growth in the action, and the weights THuB and YAcsion for two regularization losses, LHI and LAction respectively. Here, a represents our prior regarding the strength of cell birth and death in the data: larger a imposes a greater penalty on cell birth and death, thereby making easier for the model learn solutions with lower birth and death intensities. Mcanwhile, the HJB loss and the action loss, serving regularizers, are both designed ensure that the solution obtained by the algorithm has as low an action as possible- -te HJB equation being a necessary condition for the action to reach its minimum, and the inclusion of the action loss ensuring that the model learns a solution with an even smaller action under those necessary conditions.
To ensure that our algorithm generalizes well across a wide range of real-world datasets, we only used
two sets of parameters: one for the standard WFR Metric (v) 2 and one for the Modified
Metric (v)()) = 2/)15). The parameters used in each case are listed Table 6. The primary reason for using two sets is that different metrics yield different scales for the HJB loss.
Table 6: Parameter settings for standard WFR metric and modified metric.
Parameter
7HBB
Mectio
Learning Rate
Optimizer
Standard WFR Metric (vi(g) 9 Modified Metric (v2(g) = 2/15) 6.25 x 10-3
6.25 10
6.25 10
2.00
1 x 10
AdamW
6.25 x 10-2
7.00
2 x 10-5
AdamW
Sensitivity Analysis of a To demonstrate the robustness of our algorithm with respect to hyperpa- rameter selection, we first varied the penalty coefficient for growth and examined the resulting
25

---

changes in model performance. This sensitivity analysis was conducted on the 2D Mouse Blood Hematopoiesis dataset, and we performed the analysis for both the Standard WFR metric and the modified metric. The performance of the model under different values of a is shown in Table 7. The experimental results indicate that our algorithm not sensitive to a, as similar performance can be achieved across multiple different values of a. In comparison with the Standard WFR metric, however, the algorithm appears to be somewhat more sensitive to a when the modified metric used.
Table 7: On the Mouse Blood Hematopoiesis dataset, the Wasserstein distances (Wi and W2) between the predicted distributions of Var-RUOT with different a and the true distribution at each time points. Each experiment was run five times compute the mean and standard deviation.
t=
Model
W1
W2
W1
W22
Var-RUOT (Standard WFR, a 1)
0.1622:100072
1.0007
0.1280_0.012
0.1522110013
Var-RUOT (Standard WFR, a = 2)
0.1203 +0.000
0.1498+a.mae
0.1389
=0,008
0.1701 +ooms
Var-RUOT (Standard WFR, a 3)
0, 0.1402++o.00
0.1704
0.1350
0.1655 4-00132
Var-RUOT (Modified Metric, a = 5)
0.3783 0194
0.3326
0.21
000000
0.2226:10022
Var-RUOT (Modified Metric, ar = 7)
0.2953
0.31 1 7 200000
0.1917 -0.0040
0.2226 01
Var-RUOT (Modified Metric, Q 0.2737 +00095
0.3116++o0072
0.1970==..0m
0.2224 4-0.075
Sensitivity Analysis of Growth Penalty Function
In this work, we select (g) = g2/15 to illustrate a case where the conditions vg' > and ( < 0 hold. As established in Theorem 4.2, these conditions ensure that the growth rate g(æ, t) increases along the direction of the velocity field u(x () at each time step.
The specific choice of (2)(g) = g2/15 is motivated by its simplicity and, more importantly, by considerations for training stability. For a general penalty of the form (gg) g(2p)/ the first-order optimality condition a dg
do(g) = À yields an analytical solution for in terms of A:
2
2g+1
((A) 29 + Ip
(3u
x
2p
Since g is computed directly from X during training, we must avoid an excessively steep g()) curve, especially near À 0, to ensure stability. This requires the exponent (24+1)-2p be small. We achieve this by setting p = 1 (the smallest positive integer) and choosing a large value for Selecting q = 7 results in our final penalty function, vb2(g) = g2/15 which balances simplicity and stability.
2q+1
To empirically validate this choice, we conducted a sensitivity study on the 2D Mouse Blood Hematopoiesis dataset using different parameter choices for (g)). The results, shown in Table compare the Wasserstein distances (W1 and W2) at different time points.
Table 8: Sensitivity analysis on the choice of (gg).
2
Penalty Function sk(g)
W2
Wi
v'g)) = /2/ (p = 1,9 =0:==d. 4)
0.9476_nneex
1.082_1osor
1,2197.nxve
1,3414gnmm
v(g) g2/16 (p 1.9 = 7)
0,3117=0m22
0.1917=a0m0
0.1764
(gg)) 2/22 (p = 1.9 = 10)
0.4239 -usss
0.5009-00ma
0.2638+aam4
0.2194
The results demonstrate that vb(g) 22/15 achieves the best performance. Its performance is
comparable to that of (gg)) 22/21, whereas the choice of (b(g) = g2// (with a smaller q) leads to poorer reconstruction quality due to training instability. This confirms that selecting an exponent within a reasonable range (i.e., with a sufficiently large q) is effective and stable
Ablation Study of THB and 7Action In order to verify whether LUB and CAction facilitate the algorithm find solutions with lower action, we conducted ablation studies. These experiments were carried out on the EMT data and 2D Mouse Blood Hematopoiesis Data, since in this dataset the transition from the initial distribution to the terminal distribution can be achieved through relatively
26

---

simple dynamics (each particle moving in roughly straight line). Therefore, if the HJB loss and the action loss are effective, the model will learn this simple dynamics rather than a more complex one. We varied the HJB loss weight "Hun to the following values:[0, 6.25 x 10-3, 3. 1.125 x 10-2, 6.25 x 10-2, 6.25 x 10-1, 3. 3.125,. while keeping the action loss weight V fixed at We then shown both the mean W1 distances between the predicted and true distributions at four different time points and the trajectory action (as shown in Table 9 ). Similarly, we fixed THB = 1 and varied AAction over the same set of values: [0, 6.25 x 10-$ 3.125 10-2 6.25 10 10-2 6.25 10-1, 3.125], with the corresponding results illustrated Table 10. The figures indicate that as THJB and "/Action increase, the action of the learned trajectories decreases monotonically, demonstrating that both loss terms effective. However, as these weights increase, the model's ' ability to fit the distribution deteriorates.
Therefore, recommend that in practical applications, both WHB and VAttion should be set to values below 0.1, as configured in this paper.
Table 9: Comparison of W1 and Path Action for different values of the parameter "HuB on a specific dataset.
Dataset
Metric
Value of Yuun
0
6.25 x 10-3
3.125 x 10-2
6.25 x 10-2
6.25 x 1001
3.125
Mouse Blood Hematopoiesis
Wi
0,1573
0,1609
0,1550
0,1316
0.2109
0.2539
Action
3.3002
3.2165
3.2397
3.1491
3.1182
3.1200
EMT
Wi
0.2702
0.2701
0.2686
0.2719
0.2722
0.2748
Puth Action
0,4512
0,4519
0,3867
0,3641
0.2865
0.2796
Table 10: Comparison of Wi and Path Action for different values of the parameter "AAction on a specific dataset.
Dataset
Metric
Value of "Aacton
0
6.25 x 10
3.125 x 10
6.25 x 10
6,25 x 10
3.125
Mouse Blood Hematopoicsis
Wi
0,1454
0,1344
0,1674
0,1316
0.4001
0.7914
Action
3.4043
3.4497
3.2737
3.1491
2.4716
1,2059
EMT
Wi
0.2706
0,2681
0.2656
0.2774
0.3043
0.3832
Path Action
0.4192
0,4171
0,3934
0,3288
0,1686
0,0709
The Effectiveness of different types of Loss Norms
In the experiments above, we used the L2 loss for computing both CHB and LAction To investigate
the effect of the loss function type, we present the results of using the L1 loss. On the Simulation and EMT datasets, the experimental results with the L1 loss are shown in Tables 11 and 12. The results indicate that with our hyperparameter values (THE = 6.25 x 10 2 YAAciion 6.25 x 10 10-2), using either the L1 loss or the L2 loss when computing the HJB loss achieves similar reconstruction performance and Path Action values.
Table 11: Performance comparison of Var-RUOT with L and J, loss functions on Simulation Dataset. The table shows Wi and W/2 distances, as well as the final path action.
Model
W//
Wi
Wg
VUI-RUOT (1., Loss) go0e
0.1181 geom
m.1270 0.0
0.045
0.1144garoe
0.0772
0.2140gamm:
1.1105
Var-UUTT (4, Loss)
0.0489
0.1112ga-m
0.1541 0.0956
0.0589
0.1372
0 0.0664gnaa wiz
0.2188
g.124
Table 12: Performance comparison of Vur-RUOT with L2 and L2 loss functions EMT Dataset for time steps f = 1,
Model Var-UUTT (Le Loss]
W2
31/
Wy
Path Action
0.2540
0.2756
0.3544gaam
ROT (la Les)
0,2631 Lonne
0.2717-ax0
0.2796.500 wi
0,2809 2.28862mm4
0.2916uwess
0.3528
C33 Hold-one-out Experiment
27

---

Table 13: Performance comparison of Var-RUOT with L and L1 loss functions on 2D Mouse Blood Hematopoiesis Dataset for time steps 1,
t
Model
W2
W2
Wi
W2
Path Action
Var-RUOT (Le Loss)
0. 0.1459-wann +.
0. 0.1431 Hulsoe
0.1764-aamss +
Var-RUOT (I, Loss)
0.1747 1O00Ms
0.2097
0. 116900.uore
0.2102,1ouez
3.18891 counu
In order to validate whether our algorithm can learn the correct dynamical equations from a limited set of snapshot data, we conducted hold-one-out experiments on the three-gene simulated data, the EMT data, and the 2D Mouse Blood Hematopoiesis data. This experiment is designed to test the interpolation and extrapolation capabilities of the algorithm. For a dataset with n time points, we perform n experiments. In each experiment, one time point is removed from the n time points, and the model is trained using the remaining time points, Afterwards, we compute the W₁ and W2 distances between the predicted distribution and the true distribution at the missing time point. When a time point from 21,2.... .... n - 1} is removed, the model is performing interpolation; when the time point is removed, the model is performing extrapolation. The results of these experiments are shown in Table 14, Table 15, and Table 16. The experimental results indicate that our model's s interpolation performance is superior to that of TIGON and comparable to that of DeepRUOT: in the EMT data and the Mouse Blood Hematopoiesis data, our model achieves the superior extrapolation performance. To further assess the long-term extrapolation capabilities of Var-RUOT and investigate whether minor errors from early time points are amplified over the longer time period, we conducted an additional extrapolation experiment. For this, we extended the three-gene simulation dataset to nine time points. The initial five time points (1 = 0, 1, 2,3, 4) were designated for training, while the final four (I 5, 6, 7, were reserved for evaluating extrapolation performance. We compared the performance of TIGON, DeepRUOT, and Var-RUOT, with the results presented in Table 17. These results indicate that Var-RUOT maintains performance superior to TIGON and comparable to DeTpRUOT, even in this long-term extrapolation setting. A degradation in performance over longer extrapolation horizons is expected. This is primarily because we employ a non non-autonomous system, where the parameter À is function of both the state æ and time f. Consequently, the model learns the mapping X(x. t) only within the temporal domain of the training data, and its behavior beyond this range cannot be accurately inferred.
From a physical viewpoint, the dynamical equations governing the biological processes of cells can be formulated in the form of a minimum action principle (in this work, ITI RUOT Problem is a surrogate model whose action is not the true action derived from studying the biological process, but rather a simple and numerically convenient form of action). Compared to other algorithms, our method can find trajectories with lower action, i.c., It is more capable of learning dynamics that conform the prior prescribed by the action functional. These dynamics yield better extrapolation performance, which indicates that the design of the action in the RUOT problem is at least partially reasonable. From a machine learning perspective, forcing the model to learn trajectories corresponding to the minimum action serves as form of regularization that enhances the model's 's generalization capability.
In addition, we separately illustrate the learned trajectories and growth profiles on the three-gene simulated dataset after removing four different time points, as shown in Fig. and Fig. 8, respectively. The consistency in the learned results indirectly demonstrates that the model is still able to learn the correct dynamics and perform effective interpolation and extrapolation, even when snapshots at certain time points are missing. We further illustrate the interpolated and extrapolated trajectories of both the DeepRUOT and Var-RUOT algorithms on the Mouse Blood Hematopoiesis dataset, as shown in Fig. 9 and Fig. 10, respectively. This dataset comprises only three time points, t = O, 1, 2. When one time point removed, Var-RUOT tends to favor a straight-line trajectory connecting the action path), which serves
remaining two time points (since such a trajectory represents the minimum cittmumm-atioon as an effective prior and leads to a reasonably accurate interpolation. In contrast, because DeepRUOT does not explicitly incorporate the minimum-action objective into its model, the trajectories it learns tend to be more intricate and curved. These more complex trajectories might present challenges for generalization, making accurate interpolation or extrapolation more difficult.
28

---

w/o
0.
6.3
0.00
200
3.03
1.00
4.92
1.00
weight
Figure 7: Trajectories learned on the three-gene simulated dataset after individually removing t 1,2,3, 4.
w/o b-]
w/o b-2
w/o
w/o b-4
Figure 8: Growth learned on the three-gene simulated dataset after individually removing /
1, 2, 3, 4.
w/o t=1
w/o t=2
t=0
t=
t=
Figure 9: The results of the DeepRUOT algorithm the 2D Mouse Blood Hematopoiesis dataset for interpolation (with t=1 removed) and extrapolation (with t=2 removed)
29

---

Table 14: On the three-gene simulated dataset, after removing the data of each time point in turn and training on the remaining data, the Wasserstein distances (i.e., W₁ and W2z distances) between the model-predicted data for the missing time points and the ground truth are computed.
3
Model
W2
W2
Wi
W4
Wi
M4/2
W/
W2
TIGON
0,1205 suueoo
0.1679
0,0931 souoe
0.1919
0.2390scooxc
0.3369-10x8
0.2403-a2xx
0.3616
RUOT
0.1505 swoers
0,0887 200ow
0.1501
0.1 2.
0.1704:aam
0.2179
Yac-RUOOT(Our))
0,0580 anrea
0.1210gmre6
D. 104 grroos
0.229 Par
0.0943 2:
0.1 7769999
0.1401 14401ga00 g9
0.3382ga
Table 15: On the EMT dataset, after removing the data of each time point in turn and training on the remaining data, the Wasserstein distances (i.e., W1 and W2 distances) between the model-predice data for the missing time points and the ground truth are computed.
2
3
Model
W1
W2
W1
W2
W1
W2
TIGON RUOT
0.3457 10.0000
0.3733 19:000
0.3849
0.5260aommm
0.5424
a00m0 1o00m
0.3344
1o0044
1000004
0.3445 =000011
0.4947 1000019
0.5074
Var-RUOT(OOurs)
0.3018
0.3104:00m1
0.3375 =0.0077
0.3460=00002
0.4082
0.4189:amm
Table 16: On the Mouse Blood Hematopoiesis dataset, after removing the data of each time point in turn and training on the remaining data, the Wasserstein distances (i.e., W1 and W2 distances) between the modecepredicted data for the missing time points and the ground truth are computed.
t = 1
t 2
Model TIGON
W1
W2
W1
W2
0.58382200000
0.6726=0.0m
1.326420.00m
1.3928=2n.0m
RUOT
0.6235 10,0014
0.6971
10,0012
1.0723
1.1397 =000020
Var-RUOT(OOrr)
0.2696.10.05
0.3279 1.00044
0.2594
0.3016 20.006
Table 17: Long-term extrapolation performance comparison on the three-gene simulation dataset. All
models were trained three-gene gene simulation data from the first five time points (t 0, 4) and evaluated on four subsequent, unseen time points (t = 5, 6, 7, 8). The table reports performance metrics (e.g., Wi and W2 distances) against the ground truth.
Method
8
Wi
W9
W
W;
W
W2
W1
W//
TIGON
0.1932ama
0.2496ane
0, a,3437
0.401 Lete
0.5209
0.5643
0.691 B2en
0.7155
DesRRUT 0.1027 Hron4s
0.0459
0. 0.1940suae
0.2192
0.3359
0.3644
0.5047
0.5362-s1ew
VarRUOT
0.105
0.2191 aooes
0. 1866
0. 0.2604
0.2890
0,3787
a.4260=amx
0.5557=x1ev
C.4 Experiments on High Dimensional Dataset
To evaluate the effectiveness of our method on high- dimensional datasets, we first tested it on Gaussian datasets of 50 and 100 dimensions. We learned the dynamics of the data using the standard WFR metric (x)(g) = _ as well as a modified growth penalty function, ()(g)) = 2/1 which ensures vb"( < 0. The learned trajectories and growth rates are illustrated in Fig. 11. Under both choices of el(g), our method captures reasonable dynamics: the Gaussian distribution centered on the left shifts upward and downward, while the Gaussian distribution on the right exhibits growth without displacement. We also tested the performance of VarRUOT and other algorithms on a 150-dimensional Gaussian Mixture dataset, as shown in Table 18.
High Dimensional Gaussian Dataset
50D Mouse Blood Hematopoiesis and Pancreatic 3-cell Differentiation Dataset We tested our method on two high- digheniinnsii real scRNA-sep seq datasets including 50D Mouse Blood Hematopoiesis dataset and Pancreatic 3-cell Differentiation dataset. We used UMAP to reduce the dimensionality of the datasets to 2 (only for visualization). plotted the growth of each data point, and visualized the vector fields uxx, 1) on the reduced dimensions using the scvelo library. The results for the two datasets are shown in Fig. 12 and Fig. 13, respectively. As can be seen from the figures, the reduced
30

---

w/o t=1
w/o t=2
0.00
0.00
t 1.00
1.00
t 2.00
2.00
Figure 10: The results the Var-RUOT algorithm on the 2D Mouse Blood Hematopoiesis dataset for interpolation (with t=1 removed) and extrapolation (with t=2 removed)
Standard WFR Modified Metric
1000
100
1.000
Trajectories and growth learned on the 50-dimensional and 100-dimensional Gaussian datasets using the Standard WFR Metric and Modified Metric.
Figure 11:
Table 18:
Wasserstein distances (W1 and W2) between the predicted distributions of each algorithm and the true distribution at various time points on 150D Gaussian Data. Each experiment was run five times compute the mean and standard deviation.
Model
1
Path Action
Wi
W/22
Wi
W3
SH2M
7.286
7.345
9.4572200
9.558220m
PISDE
6.051 20000
6.061
2oote
7.5462000
7.59720mm
MIO Flow
6.494:aam0
6.521
7.812+80000
7.866+0000
TIGON
6.157
=.169
7.615 +0.00
7.665+0000
13.4913
DeepRUOT
6.202-1am0
6.213
7.648 inows
7.69620005
13.5992
Var RUOT
6.072-ta.0
6.089
7.596
10.1033
velocity field points from points with smaller to those with larger t, which indicates that our model can correctly learn a vector field that transfers the distribution even in high- nigh-dimensional data.
EB Dataset and NeurIPS Challenge Dataset
To further quantitatively evaluate the performance of VarRUOT on high-dimensional data, we em- ployed the 100-dimensional EB dataset and the 50-dimensional NeurlPS Challenge dataset. The
31

---

Vector Field
Growth
0.0
11
20
Figure 12: Learned vector field u(x, t) and growth on the 50D Mouse Blood Hematopoiesis dataset (data reduced using UMAP and vector field stream plots generated with the sevelo library).
Vector Field
Growth
100
20
040
Figure 13: Learned vector field u(x, 1) and growth on the Pancreatic 3-cell Differentiation dataset (data reduced using UMAP and vector field stream plots generated with the sevelo library).
quantitative results for VarRUOT against other algorithms presented in Table 19 and Table 20, respectively. On the EB dataset, Var-RUOT demonstrates distribution reconstruction performance comparable to that of TIGON and DeepRUOT, while surpassing simulation-frre approaches (OTCFM, UOTCFM) and Wasserstein Gradient Flow. For the NeurlPS Challenge dataset, Var-RUOT outper- forms all other prior methods evaluated. We also report the results for VarRUOT and competing algorithms on the unnormalized 5-dimensional EB dataset in Table 21.
Table 19: Wasserstein distances (W1 and W2) between the predicted distributions of each algorithm and the true distribution at various time points on EB Data. Each experiment was run five times to compute the mean and standard deviation.
f
31
W2
Mii
3/2
W4/
34
1141711
12,5436sz
12 22.6776mmm
14.7100gnmmm
14.8127asmo
aue
9.78724
9,89522mmm
11.1102
11,207
MIOFLOW
8.7182
2.7924 3m0
9.41392a0e
9,51801anmm
9.7547
9 36361 2:67
11.1065 2000
TIGON
8.7992
9.6497
9.7533
10. 10.1209 10.0130.1m00
11.3284 ooas
11.2452
8,5029
8803
10.0365
10.1424.3xci
11.3555
11.4501
OTCFM
9.4450 cocan
9.5280 e
1L58833-omo
14.05294:a00
14 -4.9498
19.7489
19.7489.4x00
UOTCEM
9.479- Lozaao
9.5596 2.00
11.00 116060 suee
11.3356 ouo
14,6715 10
15 15.6972 Juom
21.7653
Lozaoo
25,5683
10:000
Action Matching
12.7 12.7145m00ms
12.9813.a
-6.8206-00m
13.466344
13. 200
18.1
18.165744
WGF
9.984722000
10.04K22r
05662.10e
11.97192rmo 23.134
12.050M
13.2417
13.3035
Var-RUOT
8,999644
9.7958ga000
9.9K5 gevtes
10.2647
10.4840g.mmm
11.6323
32

---

Table 20: Wasserstein distances (Wi and W2) between the predicted distributions of each algorithm and the true distribution at various time points on NeurIPS Challenge Data. Each experiment was run five times to compute the mean and standard deviation.
Method
f
W2
W/
W2
W2
W2
SP2M
6.1876
6.2565
12.617
13.6426giinn
9.1830
10. 10.9930ganss
12. 1
13,2691 gaas
PISDE
3,6642:1ams
5,7069
5,2358
Lover
5,2779 .wee
6,2795
6.4015
8.9086
9,0298 1.0298
MIOFLOW
6.4267.1am0
6.4817
Lume
6.3141
7.2222.ecaa
9,0160,uuo
TIGON
2:9481::90
6.000
g.7e
5.8008gmre
6.55.3gmmee
6.6787ammm
8.957
g.0e87
RUOT
6.0166:1a0/
6.0579
5.7415
6.5099
6.64192mmm
8.8047 sxume
8.9103
OTCEM
2,5648mmee
3.44/gmmee
5,8151 garee
2.0938
7.34732mrme
10. 10.1416izs
10.7654_.no _nove
UOTCFM
5.5526:1.30/
5.6089
suee
5.4349 Lome
5,5086 sure
6.7482 suooe
6.6849 Losme
9.6782 couae
10.3230 .coee
Action Matching
7.3701
7,4312xc00e
8.5156 soree
8.5485 -oo0x
9.2842
9.3269
9.9665
10.0428
WGF
7
g.
2.7479g.mew
8.2521 goioe
8.9044
8.7749g 9.363620087
VaC-RUOT
4.7961=t.ms
4.9015
4,18162ues
4,2887 (X003)
6.2312
6.4587202144
9. 0.1058aumn
Table 21: Comparison of W distances (Wi) between the predicted distributions and the ground truth on the unnormalized 5-dimensional EB dataset.
Wasserstein- Distance (W)) at time
Model
f
3
t 4
MMFM
0.477
0.554
0.781
0.872
Metric FM
0.449
0.552
0.583
0.597
SF2M
0.556
0.715
0.750
0.650
MIOFlow
0.442
0.585
0,651
0.670
TIGON
0.386
0.502
0.602
0.600
DeepRUOT
0.386
0.497
0.591
0.585
UOT-FM
0.544
0.670
0.729
0.852
VARRUOT
0.416
0.486
0.509
0,511
C,5 Comparison with Tranditioal OT Solvers
To verify whether the Var-RUOT algorithm can find the correct minimum action paths for both Optimal Transport (OT) and Unbalanced Optimal Transport (UOT) problems, we constructed two normalized (balanced) and two unbalanced Gaussian Mixture datasets on 2D plane. We then compared the action computed by the Var-RUOT method (with o 0 to eliminate stochasticity) against the minimum action values calculated by traditional solvers.
To compute the ground truth for the minimum action on the normalized datasets, we used the pot library. For the unbalanced datasets, we could not find an open- solver for dynamic unbalanced optimal transport (apart from deep learning based methods like TIGON ). Therefore, we adopted static Wasserstein-Fisher-Raa (WFR) distance solver, designed as a variant of the Sinkhorn algorithm as described in (Wang et al. 2019), to serve as our baseline. The comparison results are presented in Table 22.
Table 22: Comparison of Path Action values on 2D Gaussian Mixture datasets.
Dataset
Normalized
Unbalanced
Single Gaussian
Gaussian Mixture
Single Gaussian
Gaussian Mixture
Path Action (Traditional Solver) Path Action -UOT)
0.5046
1,9062 1.8754
1.3776
2.4464 2,3909
0.4998
1.0102
The results show that for the two normalized datasets, Var-UUOT finds Path Action values comparable to those computed by the pot library. On the two unbalanced datasets, the Path Action found by Var- is even lower than that of the traditional solver. This outcome might be attributable to either (1) suboptimal solution of the baseline solver or (2) numerical roundoff errors. Nevertheless, our review of the literature did not identify alternative open-source baselines addressing this specific problem. Overall, these results demonstrate that Var-RUOT can achieve performance comparable to traditional OT solvers or WFR distance solvers for the task of finding the minimum action path, which validates its effectiveness.
33

---

D Further Discussions
D.1 Comparison with Relevant Algorithms
Several methods have been proposed to address the problem of trajectory inference for temporal snapshots of single-cell data. We list the differences between our method and previous trajectory inference algorithms in Table 23.
Table 23: Comparison of trajectory inference methods. Properties assessed are: handling Unbalanced Distributions (unequal mass); modeling Stochastic Dynamies; addressing both Unbalanced + Stochastic properties simultaneously; and formulation based on the Least Action principle via first-order conditions, not just as loss penalty. The indicates the method possesses the property. while a indicates it does not.
Method
Unbalanced Distribution
Stochastic Dynamics
Unbalanced Stochastic
Least Action
SF2M
PISDE
MIOFlow
Action Matching
TIGON
DeepRUOT
Var-RUOT (Ours)
Difference from DeepRUOT
Like DeepRUOT, Var-RUOT addresses the critical challenge of handling both unbalanced distributions and stochastic dynamics, which is essential for single-cell omics data.
The primary difference is architectural. DeepRUOT employs three interdependent neural networks to model the velocity field v(x, t), growth rate g(x, t). and log-density log p(x, t). It enforces physical constraints soft penalties in the loss function. This complex setup necessitates multi-stage pre-training procedure to ensure convergence.
In contrast, Var-RUOT uses single network to parameterize scalar field Ao (ax L). Both the velocity u (x. t) and growth rate g(x,,) t) are analytically derived from \a via first-order optimality conditions (Theorem 4)... This simpler, end-to-end design eliminates the need for pre-training and improves training stability.
Difference from Action Matching
The differences between VaO-RUOT and Action Matching could be summarized from three perspec- tives. Firstly, Var-RUOT handles both unbalanced distributions and stochastic dynamics, which is crucial for single-cell omies.
Secondly, Action Matching's training requires sampling from intermediate distributions (z()). These typically challenging obtain in single-cell omics, where data exists only sparse time-point snapshots. This data limitation can potentially hinder the ability to capture complex dynamics, whereas Var-RUOT is designed for such snapshot data, and empirical results suggest Var-RUOT's desirable performance.
Lastly, the methods employ different loss formulations. Action Matching minimizes an Action Gap loss:
[' Bn((z)
Os Ôt (1) dt.
LAM (x) [Su(x)) E (x)
On the other hand, Var-RUOT directly minimizes the residual of the Hamilton-Jacoii-elll (HJB) equation. By directly enforcing this first-order optimality condition in its loss, Rar- can find solutions with low action robustly in empirical tests.
D.2 Convergence of the Particle Method and Loss Computation
The VarRUOT algorithm involves three loss terms: LIUB, LActions and CReson Among these, CIuB and CAction computed using the Monte Carlo method. This involves sampling set of particles,
34

---

integrating along their trajectories, and then averaging the results over all particles. The convergence rate the Monte Carlo method is
where N is the number of particles used.
The reconstruction loss, LRconn is composed of two parts: Lor and LMsss Lot, calculated using the geomlos, library, represents the W/2 distance between the normalized empirical distribution and the true distribution. According the triangle inequality, W2(μ v) S W2(u" p)) + W2(u, v), the convergence of Cor as the number of particles increases depends on the rate at which the empirical measure y = * Ei ō(x - 2;)) (with x; ~ (4) approximates the true measure μ. The expected error bounded follows (Fournier and Guillin 2013):
N 1
if 2,3
EDWa SC
N 7 log + N) N-
if
if 4
E
Limitations and Broader Impacts
E.1
Limitations
The algorithm presented this paper offers new insights for solving the RUOT problem; however, still has several limitations. Firstly, although Var-RUOT parameterizes and with single neural network, and designs the loss function based on the necessary conditions for minimal action solution, since neural network optimization only finds local minima, there is still guarantee that the solution found is indeed the one with minimal action. Moreover, in practice, the HJB loss does not necessarily converge zero (A typical training curve for the HJB loss is shown in Fig. 14). Thus, effectively acts as a regularization term. This could be addressed by conducting a more detailed analysis simpler versions of the RUOT problem (for instance, Cauusian-to-Gausse Gaussian transport Dirac-to-Dirac Transport).
Furthermore, when using the modified metric, the goodness-of-fi in the distribution deteriorates, which may suggest that the u and g satisfying the optimal necessary conditions derived via the variational method are limited in transporting the initial distribution to the terminal distribution. This might reflect controllability issue control theory that warrants further investigation.
Finally, the choice of (g))) in the action is dependent on biological priors. To automate it, one could approximate 2 with neural network or derive from microscopic or mesoscopic dynamics, such as the branching Wiener process to model cell division for a more physically grounded action (Baradat and Lavenant 2021).
HJB Loss Over Epochs
0.15
0.40
0.35
0.25
0.20
0.35
0.30
Figure 14: Training curve of HJB Loss on three-gene simulation dataset.
35

---

E.2
Broader Impacts
Var-RUOT explicitly incorporates the first-order optimality conditions of the RUOT problem into both the parameterization process and the loss function. This approach enables our algorithm to find solutions with a smaller action while maintaining excellent distribution fitting accuracy. Compared to previous methods, Var-RUOT employs only single network to approximate a scalar field, which results in a faster and more stable training process. Additionally, we observe that the selection of the growth penalty function ()(g) within the WFR metric is highly correlated with the underlying biological priors. Consequently, our new algorithm provides a novel perspective on the RUOT problem.
Our approach can be extended other analogous systems. For example, in the case of simple mesoscopic particle systems where the action can be explicitly formulated such as in diffusion or chemical reaction processes, our framework can effectively infer the evolution of particle trajectories and distributions. This capability makes it applicable to tasks such as experimental data processing and interpolation. the biological or medical field, our method can be employed to predict cellular fate and to provide quantitative diagnostic results or treatment plans for certain diseases.
It should be noted that the performance of Var-RUOT largely depends the quality of the data. Datasets containing significant noise may lead the model to produce results with a slight bias. Moreover, the particular form of the action can have substantial impact on the model's outcomes, potentially affecting important biological priors. These factors could present challenges for subsequent biological analyses or clinical decision- decision-making, and care must be taken in the use and dissemination of the model-genen -enerrated interpolation results to avoid data contamination.
When applying method in biological or medical contexts, it is crucial train the model using high-quality experimental data, select an action formulation that is well-aligned with the relevant domain-specific priors, and ensure that the results are validated by domain experts. Furthermore, there is a need to enhance the interpretability of the model and to further improve training speed through methods such as simulation-free techniques. These directions represent important avenues for our future work.
36

---

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

---

Answer:
[Yes]
Justification:
We provide our proof for our theorems in Appendix A.1.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide our training procedure in Appendix A.3, and we provide code in a public repository https: htt///gtthbbb.///// com wcwm///eroooeetto///tuuuuo
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
38

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We public datasets for experiment, and provide code in a public repository https: / thub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
(oollc/guides/Coeeeee
cwtemmiisssnoooiiiiioioo onPolicy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wctpublic//aeccccccccccccccc ggwpi/gg///////eeeeee /Cae//ooeubbiiic wttsSuubbssiion.. ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes]
Justification: We provide our experimental details in Appendix and Appendix C. In Appendix C, we also provide some additional results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
We provide the mean and standard deviation values of our metrics in each experiment.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
39

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the description of compute resources in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: //t::///neeeee /ewrrps......// cc/ cc/public//eececcceeceeeeccbelle ceChieeCielliinmmmcc
Answer: [Yes]
Justification: We have all reviewed and confirmed this research conforms with the NeurIPS Code of Ethics. We have striven maintain and preserve anonymity.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss broader impacts of our work in Appendix E.2.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
40

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All datasets publicly available, and provide details of datasets in Appendix B.I.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
41

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new datasets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
42

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The method development does not involve LLMs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: http::/p..../oooeeeeeeee///22550 for what should or should not be described.
43


Reviews:
[{"review_id": "OGsfUYx1pq", "reviewer": "Reviewer_kRW9", "summary": "This paper proposes Var-RUOT, a new framework for solving the Regularized Unbalanced Optimal Transport (RUOT) problem (Def. 4.1). The authors reformulate the RUOT problem as a weighted particle system (Theorem 5.1). By applying the Hamilton–Jacobi–Bellman (HJB) equation (Theorem 4.1), they derive an action-matching loss based on the value function. In addition, they introduce an HJB loss term to enforce necessary conditions and a reconstruction loss to ensure that the distribution generated by the model remains consistent with the real data distribution. The method is evaluated on synthetic datasets (Sec. 6.1) and applied to biological problems (Sec. 6.2, 6.3).", "questions": "- Does the method scale to high-dimensional settings (e.g., more than 100 dimensions) where the marginals are Gaussian Mixture Models with many modes? It seems challenging to generalize this approach to such complex, high-dimensional cases.\n\n- What hyperparameters require careful tuning? I imagine that the ratio between the loss function might be crucial. Could the authors elaborate on this aspect or provide additional ablation studies? Also, since estimating $\\lambda$ relies on accurately minimizing the HJB loss, could the authors include a training curve for the HJB loss? Furthermore, have the authors considered using an L1 loss instead, as it might yield sharper estimates of the value surface? Any comments on this would be appreciated.", "limitations": "Please refer to Strength and Weakness section.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "- The paper is well-organized and mathematically sound, with clear examples and intuitive explanations.\n\n- To the best of my knowledge, this is the first work to apply a weighted particle method in this context. This formulation provides valuable intuition and could inspire extensions to related problems. I believe this is an innovative contribution that merits attention from the NeurIPS community.\n\n- The method demonstrates an efficient application of RUOT compared to previous works. The experimental results consistently outperform prior approaches. Moreover, the combination of toy experiments for validation and real-world biological applications highlights both the technical soundness and practical potential of the approach.\n\n\n- However, there are some limitations. Estimating the value function solely through a PINN-like loss (via the HJB equation) may lead to inaccuracies, as solving PDEs is challenging and in this case requires high precision—since the control is derived from differentiating the value function. Even small errors in estimating the value surface could cause significant deviations. Additionally, while the action-matching loss is novel, it may have limited scalability, and the need to estimate second-order derivatives with respect to the state space adds further computational burden. Despite these limitations, I still believe the paper’s contributions are valuable.", "was_revised": true, "final_justification": "I initially thought this paper should be clearly accepted, and the authors addressed all of my concerns.", "rebuttal": "Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:\n\n> Summary & Strengths\n\nWe appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT **finds paths with lower action**, **achieves stable training**, and **converges faster**. We also discuss the **biological priors** associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.\n\n>[W1] Relying solely on a PINN-like HJB loss to estimate the value function can produce significant control deviations, as even small errors cause large inaccuracies due to sensitive differentiation.\n\nThank you for your insightful question. We fully agree that since deep learning solver for the RUOT problem converts hard constraints into soft ones via the loss function, so there's no guarantee that the solution exactly satisfies the marginal constraints nor minimizes the action at every time point. Instead, the HJB loss serves as a **regularizer** that promotes a lower action value without necessarily reducing to 0. On the other hand, we point out that empirically, sensitivity analysis (lines 656–670) shows that increasing the loss weight weight $\\gamma_{\\text{HJB}}$ would decrease the action value, indicating the effectiveness of such term to achieve the decrease of action for the comnputed path. **We will add the relevant discussions in the revised manuscript**.\n\n\n>[W2] Compute Second-order derivatives causes computational burden.\n\nThank you for your comments and suggestion. We fully agree that compute second-order derivatives using PyTorch autograd **by default** will requires retaining a larger graph and thus more memory and time than first-order derivatives. Meanwhile we clarify that\n\n- Unlike previous RUOT methods such as DeepRUOT that train three networks for $\\boldsymbol{v}$, $\\rho$, and $g$, VAR-RUOT fits a single scalar field $\\lambda$, reducing the search space, improving stability, and converging faster (see Section 6.2).\n\n- Since Var-RUOT only need second-order derivatives for the $\\mathcal{L}_{\\text{HJB}}$ loss (to compute $\\nabla_x^2 \\lambda$), meaning we only require the trace of the Hessian matrix $H$. We can estimate it using the **Hutchinson method**:\n  \n  $$\n  \\text{trace}(H)=\\mathbb{E}_{\\boldsymbol{v}}[\\boldsymbol{v}^{T} H \\boldsymbol{v}], \\quad \\boldsymbol{v}\\sim N(0,I)\n  $$\n  \n  which computes the derivative along one direction, reducing cost (error $O(1/\\sqrt{m})$, with $m$ samples). We can choose to use Hutchinson early in training for speed and later switch to full autograd for accuracy. **This feature will be made as an option in our open-source code to speedup the computation**.\n\n>[Q1] Can Var-RUOT be Generalized to high-dimensional settings?\n\nThank you for your question on Var-RUOT's scalability to high-dimensional datasets. In previous manuscript, we validated its scalability on several datasets: 50 and 100 dimensional Gaussian Mixture datasets (lines 709–715), a 50-dimensional Mouse Blood Hematopoiesis dataset, and a 30-dimensional β-cell Differentiation dataset (lines 716–723). For the high-dimensional Gaussian Mixture dataset, Figure 13 qualitatively shows the learned particle trajectories and growth rates on gaussian datasets; Figures 14 and 15 display the learned velocity fields and growth rates for the real 50D and 30D datasets, demonstrating Var-RUOT's applicability.\n\nTo quantitatively evaluate scalability, we further designed simulated data at three time points (t = 0, 1, 2) where each point is generated from a mixture of three Gaussians in **150 dimensions**. We then compared various baselines on this dataset, with results as follows:\n\n\n|Method/$\\mathcal{W}_1$ Distance|t=1|t=2|PathAction|\n|---|---|---|---|\n|SF2M|7.286|9.457|-|\n|PISDE|**6.051**|7.546|-|\n|MIOFlow|6.494|7.812|-|\n|TIGON|6.157|7.615|13.4913|\n|DeepRUOT|6.202|7.648|13.5992|\n|VarRUOT|6.072|**7.544**|10.1033|\n\n\nVar-RUOT achieves comparable **reconstruction accuracy** to PISDE, slightly outperforms TIGON and DeepRUOT, and significantly outperforms other methods on the 150-dimensional Gaussian Mixture dataset. It also yields a **significantly lower path action** than TIGON/DeepRUOT, confirming its scalability in high-dimensional settings. **We will include these experiments in our revised manuscript.**\n\n\n>[Q2]Questions on hyperparameter tuning and ablation studies.\n\nThank you for your question on hyperparameter tuning. Our method requires tuning three hyperparameters: the WFR metric parameter $\\alpha$ (controlling the penalty for mass growth/decay) and two loss parameters, $\\gamma_{\\text{HJB}}$ and $\\gamma_{\\text{Action}}$. In Appendix C.2, we show that our method is robust to variations in $\\alpha$ (using the 2D Mouse Blood Hematopoiesis dataset, lines 648–655, Table 7). Similarly, sensitivity experiments on $\\gamma_{\\text{HJB}}$ and $\\gamma_{\\text{Action}}$ (lines 656–670, Figures 7–8) demonstrate that increasing either parameter reduces the Path Action. However, very high values diminish the impact of the reconstruction loss that enforces marginal constraints, which degrades performance.\n\nTo provide further insights, here we have conducted additional sensitivity tests on $\\gamma_{\\text{HJB}}$ and $\\gamma_{\\text{Action}}$ using the Mouse Blood Hematopoiesis dataset. \n\n\n- **Case with fixed $\\gamma_{\\text{Action}} = 6.25\\times {10}^{-2}$:**\n\n\n|$\\gamma_{\\text{HJB}}$|$0$|$6.25\\times10^{-3}$|$3.125\\times10^{-2}$|$6.25\\times10^{-2}$|$6.25\\times10^{-1}$|$3.125$|\n|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|$\\mathcal{W}_{1}$|0.1573|0.1609|0.1550|0.1316|0.2109|0.2539|\n|PathAction|3.3002|3.2165|3.2397|3.1491|3.1182|3.1200|\n\n\n\n- **Case with fixed $\\gamma_{\\text{HJB}}= 6.25\\times {10}^{-2}$:**\n\n\n|$\\gamma_{\\text{Action}}$|$0$|$6.25\\times10^{-3}$|$3.125\\times10^{-2}$|$6.25\\times10^{-2}$|$6.25\\times10^{-1}$|$3.125$|\n|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|$\\mathcal{W}_{1}$|0.1454|0.1344|0.1674|0.1316|0.4001|0.7914|\n|PathAction|3.4043|3.4497|3.2737|3.1491|2.4716|1.2059|\n\n\n\nThe results from these additional sensitivity experiments are consistent with those presented in the original manuscript: higher values of $\\gamma_{\\text{HJB}}$ and $\\gamma_{\\text{Action}}$ lead to a lower Path Action, yet if they are set too high, the reconstruction quality deteriorates significantly. In fact, all experiments in the main text use the same parameters, $\\gamma_{\\text{HJB}} = \\gamma_{\\text{Action}} = 6.25 \\times 10^{-2}$ (see Table 6). In practice, we recommend setting both parameters in the range of $10^{-2}$ to $10^{-1}$ and **we will make it clear in the open-source package**.\n\n>[Q3]Training curve for the HJB loss.\n\nThank you for your question. Due to rebuttal restrictions, we cannot provide an image of the HJB loss curve. The HJB loss initially rises for 10–20 epochs (since $L_{\\text{Recon}}$  dominates while parameters are near zero) and then gradually declines as $L_{\\text{Recon}}$ reaches a similar scale as $L_{\\text{HJB}}$ and $L_{\\text{Action}}$. **We will include this curve in the appendix of revised manuscript.**\n\n\n>[Q4]Can we use L1 loss instead?\n\nThank you for your valuable suggestion. In additional experiment, we have modified the HJB loss per your recommendation by replacing the quadratic penalty with a linear penalty. Then we re-ran experiments on three datasets presented in the main text for the Var-RUOT method, comparing the performance when using the original $L_{2}$ loss versus the $L_{1}$ loss.\n\n- Three Gene Simulation Dataset:\n\n|Method/$\\mathcal{W}_1$ Distance|t=1|t=2|t=3|t=4|PathAct|\n|---|---|---|---|---|---|\n|Var-RUOT($L_2$Loss)|0.0452|0.0385|0.0445|0.0572|1.1105|\n|Var-RUOT($L_1$Loss)|0.0489|0.0556|0.0589|0.0664|1.1249|\n\n\n\n- EMT Dataset:\n\n\n|Method/$\\mathcal{W}_1$ Distance|t=1|t=2|t=3|Path Action|\n|---|---|---|---|---|\n|Var-RUOT($L_{2}$Loss)|0.2540|0.2670|0.2683|0.3544|\n|Var-RUOT($L_{1}$Loss)|0.2631|0.2796|0.2809|0.3528\n\n\n\n- 2D Mouse Blood Hematopoiesis Dataset:\n\n\n|Method/$\\mathcal{W}_1$ Distance|t=1|t=2|Path Action|\n|---|---|---|---|\n|Var-RUOT($L_{2}$Loss)|0.1200|0.1431|3.1491|\n|Var-RUOT($L_{1}$Loss)|0.1747|0.1690|3.1889|\n\n\n\nThe experimental results indicate that with our chosen hyperparameter values ($\\gamma_{\\text{HJB}} = 6.25 \\times 10^{-2}$ and $\\gamma_{\\text{Action}} = 6.25 \\times 10^{-2}$), using either the $L_{1}$ loss or the $L_{2}$ loss for the HJB loss yields similar reconstruction performance and Path Action. **We will add the above analysis to the appendix during the revision of the paper.**"}, {"review_id": "ec0KEb59fG", "reviewer": "Reviewer_s4d8", "summary": "This paper introduces Var-RUOT, a new method for learning the dynamics of evolving systems using just a few snapshot observations. Traditional approaches rely on multiple neural networks and often struggle to find the most efficient, biologically realistic paths. Var-RUOT simplifies the problem by learning only a single scalar function, from which all relevant dynamics can be derived. This leads to more stable training, faster convergence, and solutions that follow the principle of least action (i.e., the most efficient evolution). The method also allows for customizing the behavior of growth or decay in the system to better reflect biological processes. Through experiments on synthetic and real biological datasets, the authors show that Var-RUOT achieves more accurate and efficient results than existing approaches.", "questions": "1. Your method relies on enforcing first-order necessary conditions for action minimization via the HJB loss. However, these conditions are not sufficient for global or even local optimality. Can you provide any theoretical or empirical evidence that minimizing the HJB residual leads to low-action trajectories beyond satisfying a PDE constraint pointwise?\n\n2. The derivation assumes that the mapping \\( \\lambda = \\alpha \\psi'(g) \\) can be inverted to recover \\( g(x,t) \\). What guarantees do you have that \\( \\lambda_\\theta(x,t) \\) stays within the image of \\( \\psi' \\) during training? How is the domain of \\( g \\) enforced or regularized when \\( \\psi \\) is non-convex or has flat regions?\n\n3. In lower-dimensional or isotropic settings, optimal transport paths with minimal action (e.g., in the WFR or Benamou-Brenier case) are analytically or numerically solvable. Why did you not include comparisons against such ground-truth geodesics to rigorously evaluate whether Var-RUOT truly finds lower-action trajectories?\n\n4. What is the quantitative effect of violating the HJB equation on the evolution of \\( \\rho(x,t) \\)? Can small residual errors compound over time and distort the learned transport path, particularly in extrapolation tasks?\n\n5. The biological interpretation of \\( \\psi''(g) \\) controlling the direction of growth is intuitive but lacks formal backing. Why is \\( \\psi(g) = g^{2/15} \\) considered a principled choice beyond its curvature sign? Have you validated this with any empirical or theoretical biological models?\n\n6. You claim that learning a single scalar field simplifies the problem. However, could this choice overly constrain the model class, especially in settings where velocity and growth cannot be jointly derived from a single potential? Are there situations where this coupling leads to underfitting?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n1. By reducing the RUOT problem to learning a single scalar field \\lambda(x,t), the method avoids redundant parameterization of velocity and growth fields, leading to a compact, interpretable, and theoretically grounded framework.\n\n2. Var-RUOT achieves faster convergence and more stable training compared to prior methods, as shown through detailed empirical comparisons across multiple datasets and metrics (e.g., action, Wasserstein distance, wall time).\n\n3. The paper connects the choice of the growth penalty function \\psi(g) to biological priors, enabling control over whether growth increases or decreases along inferred trajectories, offering practical relevance in single-cell dynamics applications.\n\n###  Weaknesses\n\n1. The paper enforces only first-order optimality conditions through the HJB residual loss, which are necessary but not sufficient for achieving true minimal action. There is no guarantee that the learned solution corresponds to a global or even local minimum of the action functional.\n\n2. While the method relies on a weighted particle system to approximate density evolution, it lacks any quantitative analysis of convergence rates or variance bounds for the empirical measure, especially under high-dimensional stochasticity.\n\n3. The method assumes that the mapping  \\lambda = \\alpha \\psi'(g) is invertible across all relevant values of \\lambda, but this is not always true for non-convex or saturating forms of \\psi. No mechanism is introduced to ensure this invertibility is respected during training.", "was_revised": true, "final_justification": "I believe this is a valuable paper and authors addressed my concerns thoroughly, therefore I increased my score to 5.", "rebuttal": "Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:\n\n> Summary & Strengths\n\nWe appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT **finds paths with lower action**, **achieves stable training**, and **converges faster**. We also discuss the **biological priors** associated with the RUOT penalty function.\n\n\n> [W1] Regarding Var-RUOT’s Guarantee for the Minimal Action Solution\n\nThank you for your comment. Since Var-RUOT is a deep learning–based RUOT solver, it guarantees only a local minimum of the action. Empirically, experiments (Table 1,2,4) along with all additional results in the rebuttal show that Var-RUOT achieves a lower action than DeepRUOT, and its Path Action is comparable to that of traditional OT solvers.\n\nWe consider guaranteeing a global minimum an interesting direction for future work and **will discuss it in the revised manuscript**. We plan to start with analytical solutions for simple cases (e.g. Dirac-to-Dirac transport) and then extend to real data by treating complex distributions as a superposition of Dirac deltas and solving the associated velocity field via a Flow Matching–like approach. \n\n\n> [W2] Convergence Rate\n\nThanks for your insightful question. From the perspective of numerical integration, we use the particle method to estimate $L_{\\text{HJB}}$ and $L_{\\text{Action}}$, which is equivalent to performing integration via Monte Carlo methods. Let the number of particles be $N$. Then, for any smooth function $f(\\boldsymbol{x})$ (In the Action Loss and HJB Loss, the integrand.), the central limit theorem tells us that the error of the integral $\\int f(\\boldsymbol{x})\\,\\mathrm{d}\\mu^N$ is of order $\\mathcal{O}\\left(\\dfrac{1}{\\sqrt{N}}\\right)$. **We will revise the paper and include this discussion on the convergence rate.** In practical applications of the Var-RUOT algorithm, our default setting is $N=512$, which ensures that the estimates of both loss have very small variance.\n\n\n> [Q1, Q4] Regarding the effect of the HJB Loss.\n\nThanks for the insightful question. Var-RUOT controls the penalty for $\\lambda(\\boldsymbol{x},t)$ violating the HJB equation by adjusting the weight $\\gamma_{\\text{HJB}}$: a higher weight increases the importance of the HJB loss, reducing the violation after training. Sensitivity analysis (Section C.2, Figure 7) shows that as **$\\gamma_{\\text{HJB}}$ increases, the Path Action decreases**, proving that a lower HJB loss leads to a smaller action solution. To address the concerns, we further provide sensitivity results for the EMT dataset and additional results for the 2D Mouse dataset:\n\n\n- EMT Dataset\n\n|$\\gamma_{\\text{HJB}}$|$0$| $6.25 \\times 10^{-3}$|$3.125 \\times 10^{-2}$| $6.25 \\times 10^{-2}$|$6.25 \\times 10^{-1}$|$3.125$|\n|-|-|-|-|-|-|-|\n|$\\mathcal{L}_{\\text{HJB}}$ |1.9016|1.3091|0.6205|0.3685|0.0223| 0.0006|\n|Path Action|0.4612|0.4584|0.3855|0.3544|0.2975|0.2808|\n\n\n- 2D Mouse Dataset\n\n|$\\gamma_{\\text{HJB}}$ |$0$| $6.25 \\times 10^{-3}$ | $3.125 \\times 10^{-2}$ | $6.25 \\times 10^{-2}$ | $6.25 \\times 10^{-1}$ | $3.125$ |\n|-|-|-|-|-|-|-|\n|$\\mathcal{L}_{\\text{HJB}}$ |1.9107|1.4605|1.0745|0.6819|0.0725|0.0074|\n|Path Action|3.3002|3.2165|3.2397|3.1491|3.1182|3.1200|\n\n\n\nThese results confirm that **incorporating the HJB loss effectively finds lower action solutions**. Based on our findings, we recommend setting $\\gamma_{\\text{HJB}}$ between $10^{-2}$ and $10^{-1}$ to balance minimal action with data reconstruction quality. We will include these discussion in our final version.\n\n\n> [W3,Q2] Reversibility Issue of $\\psi'(g)$\n\n\nThank you for your valuable comment and apologize for the insufficient discussions in previous paper. Here we discuss the two scenarios of $\\psi$ selection in more details:\n\n-  $\\psi_{1}(g)=g^{2p}$, $\\psi_{1}'(g)=2p\\cdot g^{2p-1}$, $p\\in\\mathbb{Z}^{+}$. Here, $\\psi_{1}'(g)$ is a bijection from $\\mathbb{R}$ to $\\mathbb{R}$, so $g$ is uniquely determined from $\\lambda=\\alpha\\psi_{1}'(g)$.\n\n- $\\psi_{2}(g)=g^{\\frac{2p}{2q+1}}$, $\\psi_{2}'(g)=\\frac{2p}{2q+1}g^{\\frac{2p-(2q+1)}{2q+1}}$, $p,q\\in\\mathbb{Z}^{+}$ with $2q+1>2p$. Here, $\\psi_{2}(g)$ is a bijection from $\\mathbb{R}\\setminus\\{0\\}$ to $\\mathbb{R}\\setminus\\{0\\}$ (undefined at 0). Directly solving$\\lambda=\\alpha\\psi_{2}'(g)$ yields \n$$\ng=\\left(\\frac{\\alpha}{\\lambda}\\frac{2q+1}{2p}\\right)^{\\frac{2q+1}{(2q+1)-2p}}.\n$$ \nAt $\\lambda=0$, $g(\\lambda)$ is undefined. To handle this inevitable singularity, we set a small $\\delta$ and linearly interpolate between $g(-\\delta)$ and $g(\\delta)$ for $\\lambda\\in[-\\epsilon,\\epsilon)$,i.e., $g(\\lambda) = \\dfrac{1}{2\\delta}\\left(g(\\delta) - g(-\\delta)\\right)(\\lambda + \\delta) + g(-\\delta)$ (Appendix A.2)\n\n\nAnother approach is to fit the reciprocal of $\\lambda$, denoted as $\\mu$, thereby moving the singularity at $\\lambda=0$ to $\\mu \\rightarrow \\infty$. In practice, $\\mu \\rightarrow \\infty$ does not occur, just as the singularity where $|g| \\rightarrow \\infty$ is not reached. This method avoids artificially modifying the $g(\\lambda)$ relationship and is more reasonable. **This will be our future work. We will update the paper for more thorough disucssions on $\\psi'(g)$ choice and include these detailed analysis in the appendix.**\n\n> [Q3] Comparison with Traditional OT Solvers\n\nThank you for your valuable comment. In additional experiment, here we built two 2D normalized and unbalanced datasets to test Var-RUOT’s ability to find minimal action paths. For normalized data, we used the `pot` library to obtain ground truth; for unbalanced data, we applied a Sinkhorn-based static WFR solver [(Wang Z., et al.) Wasserstein-fisher-rao document distance]. With $\\sigma=0$, Var-RUOT achieves similar—or even lower—Path Action values than traditional solvers(results are shown below), demonstrating its effectiveness. **We will include these analysis in the appendix.**\n\n\n|Dataset|Normalized|Normalized|Unbalanced|Unbalanced|\n|---|---|---|---|---|\n|-|Simple Gaussian|Gaussian Mixture|Simple Gaussian|Gaussian Mixture|\n|Path Action (Traditional Solver)|0.5046 |1.9062 |1.3776| 2.4464|\n|Path Action (Var-RUOT) |0.4998|1.8754 |1.0102| 2.3909|\n\n\n> [Q4] Extrapolation Performance\n\nWe evaluated Var-RUOT’s extrapolation capability and error propagation by extending the three-gene simulation dataset to nine time points. The first five were used for training, while the remaining four tested extrapolation. We compared TIGON, DeepRUOT, and Var-RUOT. The results are shown below:\n\n|Method/$\\mathcal{W}_{1}$  Distance|t=5|t=6|t=7|t=8|\n|---|---|---|---|---|\n|TIGON|0.1932|0.3437|0.5209|0.6913|\n|DeepRUOT|0.1027|0.1940|0.3359|0.5047|\n|VarRUOT|0.1057|0.1806|0.2890|0.4260|\n\nVar-RUOT outperforms TIGON and performs similarly to DeepRUOT in long-term extrapolation. However, as time increases, error grows because the model’s non-autonomous system learns $\\lambda(\\boldsymbol{x}, t)$ only within the training window. Beyond this range, the $t$-dependency is less accurate, leading to larger errors. Nonetheless, this non-autonomous system improves flexibility and helps prevent underfitting (see Q6). **We will include these analysis in the appendix.**\n\n\n>[Q5] Choice of $\\psi(g) = g^{2/15}$\n\nSorry for the confusion. We clarify that choosing $\\psi(g)=g^{2/15}$ primarily to **showcase an example** with $\\psi''(g)<0$. There are two reasons. \n\nFirstly, the simplest function family satisfying our conditions (lines 224–229)—that is, $\\frac{d\\psi(g)}{d|g|}>0$, $\\psi(g)=\\psi(-g)$, and $\\psi''(g)<0$—is given by\n$$\n\\psi(g)=g^{\\frac{2p}{2q+1}},\\quad p,q\\in\\mathbb{Z}^+,\\; 2p<2q+1.\n$$\n\nSecondly, for training stability, the first-order condition yields\n$$\ng(\\lambda)=\\left(\\frac{\\lambda}{\\alpha}\\,\\frac{2q+1}{2p}\\right)^{\\frac{2q+1}{2p-(2q+1)}} \\propto \\left(\\frac{1}{\\lambda}\\right)^{\\frac{2q+1}{(2q+1)-2p}}.\n$$\nTo avoid a steep $g(\\lambda)$ curve near $\\lambda=0$, we choose the smallest $p$ (i.e., $p=1$) and a large $q$ (here, $q=7$), which leads to $\\psi(g)=g^{2/15}$.\n\nHere we conduct additional experiments on the 2D Mouse Blood Hematopoiesis dataset (results are shown below) to show that while $\\psi(g)=g^{2/15}$ and $\\psi(g)=g^{2/21}$ perform similarly, however $\\psi(g)=g^{2/9}$ yields unstable training and poorer reconstruction. \n\n\n| $\\psi(g)$  / $\\mathcal{W}_1$ Distance| t=1 | t=2 | \n| --- | ----| ---- |\n| $\\psi(g)=g^{(2/9)}(p=1, q=4)$   | 0.9476 | 1.2197|\n| $\\psi(g)=g^{(2/15)}(p=1, q=7)$  | 0.2953 | 0.1917  | \n| $\\psi(g)=g^{(2/21)}(p=1, q=10)$ | 0.2858 |0.2565 | \n\n\nOverall, our experiments confirm that, within a reasonable range, different choices for the exponent on $g$ are feasible. We will include this sensitivity analysis in the appendix.\n\n> [Q6] On Whether a Single Scalar Field Causes Underfitting\n\nThanks for the insightful question. Here we clarify that:\n- Experiments show that Var-RUOT reconstructs the marginal distributions at all time points well—both on low-dimensional data (Sections 6.1 and 6.3) and high-dimensional data (Appendix C.4), the results are comparable to DeepRUOT (which uses three networks) without underfitting. Please also refer to the responses to Reviewer kRW9 [Q1] and Reviewer YDV3 [W4,5,6], which include quantitative results of Var-RUOT on high-dimensional datasets.\n- Allowing $\\lambda$ to depend on both $\\boldsymbol{x}$ and $t$ enables capturing both trajectory evolution and total mass fluctuations; if $\\lambda$ depended solely on $\\boldsymbol{x}$, it could not fit distributions showing some specific mass changing pattern (e.g., mass that first rises then falls).\n- Moreover, the DeepRUOT method employs three networks, while they are interdependent (via a Fokker–Planck loss), which makes training unstable and requires pretraining, while Var-RUOT is simpler and more stable."}, {"review_id": "5EDJpEVk7E", "reviewer": "Reviewer_5WHv", "summary": "The authors present a new method for learning high dimensional stochastic dynamics for unnormalized distributions. The authors build on the recent RUOT framework modifying the loss to ensure the learned solution satisfies the optimality condition of least action. This in turns allows them to use a single network and parametrize only a scalar field. Additionally they propose a growth penalty function which better aligns with biological priors.", "questions": "- It would be great to see a table (or something like this) clearly showing the differentiate factors and or advantages of their method compared to the many other (see above).\n- Recent works show that estimating action losses can be quite difficult and they proposed quadrature of these types of losses [3]. Do the authors encounter similar issues? \n- How and why are the authors able to eliminate the pertaining steps of the original RUOT algorithm [1] ? \n- How exactly was the g^2/15 value obtained? Was this a tuned hyper-parameter?\n\n[1] Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. Zhang. 2024.\n[3] Parametric model reduction of mean-field and stochastic systems via higher order action matching. Berman. 2024.", "limitations": "- The Limitations section is pushed to the appendix, which in my opinion is not in the spirit of the guidelines and neurips checklist.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- The motivation of trying to satisfy the condition of  least action is well motivated.\n- There is a substantial theory section.\n- The numerical results are quite strong. The method really does seem to be an improvement. The faster convergence speed is promising.\n\n**Weaknesses**\n- The overall presentation is not great with regards to what the key methodological innovations are and how they fit within the existing literature. It is quite difficult to tease out the exact differences between this paper other work such as the original RUOT paper [1] or action matching style papers [2]. This is quite important because the latter works also explicitly enforce optimality conditions WRG to least action and the loss terms look very similar to what the paper calls \"Action Loss\". \n- In particular one of the key approaches of the proposed method is to paramzeraize with a single scalar field where by the velocity is the gradient of this field and the growth function is equal to it. The title of paper might suggest that this is a new innovation, but it is in fact exactly what is done in the original action matching paper [2]. While the authors cite this paper and state \"these methods typically employ separate neural networks to parameterize the velocity and growth functions, without leveraging their optimality conditions or the inherent relationship between them\". To my knowledge this is incorrect with respect to [2].\n\n\n[1] Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. Zhang. 2024.\n[2] Action Matching: Learning Stochastic Dynamics from Samples. Neklyudov. 2023.", "was_revised": true, "final_justification": "The response was very thorough and addressed my concerns.", "rebuttal": "Thank you for carefully reviewing our paper and raising valuable questions. Below we address your concerns in a point-to-point response:\n\n> Summary & Strengths\n\nWe appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT **finds paths with lower action**, **achieves stable trainin**, and **converges faster**. We also discuss the **biological priors** associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.\n\n\n\n> [W1, W2, Q1] Differences from Previous Methods\n\nThank you for your question and sorry for the confusion. First, we sincerely apologize for not thoroughly checking our citation positions and misunderstandings caused by it. We fully agree that Action Matching is a very innovative work uses a single network to parameterize $s_{\\theta}(\\boldsymbol{x}, t)$ that inspires our research; we will review the manuscript to correct any improper citations and clarify the significance of previous methods. \n\nFollowing your great suggestion, we also made a table to clearly highlight the contribution of Var-RUOT that will **later be included in the revised Appendix**:\n\n\n|Method|Unbanlanced Distribution| Stochastic Dynamics | Unbalanced + Stochastic | Least Action|\n|---|---|---|---|---|\n|SF2M|-|+|-|-|\n|PISDE|-|+|-|+|\n|MIOFlow|-|-|-|-|\n|Action Matching|+|+|-|+|\n|TIGON|+|-|-|-|\n|DeepRUOT|+|+|+|-|\n|Var-RUOT|+|+|+|+|\n\n\n- **Unbalanced Distribution** indicates that the method can handle non-normalized distributions.\n- **Stochastic Dynamics** means the method supports stochastic dynamics.\n- **Unbalanced + Stochastic** shows the method can handle both simultaneously (e.g., DeepRUOT).\n- **Least Action** signifies that the method employs a first-order condition to minimize the action, rather than simply adding the action as a weighted loss.\n\nThe key difference between Var-RUOT and DeepRUOT is that Var-RUOT **leverages first-order optimality conditions** in its parameterization and loss design (e.g., using only a scalar field rather than three neural networks as in DeepRUOT). Compared to DeepRUOT, our approach **finds paths with lower action** and **addresses unstable training and slow convergence**, thus eliminating the need for a pre-training phase.\n\nVar-RUOT and Action Matching differ in three key ways: \n\nFirstly, Var-RUOT can **simultaneously handle imbalanced data distributions and stochastic dynamics**, both crucial for single-cell omics.\n\n\nSecondly, Action Matching's training algorithm needs to sample from intermediate distribution $q_{t}(x)$, while in single-cell omics data, we only have snapshot data at a few time points as boundary sample. So acquiring intermediate samples from $q_{t}(x)$ may be challenging due to missing evolution data in-between, potentially hindering the ability to capture complex dynamics. Empirically, our experimental results (Tables 1, 2, 4) show that Var-RUOT achieves more satisfactory results over existing results including Action Matching.\n\nLastly, Action Matching defines the Action Gap loss as:\n\n$$L_{\\text{AM}} = E_{q_{0}(x)} [s_{0}(x)] - E_{q_{1}(x)} [s_{1}(x)] + \\int_{0}^{1} E_{q_{t}(x)} \\left[ \\dfrac{1}{2} \\| \\nabla s_{t}(x) \\|^{2} + \\dfrac{\\partial s}{\\partial t}(x) \\right] \\mathrm{d} t$$ \n\nWhile Var-RUOT directly adopts the HJB Loss(Line 204-209) to ensure that the computed solution $\\lambda(\\boldsymbol{x},t)$ satisfies the HJB equation. In other words, Var-RUOT incorporates the first-order condition into the loss, enabling it to find a solution with lower action. **We will add these discussions in the revised manuscript**.\n\n\n\n\n\n> [Q2] On the Estimation of the Action Loss\n\nThank you for your question. In our experiments, we employed two methods to compute the integral losses for both the action and HJB terms. When using the Euler-Maruyama method, setting the step size too large (e.g., $\\Delta t = 0.2$) results in high variance estimates for the action and HJB losses, leading to unstable training. Therefore, we use $\\Delta t = 0.1$ in all experiments. The Stochastic Runge-Kutta method offers better convergence, allowing an integration step of $\\Delta t = 0.2$. Since our approach is simulation-based, a longer step size significantly reduces training time. In practice, we thus choose the Stochastic Runge-Kutta method and select a larger step size as long as training remains stable. We will discuss this in our final version and  cite [(Berman.) High Order Action Matching] in our paper.\n\n> [Q3] Why the Related Steps in the Original RUOT Algorithm Can Be Omitted\n\nIn the original DeepRUOT algorithm, three neural networks are employed to approximate $\\boldsymbol{v}(\\boldsymbol{x},t)$, $g(\\boldsymbol{x},t)$, and $\\log \\rho(\\boldsymbol{x},t)$, with all constraints imposed as soft constraints via loss functions. For instance, to ensure that $\\boldsymbol{v}$, $g$, and $\\rho$ satisfy the Fokker–Planck equation, DeepRUOT introduces the Fokker–Planck loss\n\n$$L_{\\text{FP}} = \\|\\partial_{t} p_{\\theta} + \\nabla_{x} \\cdot (p_{\\theta}\\boldsymbol{v})- g_{\\theta} p_{\\theta}\\| + \\lambda_{w} \\| p_{\\theta}(\\boldsymbol{x},0) - p_{0} \\|$$\n\nTraining these three interdependent networks simultaneously is challenging and often causes convergence issues. To accelerate training by obtaining a good initialization, DeepRUOT incorporates a pretraining phase: (1) training $\\boldsymbol{v}$ and $g$ using a reconstruction loss ${L}_{\\text{Recons}}$, and (2) training $p$ with Conditional Flow Matching.\n\nIn Var-RUOT, we use a single neural network $\\lambda_\\theta(\\boldsymbol{x},t)$ to approximate a scalar field. The velocity field $\\boldsymbol{u}(\\boldsymbol{x},t)$ and growth rate $g(\\boldsymbol{x},t)$ are derived analytically from $\\lambda(\\boldsymbol{x},t)$ using first-order optimality conditions (Theorem 4.1). As a result, Var-RUOT avoids handling multiple interdependent networks and eliminates the need for the pretraining phase used in DeepRUOT. This is indeed one of the main contribution of our work; ablation experiments in DeepRUOT demonstrate that these pretraining steps are indispensable (DeepRUOT[(Zhang.), DeepRUOT], Appendix B.5), and we have removed them in Var-RUOT. **We will add these discussions in the revised manuscript**.\n\n\n> [Q4] Why we Choose $g^{2/15}$\n\nThanks for your question regarding $\\psi(g)$. we choose $g^{2/15}$ primarily to **serve as an example** showing that when $\\psi'(g) > 0$ and $\\psi''(g) < 0$, at each time step $g(\\boldsymbol{x},t)$ increases along the direction of the velocity field $\\boldsymbol{u}(\\boldsymbol{x},t)$ (Theorem 4.2). As for why we adopt $\\psi_{2}(g) = g^{2/15}$ as the example, it is because: (1) It is the simplest choice that satisfies the conditions given in our paper (lines 224–229): $\\dfrac{\\mathrm{d}\\psi(g)}{\\mathrm{d}|g|} > 0$, $\\psi(g) = \\psi(-g)$, and $\\psi''(g) < 0$. (2) For training considerations: when choosing $\\psi(g)=g^{(2p)/(2q+1)}$, the first-order optimality condition allows us to derive\n$$\ng(\\lambda) = \\left(\\dfrac{\\lambda}{\\alpha} \\dfrac{2q+1}{2p}\\right)^{\\frac{2q+1}{2p - (2q+1)}} \\propto \\left(\\dfrac{1}{\\lambda}\\right)^{\\frac{2q+1}{(2q+1)-2p}}.\n$$\nWe aim for more stable training, and since $g$ is computed directly from $\\lambda$, we do not want the $g(\\lambda)$ to be overly steep (especially near $\\lambda=0$), which requires the exponent $\\dfrac{2q+1}{(2q+1)-2p}$ to be as small as possible. This means choosing $p$ as small as possible (hence $p=1$) and choosing $q$ as large as possible (we select $q=7$). Therefore, we ultimately choose the penalty function $\\psi(g)=g^{2/15}$.\n\n*Note: It is worth noting that $g(\\lambda)$ is undefined at $\\lambda=0$. In fact, for any function $\\psi(g)$ that meets the conditions outlined in lines 224–229 along with $\\psi''(g)<0$, its first derivative must be discontinuous at $g=0$, so this singularity is inevitable. To handle this issue, we select a small value $\\delta$ and linearly interpolate between $g(-\\delta)$ and $g(\\delta)$ by redefining\n$$\ng(\\lambda) = \\dfrac{1}{2\\delta}\\left(g(\\delta) - g(-\\delta)\\right)(\\lambda + \\delta) + g(-\\delta), \\quad \\lambda \\in [-\\epsilon, \\epsilon),\n$$\nto bypass the singularity (Appendix A.2). Another approach to deal with the singularity is to fit the reciprocal of $\\lambda$, denoted as $\\mu$, thereby moving the singularity at $\\lambda=0$ to $\\mu \\rightarrow \\infty$. In practice, $\\mu \\rightarrow \\infty$ does not occur, just as the singularity where $|g| \\rightarrow \\infty$ is not reached. This method avoids artificially modifying the $g(\\lambda)$ relationship and is more reasonable. This will be our future work.*\n\nTo further explain our training considerations, we experimented with two additional parameter sets on the 2D Mouse Blood Hematopoiesis dataset. The experimental results are as follows:\n\n\n|$\\psi(g)$  / $\\mathcal{W}_1$ Distance|t=1|t=2| \n|---|---| ---- |\n|$\\psi(g)=g^{(2/9)}(p=1, q=4)$|0.9476|1.2197|\n|$\\psi(g)=g^{(2/15)}(p=1, q=7)$|0.2953|0.1917| \n|$\\psi(g)=g^{(2/21)}(p=1, q=10)$|0.2858|0.2565|\n\n\n\nThe experimental results indicate that choosing $\\psi(g)=g^{2/15}$ produces performance comparable to $\\psi(g)=g^{2/21}$, while selecting $\\psi(g)=g^{2/9}$ leads to unstable training and poorer reconstruction of the distribution. Our experiments demonstrate that choosing the exponent on $g$ within a certain range is feasible. **We will update the paper to include the above analysis in the sensitivity analysis section of the Appendix.**\n\n>  On Writing Issues\n\nThank you for your suggestions regarding our writing. We will review and revise our paper to:\n- Move the Limitations section into the main text.\n- Remove the emoji from Figure 1."}, {"review_id": "Ooeo2zP3fZ", "reviewer": "Reviewer_YDV3", "summary": "The authors of the paper introduce Var-RUOT, a method to perform deep regularised unbalanced optimal transport. Var-RUOT parameterizes a scalar potential and derives the necessary conditions required by the associated field and growth penalty to minimise the isotropic time invariant RUOT problem. The potential is trained by minimising a loss function computed over simulated particles from an empirical measure. The loss has multiple components. First, a reconstruction term ensures that the model recapitulates the marginal distributions over time and preserves the total mass of individual snapshots. The second component is a Hamilton-Jacobi-Beltrami loss enforcing field optimality under the ITI-RUOT problem.\nFurthermore, the least action loss is added to regularize the dynamic optimal transport field. Finally, the paper proposes a new biological interpretation of the growth penalty function, showing that it varies in the direction of the transporting field. By controlling the second derivative of the penalty function, one can encode biological processes, such as the loss of stemness. In their experiments, the authors compare their model with existing methods on the task of reconstructing time marginals of real and simulated single-cell data. Together with distribution matching metrics the paper additionally shows that the model achieves lower path action and faster convergence compared to its counterparts, offering an efficient approach to study cellular development in cellular biology.", "questions": "* Why are the weights initialised as 1? Is there a way to instead learn them?\n* What is the ground truth for the mass objective? In other words, how do you compute M(Tk) if you only have the weights for the solution of the ODE?\n* Why did the authors pick $\\frac{2}{15}$ as an exponent of $\\phi_2$? How does the penalty function behave for different values of the exponent?\n* **Line 71**: isn’t the Schrödinger bridge always dynamic as it operates via regularisation to a reference probability path?", "limitations": "The limitations of the work have been sufficiently addressed in the Appendix.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths \n\nI generally find the paper well-written, for which I commend the authors. The math is introduced gradually and is easy to digest, which is a very positive aspect in my opinion. I also think that the problem of enforcing stronger optimality constraints to dynamical OT is relevant and well-justified, specifically in the unbalanced setting, which is very useful in computational biology. I particularly enjoyed the growth penalty explanation part, as a better interpretation of its analytical form could provide end users with no expertise in OT with a biological rationale on how to set its value to optimally model the system they are considering. \n\n## Weaknesses\n\nAt the moment, there are some concerns on my side regarding the presentation of the paper that induced me to a borderline rejection score. I would like to know more about the authors' opinions during the rebuttal phase to be able to produce a final informed assessment. \n\n* **Selection of the growth penalty:** Though I appreciate the formalisation of the growth penalty in biological terms, the results proposed by the authors show that using more sophisticated terms than the WFR standard approach is suboptimal on the snapshot reconstruction task. Therefore, I am not sure this specific aspect of the paper is ripe enough to be used as an effective contribution to the paper. I still think this aspect is interesting, and maybe additional options for the penalty functions could be considered in further iterations of the paper, or at least be shown in a setting where it is actually useful. \n* **Line 103** Here, I recommend that the authors define $\\alpha$ with a short sentence.\n* **Time index:** Overall, I would personally find it easier to understand if all the time-resolved variables and quantities in the paper had an index indicating time. Some of them already have this (for example, the density $\\rho$), but some others miss it (as, for instance, the empirical measure $\\mu$ or the weights $w$). \n* **Benchmark:** One thing that I find sub-optimal in the current presentation is the way benchmarks are conducted. First, some models are presented in Tab. 1 and 2, while many of them are dropped later in Tab. 4. I personally believe that the baselines should be the same throughout to ensure transparency. \n* **Missing baselines:** Moreover, nowadays, there are very powerful OT-based methods for trajectory inference in low dimensions, like OT-CFM and derivatives thereof. I think I would have given some priority to them as baselines to test whether stochasticity and unbalancedness help in trajectory reconstruction. Also, combinations of OT Flow Matching and unbalancedness were tackled in some existing works [1], which I think is worth comparing against. Along the same line, comparison with [2], which is presented as a special case of the current model, is beneficial. What I am trying to say is that in a research setting where simulation-based methods are overtaken by their simulation-free counterpart, comparison with the latter models is meaningful to show that the proposed approach is a valuable option. \n* **Dataset choice.** There are widespread benchmarks of methods used to match temporal snapshots in scRNA-seq. They are based on the Embryoid body dataset from [3] and the NeurIPS benchmark in [4]. They also test slightly higher dimensionalities of the input space, which I think is relevant here in a simulation-based setting. \n* **Figure presentation.** In my opinion, Fig. 3 is not very clear. I can see what the colors mean, but the paths are not clear to me. I think the presentation could be improved by making the source and the target more understandable. For example, I don't understand where one can see higher curvature (line 256) in the results from DeepRUOT. Moreover, I would add a legend explaining the difference between crosses and full dots. Moreover, in Fig. 2, Var-RUOT produces different growth patterns compared to DeepRUOT, where the growth rate first increases and then goes down again in the last time point. Is it expected? I believe this figure could benefit from more explanations. In other words, in most of the provided qualitative evidence, I cannot tell why Var-RUOT is better than DeepRUOT. \n* **Wall clock time:** Very minor point. I would add the unit of measure of wall clock time somewhere in the table caption. \n* **Description of performance:** Sometimes, the authors seem to interpret the scores in absolute terms. For example, in lines 248-250 or in lines 286-288. I think these scores are mostly comparative, so I would phrase it differently and say that the results demonstrate that Var-RUOT better models trajectories than the baselines, rather than the accurate recovery claim. \n\nI wish I could be more positive in my assessment at this stage. I am looking forward to hearing more from the authors in their responses. \n\n[1] Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z. and Theis, F., 2023. Unbalancedness in neural monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100.\n\n[2] Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q. and Makhzani, A., 2023. A computational framework for solving wasserstein lagrangian flows. arXiv preprint arXiv:2310.10649.\n\n[3] Moon, K.R., Van Dijk, D., Wang, Z., Gigante, S., Burkhardt, D.B., Chen, W.S., Yim, K., Elzen, A.V.D., Hirn, M.J., Coifman, R.R. and Ivanova, N.B., 2019. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37(12), pp.1482-1492.\n\n[4] Luecken, M.D., Burkhardt, D.B., Cannoodt, R., Lance, C., Agrawal, A., Aliee, H., Chen, A.T., Deconinck, L., Detweiler, A.M., Granados, A.A. and Huynh, S., 2021, August. A sandbox for prediction and integration of DNA, RNA, and proteins in single cells. In Thirty-fifth conference on neural information processing systems datasets and benchmarks track (Round 2).", "was_revised": true, "final_justification": "Dear AC, \n\nThank you very much for handling this paper. \n\nIn my final assessment, I increased the score assigned to the present submission from 3 to 4. The authors have addressed most of my concerns with substantial experiments, explanations, and analyses. I believe the paper is above the acceptance bar, as it introduces new theory while providing value for downstream applications. \n\nThe reason I did not score a 5 is that I believe some aspects are not fully polished, such as the exploration of the relevance of the new formulation of the non-canonical growth penalty and the lack of presentation quality in the qualitative results. \n\nI remain available to provide further justification for my score.\n\nBest regards, \n\nThe Reviewer", "rebuttal": "Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:\n\n> Summary & Strengths\n\nWe appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT **finds paths with lower action**, **achieves stable training**, and **converges faster**. We also discuss the **biological priors** associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.\n\n\n\n> [W1, Q3] On the Selection of $g^\\frac{2}{15}$\n\nThank you for your question regarding $\\psi(g)$. To clarify, we use $g^{2/15}$ **as an example** to show that when $\\psi'(g)>0$ and $\\psi''(g)<0$, $g(\\boldsymbol{x},t)$ increases along $\\boldsymbol{u}(\\boldsymbol{x},t)$ (Theorem 4.2). We chose $\\psi_2(g)=g^{2/15}$ because:  \n1. It is the simplest form  that meets $\\frac{d\\psi(g)}{d|g|}>0$, $\\psi(g)=\\psi(-g)$, and $\\psi''(g)<0$.  \n2. From the optimality condition we obtain\n$$\ng(\\lambda) = \\left(\\frac{\\lambda}{\\alpha}\\frac{2q+1}{2p}\\right)^{\\frac{2q+1}{2p-(2q+1)}} \\propto \\left(\\frac{1}{\\lambda}\\right)^{\\frac{2q+1}{(2q+1)-2p}}.\n$$\nFor a more stable training process, we keep $g(\\lambda)$ from being too steep near $\\lambda=0$. So we set $p=1$ and choose a large $q$ (i.e., $q=7$).\n\nNote that $g(\\lambda)$ is undefined at $\\lambda=0$ due to the necessary discontinuity in $\\psi'(g)$ at $g=0$. To address this, we linearly interpolate between $g(-\\delta)$ and $g(\\delta)$ (Appendix A.2). This singularity causes slowing convergence and degrading distribution reconstruction when using $\\psi_2(g)$. A more reasonable way to deal with this problem is to fit the reciprocal of $\\lambda$, which is our future work.\n\nEmpirically, on the 2D Mouse Blood Hematopoiesis dataset (results are shown below), $\\psi(g)=g^{2/15}$ performs similarly to $\\psi(g)=g^{2/21}$, while $\\psi(g)=g^{2/9}$ (with a steeper $g(\\lambda)$) leads to unstable training and poorer reconstruction. Our experiments demonstrate that choosing the exponent on $g$ within a certain range is feasible. We will **update the paper to include the above analysis in the sensitivity analysis section of the Appendix**.\n\n|$\\psi(g)$/$\\mathcal{W}_1$ Distance|t=1|t=2| \n|-|-|-|\n|$\\psi(g)=g^{(2/9)}$|0.9476|1.2197| \n|$\\psi(g)=g^{(2/15)}$|0.2953| 0.1917| \n|$\\psi(g)=g^{(2/21)}$|0.2858|0.2565|\n\nAs noted in our Limitations, choosing $\\psi(g)$ **remains open**. Future work may derive it from **physical principles** (e.g., Branching SDE, Stochastic Thermal Dynamics) or **learn it with a neural network**. **We will include these discussions in the Appendix.**\n\n> [Q1] On Weight Initialization\n\nUsing particle methods, we approximate the true measure $\\mu$ with the empirical measure $\\mu^N$. We initialize by placing a particle at every $t=0$ data point with the uniform weight 1. This yields\n$$\n\\mu^{N}(0) = \\frac{1}{N}\\sum_{i=1}^{N} w_i \\delta(x - X_{i} ) \\rightarrow \\mu(0),\n$$\nwhich justifies the rationale to set $w_i(0)=1$.\n\n> [Q2] On the Ground Truth for Mass Matching\n\nIn RUOT modeling, because cell counts vary (due to growth or death), the number of data points differs over time. If we set the total mass at $t_1=0$ as the reference number 1, then the mass at time $t_i$ as $\\frac{N_i}{N_1}$ (with $N_i$ as the count at $t_i$), which serves as the ground truth for mass matching loss. **We will clarify this in our paper.**\n\n> [Q4]On the Schrodinger Bridge Problem\n\nWe apologize for the confusion. The original Schrodinger Bridge is dynamic—minimizing the KL divergence over path space between a stochastic process and a reference one. In our paper, \"static\" means some method[(Tong, A.,et al) S2FM] first compute a static OT solution via entropy-regularized OT, then fit the velocity field using Conditional Flow Matching. **We will clarify this in the paper.**\n\n\n> [W4,5,6] Additional Baselines for Existing Datasets\n\nThanks for the great suggestion. Follow your advice, we **supplemented three baselines on existing datasets**; in addition, we provided results for **nine baselines and Var-RUOT on two new datasets**. We first **supplemented the additional baselines** on existing datasets with:\n- Wasserstein Gradient Flow (WGF in the table) [(Neklyudov, K., et al.) Wasserstein Lagrangian Flows]\n- OTCFM [(Lipman, Y., et al.) Flow Matching]\n- UOTCFM [(Eyring, L., et al.)Unbalancedness in Neural Monge Maps]\n\nSpecifically, for the 2D Mouse Blood Hematopoiesis dataset, we have added three previously unevaluated baselines: MIOFLOW, S2FM, and PISDE.\n\n> Simulation Dataset\n\n|Method/$\\mathcal{W}_1$ Distance|$t_1$|$t_2$|$t_3$|$t_4$|\n|-|-|-|-|-|\n|OTCFM|0.1035|0.2078|0.2898|0.3107|\n|UOTCFM|0.1002|0.1653|0.1711|0.4129|\n|WGF|0.4983|0.8346|0.8046|0.4493|\n|Var-RUOT|0.0452|0.0385|0.0445|0.0572|\n\n> EMT Dataset\n\n|Method/$\\mathcal{W}_1$ Distance|$t_1$|$t_2$|$t_3$|\n|-|-|-|-|\n|OTCFM|0.2557|0.2701|0.2799|\n|UOTCFM|0.2538|0.2696|0.2771|\n|WGF|0.3901|0.4381|0.2848|\n|Var-RUOT|0.2540|0.2670|0.2683|\n\n> 2D Mouse Dataset\n\n|Method/$\\mathcal{W}_1$ Distance|$t_1$|$t_2$|\n|-|-|-|\n|MIOFLOW|0.3546|0.2772|\n|S2FM|0.1706|0.1602|\n|PISDE|0.3124|0.2531|\n|OTCFM|0.3674|0.3222|\n|UOTCFM|0.3301|0.2051|\n|WGF|0.3302|0.2051|\n|Var-RUOT|0.1200|0.1917|\n\n\nNext, we show experiments on **additional datasets**. We use a 100-dimensional, 5-time-point ($t=0,1,2,3,4$) EB dataset [(Moon, K.R., et al) Visualizing structure and transitions...] and the NeurIPS challenge dataset [(Luecken, M.D., et al.) A sandbox for prediction and integration...]. For NeurIPS dataset, we extract gene expression data from donor 28483, cluster them via K-Means by pseudo time, and order the clusters (by centroid) into $t=0,1,2,3,4$. This new dataset is then used for further tests. The results are as follows:\n\n> EB Dataset\n\n|Method/$\\mathcal{W}_1$ Distance|$t_1$|$t_2$|$t_3$|$t_4$|\n|-|-|-|-|-|\n|S2FM|9.6732|11.0711|12.5436|14.7100|\n|PISDE|8.8317|9.4283|9.7872|11.1102|\n|MIOFLOW|8.7182|9.4139|9.7547|11.0080|\n|TIGON|8.7992|9.6497|10.0130|11.3284|\n|RUOT|8.8029|9.6518|10.0365|11.3555|\n|OTCFM|9.4459|11.3786|14.0529|19.7489|\n|UOTCFM|9.4794|11.6060|14.6715|21.7653|\n|ActionMatching|12.7145|16.7809|13.4663|18.1165|\n|WGF|9.9847|11.0562|11.9719|13.2417|\n|Var-RUOT|8.9158|9.7955|10.2647|11.4927|\n\n> NeurIPS Dataset\n\n|Method/$\\mathcal{W}_1$ Distance|$t_1$|$t_2$|$t_3$|$t_4$|\n|-|-|-|-|-|\n|S2FM|6.1876|12.6170|9.1830|12.5831|\n|PISDE|5.6642|5.2358|6.2793|8.9086|\n|MIOFLOW|6.4267|6.2715|7.1079|8.9035|\n|TIGON|5.9481|5.7599|6.5539|8.9573|\n|RUOT|6.0166|5.6972|6.5099|8.8047|\n|OTCFM|5.5082|5.4476|7.0935|10.1416|\n|UOTCFM|5.5526|5.4349|6.7482|9.6782|\n|ActionMatching|7.3701|8.5156|9.2842|9.9665|\n|WGF|7.1785|7.7135|8.2521|8.6629|\n|Var-RUOT|4.7961|4.1816|6.2312|9.1058|\n\nAnd we compared the path action learned by TIGON, DeepRUOT and VarRUOT:\n\n|Method/Path Action|TIGON|DeepRUOT|VarRUOT|\n|-|-|-|-|\n|EB Dataset|4.4927|5.0571|4.1536|\n|NeurIPS Dataset|10.4250|10.9898|3.6982|\n\nIn summary, the experiments show that:\n- For low-dimensional data: On the EMT dataset, Var-RUOT achieves distribution reconstruction **comparable to OTCFM and UOTCFM**, and on the three-gene simulation and 2D Mouse Blood Hematopoiesis datasets, it **outperforms the other baselines**.\n- For high-dimensional data: On the EB dataset, Var-RUOT yields results similar to TIGON and DeepRUOT, and outperforms simulation-free methods (OTCFM, UOTCFM) as well as Wasserstein Gradient Flow; on the NeurIPS dataset, it **outperforms existing methods**. Notably, compared to DeepRUOT and TIGON, Var-RUOT finds a path with **significantly lower action**.\n\nThus, the experiments confirm that Var-RUOT remains competitive on high-dimensional data, achieving good distribution reconstruction while finding lower-action paths. **We will include these discussions on new datasets & baselines in the Appendix.**\n\n\n> [W7] On the Figures\n\nThank you for your questions regarding our figures. For Figures 2 and 3:\n- Crosses represent real dataset points, and circles represent data generated by our algorithm.\n- Trajectory colors indicate particle mass, with brighter areas showing higher mass.\n\nWe will add legends and annotations for clarity.\n\nRegarding the EMT results, Figure 3 (a) shows DeepRUOT produces curved trajectories (even passing through sparse top regions), while Figure 3 (b) shows nearly straight lines connecting the start and end points. Since straight lines minimize action, Table 2 compares the Path Action of DeepRUOT and Var-RUOT, with Var-RUOT achieving significantly lower action.\n\nThe growth pattern in Var-RUOT emerges naturally from using the WFR Metric (Theorem 4.2), where the growth rate increases along the trajectory. Without a \"minimum action\" objective, many valid velocity fields and growth patterns could reconstruct the marginals.  Both methods correctly reconstruct the marginals, while Var-RUOT finds a lower-action solution, which aligns with our objective.\n\nWe clarify that our goal in presenting these figures is **not** to show that Var-RUOT has a significantly lower reconstruction loss, but to demonstrate that both methods perform similarly in marginal reconstruction. Var-RUOT’s advantages lie in it **obtains a lower-action solution**, **simpler scalar field fitting**, **faster convergence**, and the elimination of the pretraining phase in DeepRUOT.\n\n> [W2,3,8,9]Other Writing Issues\n\nThank you for your valuable suggestions. We will revise the paper to:\n- Define $\\alpha$ at line 103 as the penalty coefficient for growth/death.\n- Clarify in lines 248–250 and 286–288 that the $\\mathcal{W}_1$, $\\mathcal{W}_2$ distances and Path Action are relative indicators, with Var-RUOT performing better relative to the baseline.\n- Add the unit (seconds) for Wall Clock time in Table 3.\n- Add time index to time-dependent variables."}]

Author rebuttal:

Meta-review: This paper presents Var-RUOT, a new approach to recover the dynamics from snapshot data, aproblem central in the computational biology literature. The core innovation is to directly incorporate the first-order optimality conditions (HJB) of the RUOT problem into the models parameterization and loss. This leads to significant simplifications over the previous methods.

The work received unanimous and positive recommendation from all four reviewers who were particularly impressed by the method's theoretically grounding, the performance, and a strong rebuttal that addressed all concerns from all reviewers. For these reasons I recommend accepting this work.
Decision: Accept (poster)