Title: Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action
Venue: NeurIPS 2025
Accepted: yes
Keywords: unbalanced optimal transport

Schrödinger bridge

trajectory inference

single-cell
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

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 .

1 Introduction

Inferring continuous dynamics from finite observations is crucial when analyzing systems with many particles (Chen et al. 2018). However, in many important applications such as single-cell RNA sequencing (scRNA-seq) experiments, only a few snapshot measurements are available, which makes recovering the underlying continuous dynamics a challenging task (Ding et al. 2022). Such a task of reconstructing dynamics from sparse snapshots is commonly referred to as trajectory inference in time-series scRNA-seq modeling (Zhang et al. 2025b; Ding et al. 2022; Heitz et al. 2024; Yeo et al. 2021a; Schiebinger et al. 2019a; 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

Initial Dist.

Terminal Dist.

RUOT in Configuration Manifold

RUOT in Probability Manifold

𝝁𝑰

Previous RUOT Solver

Use 3 Neural Networks

Hard to Achieve Least Action

To Achieve Least Action,

You can use Variational Method

to derive Optimal Conditions

Only Learns One Scalar Field

𝝀(𝒙, 𝒕)

！

Represents Biological Prior

For a given timepoint

𝒕

,

𝝍′′(𝒈) > 𝟎

means

𝒈

increases

along

𝒖(𝒙, 𝒕)

𝝍 ′′ 𝒈

< 𝟎

means

𝒈

decreases

along

𝒖(𝒙, 𝒕)

dab(g)

u=

Vx入

a

dg

=入

10

0.3

(p,u.g)

inf

A

α(g(x,t)) p(x,t) dx dt.

C = 入Recon LRecon + 入HJB CHJB + LAction·

ap(x,t)

=-Vz·(o(x,t)u（x,t）)+o²V²p（x,t)+g（x,t）p（x,t).

at

( = 3.00

( = 1.00

4:100

X

velocity

y

X

t

y

t

score

y

growth

t

ODEsolver

1:2

0.2

1.5

1

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 ψ ( g ) on the particle growth or death processes in total transport energy (namely the Wasserstein-Fisher-Rao 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 et 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 et al. 2015; Song et 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 ψ ( g ) that aligns with underlying biological priors remains challenging.

To overcome these challenges, we propose Variational-RUOT (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 a 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 are 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.

Weshow how incorporating these necessary conditions into the loss function and architecture enables Var-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-Fisher-Rao 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 are 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´ c et al. 2025; Schiebinger et al. 2019b; Klein et al. 2025; Wang et al. 2025), as well as 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; Gu et al. 2025; Koshizuka and Sato 2023; Neklyudov et al. 2024; Zhang et 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 to parameterize the velocity and growth functions, without leveraging their optimality conditions or the inherent relationship between them(Zhang et al. 2025d) with few exceptions such as Action Matching (Neklyudov et al. 2023; Neklyudov et al. 2024)) . This poses challenges in achieving optimal solutions that minimize the action energy. Our work generalizes Action Matching (Neklyudov et al. 2023; Neklyudov et 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-Jacobi-Bellman (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; Neklyudov et 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 in temporal scRNA-seq 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 et al. 2024; Zhang et al. 2025a; Zhang et al. 2025c). However, these works usually adopt the default growth penalty function ψ ( g ) = 1 2 g 2 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 BenamouBrenier formulation, requires minimizing the following action functional (Benamou and Brenier 2000):

where ρ and u are subject to the continuity equation constraint:

Unbalanced Dynamical OT and Wasserstein-Fisher-Rao (WFR) metric In order to handle 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 a corresponding penalty term in the action. This leads to the optimal transport problem under the Wasserstein-Fisher-Rao (WFR) metric (Chizat et al. 2018a; Chizat et al. 2018b).

with ρ , u , and g subject to the unnormalized continuity equation constraint

Schrödinger Bridge Problem and Dynamical Formulation Schrödinger bridges aims to find the most likely way for a stochastic system to evolve from an initial distribution µ 0 to a terminal distribution µ 1 . Formally, let µ X [0 , 1] denote the probability measure induced by the stochastic process X ( t ) , 0 ≤ t ≤ 1 , and let µ Y [0 , 1] denote the probability measure induced by a given reference process Y ( t ) , 0 ≤ t ≤ 1 . The Schrödinger bridge seeks to solve min µ X [0 , 1] D KL ( µ X [0 , 1] ∥ ∥ ∥ µ Y [0 , 1] ) . In particular, if X t follows the SDE d X t = u ( X t , t ) d t + σ ( X t , t ) d W t , where W t ∈ R d is a standard Brownian motion, σ ( x , t ) ∈ R d × d is a given diffusion matrix, and the reference process is defined as d Y t = σ ( Y t , t ) d W t , then the Schrödinger bridge problem is equivalent to the following stochastic optimal control problem (Chen et al. 2016; Gentil et al. 2017):

where ρ and u are subject to the Fokker-Planck equation constraint

Regularized Unbalanced Optimal Transport If we consider both unnormalized probability densities and stochasticity simultaneously, we arrive at 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:

where ψ : R → [0 , + ∞ ) is a growth penalty function, α is a penalty coefficient, and the quantities ρ , u and g are subject to the following constraint, which is an unnormalized continuity equation:

Remark 3.1. When modeling the particle dynamics in these OT variants, one can use either autonomous systems (where v or g does not explicitly depend on t ) or non-autonomous systems (where v or g explicitly depends on t ). Here, we follow methods such as DeepRUOT(Zhang et al. 2025a) to choose non-autonomous systems. Indeed non-autonomous 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 of isotropic time-invariant diffusion, i.e., a ( x , t ) = σ 2 I . We refer to RUOT problem in this scenario as the isotropic time-invariant RUOT problem .

Definition 4.1 (Isotropic Time-Invariant (ITI) RUOT Problem) . Consider the following minimumaction problem with the action functional given by

Here, ψ : R → [0 , + ∞ ) is the growth penalty function, and the triplet ( ρ, u , g ) is subject to the constraint of the Fokker-Planck equation

Additionally, p satisfies the initial and terminal conditions ρ ( · , 0) = µ 0 , ρ ( · , 1) = µ 1 .

In particular, if ψ ( g ( x , t )) = 1 2 g 2 ( 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 T to attain a minimum are

Here, λ ( x , t ) is a scalar field. The proof of this theorem can be found in Appendix A.1.1.

Remark 4.1. Substituting the necessary conditions satisfied by u and g into the Fokker-Planck equation, the evolution of the probability density ρ ( x, t ) is determined by ∂ρ ( x ,t ) ∂t = -∇ x · ( ρ ( x , t ) ∇ x λ ( x , t ) ) + 1 2 σ 2 ∇ 2 x ρ ( x , t ) + ( ψ ′ ) -1 ( λ ( x ,t ) α ) ρ ( x , t ) , where ψ ′ = d ψ ( g ) d g , and ( ψ ′ ) -1 denotes the inverse function of ψ ′ .

Remark 4.2. If we choose the growth penalty function to take the form used in the WFR metric, i.e., ψ ( g ) = 1 2 g 2 , and set α = 1 , σ = 0 , then the above optimal necessary conditions immediately degenerate to u = ∇ x λ, g = λ, ∂λ ∂t + 1 2 ∥∇ x λ ∥ 2 + 1 2 λ 2 = 0 , which is same as the form derived in (Neklyudov et al. 2024) under the WFR metric. If we let g = 0 and ψ (0) = 0 it becomes u = ∇ x λ, ∂λ ∂t + 1 2 ∥∇ x λ ∥ 2 + 1 2 σ 2 ∇ 2 x λ = 0 , which is same as the form derived in (Neklyudov et al. 2024; Zhang et al. 2024b; Chen et 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 , t ) . Moreover, since the initial density ρ ( · , 0) is known, once the necessary conditions are satisfied, the evolution equation (i.e., the Fokker-Planck equation) is completely determined by λ ( x , t ) . Thus, the scalar field λ ( x , t ) fully determines the system's evolution, and we only need to solve for one λ ( x , t ) , which simplifies the problem.

However, these necessary conditions introduce a coupling between u ( x , t ) and g ( x , t ) , and this coupling could contradict biological prior knowledge. In biological data, it is generally believed that cells located at the upstream of a trajectory are stem cells with the highest proliferation and differentiation capabilities, and thus the corresponding g 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 a given time t depends on the form of the growth penalty function.

Theorem 4.2 (The relationship between u and g ; Biological prior) . At a fixed time t , if d ψ 2 ( g ) d g 2 > 0 , then g ( x , t ) ascends in the direction of the velocity field u ( x , t ) (i.e., u ( x , t ) T ( ∇ x g ( x , t )) > 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 g value, it is necessary to ensure that that d 2 ψ ( g ) d g 2 < 0 .

5 Solving ITI ROUT Problem Through Neural Network

Given samples from distributions ρ t at K discrete time points, t ∈ { T 1 , · · · , T K } , we aim to recover the continuous evolution process of the distributions by solving the ITI RUOT problem, that is, by minimizing the action functional while ensuring that ρ ( x , t ) matches the distributions ρ t at the corresponding time points. Since the values of u ( x , t ) and g ( x , t ) , as well as the evolution of ρ ( x , t ) over time, are fully determined by the scalar field λ ( x , t ) in variational form (Appendix A.1.1), we approximate this scalar field using a single neural network. Specifically, we parameterize λ ( x , t ) as λ θ ( x , t ) , where θ 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.

where u : R d × [0 , T ] → R d is a time-varying vector field, g : R d × [0 , T ] → R is a growth rate function, σ : [0 , T ] → [0 , + ∞ ) is a time-varying diffusion coefficient, and W t is an N -dimensional standard Brownian motion with independent components in each coordinate. The initial conditions are X 0 i ∼ ρ ( x , 0) and w i (0) = 1 . In the limit as N → ∞ , the empirical measure µ N ( x , t ) = 1 N ∑ N i =1 w i ( t ) δ ( x -X t i ) converges to the solution of the following Fokker-Planck equation:

with the initial condition ρ ( x , 0) = ρ 0 ( x ) .

The proof is provided in Appendix A.1.3. This theorem implies that we can approximate the evolution of ρ ( x , t ) by simulating N particles, where each particle's weight w i is governed by an ODE and its position X i is governed by an SDE. The evolution of the empirical measure µ N ( x , t ) thereby approximates the evolution of ρ ( x , t ) .

5.2 Reformulating the Loss in Weighted Particle Form

The total loss function consists of three components such that L = L Recon + γ HJB L HJB + γ Action L Action . Here, L Recon ensures that the distribution generated by the model closely matches the true data distribution, L HJB enforces that the learned λ θ ( x , t ) satisfies the HJB equation in the necessary conditions, and L Action minimizes the action as much as possible.

Reconstruction Loss Minimizing the reconstruction loss guarantees that the distribution generated by the model is consistent with the real data distribution. Since in the ITI RUOT problem the probability density ρ ( x , t ) is not normalized, we need to match both the total mass and the discrepancy between the two distributions. Our reconstruction loss is given by L Recon = γ Mass L Mass + L OT , where, at time point k , the true mass is ∫ R d ρ ( x , T k ) d x = M ( T k ) , and the weight of particle i is w i ( T k ) . Thus, the total mass of the model-generated distribution is ˆ M ( T k ) = 1 ∑ N w i ( T k ) .

N i =1 The mass reconstruction loss is then defined as L Mass = ∑ K k =1 ( M ( T k ) -ˆ M ( T k ) ) 2 . Let the true distribution at time point k be ρ ( x , T k ) . Its normalized version is given by ˜ ρ ( x , T k ) = ρ ( x ,T k ) ∫ R d ρ ( x ,T k ) d x , while the normalized model-generated distribution is ˆ ˜ ρ ( x , T k ) = 1 N ∑ N i =1 w i ( T k ) δ ( x -X i ) ˆ M ( T k ) . The distribution reconstruction loss is then defined as L OT = ∑ K k =1 W 2 ( ˜ ρ ( · , T k ) , ˆ ˜ ρ ( · , T k ) ) , where γ Mass is a hyperparameter that controls the importance of the mass reconstruction loss.

HJB Loss Minimizing the HJB loss ensures that the learned λ θ ( x , t ) obeys the HJB equation constraints specified in the necessary conditions. Since the gradient operator in the HJB equation is a local operator, we compute the HJB loss by integrating the extent to which λ θ ( x , t ) violates the HJB equation along the trajectories. When using N particles, the HJB loss is given by: L N HJB =

g θ is obtained from the necessary condition α d ψ ( g θ ) d g θ = λ θ .

Remark

∫

T

K

∫

0

R

d

ˆ

ρ

(

x

5.1.

, t

)

The

∂λ

(

x

,t

)

∂t

+

1

2

expectation

∥∇

x

λ

∥

+

2

1

2

σ

of

∇

x

λ

+

HJB

g

(

x

, t

)

-

Loss

αψ

(

g

)

is

d

x

d

t

E

[

L

N

HJB

where

]

ˆ

ρ

(

x

, t

)

=

=

ρ ( x , t ) ∫ R d ρ ( x , t )d x is the probability density by normalizing ρ ( x , t ) . The proof is left in Appendix A.1.4.

2

(

2

)

2

Action Loss Since the variational method provides only the necessary conditions for achieving minimal action rather than sufficient ones, we need to incorporate the action into the loss so that the action is minimized as much as possible. The action loss is also computed via simulating weighted particles. When using N particles, it is given by L N Action =

Remark 5.2. The expectation of action loss is exactly the action defined in the ITI RUOT problem (Definition 4.1): E [ L N Action ] = 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 λ θ . 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: if ψ ′′ ( g ) > 0 , then at any given time t , g increases in the direction of the velocity field, and vice versa. Therefore, we choose two representative forms of ψ ( g ) for our solution. Given that ψ ( g ) penalizes nonzero g , it should satisfy the following properties: (1) The further g deviates from 0 , the larger ψ ( g ) becomes, i.e., d ψ ( g ) d | g | > 0 . (2) Birth and death are penalized equally when prior knowledge is absent, i.e., ψ ( g ) = ψ ( -g ) .

Case 1: ψ ′′ ( g ) > 0 In the case where ψ ′′ ( g ) > 0 , a typical form that meets the requirements is ψ ( g ) = Cg 2 p , p ∈ Z + , C > 0 . We select the form used in the WFR Metric, namely, ψ 1 ( g ) = 1 2 g 2 . The optimality conditions are presented in Appendix A.2.

Case 2: ψ ′′ ( g ) < 0 For the case where ψ ′′ ( g ) < 0 , a typical form that meets the conditions is ψ ( g ) = C g (2 p ) / (2 q +1) , where p, q ∈ Z + and 2 p < 2 q + 1 . In order to obtain a smoother g ( λ ) relationship from the necessary conditions, and as a illustrative example , we choose ψ 2 ( g ) = g 2 / 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 ) = 1 2 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.1. The trajectories learned by DeepRUOT and Var-RUOT are illustrated in Fig. 2, and the W 1 and W 2 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 Var-RUOT achieves a lower path action while maintaining distribution matching accuracy. To further assess the performance of Var-RUOT on high-dimensional data, we also conducted experiments on an Epithelial Mesenchymal Transition (EMT) dataset(Sha et al. 2024; Cook and Vanderhyden 2020). This dataset was reduced to a 10-dimensional feature space, and the trajectories obtained after applying PCA for dimensionality reduction are shown in Fig. 3. Both Var-RUOT and DeepRUOT learn dynamics that can transform the distribution at t = 0 into the distributions at t = 1 , 2 , 3 . Var-RUOT learns the nearly straight-line trajectory corresponding to the minimum action, whereas DeepRUOT learns a curved trajectory. The results of W 1 , W 2 distance and action, summarized in Table 2, Var-RUOT also learns trajectories with smaller action while achieving matching accuracy comparable to that of other algorithms.

t

= 1

t

= 2

t

= 3

t

= 4

Path Action

Model

W

1

W

2

W

1

W

2

W

1

W

2

W

1

W

2

SF2M (Tong et al. 2024b)

0.1914

±

0.0051

0.3253

±

0.0059

0.4706

±

0.0200

0.7648

±

0.0059

0.7648

±

0.0260

1.0750

±

0.0267

2.1879

±

0.0451

2.8830

±

0.0741

-

PISDE (Jiang and Wan 2024)

0.1313

±

0.0023

0.3232

±

0.0013

0.2311

±

0.0015

0.5356

±

0.0015

0.4103

±

0.0006

0.7913

±

0.0035

0.5418

±

0.0015

0.9579

±

0.0037

-

MIO Flow (Huguet et al. 2022)

0.1290

±

0.0000

0.2087

±

0.0000

0.2963

±

0.0000

0.4565

±

0.0000

0.6461

±

0.0000

1.0165

±

0.0000

1.1473

±

0.0000

1.7827

±

0.0000

-

Action Matching (Neklyudov et al. 2023)

0.3801

±

0.0000

0.5033

±

0.0000

0.5028

±

0.0000

0.5637

±

0.0000

0.6288

±

0.0000

0.6822

±

0.0000

0.8480

±

0.0000

0.9034

±

0.0000

1.5491

OTCFM (Tong et al. 2024a)

0.1035

±

0.0000

0.3043

±

0.0000

0.2078

±

0.0000

0.4923

±

0.0000

0.2898

±

0.0000

0.5867

±

0.0000

0.3107

±

0.0000

0.4358

±

0.0000

-

UOTCFM(Eyring et al. 2024a)

0.1002

±

0.0000

0.2911

±

0.0000

0.1653

±

0.0000

0.3578

±

0.0000

0.1711

±

0.0000

0.2620

±

0.0000

0.4129

±

0.0000

0.7583

±

0.0000

-

WLF(Neklyudov et al. 2024)

0.4983

±

0.0000

0.5273

±

0.0000

0.8346

±

0.0000

0.8357

±

0.0000

0.8046

±

0.0000

0.8815

±

0.0000

0.4493

±

0.0000

0.8571

±

0.0000

-

TIGON (Sha et al. 2024)

0.0519

±

0.0000

0.0731

±

0.0000

0.0763

±

0.0000

0.1559

±

0.0000

0.1387

±

0.0000

0.2436

±

0.0000

0.1908

±

0.0000

0.2203

±

0.0000

1.2442

DeepRUOT (Zhang et al. 2025a)

0.0569

±

0.0019

0.1125

±

0.0033

0.0811

±

0.0037

0.1578

±

0.0079

0.1246

±

0.0040

0.2158

±

0.0081

0.1538

±

0.0056

0.2588

±

0.0088

1.4058

Var-RUOT (Ours)

0.0452

±

0.0024

0.1181

±

0.0064

0.0385

±

0.0022

0.1270

±

0.0121

0.0445

±

0.0033

0.1144

±

0.0160

0.0572

±

0.0034

0.2140

±

0.0067

1.1105

±

0.0515

(a) DeepRUOT

Growth

Trajectory

Growth

Trajectory

(b) Var-RUOT

t =0.00

0.5

t

t = 1.00

t

=

t =2.00

0.4

t

3

t=3.00

t =4.00

0.3

0.2

0.1

0.0

Growth

6.2 Var-RUOT Stabilizes and Accelerates Training Process

To demonstrate that Var-RUOT 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 a specified threshold (set to 0.30 in this study). Each training session was capped at a maximum of 500 epochs. If an algorithm'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.1, we further demonstrate our training speed and stability by plotting the loss decay curves.

(a) DeepRUOT

Growth

Trajectory

Growth

Trajectory

(b) Var-RUOT

t=0

t = 1

0.3

t=

2

t

=

3

0.2

0.1

t=0.00

0.0

t = 1.00

t = 2.00

t = 3.00

0.1

Growth

6.3 Different ψ ( g ) Represents Different Biological Prior

To illustrate that the choice of ψ ( g ) represents different biological priors, we present the learned dynamics under two selections of ψ ( g ) . We apply our algorithm on 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., ψ 1 ( g ) = 1 2 g 2 , from which it can be observed that at time points t = 0 , 1 , 2 , along the direction of the drift vector field u ( x , t ) , g ( x , t ) gradually increases. In Fig. 4(b), on the other hand, the alternative selection ψ 2 ( g ) = g 2 / 15 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 to those obtained using the WFR metric. Therefore we do not report its action here.

In addition to the three experiments presented in 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 a minimal action. Additionally, we explored alternative modeling choices, including a comparison

Growth

(a) Standard WFR Metric

𝝍 𝒈

= 𝟏

𝟐

𝒈 𝟐

𝒕 = 𝟎

Growth

𝒕 = 𝟏

Growth

𝒕 = 𝟐

Growth

Growth

𝝍 𝒈

= 𝒈 𝟐/𝟏𝟓

Trajectory

𝒕 = 𝟎

Growth

𝒕 = 𝟏

Growth

𝒕 = 𝟐

Growth

Trajectory

0.12

0.10

- 0.08

0.06

t=0.00

t = 1.00

t =2.00

0.04

0.56

0.54

0.52

0.50

0.48

0.46

t=0.00

0.44

t =2.00

t = 1.00

0.42

between using L 1 and L 2 norms for the loss terms, and an analysis of the impact of different growth penalty functions beyond the ψ 2 ( g ) = g 2 15 example where ψ ′′ ( g ) < 0 . 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 can effectively perform both interpolation and extrapolation, with the learned minimum-action dynamics leading to more accurate extrapolation outcomes (Appendix C.3). Furthermore, we carried out experiments on 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 Var-RUOT against that from a traditional OT solver(Appendix C.5). The experimental results show that the path actions obtained by both methods are similar, which helps to validate our approach.

7 Conclusion

In this paper, we propose a 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 it 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 new insights for solving the RUOT problem, it is subject to several key limitations. Firstly, because the algorithm is based on neural network optimization, it can only converge to a local minimum rather than guaranteeing the attainment of the global minimum of the action. In practice, the HJB loss functions more as a regularization term than as a hard constraint that converges to zero. In addition, when using the modified metric, the goodness-of-fit 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 fundamental behavior and limitations, and by automating the selection of ψ ( g ) either by approximating it with a neural network or by deriving it from first-principle-based microscopic dynamics, such as a branching Wiener (Baradat and Lavenant 2021). We discussed more limitations of our work and potential broader impacts in Appendix E.1.

Acknowledgments and Disclosure of Funding

This work was supported by the National Key R&D Program of China (No. 2021YFA1003301 to T.L.), National Natural Science Foundation of China (NSFC No. 12288101 to T.L. & P.Z., and 8206100646, T2321001 to P.Z.) and The Fundamental Research Funds for the Central Universities, Peking University. We acknowledge the support from the High-performance Computing Platform of Peking University for computation.

References

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. 'Stochastic interpolants: A unifying framework for flows and diffusions'. In: arXiv preprint arXiv:2303.08797 (2023).

Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. 'Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold'. In: The Thirteenth International Conference on Learning Representations . 2025.

Aymeric Baradat and Hugo Lavenant. 'Regularized unbalanced optimal transport as entropy minimization with respect to branching brownian motion'. In: arXiv preprint arXiv:2111.01666 (2021).

Jean-David Benamou and Yann Brenier. 'A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem'. In: Numerische Mathematik 84.3 (2000), pp. 375393.

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

Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. 'The schrödinger bridge between gaussian measures has a closed form'. In: International Conference on Artificial Intelligence and Statistics . PMLR. 2023a, pp. 5802-5833.

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 singlecell perturbation responses using neural optimal transport'. In: Nature methods 20.11 (2023b), pp. 1759-1768.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. 'Neural ordinary differential equations'. In: Advances in neural information processing systems 31 (2018).

Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. 'Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory'. In: International Conference on Learning Representations . 2022a.

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.

Yongxin Chen, Tryphon T 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.

Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. 'An interpolating distance between optimal transport and Fisher-Rao metrics'. In: Foundations of Computational Mathematics 18 (2018a), pp. 1-44.

Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. 'Unbalanced optimal transport: Dynamic and Kantorovich formulations'. In: Journal of Functional Analysis 274.11 (2018b), pp. 3090-3123.

Lénaıc Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. 'Trajectory inference via mean-field langevin in path space'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 16731-16742.

Shui-Nee Chow, Wuchen Li, and Haomin Zhou. 'Wasserstein hamiltonian flows'. In: Journal of Differential Equations 268.3 (2020), pp. 1205-1219.

David P Cook and Barbara C Vanderhyden. 'Context specificity of the EMT transcriptional response'. In: Nature communications 11.1 (2020), p. 2142.

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.

Jun Ding, Nadav Sharon, and Ziv Bar-Joseph. 'Temporal modelling using single-cell transcriptomics'. In: Nature Reviews Genetics 23.6 (2022), pp. 355-368.

Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian Theis. Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation . 2024a. arXiv: 2311.15100 [cs.CV] .

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 on Learning Representations . 2024b.

Nicolas Fournier and Arnaud Guillin. On the rate of convergence in Wasserstein distance of the empirical measure . 2013. arXiv: 1312.2128 [math.PR] .

Ivan Gentil, Christian Léonard, and Luigia Ripani. 'About the analogy between optimal transport and minimal entropy'. In: Annales de la Faculté des sciences de Toulouse: Mathématiques . Vol. 26. 3. 2017, pp. 569-600.

Anming Gu, Edward Chien, and Kristjan Greenewald. 'Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior'. In: The Thirteenth International Conference on Learning Representations . 2025.

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

Jonathan Ho, Ajay Jain, and Pieter Abbeel. 'Denoising diffusion probabilistic models'. In: Advances in neural information processing systems 33 (2020), pp. 6840-6851.

Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. 'Manifold interpolating optimal-transport flows for trajectory inference'. In: Advances in neural information processing systems 35 (2022), pp. 29705-29718.

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: 13674811.

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.

Takeshi Koshizuka and Issei Sato. 'Neural Lagrangian Schrödinger Bridge: Diffusion Modeling for Population Dynamics'. In: The Eleventh International Conference on Learning Representations . 2023.

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.

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 A 34.4 (2014), pp. 1533-1574.

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

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. 'Flow Matching for Generative Modeling'. In: The Eleventh International Conference on Learning Representations . 2023.

Xingchao Liu, Chengyue Gong, and Qiang Liu. 'Flow straight and fast: Learning to generate and transfer data with rectified flow'. In: arXiv preprint arXiv:2209.03003 (2022).

Malte Luecken, Daniel Burkhardt, Robrecht Cannoodt, Christopher Lance, Aditi Agrawal, Hananeh Aliee, Ann Chen, Louise Deconinck, Angela Detweiler, Alejandro Granados, Shelly Huynh, Laura Isacco, Yang Kim, Dominik Klein, BONY DE KUMAR, Sunil Kuppasani, Heiko Lickert, Aaron McGeever, 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.

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

Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. 'Visualizing structure and transitions in high-dimensional biological data'. In: Nature biotechnology 37.12 (2019), pp. 1482-1492.

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. 25858-25889.

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.

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.

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.

Qiangwei Peng, Peijie Zhou, and Tiejun Li. 'stVCR: Reconstructing spatio-temporal dynamics of cell development using optimal transport'. In: bioRxiv (2024), pp. 2024-06.

Katarina Petrovi´ c, 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.

Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Y. Chen, Aviv Regev, and Romain Lopez. 'Modeling Complex System Dynamics with Flow Matching Across Time and Conditions'. In: The Thirteenth International Conference on Learning Representations . 2025.

Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. 'Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming'. In: Cell 176.4 (2019a), pp. 928-943.

Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. 'Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming'. In: Cell 176.4 (2019b), pp. 928-943.

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.

Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. 'Diffusion Schrödinger bridge matching'. In: Advances in Neural Information Processing Systems 36 (2024).

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. 'Deep unsupervised learning using nonequilibrium thermodynamics'. In: International conference on machine learning . PMLR. 2015, pp. 2256-2265.

Jiaming Song, Chenlin Meng, and Stefano Ermon. 'Denoising diffusion implicit models'. In: arXiv preprint arXiv:2010.02502 (2020).

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 'Score-Based Generative Modeling through Stochastic Differential Equations'. In: International Conference on Learning Representations . 2021.

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.

Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. 'Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics'. In: International conference on machine learning . PMLR. 2020, pp. 9526-9536.

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.

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.

Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. 'Trajectory inference for a branching SDE model of cell differentiation'. In: arXiv preprint arXiv:2307.07687 (2023).

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 β -cell differentiation'. In: Nature 569.7756 (2019), pp. 368-373.

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.

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

Zihao Wang, Datong Zhou, Yong Zhang, Hao Wu, and Chenglong Bao. Wasserstein-FisherRao Document Distance . 2019. arXiv: 1904.10294 [cs.LG] .

Caleb Weinreb, Alejo Rodriguez-Fraticelli, Fernando D Camargo, and Allon M Klein. 'Lineage tracing on transcriptional landscapes links state to fate during differentiation'. In: Science 367.6479 (2020), eaaw3381.

Hao Wu, Shu Liu, Xiaojing Ye, and Haomin Zhou. 'Parameterized wasserstein hamiltonian flow'. In: SIAM Journal on Numerical Analysis 63.1 (2025), pp. 360-395.

Liu Yang, Constantinos Daskalakis, and George E Karniadakis. 'Generative ensemble regression: Learning particle dynamics from observations of ensembles with physics-informed deep generative models'. In: SIAM Journal on Scientific Computing 44.1 (2022), B80-B99.

Maosheng Yang. 'Topological Schrödinger Bridge Matching'. In: The Thirteenth International Conference on Learning Representations . 2025.

Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. 'Generative modeling of singlecell time series with PRESCIENT enables prediction of cell trajectories with interventions'. In: Nature communications 12.1 (2021a), p. 3222.

Grace Hui Ting Yeo, Sachit D Saksena, and David K Gifford. 'Generative modeling of singlecell time series with PRESCIENT enables prediction of cell trajectories with interventions'. In: Nature communications 12.1 (2021b), p. 3222.

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

Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. 'scNODE: generative model for temporal single cell transcriptomic data prediction'. In: Bioinformatics 40.Supple-ment_2 (2024a), pp. ii146-ii154. ISSN: 1367-4811.

Peng Zhang, Ting Gao, Jin Guo, and Jinqiao Duan. 'Action Functional as Early Warning Indicator in the Space of Probability Measures'. In: arXiv preprint arXiv:2403.10405 (2024b).

Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. 'Optimal transport analysis reveals trajectories in steady-state systems'. In: PLoS computational biology 17.12 (2021), e1009466.

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.

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.

Zhenyi Zhang, Zihan Wang, Yuhao Sun, Tiejun Li, and Peijie Zhou. 'Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrödinger Bridge'. In: Advances in Neural Information Processing Systems (2025c).

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

Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. 'Denoising Diffusion Bridge Models'. In: The Twelfth International Conference on Learning Representations . 2024.

Qunxi Zhu, Bolin Zhao, Jingdong Zhang, Peiyang Li, and Wei Lin. 'Governing equation discovery of a complex system from snapshots'. In: arXiv preprint arXiv:2410.16694 (2024).

A Technical Details

A.1 Proof of Theorems

A.1.1 Proof for Theorem 4.1

Theorem A.1. The RUOT problem with isotropic and time-invariant diffusion intensity is formulated as:

In this problem, the necessary conditions for the action A to achieve a minimum are given by:

where λ ( x , t ) is a scalar field.

Proof. In order to incorporate the constraints of the Fokker-Planck equation, we construct an augmented action functional:

We take variations with respect to u , g , and ρ . 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 → u + δ u . The variation of the augmented action is

Here, S ∞ denotes the boundary at infinity in R d and d S is the surface element. Based on the assumption that

and using the arbitrariness of δ u , we obtain the optimality condition

Step 2: Variation with respect to g .

Let g → g + δg , then the variation of the augmented action becomes

Since δg is arbitrary, we immediately obtain the optimality condition

Step 3: Variation with respect to ρ .

Let ρ → ρ + δρ . Then the variation of the augmented action is given by

Since δρ is arbitrary, the corresponding optimality condition is

Substituting the previously obtained condition u = ∇ x λ , we arrive at the final optimality condition:

A.1.2 Proof for Theorem 4.2

Theorem A.2. The choice of ψ ( g ) affects whether g ascends or descends along the direction of the velocity field u at a given time. Specifically, at a fixed time t , if

then g ( x , t ) ascends in the direction of the velocity field u ( x , t ) (i.e., u ( x , t ) T ( ∇ x g ( x , t )) > 0 ); otherwise, it descends.

Proof. Let we have

Using the optimality condition for g from Appendix A.1.1,

taking the gradient with respect to x on both sides yields

The condition for g to increase along the velocity field is that the inner product between ∇ x g ( x , t ) and u ( x , t ) is positive everywhere. Using the optimality condition for the velocity,

Since α > 0 , the condition

is equivalent to requiring that

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 i is given by X t i ∈ R d and its weight by w i ( t ) > 0 . The dynamics of each particle are described by

where u : R d × [0 , T ] → R d is a time-varying vector field, g : R d × [0 , T ] → R is a growth rate function, σ : [0 , T ] → [0 , + ∞ ) is a time-varying diffusion coefficient, and W t is an N -dimensional standard Brownian motion with independent components in each coordinate. The initial conditions are X 0 i ∼ ρ ( x , 0) and w i (0) = 1 . In the limit as N →∞ , the empirical measure

converges to the solution of the following Fokker-Planck equation:

with the initial condition ρ ( x , 0) = ρ 0 ( x ) .

Proof. Consider a smooth test function ϕ : R d → R . We study the evolution of the expectation

By applying Itô's formula, we have

Using Itô's formula to compute d ϕ ( X t i ) , we obtain

Since d w i ( t ) = g ( X t i , t ) w i ( t ) d t contains no stochastic term (i.e., there is no d W ), the term d w i ( t ) d ϕ ( X t i ) is of higher order and can be neglected. Therefore, we have:

Next, we compute

Thus, in the limit as N →∞ , and let ρ ( x , t ) = µ ∞ ( x , t ) , we have

By integrating by parts, we obtain and

Hence, we deduce that

Since ϕ ( x ) is arbitrary, we obtain the Fokker-Planck equation:

A.1.4 Proposition : the Expectation of HJB Loss

Proposition A.1. Consider the following HJB loss:

where

The expectation of HJB loss is

where ˆ ρ ( x , t ) = ρ ( x , t ) ∫ d R ρ ( x , t )d x is the normalized probability density.

Proof. Taking the expectation of L N HJB is equivalent to drawing N particles each time to obtain L N HJB , repeating this process infinitely many times, and computing the average of the L N HJB obtained in each instance. Since the particles are independent, this operation is directly equivalent to taking the number of particles N →∞ , thus:

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:

where

The expectation of action loss is equal to the action defined in the RUOT formulation, namely,

Proof. Taking the expectation of L N Action is equivalent to drawing N particles each time to obtain L N Action , repeating this process infinitely many times, and computing the average of the L N Action obtained in each instance. Since the particles are independent, this operation is directly equivalent to taking the number of particles N →∞ , thus:

In the final equality of the proof, we employed the previously proven Appendix A.1.3.

A.2 Optimality Conditions Under Different ψ ( g )

In our Experiment, we use two different ψ ( g ) as examples. When ψ 1 ( g ) = 1 2 g 2 , the optimality conditions are :

When ψ 2 ( g ) = g 2 / 15 , the optimality conditions are :

Note that in this case the function g ( λ ) exhibits a singularity at λ = 0 . In fact, given the two properties we imposed on ψ ( g ) ( d ψ ( g ) d | g | > 0 and ψ ( g ) = ψ ( -g ) ) along with the constraint ψ ′′ ( g ) < 0 , it follows that ψ ′ ( g ) must be discontinuous at 0 , and hence g ( λ ) = ( ψ ′ ) -1 ( λ ( x ,t ) α ) necessarily has a singularity at λ = 0 . For the sake of training stability, we slightly modify g ( λ ) to remove this singularity. We redefine g ( λ ) as:

where δ is a small positive constant. In our computations, we set δ = 0 . 1 . Another possible treatment for this singularity involves fitting the reciprocal of λ , µ , thereby canceling the singularity at λ = 0 . In this case, g = ( 2 αµ 15 ) 15 13 , and thus the singularity near zero points no longer exists. We leave how to eliminate this singularity more systematically in numerical computation as the future work.

A.3 Training Algorithm and Computation of Losses

For the algorithms mentioned above, we use the particle method to compute the losses. In Appendix D.2, we discuss the convergence rates of these losses: both L HJB and L Action have a convergence rate of O ( 1 √ N ) , while for data with high dimensionality ( d ≫ 1 ), the convergence

Algorithm 1 Training Var-RUOT

Require: Datasets D 1 , . . . , D K , batch size N , training epochs N Epoch, initialized network λ θ ( x , t ) . Ensure: Trained scalar field λ θ (x , t ) .

1: for i = 1 to N Epoch do

2: From the data at the first time point, D 1 , sample N particles and set all their weights to w i (0) = 1 , for i = { 1 , 2 , · · · , N } .

rate of L OT is O ( N -2 d ) . To balance training speed, memory consumption, and the algorithm's convergence rate, we recommend using approximately N = 10 3 to 10 4 particles to compute L HJB, L Action, and L Recon. Furthermore, to ensure high accuracy in the integration for L HJB and L Action , we advise using a high-order integrator (Berman et al. 2024) (e.g., a Stochastic Runge-method (Li et al. 2020) ) and setting the integration step size to approximately ∆ t = 0 . 1 . L Mass measures the discrepancy between the mass learned by the model and the true mass. The true mass at time point i (for i > 1 ) is defined as the ratio of the number of data points at that time, N i , to the number of data points at the initial time, N 1 .

In particular, computing L HJB requires second-order derivatives to calculate the Laplacian ∇ 2 x λ ( x , t ) . This step can be computationally expensive. However, we only need the trace of 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 ( H ) = E v [ v T H v ] , where the vector v is drawn from a distribution with zero mean and unit covariance, such as N (0 , I ) .

B Experiential Details

B.1 Additional Information for Datasets

Simulation Dataset In the main text, we utilize a simulated dataset derived from a three-gene 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:

where X i ( t ) represents the gene expression levels of the i th cell at time t . The coefficients α i , γ i , and β control the strengths of self-activation, inhibition, and the external stimulus, respectively.

The parameters δ i indicate the rates of gene degradation, and the terms η i ξ t account for stochastic influences using additive white noise.

The probability of cell division is linked to the expression level of X 2 and is given by

When a cell divides, the resulting daughter cells are created with each gene perturbed by an independent random noise term, η d N (0 , 1) , around the parent cell's gene expression profile ( X 1 ( t ) , X 2 ( t ) , X 3 ( t )) . Detailed hyper-parameters are provided in Table 5. The initial population of cells is independently sampled from two normal distributions, N ([2 , 0 . 2 , 0] , 0 . 1) and N ([0 , 0 , 2] , 0 . 1) . At every time step, any negative expression values are set to zero.

Other Datasets Used in Main Text In addition to the three-gene 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 TGFB1-induced epithelial-mesenchymal transition (EMT). It comprises data from four distinct time points, containing a 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 et al. 2020; Sha et al. 2024) and was reduced to 2-dimensional space using nonlinear dimensionality reduction .

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 high-dimensional data), we constructed two 2D Gaussian datasets (a balanced version and an unbalanced version) and three high-dimensional Gaussian datasets.The 2D Gaussian dataset consists of two time points, where the distribution at each time point is a mixture of multiple Gaussian distributions, as illustrated in 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 DeepRUOT (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,302 cells collected at three time points. We reduced its dimensionality to 50 using PCA, and a subset of this processed data was used for the experiments in our main text. The second is the Pancreatic β -cell Differentiation dataset from (Veres et al. 2019), consisting of 51,274 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

(c) 150dim Gaussian Mixture

(b) 50dim, 100dim Gaussian Mixture

(a) 2dim Gaussian Mixture

t=0.00

t*1.00

1=1.00

(=0.00

t1.00

t=0.00

0.5

0.8

-0.5

PC1

PC1

0.0

0.5

PC1

3D culture throughout a 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 et al. 2021). From this, we utilized the gene expression data for the donor with DonorID 28483. We extracted the 1,079 cells from this donor that were annotated with pseudotime values. To create discrete time points, we performed K-Means clustering on these cells based on their pseudotime. The resulting clusters were then ordered by their centroid pseudotime values and assigned discrete labels t = 0 , 1 , 2 , 3 , 4 . This data was also standardized on a per-feature basis.

B.2 Evaluation Metrics

To assess the fitting accuracy of the learned dynamics to the data distribution, we compute the W 1 and W 2 distances between the data points generated by the model and the real data points. They are defined as and

We compute these two metrics using the emd function from the pot library.

To evaluate the action of 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:

where

with initial conditions X 0 i ∼ ρ ( x , 0) and w i (0) = 1 . We run our model 5 times on each dataset, to calculate the mean and standard deviation of W 1 , W 2 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 0 . 10 . 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 a 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 σ directly affects the results (both the OT loss and the path action). Therefore, when running the RUOT and our Var-RUOT models on 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 on the EMT dataset, our algorithm converges roughly 20 times faster than the fastest alternative (TIGON).

(b) Training Process on Dataset EMT

Var-RUOT

TIGON

RUOT with

3 Epoch Pretraining

RUOT w/o Pretraining

About 20x Faster

Var-RUOT

RUOT with

3 Epoch Pretraining

TIGON

RUOT w/o Pretraining

About 10x Faster

(a) Training Process on Dataset Simulation Gene

10-1

10

100

10-1

10-1

+

0

100

200

300

400

500

0

100

200

300

400

500

Step

Step

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 λ ( x , t ) is a fully connected network augmented with layer normalization and residual connections. It consists of 2 hidden layers, each with 512 dimensions. In our algorithm, the main hyperparameters that need tuning include the penalty coefficient α for growth in the action, and the weights γ HJB and γ Action for the two regularization losses, L HJB and L Action, respectively. Here, α represents our prior regarding the strength of cell birth and death in the data: a larger α imposes a greater penalty on cell birth and death, thereby making it easier for the model to learn solutions with lower birth and death intensities. Meanwhile, the HJB loss and the action loss, serving as regularizers, are both designed to ensure that the solution obtained by the algorithm has as low an action as possible-the 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 ( ψ ( g ) = 1 2 g 2 ) and one for the Modified Metric ( ψ ( g ) = g 2 / 15 ). The parameters used in each case are listed in Table 6. The primary reason for using two sets is that different metrics yield different scales for the HJB loss.

Sensitivity Analysis of α To demonstrate the robustness of our algorithm with respect to hyperparameter selection, we first varied the penalty coefficient α for growth and examined the resulting

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 α is shown in Table 7. The experimental results indicate that our algorithm is not sensitive to α , as similar performance can be achieved across multiple different values of α . In comparison with the Standard WFR metric, however, the algorithm appears to be somewhat more sensitive to α when the modified metric is used.

Table 7: On the Mouse Blood Hematopoiesis dataset, the Wasserstein distances ( W 1 and W 2 ) between the predicted distributions of Var-RUOT with different α and the true distribution at each time points. Each experiment was run five times to compute the mean and standard deviation.

Sensitivity Analysis of Growth Penalty Function

In this work, we select ψ 2 ( g ) = g 2 / 15 to illustrate a case where the conditions ψ ′ ( g ) > 0 and ψ ′′ ( g ) < 0 hold. As established in Theorem 4.2, these conditions ensure that the growth rate g ( x , t ) increases along the direction of the velocity field u ( x , t ) at each time step.

The specific choice of ψ 2 ( g ) = g 2 / 15 is motivated by its simplicity and, more importantly, by considerations for training stability. For a general penalty of the form ψ ( g ) = g (2 p ) / (2 q +1) , the first-order optimality condition α d ψ ( g ) d g = λ yields an analytical solution for g in terms of λ :

Since g is computed directly from λ during training, we must avoid an excessively steep g ( λ ) curve, especially near λ = 0 , to ensure stability. This requires the exponent 2 q +1 (2 q +1) -2 p to be small. We achieve this by setting p = 1 (the smallest positive integer) and choosing a large value for q . Selecting q = 7 results in our final penalty function, ψ 2 ( g ) = g 2 / 15 , which balances simplicity and stability.

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 8, compare the Wasserstein distances ( W 1 and W 2 ) at different time points.

The results demonstrate that ψ ( g ) = g 2 / 15 achieves the best performance. Its performance is comparable to that of ψ ( g ) = g 2 / 21 , whereas the choice of ψ ( g ) = g 2 / 9 (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 γ HJB and γ Action In order to verify whether L HJB and L Action 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

simple dynamics (each particle moving in a 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 γ HJB to the following values: [0 , 6 . 25 × 10 -3 , 3 . 125 × 10 -2 , 6 . 25 × 10 -2 , 6 . 25 × 10 -1 , 3 . 125] , while keeping the action loss weight γ Action fixed at 1 . We then shown both the mean W 1 distances between the predicted and true distributions at four different time points and the trajectory action (as shown in Table 9 ). Similarly, we fixed γ HJB = 1 and varied λ Action over the same set of values: [0 , 6 . 25 × 10 -3 , 3 . 125 × 10 -2 , 6 . 25 × 10 -2 , 6 . 25 × 10 -1 , 3 . 125] , with the corresponding results illustrated in Table 10. The figures indicate that as γ HJB and γ Action increase, the action of the learned trajectories decreases monotonically, demonstrating that both loss terms are effective. However, as these weights increase, the model's ability to fit the distribution deteriorates. Therefore, we recommend that in practical applications, both γ HJB and γ Action should be set to values below 0 . 1 , as configured in this paper.

The Effectiveness of different types of Loss Norms

In the experiments above, we used the L 2 loss for computing both L HJB and L Action. To investigate the effect of the loss function type, we present the results of using the L 1 loss. On the Simulation and EMT datasets, the experimental results with the L 1 loss are shown in Tables 11 and 12. The results indicate that with our hyperparameter values ( γ HJB = 6 . 25 × 10 -2 , γ Action = 6 . 25 × 10 -2 ), using either the L 1 loss or the L 2 loss when computing the HJB loss achieves similar reconstruction performance and Path Action values.

C.3 Hold-one-out Experiment

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 1 and W 2 distances between the predicted distribution and the true distribution at the missing time point. When a time point from { 1 , 2 , · · · , n -1 } is removed, the model is performing interpolation; when the time point n 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 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 ( t = 0 , 1 , 2 , 3 , 4 ) were designated for training, while the final four ( t = 5 , 6 , 7 , 8 ) 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 DeepRUOT, 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-autonomous system, where the parameter λ is a function of both the state x and time t . Consequently, the model learns the mapping λ ( 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.e., It is more capable of learning dynamics that conform to 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 a form of regularization that enhances the model'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. 7 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 = 0 , 1 , 2 . When one time point is removed, Var-RUOT tends to favor a straight-line trajectory connecting the remaining two time points (since such a trajectory represents the minimum-action path), which serves 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.

w/o t=1

w/o t=3

w/o t=4

w/o t=2

w/o t=3

w/o t=4

weight

t =0.00

t =0.00

t =0.00

t=0.00

4.5

t =1.00

t = 2.00

t = 1.00

t = 2.00

t =1.00

t =2.00

t =1.00

t = 2.00

4.0

t =3.00

t =3.00

t =3.00

t =3.00

t = 4.00

t =4.00

t = 4.00

t = 4.00

3.5

3.0

2.5

2.0

1.5

1.0

w/o t=1

w/o t=3

w/o t=4

w/o t=2

w/o t=3

w/o t=4

growth

0.4

E°0 -

0.1

0.0

w/o t=2

w/o t=1

t=0

t=0

t：

=

1

t=

1

t

=2

t

=

2

C.4 Experiments on High Dimensional Dataset

High Dimensional Gaussian Dataset To evaluate the effectiveness of our method on highdimensional 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 ( ψ ( g ) = 1 2 g 2 ) as well as a modified growth penalty function, ψ ( g ) = g 2 / 15 , which ensures ψ ′′ ( g ) < 0 . The learned trajectories and growth rates are illustrated in Fig. 11. Under both choices of ψ ( 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.

50D Mouse Blood Hematopoiesis and Pancreatic β -cell Differentiation Dataset We tested our method on two high-dimensional real scRNA-seq datasets including 50D Mouse Blood Hematopoiesis dataset and Pancreatic β -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 u ( x , t ) 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

w/o t=2

w/o t=1

t=0.00

t=0.00

t =1.00

t =1.00

Standard WFR

50 Dim

100 Dim

t=0.00

t = 0.00

12.5

t = 1.00

t = 1.00

10.0

.5

.0

2.5

0'0

S'2-

5.0

t = 0.00

t = 0.00

t = 1.00

t = 1.00

velocity field points from points with smaller t to those with larger t , which indicates that our model can correctly learn a vector field that transfers the distribution even in high-dimensional data.

EB Dataset and NeurIPS Challenge Dataset

To further quantitatively evaluate the performance of VarRUOT on high-dimensional data, we employed the 100-dimensional EB dataset and the 50-dimensional NeurIPS Challenge dataset. The

Vector Field

Growth

8

0.0

1.0

2.0

Vector Field

Growth

0.0

1.0

2.0

3.0

4.0

5.0

6.0

7.0

quantitative results for VarRUOT against other algorithms are 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-free approaches (OTCFM, UOTCFM) and Wasserstein Gradient Flow. For the NeurIPS Challenge dataset, Var-RUOT outperforms 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.

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 a 2D plane. We then compared the action computed by the Var-RUOT method (with σ = 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-source solver for dynamic unbalanced optimal transport (apart from deep learning-based methods like TIGON ). Therefore, we adopted a static Wasserstein-Fisher-Rao (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.

The results show that for the two normalized datasets, Var-RUOT finds Path Action values comparable to those computed by the pot library. On the two unbalanced datasets, the Path Action found by Var-RUOT 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.

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.

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 ρ ( x , t ) . It enforces physical constraints as soft penalties in the loss function. This complex setup necessitates a multi-stage pre-training procedure to ensure convergence.

In contrast, Var-RUOT uses a single network to parameterize a scalar field λ θ ( x , t ) . Both the velocity u ( x , t ) and growth rate g ( x , t ) are analytically derived from λ θ ( x , t ) via first-order optimality conditions (Theorem 4.1). This simpler, end-to-end design eliminates the need for pre-training and improves training stability.

Difference from Action Matching

The differences between Var-RUOT and Action Matching could be summarized from three perspectives. Firstly, Var-RUOT handles both unbalanced distributions and stochastic dynamics, which is crucial for single-cell omics.

Secondly, Action Matching's training requires sampling from intermediate distributions q t ( x ) . These are typically challenging to obtain in single-cell omics, where data exists only as 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:

On the other hand, Var-RUOT directly minimizes the residual of the Hamilton-Jacobi-Bellman (HJB) equation. By directly enforcing this first-order optimality condition in its loss, Var-RUOT 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: L HJB, L Action, and L Recon. Among these, L HJB and L Action are computed using the Monte Carlo method. This involves sampling a set of particles,

integrating along their trajectories, and then averaging the results over all particles. The convergence rate of the Monte Carlo method is O ( 1 √ N ) , where N is the number of particles used.

The reconstruction loss, L Recon, is composed of two parts: L OT and L Mass. L OT, calculated using the geomloss library, represents the W 2 distance between the normalized empirical distribution and the true distribution. According to the triangle inequality, W 2 ( µ N , ν ) ≤ W 2 ( µ N , µ ) + W 2 ( µ, ν ) , the convergence of L OT as the number of particles increases depends on the rate at which the empirical measure µ N = 1 N ∑ N i =1 δ ( x -x i ) (with x i ∼ µ ) approximates the true measure µ . The expected error is bounded as follows (Fournier and Guillin 2013):

E Limitations and Broader Impacts

E.1 Limitations

The algorithm presented in this paper offers new insights for solving the RUOT problem; however, it still has several limitations. Firstly, although Var-RUOT parameterizes u and g with a single neural network, and designs the loss function based on the necessary conditions for the minimal action solution, since neural network optimization only finds local minima, there is still no guarantee that the solution found is indeed the one with minimal action. Moreover, in practice, the HJB loss does not necessarily converge to zero (A typical training curve for the HJB loss is shown in Fig. 14). Thus, it effectively acts as a regularization term. This could be addressed by conducting a more detailed analysis on simpler versions of the RUOT problem (for instance, Gaussian-to-Gaussian transport or Dirac-to-Dirac Transport).

Furthermore, when using the modified metric, the goodness-of-fit 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 a controllability issue in 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 ψ with a neural network or derive it 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.45

0.40

0.35

0.30

0.20

0.15

0.10

0.05

0

25

50

75

100

125

150

175

200

Epoch

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 a 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 to 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. In 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 on 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 a substantial impact on the model's outcomes, potentially affecting important biological priors. These factors could present challenges for subsequent biological analyses or clinical decision-making, and care must be taken in the use and dissemination of the model-generated interpolation results to avoid data contamination.

When applying our method in biological or medical contexts, it is crucial to 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.

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.1 provides the proof of the important theorem of the claim; our numerical results are presented in Section 6.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss our limitations in Section 7 and Appendix E.1.

Guidelines:

The answer NA means that the paper has no 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 (e.g., independence assumptions, noiseless settings, model well-specification, 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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

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

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs 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 to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide our training procedure in Appendix A.3, and we provide our code in a public repository https://github.com/ZerooVector/VarRUOT .

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken 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, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or 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., 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 all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification:We use public datasets for our experiment, and we provide our code in a public repository https://github.com/ZerooVector/VarRUOT .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide our experimental details in Appendix B 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 in the core of the paper to a 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, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide the mean and standard deviation values of our metrics in each experiment.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have all reviewed and confirmed this research conforms with the NeurIPS Code of Ethics. We have striven to maintain and preserve anonymity.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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.

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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 are publicly available, and we provide details of datasets in Appendix B.1.

Guidelines:

The answer NA means that the paper does not use 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 is used and, if 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as 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 in the supplemental material is fine, but if the main contribution 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: 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, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The method development does not involve LLMs as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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).\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe limitations of the work have been sufficiently addressed in the Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- The Limitations section is pushed to the appendix, which in my opinion is not in the spirit of the guidelines and neurips checklist.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS AND WEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nPlease refer to Strength and Weakness section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}