Original: N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE : D IFFUSION M ODELING FOR P OPULATION D YNAMICS
Revised: N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE : D IFFUSION M ODELING FOR P OPULATION D YNAMICS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Population dynamics is the study of temporal and spatial variation in the size of populations of organisms and is a major part of population ecology.
Revised: Population dynamics is the study of temporal and spatial variation in the size of populations of organisms and is a major part of population ecology.

Original: One of the main difficulties in analyzing population dynamics is that we can only obtain observation data with coarse time intervals from fixed-point observations due to experimental costs or measurement constraints.
Revised: One of the main difficulties in analyzing population dynamics is that we can only obtain observation data with coarse time intervals from fixed-point observations due to experimental costs or measurement constraints.

Original: Recently, modeling population dynamics by using continuous normalizing flows (CNFs) and dynamic optimal transport has been proposed to infer the sample trajectories from a fixed-point observed population.
Revised: Recently, modeling population dynamics by using continuous normalizing flows (CNFs) and dynamic optimal transport has been proposed to infer the sample trajectories from a fixed-point observed population.

Original: While the sample behavior in CNFs is deterministic, the actual sample in biological systems moves in an essentially random yet directional manner.
Revised: While the sample behavior in CNFs is deterministic, the actual sample in biological systems moves in an essentially random yet directional manner.

Original: Moreover, when a sample moves from point A to point B in dynamical systems, its trajectory typically follows the principle of least action in which the corresponding action has the smallest possible value.
Revised: Moreover, when a sample moves from point A to point B in dynamical systems, its trajectory typically follows the principle of least action in which the corresponding action has the smallest possible value.

Original: To satisfy these requirements of the sample trajectories, we formulate the Lagrangian Schr ¨ odinger bridge (LSB) problem and propose to solve it approximately by modeling the advection-diffusion process with regularized neural SDE.
Revised: To satisfy these requirements of the sample trajectories, we formulate the Lagrangian Schr ¨ odinger bridge (LSB) problem and propose to solve it approximately by modeling the advection-diffusion process with regularized neural SDE.

Original: We also develop a model architecture that enables faster computation of the loss function.
Revised: We also develop a model architecture that enables faster computation of the loss function.

Original: Experimental results show that the proposed method can efficiently approximate the population-level dynamics even for high-dimensional data and that using the prior knowledge introduced by the Lagrangian enables us to estimate the sample-level dynamics with stochastic behavior.
Revised: Experimental results show that the proposed method can efficiently approximate the population-level dynamics even for high-dimensional data and that using the prior knowledge introduced by the Lagrangian enables us to estimate the sample-level dynamics with stochastic behavior.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The population dynamics of time-evolving individuals appears in various scientific fields, such as cell population in biology (Schiebinger et al., 2019;
Revised: The population dynamics of time-evolving individuals appears in various scientific fields, such as cell population in biology (Schiebinger et al., 2019;

Original: Yang & Uhler, 2018), air in meteorology (Fisher et al., 2009), and healthcare statistics (Manton et al., 2008) in medicine.
Revised: Yang & Uhler, 2018), air in meteorology (Fisher et al., 2009), and healthcare statistics (Manton et al., 2008) in medicine.

Original: However, tracking individuals over a long period is often difficult due to experimental costs.
Revised: However, tracking individuals over a long period is often difficult due to experimental costs.

Original: Furthermore, it can sometimes be impossible to track the time evolution.
Revised: Furthermore, it can sometimes be impossible to track the time evolution.

Original: For example, since single-cell RNA sequencing (scRNA-seq) destroys all measured cells, we cannot analyze the behavior of individual cells over time in cell transcriptome measurements.
Revised: For example, since single-cell RNA sequencing (scRNA-seq) destroys all measured cells, we cannot analyze the behavior of individual cells over time in cell transcriptome measurements.

Original: Instead, we only obtain individual samples from crosssectional populations without alignment across time steps at a few distinct time points.
Revised: Instead, we only obtain individual samples from crosssectional populations without alignment across time steps at a few distinct time points.

Original: Under these constraints on data measurements, our goal is to better understand the time evolution of samples in the populations.
Revised: Under these constraints on data measurements, our goal is to better understand the time evolution of samples in the populations.

Original: Existing methods attempt to estimate population-level dynamics following the Wasserstein gradient flow using a recurrent neural network (RNN)
Revised: Existing methods attempt to estimate population-level dynamics following the Wasserstein gradient flow using a recurrent neural network (RNN)

Original: (Hashimoto et al., 2016) or the Jordan-Kinderlehrer-Otto (JKO) flow (Bunne et al., 2021).
Revised: (Hashimoto et al., 2016) or the Jordan-Kinderlehrer-Otto (JKO) flow (Bunne et al., 2021).

Original: Recent studies have attempted to interpolate the trajectories of individual samples between cross-sectional populations at multiple time points by using optimal transport (OT) (Schiebinger et al., 2019; Yang & Uhler, 2018), or CNF (Tong et al., 2020).
Revised: Recent studies have attempted to interpolate the trajectories of individual samples between cross-sectional populations at multiple time points by using optimal transport (OT) (Schiebinger et al., 2019; Yang & Uhler, 2018), or CNF (Tong et al., 2020).

Original: Using a
Revised: Using a

Original: CNF generates continuous-time non-linear sample trajectories from multiple time points.
Revised: CNF generates continuous-time non-linear sample trajectories from multiple time points.

Original: In addition,
Revised: In addition,

Original: Tong et al.
Revised: Tong et al.

Original: (2020) proposed a regularization for CNF that encourages a straight trajectory on the basis of the OT theory.
Revised: (2020) proposed a regularization for CNF that encourages a straight trajectory on the basis of the OT theory.

Original: Since the probability distribution transformation based on ordinary differential equations (ODEs) is used in CNF, the behavior of each sample is described by its initial condition in a completely deterministic manner.
Revised: Since the probability distribution transformation based on ordinary differential equations (ODEs) is used in CNF, the behavior of each sample is described by its initial condition in a completely deterministic manner.

Original: However, samples in population are known to move stochastically and diffuse in nature, e.g ., biological system (Kolomgorov et al., 1937).
Revised: However, samples in population are known to move stochastically and diffuse in nature, e.g ., biological system (Kolomgorov et al., 1937).

Original: To handle the stochastic and complex behavior of individual samples, we propose to model the advection-diffusion processes by using SDEs to describe the time evolution of the sample.
Revised: To handle the stochastic and complex behavior of individual samples, we propose to model the advection-diffusion processes by using SDEs to describe the time evolution of the sample.

Original: Furthermore, on the basis of the principle of least action , we estimate the sample trajectories that minimize action , defined by the time integral of the Lagrangian determined from the prior knowledge.
Revised: Furthermore, on the basis of the principle of least action , we estimate the sample trajectories that minimize action , defined by the time integral of the Lagrangian determined from the prior knowledge.

Original: We formulate this problem as the Lagrangian Schr¨odinger bridge (LSB) problem, which is a special case of the stochastic optimal transport (SOT) problem, and propose an approximate solution method neural Lagrangian Schr¨odinger bridge (NLSB).
Revised: We formulate this problem as the Lagrangian Schr¨odinger bridge (LSB) problem, which is a special case of the stochastic optimal transport (SOT) problem, and propose an approximate solution method neural Lagrangian Schr¨odinger bridge (NLSB).

Original: In NLSB, we train regularized neural SDE (Li et al., 2020; Tzen & Raginsky, 2019a;b) by minimizing the Wasserstein loss between the ground-truth and the predicted population.
Revised: In NLSB, we train regularized neural SDE (Li et al., 2020; Tzen & Raginsky, 2019a;b) by minimizing the Wasserstein loss between the ground-truth and the predicted population.

Original: The Lagrangian design defining regularization allows the sample-level dynamics to reflect various prior knowledge such as OT, manifold geometry, and local velocity arrows proposed by Tong et al.
Revised: The Lagrangian design defining regularization allows the sample-level dynamics to reflect various prior knowledge such as OT, manifold geometry, and local velocity arrows proposed by Tong et al.

Original: In addition, we parameterize a potential function instead of the drift function.
Revised: In addition, we parameterize a potential function instead of the drift function.

Original: Adopting the model architecture of the potential function from OT-Flow (Onken et al., 2021) will speed up the computation of the potential function’s gradient and the regularization term.
Revised: Adopting the model architecture of the potential function from OT-Flow (Onken et al., 2021) will speed up the computation of the potential function’s gradient and the regularization term.

Original: As a result, we capture the population-level dynamics as well as or better than conventional methods, and can more accurately predict the sample trajectories.
Revised: As a result, we capture the population-level dynamics as well as or better than conventional methods, and can more accurately predict the sample trajectories.

Original: In short, our contributions are summarized as follows.
Revised: In short, our contributions are summarized as follows.

Original: We formulate the LSB problem to estimate the stochastic sample trajectory according to the principle of least action .
Revised: We formulate the LSB problem to estimate the stochastic sample trajectory according to the principle of least action .

Original: We propose NLSB to approximate the LSB problem practically by modeling the advectiondiffusion process with regularized neural SDE on the basis of the prior knowledge introduced by the Lagrangian.
Revised: We propose NLSB to approximate the LSB problem practically by modeling the advectiondiffusion process with regularized neural SDE on the basis of the prior knowledge introduced by the Lagrangian.

Original: Our model architecture of the potential function enables its gradient and diagonal components of the Hessian to be computed explicitly, resulting in fast computation.
Revised: Our model architecture of the potential function enables its gradient and diagonal components of the Hessian to be computed explicitly, resulting in fast computation.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In Section 2.1, we first introduce the method of combining CNF and the OT theory, the basis of the proposed method.
Revised: In Section 2.1, we first introduce the method of combining CNF and the OT theory, the basis of the proposed method.

Original: Next, we explain our newly introduced theories and techniques for modeling population dynamics: the diffusion modeling using neural SDE (Section 2.2) and the SOT theory (Section 2.3).
Revised: Next, we explain our newly introduced theories and techniques for modeling population dynamics: the diffusion modeling using neural SDE (Section 2.2) and the SOT theory (Section 2.3).

Original: 2. F LOWS R EGULARIZED BY O PTIMAL T RANSPORT
Revised: 2. F LOWS R EGULARIZED BY O PTIMAL T RANSPORT

Original: CNFs (Chen et al., 2018) are a method for learning the continuous transformation between two distributions p and q by modeling the ordinary differential equation (ODE):
Revised: CNFs (Chen et al., 2018) are a method for learning the continuous transformation between two distributions p and q by modeling the ordinary differential equation (ODE):

Original: [Equation]
Revised: [Equation]

Original: subject to x ( t
Revised: subject to x ( t

Original: [Equation]
Revised: [Equation]

Original: ∼ p
Revised: ∼ p

Original: [Equation]
Revised: [Equation]

Original:  where v θ is the velocity model with learnable parameters θ .
Revised: , where v θ is the velocity model with learnable parameters θ .

Original: Several regularizations of CNFs leading to straight trajectories have been proposed on the basis of the OT theory.
Revised: Several regularizations of CNFs leading to straight trajectories have been proposed on the basis of the OT theory.

Original: The likelihood maximization problem of regularized CNF is derived by replacing the terminal constraint of the Brenier-Benamou formulation (Benamou & Brenier, 2000) with
Revised: The likelihood maximization problem of regularized CNF is derived by replacing the terminal constraint of the Brenier-Benamou formulation (Benamou & Brenier, 2000) with

Original: Kullback–Leibler (KL) divergence.
Revised: Kullback–Leibler (KL) divergence.

Original: RNODE (Finlay et al., 2020), OT-Flow (Onken et al., 2021), and
Revised: RNODE (Finlay et al., 2020), OT-Flow (Onken et al., 2021), and

Original: TrajectoryNet (Tong et al., 2020) introduced a regularization ˜ R e in Eq.
Revised: TrajectoryNet (Tong et al., 2020) introduced a regularization ˜ R e in Eq.

Original: Potential Flow (Yang &
Revised: Potential Flow (Yang &

Original: Karniadakis, 2020) and OT-Flow modeled the potential function Φ that satisfies v = −∇ Φ instead of modeling the velocity function v .
Revised: Karniadakis, 2020) and OT-Flow modeled the potential function Φ that satisfies v = −∇ Φ instead of modeling the velocity function v .

Original: They also proposed an additional OT-based regularization ˜ R h derived from the Hamilton–Jacobi–Bellman (HJB) equation (Evans, 1983) satisfied by the potential function shown in Eq.
Revised: They also proposed an additional OT-based regularization ˜ R h derived from the Hamilton–Jacobi–Bellman (HJB) equation (Evans, 1983) satisfied by the potential function shown in Eq.

Original: These OT-based regularizations have resulted in faster CNFs (Finlay et al., 2020; Onken et al., 2021) and improved the modeling of cellular dynamics (Tong et al., 2020).
Revised: These OT-based regularizations have resulted in faster CNFs (Finlay et al., 2020; Onken et al., 2021) and improved the modeling of cellular dynamics (Tong et al., 2020).

Original: [Equation]
Revised: [Equation]

Original: whereΩ is the bounded domain and ρ t is the law of the sample x at the time t .
Revised: where ρ t is the law of the sample x at the time t .

Original: 2. N EURAL SDE AND A PPLICATIONS TO D IFFUSION M ODELING
Revised: 2. N EURAL SDE AND A PPLICATIONS TO D IFFUSION M ODELING

Original: SDE has been used to model real-world random phenomena in a wide range of areas, such as chemistry, biology, mechanics, economics and finance (Higham, 2001).
Revised: SDE has been used to model real-world random phenomena in a wide range of areas, such as chemistry, biology, mechanics, economics and finance (Higham, 2001).

Original: As an extension of neural
Revised: As an extension of neural

Original: ODEs, neural SDE has been proposed to model drift and diffusion functions with neural networks (NNs), as follows.
Revised: ODEs, neural SDE has been proposed to model drift and diffusion functions with neural networks (NNs), as follows.

Original: [Equation]
Revised: [Equation]

Original: where { X t } t ∈ [ t 0 ,t K − 1 ] is a continuous R d -valued stochastic process, f θ : R d × [ t 0 , t K − 1 ] (cid:55)→ R d is a drift function, g ϕ : R d × [ t 0 , t K − 1 ] (cid:55)→ R d × m is a diffusion function and { W t } t ∈ [ t 0 ,t K − 1 ] is an m -dimensional Wiener process.
Revised: where { X t } t ∈ [ t 0 ,t K − 1 ] is a continuous R d -valued stochastic process, f θ : R d × [ t 0 , t K − 1 ] (cid:55)→ R d is a drift function, g ϕ : R d × [ t 0 , t K − 1 ] (cid:55)→ R d × m is a diffusion function and { W t } t ∈ [ t 0 ,t K − 1 ] is an m -dimensional Wiener process.

Original: Score-based generative models (SGM) (Song et al., 2020), which use score matching to learn the reverse diffusion process of generating images from noise as SDE, have demonstrated the ability to produce high-quality data.
Revised: Score-based generative models (SGM) (Song et al., 2020), which use score matching to learn the reverse diffusion process of generating images from noise as SDE, have demonstrated the ability to produce high-quality data.

Original: Other recent studies on diffusion modeling using SDE (De Bortoli et al., 2021;
Revised: Other recent studies on diffusion modeling using SDE (De Bortoli et al., 2021;

Original: Wang et al., 2021; Vargas et al., 2021; Chen et al., 2021a) have proposed learning SDE solutions to the SB problem.
Revised: Wang et al., 2021; Vargas et al., 2021; Chen et al., 2021a) have proposed learning SDE solutions to the SB problem.

Original: Vargas et al.
Revised: Vargas et al.

Original: and De Bortoli et al.
Revised: and De Bortoli et al.

Original: (2021) solved the SB problem by combining iterative proportional fitting (IPF) with mean-matching regression of the SDE drift function using Gaussian process (GP) and NN, respectively.
Revised: (2021) solved the SB problem by combining iterative proportional fitting (IPF) with mean-matching regression of the SDE drift function using Gaussian process (GP) and NN, respectively.

Original: 2. S TOCHASTIC O PTIMAL T RANSPORT
Revised: 2. S TOCHASTIC O PTIMAL T RANSPORT

Original: Mikami (2008) generalized the OT problem and defined the SOT problem as a random mechanics problem determined by the principle of least action .
Revised: Mikami (2008) generalized the OT problem and defined the SOT problem as a random mechanics problem determined by the principle of least action .

Original: The SOT problem with the endpoint marginals fixed to µ 0 and µ 1 is represented as
Revised: The SOT problem with the endpoint marginals fixed to µ 0 and µ 1 is represented as

Original: [Equation]
Revised: [Equation]

Original: where L ( t, x , u ) is continuous and convex in u and A is the set of all R d -valued, continuous semimartingales { X t } t 0 ≤ t ≤ t 1 on a complete filtered probability space such that there exists a Borel measurable drift function β X ( X t , t ) for which satisfies several conditions (see (Mikami, 2021)).
Revised: where L ( t, x , u ) is continuous and convex in u and A is the set of all R d -valued, continuous semimartingales { X t } t 0 ≤ t ≤ t 1 on a complete filtered probability space such that there exists a Borel measurable drift function β X ( X t , t ) for which satisfies several conditions (see Appendix A.3).

Original: The function L , called the Lagrangian, is the transport cost defined on the space-time of the system and allows us to describe phenomena consistently regardless of the choice of coordinate system.
Revised: The function L , called the Lagrangian, is the transport cost defined on the space-time of the system and allows us to describe phenomena consistently regardless of the choice of coordinate system.

Original: Mikami (2008) also introduced another variational version of the SOT problem for a flow of marginal distributions which satisfies the Fokker–Planck (FP) equation, which is a advection-diffusion equation describing the time evolution of the probability density function.
Revised: Mikami (2008) also introduced another variational version of the SOT problem for a flow of marginal distributions which satisfies the Fokker–Planck (FP) equation, which is a advection-diffusion equation describing the time evolution of the probability density function.

Original: [Equation]
Revised: [Equation]

Original: satisfies the FP eq
Revised: satisfies the [Equation] eq

Original: [Equation]
Revised: [Equation]

Original: where ρ t and p t are the law and density of the random variable X t , respectively.
Revised: where ρ t and p t are the law and density of the random variable X t , respectively.

Original: Proposition 1 from (Mikami, 2021) states the optimal values of Eqs.
Revised: Proposition 1 from (Mikami, 2021) states the optimal values of Eqs.

Original: (4) and (5) match, i.e .
Revised: (4) and (5) match, i.e .

Original: V ( µ 0 , µ 1 ) = v( µ 0 , µ 1 ) .
Revised: V ( µ 0 , µ 1 ) = v( µ 0 , µ 1 ) .

Original: When L ( t, x , u ) = 12 || u || 2 , the problem 4 and 5 are regarded as the special case of Schr¨odinger bridge (SB) (Jamison, 1975) problem. 
Revised: When L ( t, x , u ) = 12 || u || 2 , these two SOT problems are regarded as the special case of Schr¨odinger bridge (SB) (Jamison, 1975; Chen et al., 2021b; Dai Pra, 1991; L´eonard, 2013; Mikami, 1990) problem. See Appendix A for details on the SOT theory.

Original:  P ROPOSED M ETHOD
Revised:  P ROPOSED M ETHOD

Original: In this section, we first introduce Lagrangian-Schr¨odinger bridge (LSB) problem.
Revised: In this section, we first introduce Lagrangian-Schr¨odinger bridge (LSB) problem.

Original: We next present an approximate solution for the LSB problem: neural Lagrangian Schr¨odinger bridge (NLSB).
Revised: We next present an approximate solution for the LSB problem: neural Lagrangian Schr¨odinger bridge (NLSB).

Original: Finally, we give specific examples of our method for specific use cases.
Revised: Finally, we give specific examples of our method for specific use cases.

Original: Figure 2 illustrates the position of the
Revised: Figure 2 illustrates the position of the

Original: LSB problem in the transport problem and the relationship between the NLSB and existing methods.
Revised: LSB problem in the transport problem and the relationship between the NLSB and existing methods.

Original: 3. L AGRANGIAN S CHR ¨ ODINGER B RIDGE
Revised: 3. L AGRANGIAN S CHR ¨ ODINGER B RIDGE

Original: We consider the problem of estimating dynamics at both the population and individual sample levels by using probability densities at two known end points.
Revised: We consider the problem of estimating dynamics at both the population and individual sample levels by using probability densities at two known end points.

Original: To solve this problem, we make two assumptions about the target system model.
Revised: To solve this problem, we make two assumptions about the target system model.

Original: The stochastic behavior of individual samples yields the population diffusion phenomenon.
Revised: The stochastic behavior of individual samples yields the population diffusion phenomenon.

Original: Individual samples are encouraged to move according to the principle of least action .
Revised: Individual samples are encouraged to move according to the principle of least action .

Original: The principle of least action states that when a sample moves from point A to point B, its trajectory is the one that have least action, which is known as a fundamental principle in dynamical systems.
Revised: The principle of least action states that when a sample moves from point A to point B, its trajectory is the one that has least action, which is known as a fundamental principle in dynamical systems.

Original: We propose to explicitly model diffusion phenomena as well as drift by using SDEs for the time evolution of the sample and to explore the sample paths that minimize action defined by the time integral of the Lagrangian.
Revised: We propose to explicitly model diffusion phenomena as well as drift by using SDEs for the time evolution of the sample and to explore the sample paths that minimize action defined by the time integral of the Lagrangian.

Original: Later in Section 3.3, we will give some examples of prior information described by the Lagrangian.
Revised: Later in Section 3.3, we will give some examples of prior information described by the Lagrangian.

Original: We formulate this problem as a special case of the SOT problem and call it the Lagrangian-Schr¨odinger Bridge (LSB) problem . 
Revised: We formulate this problem as a special case of the SOT problem and call it the Lagrangian-Schr¨odinger Bridge (LSB) problem . The position of the LSB problem in the

Revised: SOT problem is clarified in Appendices A.3 and A.4.

Original: Definition 1 (LSB problem) .
Revised: Definition 1 (LSB problem) .

Original: [Equation]

Original: subject to d

Original: [Equation]
Revised: [Equation]

Original: where ρ t is the law of the random variable X t , depending on the functions f θ and g ϕ .
Revised: where ρ t is the law of the random variable X t , depending on the functions f θ and g ϕ .

Original: The connection between the LSB problem and the SB problem formulated as stochastic optimal control (SOC) problem (Chen et al., 2021b; Dai Pra, 1991; L ´ eonard, 2013; Mikami, 1990) is discussed in Appendix C.

Original: 3. N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE
Revised: 3. N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE

Original: In the setting of our paper, it is difficult to solve the LSB problem because the exact endpoint constraints µ 0 and µ 1 are unknown and only samples from them are available.
Revised: In the setting of our paper, it is difficult to solve the LSB problem because the exact endpoint constraints µ 0 and µ 1 are unknown and only samples from them are available.

Original: Therefore, we propose a practical solution method for the LSB problem using neural SDEs (Eq.
Revised: Therefore, we propose a practical solution method for the LSB problem using neural SDEs (Eq.

Original: (3)) with regularization.
Revised: (3)) with regularization.

Original: We propose to approximate the LSB problem by learning neural SDEs with the gradients of the loss (See Appendix A for details.): 
Revised: We propose to approximate the LSB problem by learning neural SDEs with the gradients of the loss (See Appendix A.5 for details.): [Equation]

Original: [Equation]
Revised: [Equation]

Original: where D  i,j ( · ; ϕ ) is the entry in the i -th row and j -th column of the diffusion coefficient matrix 
Revised: where D is the distribution discrepancy measure, D i,j ( · ; ϕ ) is the entry in the i -th row and j -th column of the diffusion coefficient matrix D ϕ

Original: D ϕ = 12 g ϕ g ⊤ ϕ  and Φ θ is the potential function satisfying f θ =[ ∇ u L ] − 1 (  −∇ x Φ θ ) .
Revised: = 12 g ϕ g ⊤ ϕ , H is the Hamiltonian defined by H ( t, x , z ) : = ⟨ z , f θ ⟩ − L ( t, x , f θ ) , and Φ θ is the potential function satisfying f θ = ∇ z H ( t, x , −∇ x Φ θ ( x , t )) .

Original: We briefly explain the action cost R e and the HJB regularization R h .
Revised: We briefly explain the action cost R e and the HJB regularization R h .

Original: Equation 7 is actually computed by Eq.
Revised: Equation 7 is actually computed by Eq.

Original: First,similar to the OT-based regularization (Finlay et al., 2020; Onken et al., 2021; Tong et al., 2020), we relax the constraint at time t 1 with the Wasserstein-2 cost, i.e . 
Revised: First, we relax the constraint at time t 1 , i.e . R

Original: R e
Revised: e

Original: + W 2 ( µ 1 , ρ t 1 ) . The use of Wasserstein-2 cost instead of KL-divergence is based on existing studies on population dynamics (Bunne et al., 2021; Hashimoto et al., 2016).
Revised: + D ( µ 1 , ρ t 1 ) , by using Wasserstein distance (Bunne et al., 2021; Hashimoto et al., 2016), KL-divergence (Finlay et al., 2020; Onken et al., 2021; Tong et al., 2020), or a combination of these (Lavenant et al., 2021) for the discrepancy measure D .

Original: This objective function can be reinterpreted with W 2 as the main cost and the action cost R e as the regularization.
Revised: This objective function can be reinterpreted with D ( µ 1 , ρ t 1 ) as the data-fitting term and the action cost R e as the regularization.

Original: Second, we exploit SOT theory by incorporating further structure into the modeling.
Revised: Second, we exploit SOT theory by incorporating further structure into the modeling.

Original: Similar to OT-Flow (Onken et al., 2021), we model the drift functions by using the potential function f θ = 
Revised: Similar to OT-Flow (Onken et al., 2021), we model the drift functions by using the potential function f θ = ∇ z H ( t, x , −∇ x

Original: [ ∇ u L ] − 1 ( −∇ x Φ θ ) and encourage the potential function Φ θ to satisfy the (stochastic)
Revised: Φ θ ( x , t )) and encourage the potential function Φ θ to satisfy the (stochastic)

Original: HJB equation (Yong & Zhou, 1999) by minimizing the PDE loss R h .
Revised: HJB equation (Yong & Zhou, 1999) by minimizing the PDE loss R h .

Original: The relation between the drift and the potential function is an analogue of Hamilton’s equations of motion.
Revised: The relation between the drift and the potential function is an analogue of Hamilton’s equations of motion.

Original: The HJB equation represents Bellman’s principle of optimality in continuous-time optimization.
Revised: The HJB equation represents Bellman’s principle of optimality in continuous-time optimization.

Original: 3. E XAMPLES OF L AGRANGIAN IN N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE
Revised: 3. E XAMPLES OF L AGRANGIAN IN N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE

Original: In this section, we provide the three Lagrangian examples of the NLSB and their use cases.
Revised: In this section, we provide the three Lagrangian examples of the NLSB and their use cases.

Original: The examples demonstrate that the Lagrangian design allows a variety of prior knowledge to be reflected in the sample trajectories.
Revised: The examples demonstrate that the Lagrangian design allows a variety of prior knowledge to be reflected in the sample trajectories.

Original: Potential-free system.
Revised: Potential-free system.

Original: Without a specific external force on the individual sample, the principle of least action typically indicates that when a sample moves from point A to point B, it tries to minimize energy by reducing the travel distance.
Revised: Without a specific external force on the individual sample, the principle of least action typically indicates that when a sample moves from point A to point B, it tries to minimize energy by reducing the travel distance.

Original: This means that the drift is encouraged to be straight and the stochastic movement to be small, i.e ., the Lagrangian is formulated by L ( t, x , u )
Revised: This means that the drift is encouraged to be straight and the stochastic movement to be small, i.e ., the Lagrangian is formulated by L ( t, x , u )

Original: =
Revised: =

Original: 12 || u
Revised: 12 || u

Original: || 2 , and the drift function is given as f θ =
Revised: || 2 , and the drift function is given as f θ =

Original: −∇ x Φ θ ( x , t ) .
Revised: −∇ x Φ θ ( x , t ) .

Original: Cellular system.
Revised: Cellular system.

Original: Tong et al.
Revised: Tong et al.

Original: (2020) proposed to introduce the prior knowledge of manifold geometry and local velocity arrows such as RNA-velocity for modeling cellular systems.
Revised: (2020) proposed to introduce the prior knowledge of manifold geometry and local velocity arrows such as RNA-velocity for modeling cellular systems.

Original: In the NLSB, these prior knowledge can be handled consistently by designing the Lagrangian.
Revised: In the NLSB, these prior knowledge can be handled consistently by designing the Lagrangian.

Original: First, we introduce a density-based penalty to constrain the sample trajectories on the data manifold.
Revised: First, we introduce a density-based penalty to constrain the sample trajectories on the data manifold.

Original: We estimate the density function U ( x , t ) from the data, e.g ., by using Gaussian mixture models (GMM) and add it to the Lagrangian.
Revised: We estimate the density function U ( x , t ) from the data, e.g ., by using Gaussian mixture models (GMM) and add it to the Lagrangian.

Original: Next, we redefine the velocity regularization, which is formulated as cosine similarity maximization by Tong et al.
Revised: Next, we redefine the velocity regularization, which is formulated as cosine similarity maximization by Tong et al.

Original: (2020), as a squared error minimization
Revised: (2020), as a squared error minimization

Original: and add it to the Lagrangian as well.
Revised: and add it to the Lagrangian as well.

Original: Therefore, the Lagrangian for the cellular system is defined by
Revised: Therefore, the Lagrangian for the cellular system is defined by

Original: [Equation]
Revised: [Equation]

Original: where v ( x , t ) is the reference velocity of the cell at the position x and the time t .
Revised: where v ( x , t ) is the reference velocity of the cell at the position x and the time t .

Original: The drift function is obtained by f θ = 1 2 ( v ( x , t ) − ∇ x Φ θ ( x , t )) .
Revised: The drift function is obtained by f θ = 12 ( v ( x , t ) − ∇ x Φ θ ( x , t )) .

Original: Random dynamical system.
Revised: Random dynamical system.

Original: Let m be the mass of the sample, U ( x ) be the potential energy of the system, and L ( t, x , u )= m 2 u ( x , t )
Revised: Let m be the mass of the sample, U ( x ) be the potential energy of the system, and L ( t, x , u )

Original: ⊤ u ( x , t )  − U ( x ) be the Lagrangian in the Newtonian dynamical system.
Revised: = m 2 ∥ u ( x , t ) ∥ 2 − U ( x ) be the Lagrangian in the Newtonian dynamical system.

Original: Then, the individual samples are encouraged to follow a stochastic analogue of the equations of motion of the Newtonian mechanics.
Revised: Then, the individual samples are encouraged to follow a stochastic analogue of the equations of motion of the Newtonian mechanics.

Original: In practice, the potential function can be used to roughly incorporate information such as obstacles and regions where samples cannot or are likely to exist.
Revised: In practice, the potential function can be used to roughly incorporate information such as obstacles and regions where samples cannot or are likely to exist.

Original: Practical examples are shown in Figure 1 and Appendix E.
Revised: Practical examples are shown in Figs. 1 and 10.

Original:  I MPLEMENTATION D ETAILS OF N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE 4. T RAINING FOR N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE
Revised:  I MPLEMENTATION D ETAILS OF N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE 4. T RAINING FOR N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE

Original: In this section, we describe a practical learning method for neural SDEs by minimizing loss 7 on the training data, as described in Algorithm 1.
Revised: In this section, we describe a practical learning method for neural SDEs by minimizing loss in Eq. (7) on the training data, as described in Algorithm 1.

Original: First, for the data, only individual samples from a cross-sectional population with no alignment across time steps at K separate time points are available.
Revised: First, for the data, only individual samples from a cross-sectional population with no alignment across time steps at K separate time points are available.

Original: Let T = { t 0 , . . .
Revised: Let T = { t 0 , . . .

Original: , t K − 1 } be a set of time points and denote the data set at time t i as X t i .
Revised: , t K − 1 } be a set of time points and denote the data set at time t i as X t i .

Original: Next, we describe how we calculate the loss 7 using the training data {X t i } t i ∈ T .
Revised: Next, we describe how we calculate the loss in Eq. (7) using the training data {X t i } t i ∈ T .

Original: The first term, the Wasserstein cost, is computed using the L2-Sinkhorn divergence W ϵ between the observed data and the predicted sample at all observation points in T except the initial point at time t 0 .
Revised: We compute the distribution discrepancy using the L2-Sinkhorn divergence W ϵ between the observed data and the predicted sample at all observation points in T except the initial point at time t 0 .

Original: The Sinkhorn divergence can efficiently approximate the Wasserstein distance and solve the entropic bias problem when using the Sinkhorn algorithm, i.e . µ = ν ⇔ W ϵ ( µ, ν ) = 0 .
Revised: The Sinkhorn divergence can efficiently approximate the Wasserstein distance and solve the entropic bias problem when using the Sinkhorn algorithm, i.e . µ = ν ⇔ W ϵ ( µ, ν ) = 0 .

Original: The prediction sample is obtained by numerically simulating neural SDE from the training data at one previous time point by a standard SDE solver such as the Euler-Maruyama method.
Revised: The prediction sample is obtained by numerically simulating neural SDE from the training data at one previous time point by a standard SDE solver such as the Euler-Maruyama method.

Original: The second and third terms, R e and R h , are also approximated on the sample paths obtained from the numerical simulations of SDE.
Revised: The second and third terms, R e and R h , are also approximated on the sample paths obtained from the numerical simulations of SDE.

Original: Therefore, the computational cost of empirical R e and R h is not high and can be further accelerated by the model architecture described in the next section.
Revised: Therefore, the computational cost of empirical R e and R h is not high and can be further accelerated by the model architecture described in the next section.

Original: To summarize, we simultaneously solve K − 1 approximated LSB problems by minimizing the subsequent loss for the model parameters θ, ϕ .
Revised: To summarize, we simultaneously solve K −approximated LSB problems by minimizing the subsequent loss for the model parameters θ, ϕ .

Original: [Equation]
Revised: [Equation]

Original: where µ k , ρ θ,ϕt k are the ground-truth and predicted probability measures at time t k expressed by data samples, respectively.
Revised: where µ k , ρ θ,ϕt k are the ground-truth and predicted probability measures at time t k expressed by data samples, respectively.

Original: ˆ R θe ( t k − 1 , t k ) and ˆ R θ,ϕh ( t k − 1 , t k ) are empirical quantities computed on the simulated sample paths from t k − 1 to t k with data ∀ x ( t k − 1 )
Revised: ˆ R θe ( t k − 1 , t k ) and ˆ R θ,ϕh ( t k − 1 , t k ) are empirical quantities computed on the simulated sample paths from t k − 1 to t k with data ∀ x ( t k − 1 )

Original: ∈
Revised: ∈

Original: X t k − 1 as the initial value.
Revised: X t k − 1 as the initial value.

Original: The weight coefficients λ e ( t k − 1 , t k ) and λ h ( t k − 1 , t k ) are tuned for each interval [ t k − 1 , t k ] respectively.
Revised: The weight coefficients λ e ( t k − 1 , t k ) and λ h ( t k − 1 , t k ) are tuned for each interval [ t k − 1 , t k ] respectively.

Original: [Equation]
Revised: [Equation]

Original:  S ELECTION FOR S PEEDUP
Revised: A RCHITECTURE S ELECTION FOR S PEEDUP

Original: We adopt the model of the potential function Φ θ proposed by OT-Flow (Onken et al., 2021) because it has two advantages in our framework.
Revised: We adopt the model of the potential function Φ θ proposed by OT-Flow (Onken et al., 2021) because it has two advantages in our framework.

Original: First, it can compute the gradient ∇ x Φ explicitly, which enables us to calculate the drift function easily.
Revised: First, it can compute the gradient ∇ x Φ explicitly, which enables us to calculate the drift function easily.

Original: Second, the model is designed for the fast and exact computation of the diagonal component of the potential function’s Hessian.
Revised: Second, the model is designed for the fast and exact computation of the diagonal component of the potential function’s Hessian.

Original: Moreover, when we assume that the diffusion model’s output is a diagonal matrix, i.e . g ϕ
Revised: Moreover, when we assume that the diffusion model’s output is a diagonal matrix, i.e . g ϕ

Original: ( x , t ) ∈ R d × d
Revised: ( x , t ) ∈ R d × d

Original: and g i,j = 0 ( i ̸ = j ) , the function D ϕ is also a diagonal matrix, and the (cid:80) i,j in the second term of R h shown in Eq.
Revised: and g i,j = 0 ( i ̸ = j ) , the function D ϕ is also a diagonal matrix, and the (cid:80) i,j in the second term of R h shown in Eq.

Original: (9) turns into the sum of the diagonal components only.
Revised: (9) turns into the sum of the diagonal components only.

Original: Combining these two tricks enables speeding up the computation of R h , which requires the expensive computation of the Hessian of the potential function Φ .
Revised: Combining these two tricks enables speeding up the computation of R h , which requires the expensive computation of the Hessian of the potential function Φ .

Original: While this model architecture trick was originally used for speeding up the computation of the Jacobian term in neural ODE maximum likelihood training, we propose  to use it as a technique to speed up the computation of R h .
Revised: While this model architecture trick was originally used for speeding up the computation of the Jacobian term in neural ODE maximum likelihood training, we propose for the first time to use it as a technique to speed up the computation of R h in neural SDE.

Original: We use a two-layer fully connected neural network (FCNN) for the diffusion function.
Revised: We use a two-layer fully connected NN for the diffusion function.

Original: We also adopt the device used by Kidger et al.
Revised: We also adopt the device used by Kidger et al.

Original: (2021) in which the activation function tanh is used after the final layer to prevent the output of the diffusion function from becoming excessively large.
Revised: (2021) in which the activation function tanh is used after the final layer to prevent the output of the diffusion function from becoming excessively large.

Original: For details of the model architecture, see Appendix B.
Revised: See Appendix B for details on the model of Φ .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluated our methods on two datasets.
Revised: We evaluated our methods on two datasets.

Original: First, we used artificial synthetic data generated from one-dimensional SDEs, where the predicted trajectory and uncertainty can be compared with the ground-truth and easily evaluated by visualization.
Revised: First, we used artificial synthetic data generated from one-dimensional SDEs, where the predicted trajectory and uncertainty can be compared with the ground-truth and easily evaluated by visualization.

Original: We set the Lagrangian for the potential-free system, i.e . L ( t, x , u ) =
Revised: We set the Lagrangian for the potential-free system, i.e . L ( t, x , u ) =

Original: 12 || u
Revised: 12 || u

Original: || 2 .
Revised: || 2 .

Original: Second, we used the evolution of single-cell populations obtained from a developing human embryo system.
Revised: Second, we used the evolution of single-cell populations obtained from a developing human embryo system.

Original: We used the Lagrangian for the cellular system and compared several combinations of the regularization terms.
Revised: We used the Lagrangian for the cellular system and compared several combinations of the regularization terms.

Original: In Tables 1 and 2, “E” is the energy term, “D” is the density term, and “V” is the velocity term.
Revised: In Tables 1 and 2, “E” is the energy term, “D” is the density term, and “V” is the velocity term.

Original: The density term U ( x , t ) is the log-likelihood function of the data estimated by GMM.
Revised: The density term U ( x , t ) is the log-likelihood function of the data estimated by GMM.

Original: We compared our methods against standard neural SDE, OTFlow (Onken et al., 2021),
Revised: We compared our methods against standard neural SDE, OTFlow (Onken et al., 2021),

Original: TrajectoryNet (Tong et al., 2020) and an iterative proportional maximum likelihood (IPML) algorithm
Revised: TrajectoryNet (Tong et al., 2020) and an iterative proportional maximum likelihood (IPML) algorithm

Original: (Vargas et al., 2021).
Revised: (Vargas et al., 2021).

Original: The drift and diffusion models of neural SDE are the same as NLSB.
Revised: The drift and diffusion models of neural SDE are the same as NLSB.

Original: The velocity model of TrajectoryNet includes the concatsquash layers used in Grathwohl et al.
Revised: The velocity model of TrajectoryNet includes the concatsquash layers used in Grathwohl et al.

Original: The base models of OT-Flow and TrajectoryNet were trained with the standard neural ODEs scheme.
Revised: The base models of OT-Flow and TrajectoryNet were trained with the standard neural ODEs scheme.

Original: +OT represents a model trained with the OT-based regularization defined by Eqs.
Revised: +OT represents a model trained with the OT-based regularization defined by Eqs.

Original: (1) and (2).
Revised: (1) and (2).

Original: We used only ˜ R e for TrajectoryNet; we used both ˜ R e and ˜ R h for
Revised: We used only ˜ R e for TrajectoryNet; we used both ˜ R e and ˜ R h for

Original: OT-Flow.
Revised: OT-Flow.

Original: We set the interval-dependent coefficients ˜ λ
Revised: We set the interval-dependent coefficients ˜ λ

Original: e , ˜ λ
Revised: e , ˜ λ

Original: h for the OT-based regularization as well as Eq.
Revised: h for the OT-based regularization as well as Eq.

Original: The drift models of IPML (GP) and IPML (NN) are sparse GP and the same networks as NLSB, respectively.
Revised: The drift models of IPML (GP) and IPML (NN) are sparse GP and the same networks as NLSB, respectively.

Original: The diffusion coefficients of IPML were tuned as hyperparameters.
Revised: The diffusion coefficients of IPML were tuned as hyperparameters.

Original: See Appendix D for more details on hyperparameters, learning rates, and optimizers.
Revised: See Appendix C for more details on hyperparameters, learning rates, and optimizers.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 5. S YNTHETIC P OPULATION D YNAMICS : O RNSTEIN –U HLENBECK P ROCESS
Revised: 5. S YNTHETIC P OPULATION D YNAMICS : O RNSTEIN –U HLENBECK P ROCESS

Original: Data .
Revised: Data .

Original: We used a time-dependent one-dimensional Ornstein–Uhlenbeck (OU) process defined by the following SDE:
Revised: We used a time-dependent one-dimensional Ornstein–Uhlenbeck (OU) process defined by:

Original: [Equation]
Revised: [Equation]

Original: where µ = 0 .
Revised: where µ = 0 .

Original: 4 , θ = 0 .
Revised: 4 , θ = 0 .

Original: 1 , σ = 0 . 8 , and t K − 1 = 4 .
Revised: 1 , σ = 0 . 8 , and t K − 1 = 4 .

Original: First, we simulated the several trajectories from t = 0 to 4 and then extracted only the data at the time of T = { 0 , 1 , 2 , 3 , 4 } as snapshots for training.
Revised: First, we simulated the several trajectories from t = 0 to 4 and then extracted only the data at the time of T = { 0 , 1 , 2 , 3 , 4 } as snapshots for training.

Original: We generated 2048 and 512 samples for each time point as training and validation data, respectively.
Revised: We generated 2048 and 512 samples for each time point as training and validation data, respectively.

Original: Performance metrics .
Revised: Performance metrics .

Original: We evaluated the estimation performance of the dynamics in the time interval [ t 0 , t K − 1 ] by using two metrics on test data: marginal distribution discrepancy (MDD) and conditional distribution discrepancy (CDD).
Revised: We evaluated the estimation performance of the dynamics in the time interval [ t 0 , t K − 1 ] by using two metrics on test data: marginal distribution discrepancy (MDD) and conditional distribution discrepancy (CDD).

Original: A smaller MDD between µ t and ρ t calculated on the basis of the Wasserstein-2 distanceW ( µ t , ρ t ) indicates better prediction of population-level dynamics at time t .
Revised: A smaller MDD between µ t and ρ t calculated on the basis of the Wasserstein-2 distance indicates better prediction of population-level dynamics at time t .

Original: We calculated MDD at 12 equally spaced time points and the earth mover’s distance with L2 cost (EMD-L2) between 1000 samples generated from the ground-truth SDE and predicted by the model.
Revised: We calculated MDD at 12 equally spaced time points and the earth mover’s distance with L 2 cost (EMD-L2) between 1000 samples generated from the ground-truth SDE and predicted by the model.

Original: When evaluating the SDE-based method, we ran 100 simulations from the same initial values, computing the MDD value each time and computing their mean and variance.
Revised: When evaluating the SDE-based method, we ran 100 simulations from the same initial values, computing the MDD value each time and computing their mean and variance.

Original: A smaller CDD between µ x ( t ) | x ( t 0 ) and ρ x ( t ) | x ( t 0 ) using the Wasserstein-2 distance indicates better prediction of the time evolution of the initial sample x
Revised: A smaller CDD between µ x ( t ) | x ( t 0 ) and ρ x ( t ) | x ( t 0 ) using the Wassersteindistance indicates better prediction of the time evolution of the initial sample x

Original: ( t 0 ) .
Revised: ( t 0 ) .

Original: In short, it is a metric for evaluating the time evolution of at the individual sample level.
Revised: In short, it is a metric for evaluating the time evolution of at the individual sample level.

Original: In the actual CDD calculation, we first prepared 1000 samples x ( t 0 ) at the initial time point.
Revised: In the actual CDD calculation, we first prepared 1000 samples x ( t 0 ) at the initial time point.

Original: We then generated 100 trajectories from each initial sample by the trained model and the ground-truth SDE, and calculated the EMD-L2 for the samples from µ x ( t ) | x ( t 0 ) and
Revised: We then generatedtrajectories from each initial sample by the trained model and the ground-truth SDE, and calculated the EMD-L2 for the samples from µ x ( t ) | x ( t 0 ) and

Original: ρ x ( t ) | x ( t 0 ) at 12 equally spaced time points.
Revised: ρ x ( t ) | x ( t 0 ) at 12 equally spaced time points.

Original: Results .
Revised: Results .

Original: The evaluation results are shown in Figure 3, and the visualization of trajectories is shown in Figure 4.
Revised: The evaluation results are shown in Fig. 3, and the visualization of trajectories is shown in Fig.

Original: Figures 3 show that NLSB outperforms neural SDE and is comparable to other ODEbased methods in estimating populations with small variance.
Revised: Figure 3 shows that NLSB and IPML outperform neural SDE and is comparable to other ODE-based methods in estimating populations with small variance.

Original: In contrast, the SDE-based methods outperform ODE-based methods when estimating populations with a large variance.
Revised: In contrast, the SDE-based methods outperform ODE-based methods when estimating populations with a large variance.

Original: That indicates that NLSB  can estimate population-level dynamics even when the population variance is large or small.
Revised: That indicates that NLSB and IPML can estimate population-level dynamics even when the population variance is large or small.

Original: Furthermore, NLSB has a smaller CDD value than neural SDE and is comparable to IPML.
Revised: Furthermore, NLSB and IPML have a smaller CDD value than neural SDE.

Original: Figure 4b shows that the average behavior of samples E [ˆ X ( t ) | X (0)] estimated by neural SDE is different from that of the ground-truth SDE (see Figure 4a), especially in the interval [0 , 1] .
Revised: Figure 4b shows that the average behavior of samples E [ X ( t ) | X (0)] estimated by neural SDE is different from that of the ground-truth SDE (see Fig. 4a), especially in the interval [0 , 1] .

Original: In contrast, the prediction by NLSB  in Figure 4c is much closer to the ground-truth.
Revised: In contrast, the predictions by NLSB and IPML in Figs. 4c, 6d and 6e are much closer to the ground-truth.

Original: These results show that the prior knowledge introduced by the Lagrangian helps to estimate the sample-level dynamics. 
Revised: These results show that the prior knowledge of the potential-free system helps to estimate the sample-level dynamics. See Appendix D.1 for further results and analysis.

Original: 5. S INGLE -C ELL P OPULATION D YNAMICS
Revised: 5. S INGLE -C ELL P OPULATION D YNAMICS

Original: Data .
Revised: Data .

Original: We evaluated on embryoid body scRNA-seq data (Moon et al., 2019).
Revised: We evaluated on embryoid body scRNA-seq data (Moon et al., 2019).

Original: This data shows the differentiation of human embryonic stem cells from embryoid bodies into diverse cell lineages, including mesoderm, endoderm, neuroectoderm, and neural crest, over 27 days.
Revised: This data shows the differentiation of human embryonic stem cells from embryoid bodies into diverse cell lineages, including mesoderm, endoderm, neuroectoderm, and neural crest, over 27 days.

Original: During this period, cells were collected at five different snapshots ( t 0 : day 0 to 3 , t 1 : day 6 to 9 , t 2 : day 12 to 15 , t 3 : day18 to21 , t 4 : day 24 to27 ).
Revised: During this period, cells were collected at five different snapshots ( t 0 : day 0 to 3 , t 1 : day 6 to 9 , t 2 : day 12 to 15 , t 3 : day to , t 4 : day 24 to

Original: The collected cells were then measured by scRNAseq, filtered at the quality control stage, and mapped to a low-dimensional feature space using a principal component analysis (PCA).
Revised: The collected cells were then measured by scRNAseq, filtered at the quality control stage and mapped to a low -dimensional feature space using a principal component analysis (PCA).

Original: For details, see Appendix E.2 in (Tong et al.,2020).
Revised: For details, see Appendix E.2 in (Tong et al.,

Original: We split the dataset into train, validation( ∼ 8 . 5% ) and test data ( ∼ 15% ).
Revised: We split the dataset into train, validation( ∼ 8 . 5% ) and test data ( ∼ 15% )

Original: Performance metrics .
Revised: Performance metrics .

Original: Unlike the experiment described in Section 5.1, there are no ground-truth trajectories in the real data.
Revised: Unlike the experiment described in Section 5.1, there are no ground-truth trajectories in the real data.

Original: Thus, MDD can be calculated only at the time of observation, and CDD between the ground truth and predicted trajectories cannot be calculated.
Revised: Thus, MDD can be calculated only at the time of observation, and CDD between the ground truth and predicted trajectories cannot be calculated.

Original: To evaluate the sample-level dynamics, the model was trained on the full data and the data without only one intermediate snapshot, respectively.
Revised: To evaluate the sample-level dynamics, the model was trained on the full data and the data without only one intermediate snapshot, respectively.

Original: We then calculated CDD between the predicted trajectories by those models.
Revised: We then calculated CDD between the predicted trajectories by those models.

Original: Let D − t
Revised: Let D − t

Original: k be the training data without a snapshot at time t k .
Revised: k be the training data without a snapshot at time t k .

Original: Larger CDD between them indicates the prediction of the sample-level dynamics is not robust to the exclusion of intermediate snapshots, representing poorer performance in interpolating the sample-level dynamics.
Revised: Larger CDD between them indicates the prediction of the sample-level dynamics is not robust to the exclusion of intermediate snapshots, representing poorer performance in interpolating the sample-level dynamics.

Original: When evaluating the SDE-based methods, we calculated the mean and standard deviation of 100 MDD scores.
Revised: When evaluating the SDE-based methods, we calculated the mean and standard deviation of 100 MDD scores.

Original: All performance metrics are calculated on the test data.
Revised: All performance metrics are calculated on the test data.

Original: Results .
Revised: Results .

Original: Table 1 and Figure 5 show that the NLSB can predict population-level dynamics with better performance and can be trained in a shorter time against all existing ODE-based methods as the data become higher dimensional.
Revised: Table 1 and Figure 5 show that the NLSB can predict population-level dynamics with better performance and can be trained in a shorter time against all existing ODE-based methods as the data become higher dimensional.

Original: In particular, the SDE-based methods significantly outperform the ODE-based ones in predicting the transitions where the samples from t 1 to t 2 and from t 3 to tare highly diffuse, indicating that the explicit modeling of diffusion is effective.
Revised: In particular, the SDE-based methods significantly outperform the ODE-based ones in predicting the transitions where the samples from t 1 to t 2 and from t 3 to t 4 are highly diffuse, indicating that the explicit modeling of diffusion is effective.

Original: (see Figure 8 in Appendix E).
Revised: (see Fig.in Appendix D).

Original: Overall, the standard deviation of the MDD values is smaller for NLSB than for neural SDE, demonstrating less variation in the approximation accuracy of the marginal distribution.
Revised: Overall, the standard deviation of the MDD values is smaller for NLSB than for neural SDE, demonstrating less variation in the approximation accuracy of the marginal distribution.

Original: Tableshows that NLSB estimates the sample-level dynamics robustly with and without population at the intermediate time point than neural SDE with some exceptions.
Revised: Table 2 shows that NLSB estimates the sample-level dynamics robustly with and without population at the intermediate time point than neural SDE with some exceptions.

Original: Especially, energy regularization is the most stable and effective.
Revised: Especially, energy regularization is the most stable and effective.

Original: This result suggests that the LSB-based regularization helps interpolate the sample-level dynamics.
Revised: This result suggests that the LSB-based regularization helps interpolate the sample-level dynamics.

Original: Seethe Appendix E for a visualization of trajectories with uncertainty.
Revised: See Appendix D.2 for further results and analysis.

Original:  D ISCUSSION AND C ONCLUSION
Revised:  D ISCUSSION AND C ONCLUSION

Original: In this work, we proposed a novel framework for estimating population dynamics that reflect prior knowledge about the target system.
Revised: In this work, we proposed a novel framework for estimating population dynamics that reflect prior knowledge about the target system.

Original: Unlike existing methods with OT (Schiebinger et al., 2019;
Revised: Unlike existing methods with OT (Schiebinger et al., 2019;

Original: Yang & Uhler, 2018), or CNF (Tong et al., 2020) for a biological system, we explicitly modeled the diffusion phenomena by using SDEs for the samples with stochastic behavior.
Revised: Yang & Uhler, 2018), or CNF (Tong et al., 2020) for a biological system, we explicitly modeled the diffusion phenomena by using SDEs for the samples with stochastic behavior.

Original: This allowed us to handle the uncertainty of the trajectory (Fig. 7) and to successfully capture the diffuse transitions 
Revised: This allowed us to handle the uncertainty of the trajectory (Fig. 7) and to successfully capture the diffuse transitions (Table 1 and Fig. In contrast, Vargas et al. (2021) and Bunne et al. (2022) estimated the sample

Original: [Table]
Revised: [Table]

Original:  points pe predicted  generated 
Revised: ated at 7 equally spaced time points within time period [ t k − 1 , t k +1 ] around excluded time point t k for each k = 1 , 2 and 3 . The CDD value at the time t ∈ [ t k − 1 , t k +1 ] was computed between the two groups of predicted samples generated from the previous samples at t k − 1 using the model trained on all data and the data D − t k .

Original: Mean CDD ↓ [ t 0 , t 2 ] [ t 1 , t 3 ] [ t 2 , t 4 ]
Revised: Mean CDD ↓ [ t 0 , t 2 ] [ t 1 , t 3 ] [ t 2 , t 4 ]

Original: NLSB (E) 0 .0 .0 .
Revised: NLSB (E) 0 .0 .0 .

Original: NLSB (D) 1 .0 .0 .
Revised: NLSB (D) 1 .0 .0 .

Original: NLSB (V) 1 .0 .0 .
Revised: NLSB (V) 1 .0 .0 .

Original: NLSB (E+D+V) 0 .0 .0 .
Revised: NLSB (E+D+V) 0 .0 .0 .

Original: Neural SDE 1 .0 .0 .
Revised: Neural SDE 1 .0 .0 .

Original: IPML (GP) 0 .0 .0 .
Revised: IPML (GP) 0 .0 .0 .

Original: IPML (NN) 0 .0 .0 .
Revised: IPML (NN) 0 .0 .0 .

Original: . 1 and Fig. In contrast, Vargas et al. (2021) and Bunne et al. (2022) estimated the sample trajectories of biological systems using the SDE solution to the SB problem.
Revised: trajectories of biological systems using the SDE solution to the SB problem.

Original: Vargas et al.
Revised: Vargas et al.

Original: (2021) proposed IPML algorithm to solve the SB problem.
Revised: (2021) proposed IPML algorithm to solve the SB problem.

Original: Bunne et al.
Revised: Bunne et al.

Original: (2022) proposed GSB-Flow, in which two SB problems are solved sequentially.
Revised: (2022) proposed GSB-Flow, in which two SB problems are solved sequentially.

Original: The first SB problem with Gaussian marginals was solved in closed form, and the next SB problem using the closed-form solution as the reference path was solved by maximum likelihood training.
Revised: The first SB problem with Gaussian marginals was solved in closed form, and the next SB problem using the closed-form solution as the reference path was solved by maximum likelihood training.

Original: Compared with these methods, our method handled a wider class of SDEs, and the diffusion function of the SDE was learned from the data using a backpropagation.
Revised: Compared with these methods, our method handled a wider class of SDEs, and the diffusion function of the SDE was learned from the data using a backpropagation.

Original: Moreover, designing the Lagrangian enables us to flexibly incorporate prior knowledge about the target system into the model.
Revised: Moreover, designing the Lagrangian enables us to flexibly incorporate prior knowledge about the target system into the model.

Original: Our Lagrangian-based regularization also treated the biological constraints proposed by Tong et al.
Revised: Our Lagrangian-based regularization also treated the biological constraints proposed by Tong et al.

Original: and Maoutsa & Opper (2021) in a unified manner.
Revised: and Maoutsa & Opper (2021) in a unified manner.

Original: We demonstrated that NLSB can efficiently estimate the population-level dynamics with better performance than existing methods even for high-dimensional data and that the prior knowledge introduced by the Lagrangian is useful to estimate the sample-level dynamics.
Revised: We demonstrated that NLSB can efficiently estimate the population-level dynamics with better performance than existing methods even for high-dimensional data and that the prior knowledge introduced by the Lagrangian is useful to estimate the sample-level dynamics.

Original: Our method is limited in that it cannot model reaction phenomena such as cell birth and death, and restriction to the diagonal diffusion matrix (Section 4.2) may cause the model to be less expressive.
Revised: Our method is limited in that it cannot model reaction phenomena such as cell birth and death, and restriction to the diagonal diffusion matrix (Section 4.2) may cause the model to be less expressive.

Original: Future work includes developing methods that can handle reaction phenomena with advection and diffusion.
Revised: Future work includes developing methods that can handle reaction phenomena with advection and diffusion.

Original: A CKNOWLEDGMENTS
Revised: A CKNOWLEDGMENTS

Original: We thank the lab members for their discussions and inspiration for our research. 
Revised: We thank the lab members for their discussions and inspiration for our research. A T HEORY OF S TOCHASTIC O PTIMAL T RANSPORT

Revised: The SOT problem is the problem of finding a stochastic process that minimizes the expected cost under fixed marginal distributions at several time points.

Revised: The SOT problem is a stochastic analog of the OT problem, especially related to the dynamic formulation of the OT problem introduced by

Revised: Benamou & Brenier (2000).

Original: A D ERIVATION OF N EURAL L AGRANGIAN S CHR ¨ ODINGER B RIDGE
Revised: It is also considered as a stochastic optimal control (SOC) problem with additional terminal constraint. Furthermore, the classical Schr¨odinger Bridge (SB) problem is a special case of the SOT problem and is related to Nelson’s stochastic mechanics, a reformulation of quantum mechanics using diffusion processes. An overview of the relationships among these problems is illustrated in Fig. We first explain the dynamic formulation of the OT problem in Appendix A.1, the SB problem in Appendix A.2 and then describe the SOT problem given two fixed marginal distributions as a generalization of the SB problem in Appendix A.3.

Revised: Next, we discuss the LSB problem, which is the main focus of our paper among the SOT problems in Appendix A.4, and finally, we discuss the relationship between NLSB and OT-Flow in Appendix A.5.

Revised: A. D YNAMIC F ORMULATION OF O PTIMAL T RANSPORT P ROBLEM

Revised: Benamou & Brenier (2000) redefined the OT problem with a quadratic cost in a continuum mechanics framework.

Revised: They introduced a time parameter t ∈ [0 , 1] and considered the transport process from µ to ν as the advection in the time interval [0 , 1] .

Original: We consider the LSB problem constrained by the Fokker-Planck (FP) equation corresponding to Ito
Revised: The OT problem was then reformulated as a minimization problem with respect to the time-varying velocity vector field v t and the time evolution of the distribution p t advected by v t . Theorem 2 (Brenier-Benamou formulation (Eularian formalism); (Benamou & Brenier, 2000)) . [Equation] subject to ∂

Original: SDE (Eq. (3)) over the bounded domain Ω ⊂ R d .
Revised: [Equation] where p and q are the densities of probability measures µ and ν , respectively. The first condition in Eq. (13) is known as the continuity equation and represents the conservation of mass in time evolution.

Original: minimize θ,ϕ
Revised: We can consider an equivalent formulation of Eq. (12) by introducing Lagrangian coordinates X ( t, x ) . Lagrangian coordinates X ( t, x ) represent the position at time t of the particle whose initial position is x , i.e . X (0 , x ) = x .

Revised: Theorem 3 (Brenier-Benamou formulation (Lagrangian formalism); (Benamou & Brenier, 2000)) .

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

Revised: The right-hand side of Eq.

Revised: (14) can be viewed as a problem of finding the shortest path (a.k.a. geodesic) for each particle in the sense of Euclidean space between probability distributions specified at times t = 0 , 1 .

Revised: In continuum mechanics, Lagrangian formalism (Eq. (14)) describes the motion of each individual particle, while Eulerian formalism (Eq.

Revised: (12)) focuses on the global property of all particles.

Revised: Benamou & Brenier (2000) also showed the optimality conditions of the dynamic formulation.

Revised: They introduced the Lagrangian multiplier of the constraint of Eq.

Revised: (13) and obtained the saddle point conditions using the variational method.

Revised: Theorem 4 (Optimality conditions for the dynamic formulation; (Benamou & Brenier, 2000)) .

Revised: There exists a space-time dependent potential function Φ t : R d × [0 , 1] (cid:55)→ R which satisfies:

Revised: [Equation]

Revised: where f ∗ and g ∗ are Kantorovich potentials.

Revised: Equation 16 is Hamilton’s equation of motion with the Hamiltonian defined by H ( p , x ) :

Revised: = sup v (cid:8) ⟨ p , v t ( x ) ⟩ − 12 ∥ v t ( x ) ∥ 2 (cid:9) and the momentum as p :=

Revised: −∇ x Φ t and equation 17 is Hamilton-Jacob-Bellman (HJB) equation.

Revised: Equations 16 and 18 indicate that the particle, whose initial position is x 0 ∈ X , moves straight ahead at the constant velocity v 0 =

Revised: −∇ x Φ 0 ( x 0 )

Revised: =

Revised: −∇ x f ∗ ( x 0 ) during t ∈ [0 , 1] .

Revised: Even without using the variational method, Hamilton’s equation of motion (Eq.

Revised: and the HJB equation (Eq.

Revised: (17)) can be also derived from the Bellman’s principle of optimality.

Revised: The existence of a potential function satisfying the equation 16 is guaranteed from the Pontryagin Maximum Principle (Evans, 1983; 2010).

Revised: A. S CHR ¨ ODINGER B RIDGE

Revised: Definition 5 (SB problem; (Jamison, 1975)) .

Revised: Let Ω

Revised: =

Revised: C

Revised: ([0 , 1] , R d ) be the space of R d -valued continuous functions on time interval [0 , 1] .

Revised: Denote by P (Ω) the probability measures space on the path pace Ω .

Revised: The SB problem is defined by

Revised: [Equation]

Revised: where µ 0 , ν 1 ∈ P ( R d ) are the probability measures at the time 0 and 1 , respectively, and the relative entropy D KL = (cid:82) log (cid:0)

Revised: d Q d P (cid:1)

Revised: d Q if Q ≪ P , and D KL = ∞ otherwise.

Revised: The case of no prior dynamics is considered classically, i.e ., the reference path measure P is the

Revised: Brownian diffusion SDE d X t = √ 2 ϵ d W t , where W t is the standard Wiener process.

Revised: We refer to this setting as the classical SB problem, following the reference (Caluya & Halder, 2021).

Revised: Chen et al.

Revised: (2021b) derived the optimality condition for problem Eq.

Revised: (19), which is characterized by the forward and backward time-harmonic equations.

Revised: Theorem 6.

Original: 12) where p t is the probability density of  the probability measure ρ t .
Revised: Let Ψ( x , t ) and ˆΨ( x , t ) be the solutions to the following PDEs: [Equation] where p 0 and p T are the probability density of µ 0 and µ T , respectively. Then, the solution to the SB problem (Eq.

Revised: (19) ) can be described by the following forward or backward SDE:

Revised: [Equation]

Revised: Dai Pra (1991) considered a dynamic formulation of the classic SB problem (Eq.

Revised: by interpreting it as a SOC problem with the additional terminal constraint.

Revised: Theorem 7 (Dynamic formulation; (Dai Pra, 1991)) .

Revised: Let u ∗ = 2 ϵ ∇ x log Ψ , where Ψ satisfies the

Revised: SB optimality (Eq.

Revised: Then, u ∗ is the minimizer of the following optimization problem:

Revised: [Equation]

Revised: subject to

Revised: [Equation]

Revised: L ´ eonard (2013) formulated the SB problem in Eq.

Revised: (19) into a variational SOC problem equivalent to the above problem (Eq.

Revised: Theorem 8 (Dynamic formulation (Eularian formalism); (L´eonard, 2013)) .

Revised: Let u ∗ t ( x ) =

Revised: ∇ x log E [Ψ( X T , T ) | X t = x ] , where Ψ satisfies the SB optimality (Eq.

Revised: Then, u ∗ t is the minimizer of the following optimization problem:

Revised: [Equation]

Revised: where p t is the probability density of the probability measure ρ t .

Revised: The first condition in Eq.

Revised: (25) is Fokker-Planck (FP) equation.

Revised: The two versions of the dynamic formulation (Eq.

Revised: (22) and Eq.

Original: The domain Ω satisfies the no-flux condition:
Revised: (24)) are equivalent, i.e . V S,ϵ ( µ 0 , µ 1 ) = v S,ϵ ( µ 0 , µ 1 ) . Furthermore, the dynamic solution of the OT problem (Eq. (16)) is obtained as the zero-noise limit (Mikami, 2004;

Revised: L´eonard, 2012) of the classical SB problem.

Revised: Theorem 9 (Zero-noise limit of the classical SB problem; (Mikami, 2004)) .

Revised: Let X ϵ ( t ) be the solution of the SB problem (Eq.

Revised: Suppose that µ 0 , µ 1 ∈ P ( R d ) have finite second moments and the density function p 0 ( x ) :

Revised: = µ 0 (d x ) /

Revised: d x exists.

Revised: Then, the following holds

Revised: [Equation]

Revised: [Equation]

Revised: and there exists a convex function φ satisfying

Original: [Equation]
Revised: [Equation]

Revised: The map ∇ x φ is the optimal transport map of the OT problem with a quadratic cost.

Revised: The connection between the classical SB problem and the dynamic formulation of the OT problem with Fisher information regularization (Chen et al., 2016) is also well-known.

Revised: A. S TOCHASTIC O PTIMAL T RANSPORT WITH T WO E NDPOINT M ARGINALS

Revised: Mikami (2008) generalized the OT problem and defined the SOT problem as a random mechanics problem determined by the principle of least action .

Revised: The SOT problem with the endpoint marginals fixed to µ 0 and µ 1 is given by the following.

Revised: Definition 10 (SOT problem; (Mikami, 2021)) .

Revised: Let L be a continuous function, the Lagrangian and let u (cid:55)→ L ( t, x , u ) be convex.

Revised: The SOT problem with two endpoint marginals is defined by

Revised: [Equation]

Revised: where A is the set of all R d -valued, continuous semimartingales { X t } 0 ≤ t ≤ 1 on a complete filtered probability space such that there exists a Borel measurable drift function β X ( X , t ) for which satisfies the following conditions:

Revised: [Equation]

Revised: Definition 11 (SOT problem for marginal flows; (Mikami, 2008)) .

Revised: [Equation]

Revised: where p 0 and p 1 are the densities of probability measures ρ 0 and ρ 1 , respectively.

Revised: The minimizer of the SOT problem for marginal flows (Eq.

Revised: (29)) is obtained by f ∗ ( x , t ) =

Revised: E [ β X ( X , t ) | ( t, X t = x )] .

Revised: We introduce assumptions from (Mikami, 2021).

Revised: (A-1)

Revised: L ∈ C 1 ( R d × R d ; [0 , ∞ ]) .

Original: (13) where n is the unit outward normal vector to the boundary of Ω .
Revised: u (cid:55)→ L ( t, x , u ) is strictly convex. L ( t, x , u ) / (1 + L ( t, y , u ) and |∇ x

Revised: L

Revised: ( t, x , u )

Revised: |

Revised: + L ( t,

Revised: x

Revised: , u ))

Revised: are bounded

Revised: on

Revised: t

Revised: ∈

Revised: and x , y , u ∈ R d .

Revised: sup

Revised: x ∈ R d |∇ x L ( t, x , u ) | 1 is

Revised: locally bounded

Revised: lim | u |→∞ inf L ( t, x , u ) / | u | = ∞ .

Revised: (A-2) ∇ 2 u L ( t, x , u ) is bounded uniformly nondegenerate on [0 , 1] × R d × R d .

Revised: Theorem 12 (Trevisan’s Superposition Principle; (Trevisan, 2016)) .

Revised: Assume that there exists f : R d × [0 , 1] (cid:55)→ R d and { ρ t } 0 ≤ t ≤ 1 ⊂ P ( R d ) such that f satisfies the FP equation,

Revised: i.e . f ∈ A (cid:16) { ρ t } 0 ≤ t

Revised: [Equation]

Revised: Then, there exists a semimartingale { X t } 0 ≤ t ≤ 1 for which the following holds:

Revised: [Equation]

Revised: From Theorem 12, it can be easily shown that V ( µ 0 , µ 1 ) =

Revised: v( µ 0 , µ 1 ) and there exist minimizers X ∗ of the SOT problem (Eq.

Revised: for which β X ∗ ( X ∗ , t ) =

Revised: f ∗ ( X ∗ t , t ) .

Revised: Theorem 13 (Duality theorem; (Mikami, 2021)) .

Revised: Suppose that (A-1) holds.

Revised: Then, for any µ 0 , µ 1 ∈

Revised: P ( R d ) ,

Revised: [Equation]

Revised: where C ∞ b ( R d ) is the set of all infinitely differentiable functions on R d , which have bounded continuous derivative and Φ( x , t ; f ) is the viscosity solution to the HJB equation:

Revised: [Equation]

Revised: The Hamiltonian H is defined by H ( t, x , z ) :

Revised: = sup u {⟨ z , u ⟩ − L ( t, x , u ) } .

Revised: Theorem 14 ((Mikami & Thieullen, 2006)) .

Revised: Suppose that (A-1) and (A-2) hold and V ( µ 0 , µ 1 ) is finite.

Revised: Then, there exists a minimizer X ∗ ∈

Revised: A of V ( µ 0 , µ 1 ) given by

Revised: [Equation]

Revised: For any maximizing sequence { Φ n } n ≥ 1 of Eq.

Revised: (34) , there exists a subsequence { n k } k ≥ 1 such that

Revised: [Equation]

Revised: A. L AGRANGIAN S CHR ¨ ODINGER B RIDGE P ROBLEM

Revised: We consider the LSB problem constrained by the FP equation corresponding to Ito SDE (Eq.

Revised: (3)) over the Euclidean space R d .

Revised: Definition 15 (LSB problem constrained by the FP equation) .

Revised: [Equation]

Revised: The the LSB problem in Eq.

Revised: (39) is a more general problem that does not fix D i,j in the SOT problem for marginal flows (Eq.

Revised: Thus, a solution to the LSB problem in Eq.

Revised: (39) clearly exists on the basis of the existence of a solution to the SOT problem in Eq.

Revised: We solve the following relaxed

Revised: LSB problem, where the terminal constraint are replaced by the soft constraint.

Revised: Definition 16 (Relaxed LSB problem constrained by the FP equation) .

Revised: [Equation]

Original: {+where G is the terminal cost introduced by relaxing the constraint ρ 1 = µ 
+}
Revised: {+where G is the terminal cost introduced by relaxing the constraint ρ 1 = µ 
+}

Revised: and δδρ 1 is the variational derivative with respect to ρ 1 .

Original: We will solve the LSB problem as shown in Eq.
Revised: We derive the optimality conditions for the LSB problems in Eqs.

Revised: (39) and (42) using variational method in the following theorem.

Revised: Theorem 17 (Optimality conditions for the LSB problem) .

Revised: There exists a space-time dependent potential function Φ: R d × [0 , 1] (cid:55)→ R , which satisfies:

Revised: [Equation]

Revised: Equation 45 and 46 are the optimality conditions for both the relaxed LSB problem (Eq. (42) and the

Original: (6) by reformulating it as a saddle point problem for ( p t , m t ) :
Revised: LSB problem in Eq. Proof. We derive the optimality conditions for the relaxed LSB problem as shown in Eq. (42) by reformulating it as a saddle point problem for ( p t , m t ) :

Original: = ( p t , p t f t ) .
Revised: = ( p t , p t f t ) .

Original: Let L be the the Lagrangian with the time-space-dependent Lagrange multiplier Φ( x , t ) .
Revised: Let L be the the Lagrangian with the time-space-dependent Lagrange multiplier Φ( x , t ) .

Original: [Equation]
Revised: [Equation]

Original: The term (a) is transformed by performing a partial integral over t :(cid:90) t
Revised: The term (a) is transformed by performing a partial integral over t :

Original: [Equation]
Revised: [Equation]

Original: The term (b) is simplifiedby the Stokes’ theorem and the no-flux condition, as shown in Eq. (13): (cid:90)
Revised: The term (b) is simplified as follows.

Original: [Equation]
Revised: [Equation]

Original: The term (c) is transformed by performing a partial integral over x :
Revised: The term (c) is transformed by performing a partial integral over x :

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: we can rewrite the LSB problem as
Revised: we can rewrite the LSB problem as

Original: [Equation]
Revised: [Equation]

Original: The saddle point ( p ∗ , m ∗ , Φ ∗ ) of the problem in Eq.
Revised: The saddle point ( p ∗ , m ∗ , Φ ∗ ) of the problem in Eq.

Original: (14) satisfies the following conditions:
Revised: (47) satisfies the following conditions:

Original: [Equation]
Revised: [Equation]

Revised: where z ∗ := ∇ f L ( t, x , f ∗ ) and the Hamiltonian is defined by H ( t, x , z ) :

Revised: = sup u {⟨ z , u ⟩ − L ( t, x , u ) } .

Revised: The Lagrangian satisfies L ( t, x , u ) =

Original: Therefore, the optimal drift function is given byf ∗ = m ∗ p ∗ =
Revised: sup z {⟨ u , z ⟩ − H ( t, x , z ) } . Therefore [Equation] the optimal drift function is given by

Original: [ ∇ u L ] − 1 ( −∇ x Φ) and the potential function Φ is the solution of the HJB equation: 
Revised: [Equation] [Equation] The potential function Φ is the solution of the HJB equation: [Equation]

Revised: Equation 53 and 54 are also optimal conditions even for the LSB problem in Eq.

Revised: (39) where the terminal constraint is strictly satisfied, i.e . G ( x ) = 0 .

Revised: From Theorem 12, we can show that there exists a semimartingale { X t } 0 ≤ t ≤ 1 for which the following holds:

Original: [Equation]
Revised: [Equation]

Original: B M ODEL A RCHITECTURE
Revised: where ( f ∗ , ρ ∗ t ) is the minimizer of the LSB problems in Eqs. (39) and (42)) and g ∗ is the diffusion function determined from the minimizer ( f ∗ , ρ ∗ t ) . Therefore, the LSB problems constrained by the

Revised: FP equation (Eqs.

Revised: (39) and (42))

Revised: and the original LSB problem defined by Eq.

Revised: (6) are equivalent.

Revised: A. D ERIVATION OF NLSB AND ITS CONNECTION TO OT-F LOW

Revised: The loss function of the NLSB shown in Eq.

Revised: (7) is the objective function of the relaxed LSB problem (Eq. (42)) with additional PDE loss applied to satisfy the optimality condition, HJB equation.

Revised: Theorem 17 justifies the minimization of HJB-PDE loss and a parameterization that strictly satisfies Eq.

Revised: (45) during training when the terminal constraint is not satisfied, i.e . G ( x ) ̸ = 0 .

Revised: For computational efficiency, we practically define the HJB-PDE loss R h in a weak form as shown in Eq.

Revised: (9) and evaluate R h on the path of the simulated SDE.

Revised: NLSB and OT-Flow (Onken et al., 2021) have strong theoretical connections.

Revised: NLSB is the practical solution to the LSB problem (Eq.

Revised: (6)) relaxed by the distribution discrepancy measure

Revised: D using neural SDE, while OT-Flow is the solution to the Brenier-Benamou formulation of the OT problem (Eq.

Revised: (12)) relaxed by the KL-divergence using neural ODE.

Revised: The optimality conditions for the LSB problem shown in Theorem 17 are analogues of the optimality conditions for the Brenier-Benamou problem shown in Theorem 4.

Revised: In NLSB and OT-Flow, the potential function is modeled using

Revised: NN and optimized by both action cost (Eqs.

Revised: (1) and (8)) and HJB-PDE loss

Revised: Eqs.

Revised: (2) and (9)) based on the respective optimality conditions.

Revised: Theoretically, the LSB problem with the Lagrangian

Revised: L ( t, x , u ) = 12 ∥ u ∥ 2 and the diffusion function g = √ 2 ϵ is reduced to the classical SB problem (Eq.

Revised: and the solution to the Brenier-Benamou problem is recovered from the zero-noise limit of the solution to the classical SB problem (Theorem 9).

Original: The model structure of the potential function proposed in OT-Flow (Onken et al., 2021) is shown below.
Revised: The model structure of the potential function proposed in OT-Flow (Onken et al., 2021) is shown below.

Original: [Equation]
Revised: [Equation]

Original: where s = ( x , t ) ∈
Revised: where s = ( x , t ) ∈

Original: R
Revised: R

Original: d
Revised: d

Original: +1 is a input vector, w , K
Revised: +1 is a input vector, w , K

Original: 0 ∈
Revised: 0 ∈

Original: R m × ( d +1) , K i [Equation] 1 ≤ i ≤ M ) ∈
Revised: R m × ( d +1) , K i (1 ≤ i ≤ M ) ∈

Original: R m × m

Original: [Equation]

Original: i

Original: [Equation] i ≤ M ) ∈ R m × m , A , b , and c are learnable parameters,  h is a fixed step size, and the activation function is defined by σ ( x ) =
Revised: R m × m , b i (0 ≤ i ≤ M ) ∈ R m × m , A , b , and c are learnable parameters, m is the number of dimensions of the hidden representation vector, h is a fixed step size, and the activation function is defined by σ ( x ) =

Original: log(exp( x ) + exp( − x )) .
Revised: log(exp( x ) + exp( − x )) .

Original: The gradient of the potential function is described .
Revised: The gradient of the potential function is described .

Original: [Equation]
Revised: [Equation]

Original: We can write σ ′ down as tanh since σ is defined as above.
Revised: We can write σ ′ down as tanh since σ is defined as above.

Original: Finally, the diagonal components of the potential function’s Hessian is shown below.
Revised: Finally, the diagonal components of the potential function’s Hessian is shown below.

Original: [Equation]
Revised: [Equation]

Original: where ⊙ is the element-wise product [Equation] ∇ s u i can be obtained by using the following update equation:
Revised: where ⊙ is the element-wise product, ∇ s u i can be obtained by using the following update equation:

Original: [Equation]
Revised: [Equation]

Original: The diagonal component of the potential function Hessian can be computed at a computation cost of
Revised: The diagonal component of the potential function Hessian can be computed at a computation cost of

Original: O ( m 2 d ) .
Revised: O ( m 2 d ) .

Original: C S CHR ¨ ODINGER B RIDGE AND L AGRANGIAN S CHR ¨ ODINGER B RIDGE

Original: We briefly explain a connection between the SB problem and the LSB problem. Definition 2 (SB problem; (Jamison, 1975)) .
Revised: The complexity O ( m 2 d ) indicates that there is a trade-off between the expressive power of the DNN and the computational cost of the Hessian.

Original: Let Ω = C
Revised: OT-Flow modeled the potential function Φ θ that satisfies v θ = −∇ x

Original: ([0 , T ] , R d ) be the space of R d -valued continuous functions on time interval [0 , T ] . Denote by P (Ω) the probability measures space on the path pace Ω . The SB problem is defined by
Revised: Φ θ instead of directly modeling the velocity function v of the neural ODE. The continuous transformation from x (0) to x (1) is described by

Original: [Equation]
Revised: [Equation]

Original: where µ, ν ∈ P ( R d ) and the relative entropy D KL = (cid:82) log (cid:0) d Q d P (cid:1) d Q if Q ≪ P , and D KL = ∞ otherwise. The case of no prior dynamics is considered classically, i.e ., the reference path measure P is the Brownian diffusion SDE d X t = √ 2 ϵ d W t , where W t is  the standard Wiener process.
Revised: OT-Flow is trained by likelihood maximization in the well-known CNF framework.

Original: Chen et al. (2021b) derived the optimality condition for problem 16, which is characterized by the forward and backward time-harmonic equations. Theorem 3. Let Ψ( x , t ) and ˆΨ( x , t ) be the solutions to the following PDEs:
Revised: The likelihood computation is performed as follows.

Original: [Equation]
Revised: [Equation]

Original: where p 0 and p T are the probability density of µ 0 and µ T , respectively.

Original: Then, the solution to the SB problem (16) can be described by the following forward or backward SDE:

Original: [Equation]
Revised: [Table]

Original: Dai Pra (1991) considered a dynamic formulation of the SB problem (16) by interpreting it as a SOC problem with the additional terminal constraint. Theorem 4. Let u ∗ = 2 ϵ ∇ x log Ψ , where Ψ satisfies the SB optimality (17).
Revised: where p 1 is the density of the target distribution and p 0 is the density of the prior distribution, which is usually a Gaussian or Laplace distribution. The fast and exact computation of the diagonal component of the potential function’s Hessian is useful for the computation of Eq.

Original: Then, u ∗ is the minimizer of thefollowing optimization problem: [Equation]
Revised: In contrast, we propose for the first time to use this technique to speed up the computation of the

Original: where U = (cid:8) u : [0 , T ] × R d (cid:55)→ R d | ⟨ u , u ⟩ < ∞ (cid:9) . L ´ eonard (2013) formulated the SB problem16 into a variational SOC problem equivalent to the above problem 18.
Revised: HJB-PDE loss R h in the NLSB under the assumption that the diffusion model’s output is a diagonal matrix.

Original: Theorem 5. Let v ∗ t ( x ) = ∇ x log E [Ψ( X T , T ) | X t = x ] , where Ψ satisfies the SB optimality (17). Then, v ∗ t is the minimizer of the following optimization problem:
Revised: The computed loss R h is given by

Original: [Equation]
Revised: [Equation]

Original: where p t is the probability density of the probability measure ρ t .

Original: The generalizations of problems 18 and 19 correspond to problems 4 and 5, respectively.

Original: When the Lagrangian L ( t, x , u ) = 12 || u || 2 and the diffusion function is fixed, the LSB problem 6 is also reduced to the problem 19.

Original: D E XPERIMENTAL D ETAILS
Revised: C E XPERIMENTAL D ETAILS

Original: In our implementation, we modified code in the TrajectoryNet 1 , OT-FLow 2 and IPML 3 repositories, which were released under the MIT license.
Revised: In our implementation, we modified code in the TrajectoryNet 1 , OT-FLow 2 and IPML 3 repositories, which were released under the MIT license.

Original: The scRNA-seq data used in Section 5.2 are licensed under Creative Commons Attribution 4.0 International license.
Revised: The scRNA-seq data used in Section 5.2 are licensed under Creative Commons Attribution 4.0 International license.

Original: Our experimental environment consists of an Intel Xeon Plantinum 8360Y (36-core) CPU and a single NVIDIA A100 GPU.
Revised: Our experimental environment consists of an Intel Xeon Plantinum 8360Y (36-core) CPU and a single NVIDIA A100 GPU.

Original: Code is available in the supplementary material.
Revised: Code is available in the supplementary material.

Original: D. T RAINING C ONFIGURATION FOR E XPERIMENTS ON S YNTHETIC D ATA
Revised: C. T RAINING C ONFIGURATION FOR E XPERIMENTS ON S YNTHETIC D ATA

Original: We describe the details of the training setup for the experiments on synthetic data conducted in Section 5.1.https:// github
Revised: We describe the details of the training setup for the experiments on synthetic data conducted in Section 5.1.

Original: com/EmoryMLIP/OT-Flowhttps://github.com/AforAnonyMeta/

Original: [Table]

Original: C OMMON SETTINGS FOR NLSB, NEURAL SDE, T RAJECTORY N ET AND OT-F LOW
Revised: C OMMON SETTINGS FOR NLSB, NEURAL SDE, T RAJECTORY N ET AND OT-F LOW

Original: We generated 2048 and 512 samples for each time point as training and validation data, respectively.
Revised: We generated 2048 and 512 samples for each time point as training and validation data, respectively.

Original: We trained all models with a batch size of 512 for each time point and used the early stopping method, which monitors the EMD-L2 value on validation data.
Revised: We trained all models with a batch size of 512 for each time point and used the early stopping method, which monitors the EMD-L2 value on validation data.

Original: We used Adam optimizer to optimize all learnable parameters with a learning rate of 0 .
Revised: We used Adam optimizer to optimize all learnable parameters with a learning rate of 0 .

Original: 001 and the decay rate of β 1 = 0 .
Revised: 001 and the decay rate of β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: We searched all weight coefficients of regularization terms in [0 . 0 , 0 . 5] and selected those with the largest possible coefficients among those with sufficiently small EMD-L2 values on the validation data.
Revised: We searched all weight coefficients of regularization terms in [0 . 0 , 0 . 5] and selected those with the largest possible coefficients among those with sufficiently small EMD-L2 values on the validation data.

Original: S ETTINGS FOR NLSB AND NEURAL SDE
Revised: S ETTINGS FOR NLSB AND NEURAL SDE

Original: For the potential model Φ θ ( s ) [Table] we set the number of ResNet layers M = 2 , the step size h = 1 . 0 , the rank of matrix rank( A ) = [Table]
Revised: For the potential model Φ θ ( s ) , we set the number of ResNet layers M = 2 , the step size h = 1 . 0 , the rank of matrix rank( A ) = 10 , and the dimension of the hidden vector z to 2 .

Original: and the dimension of the hidden vector z to 2 .

Original: We used a two-layer
Revised: We used a two-layer

Original: FCNN of a hidden dimension 16 for the diffusion function.
Revised: FCNN of a hidden dimension 16 for the diffusion function.

Original: The activations functions were LipSwish.
Revised: The activations functions were LipSwish.

Original: We trained the neural SDE in the framework of the torchsde library 4 and used the Euler-Maruyama method with the constant step size of 0 .
Revised: We trained the neural SDE in the framework of the torchsde library 4 and used the Euler-Maruyama method with the constant step size of 0 .

Original: 01 as an SDE solver.
Revised: 01 as an SDE solver.

Original: Backpropagation was performed without using the adjoint method.
Revised: Backpropagation was performed without using the adjoint method.

Original: In NLSB, we used the Lagrangian for the potential-free system
Revised: In NLSB, we used the Lagrangian for the potential-free system

Original: L ( t, x , u ) =
Revised: L ( t, x , u ) =

Original: 12 || u
Revised: 12 || u

Original: || 2 .
Revised: || 2 .

Original: We set the weight coefficients λ e , λ h shown in Table 3.
Revised: We set the weight coefficients λ e , λ h shown in Table [Table]

Original: S ETTINGS FOR T RAJECTORY N ET
Revised: S ETTINGS FOR T RAJECTORY N ET

Original: We used three concatsquash layers of a hidden dimension 16 with hyperbolic tangent activations to model the velocity function v .
Revised: We used three concatsquash layers of a hidden dimension 16 with hyperbolic tangent activations to model the velocity function v .

Original: A concatsquash layer cs was defined in the released code of FFJORD (Grathwohl et al., 2018) as:
Revised: A concatsquash layer cs was defined in the released code of FFJORD (Grathwohl et al., 2018) as:

Original: [Equation]
Revised: [Equation]

Original: where σ is the sigmoid function, and W x , W t , W b , b x , b t , b b are all learnable parameters.
Revised: where σ is the sigmoid function [Table] and W x [Table]

Revised: W t

Revised: [Table]

Revised: W b

Revised: [Table]

Revised: b x

Revised: [Table]

Original:  We trained the model with torchdiffeq library 5 and used the dopri5 solver with both absolute and relative tolerances set to 10 − 5 .
Revised: b t [Table] b b are all learnable parameters. We trained the model with torchdiffeq library 5 and used the dopri5 solver with both absolute and relative tolerances set to 10 − 5 .

Original: In the training for TrajectoryNet with OT-based regularization, we set the weight coefficients shown in Table 3.
Revised: In the training for TrajectoryNet with OT-based regularization [Table] we set the weight coefficients shown in Table 3.

Original: S ETTINGS FOR OT-F LOW
Revised: S ETTINGS FOR OT-F LOW

Original: We used the same potential model Φ in NLSB and neural SDE with the same hyperparameters.
Revised: We used the same potential model Φ in NLSB and neural SDE with the same hyperparameters.

Original: We trained the model with torchdiffeq library.
Revised: We trained the model with torchdiffeq library.

Original: We used the Euler method as the ODE solver with a constant step size of 0 .
Revised: We used the Euler method as the ODE solver with a constant step size of 0 .

Original: 01 and both absolute and relative tolerances set to 10 − 5 .
Revised: 01 and both absolute and relative tolerances set to 10 − 5 .

Original: In the training for
Revised: In the training for

Original: OT-Flow with OT-based regularization, we set the weight coefficients shown in Table 3.
Revised: OT-Flow with OT-based regularization, we set the weight coefficients shown in Table 3.

Original: S ETTINGS FOR IPML
Revised: S ETTINGS FOR IPML

Original: We generated 2048 and 512 samples for each time point as training and validation data, respectively.
Revised: We generated 2048 and 512 samples for each time point as training and validation data, respectively.

Original: The drift model of IPML (GP) is sparse GP with the exponential kernel.
Revised: The drift model of IPML (GP) is sparse GP with the exponential kernel.

Original: We selected 100 inducing points using the K-means algorithm.
Revised: We selected 100 inducing points using the K-means algorithm.

Original: The drift model of IPML (NN) is the same network as for NLSB, including hyperparameters.
Revised: The drift model of IPML (NN) is the same network as for NLSB, including hyperparameters.

Original: We solved the SB problem by using IPF algorithm with 15 iterations.We
Revised: We solved the SB problem by using IPF algorithm with 15 iterations.

Original: [Table] used the Euler-Maruyama method as the SDE solver and set the diffusion coefficients [Table] as follows.
Revised: We used the Euler-Maruyama method as the SDE solver and set the diffusion coefficients, as follows.

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Table]

Original: [Equation] [Equation] [Equation] Other experimental settings have not been changed from (Vargas et al., 2021).
Revised: Other experimental settings have not been changed from (Vargas et al., 2021).

Original: D. T RAINING C ONFIGURATION FOR E XPERIMENTS ON S C RNA- SEQ D ATA
Revised: C. T RAINING C ONFIGURATION FOR E XPERIMENTS ON S C RNA- SEQ D ATA

Original: We describe the details of the training setup for the experiments on scRNA-seq data conducted in Section 5.2.
Revised: We describe the details of the training setup for the experiments on scRNA-seq data conducted in Section 5.2.

Original: C OMMON SETTINGS FOR NLSB, NEURAL SDE, T RAJECTORY N ET AND OT-F LOW
Revised: C OMMON SETTINGS FOR NLSB, NEURAL SDE, T RAJECTORY N ET AND OT-F LOW

Original: We reused the pre-processed data available in the released repository of TrajectoryNet 1 and split the data into 200 samples ( ∼ 8 .
Revised: We reused the pre-processed data available in the released repository of TrajectoryNet 1 and split the data into 200 samples ( ∼ 8 .

Original: 5% ) of validation data, 350 samples ( ∼ 15% ) of test data, and the rest as train data for each time point.
Revised: 5% ) of validation data, 350 samples ( ∼ 15% ) of test data, and the rest as train data for each time point.

Original: We trained all models with a batch size of 1000 for each time point and used the early stopping method, which monitors the EMD-L2 value on the validation data.
Revised: We trained all models with a batch size of 1000 for each time point and used the early stopping method, which monitors the EMD-L2 value on the validation data.

Original: We used Adam optimizer for training all methods with learning rate 0 .
Revised: We used Adam optimizer for training all methods with learning rate 0 .

Original: 001 and the decay rate β 1 = 0 .
Revised: 001 and the decay rate β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: We searched all weight coefficients of regularization terms in [0 . 0 , 0 . 5] and selected those with the largest possible coefficients among those with sufficiently small EMD-L values on the validation data.
Revised: We searched all weight coefficients of regularization terms in [0 . 0 , 0 . 5] and selected those with the largest possible coefficients among those with sufficiently small EMD-L values on the validation data.

Original: S ETTINGS FOR NLSB
Revised: S ETTINGS FOR NLSB

Original: All settings for the potential model, the diffusion function model, and the SDE solver were the same as in Appendix D.1.
Revised: All settings for the potential model, the diffusion function model, and the SDE solver were the same as in Appendix C.1.

Original: We set the noise type to diagonal.
Revised: We set the noise type to diagonal.

Original: We set the weight coefficients of regularization terms shown in Table 4.
Revised: We set the weight coefficients of regularization terms shown in Table 4.

Original: For the calculation of the density regularization term, the time-dependent density function U ( x , t ) was defined by
Revised: For the calculation of the density regularization term, the time-dependent density function U ( x , t ) was defined by

Original: [Equation]
Revised: [Equation]

Original: where c is a hyperparameter to change the scale, and µ ( t ) m , Σ ( t ) m are mean and variance parameters of the mixed Gaussian distribution, respectively.
Revised: where c is a hyperparameter to change the scale, and µ ( t ) m , Σ ( t ) m are mean and variance parameters of the mixed Gaussian distribution, respectively.

Original: When t ∈ [ t k , t k +1 ] , the parameters Θ t were estimated with the data at t k and t k +1 by GMM and the number of mixture components M t was determined by the value of Bayesian information criterion (BIC).
Revised: When t ∈ [ t k , t k +1 ] , the parameters Θ t were estimated with the data at t k and t k +1 by GMM and the number of mixture components M t was determined by the value of Bayesian information criterion (BIC).

Original: The hyperparameter c was searched among { 0 . 1 , 1 . 0 , 10 . 0 } and we set c = 10 . 0 for NLSB (D) and c = 0 .
Revised: The hyperparameter c was searched among { 0 . 1 , 1 . 0 , 10 . 0 } and we set c = 10 . 0 for NLSB (D) and c = 0 .

Original: 1 for NLSB (E+D+V).
Revised: 1 for NLSB (E+D+V).

Original: For the calculation of the velocity regularization term, we used the same reference velocity as those used in
Revised: For the calculation of the velocity regularization term, we used the same reference velocity as those used in

Original: TrajectoryNet (Tong et al., 2020).
Revised: TrajectoryNet (Tong et al., 2020).

Original: S ETTINGS FOR NEURAL SDE, T RAJECTORY N ET AND OT-F LOW
Revised: S ETTINGS FOR NEURAL SDE, T RAJECTORY N ET AND OT-F LOW

Original: We used the velocity model of TrajectoryNet with 64 hidden dimensions.
Revised: We used the velocity model of TrajectoryNet with 64 hidden dimensions.

Original: All other settings such as potential models, diffusion function models, SDE solvers, and ODE solvers were the same as in Appendix D.1.
Revised: All other settings such as potential models, diffusion function models, SDE solvers, and ODE solvers were the same as in Appendix C.1.

Original: In the training for TrajectoryNet and OT-Flow with OT-based regularization, we set the weight coefficients shown in Table 4.
Revised: In the training for TrajectoryNet and OT-Flow with OT-based regularization, we set the weight coefficients shown in Table 4.

Original: S ETTINGS FOR IPML
Revised: S ETTINGS FOR IPML

Original: We set the diffusion coefficients to g ( t ) = 0 .
Revised: We set the diffusion coefficients to g ( t ) = 0 .

Original: 5( t ∈ [0 . 0 , 1 . 0]) , 1 . 

Revised: 5 ( t ∈ [0 . 0 , 1 . 0]) , 1 .

Original: ( t ∈ [1 . 0 , 4 . 0]) .
Revised: ( t ∈ [1 . 0 , 4 . 0]) .

Original: All other settings were the same as in Appendix D.1.
Revised: All other settings were the same as in Appendix C.1.

Original: D. T RAINING C ONFIGURATION FOR A DDITIONAL E XPERIMENTS ON S YNTHETIC D ATA
Revised: C. T RAINING C ONFIGURATION FOR A DDITIONAL E XPERIMENTS ON S YNTHETIC D ATA

Original: We describe the details of the training setup for the addtional experiments on synthetic data conducted in Section E.3.
Revised: We describe the details of the training setup for the addtional experiments on synthetic data conducted in Appendix D.3.

Original: D ATA
Revised: D ATA

Original: We used two-dimensional uniform distributions U 0 and U 1 for the endpoints at time t = 0 and t = 1 .
Revised: We used two-dimensional uniform distributions U 0 and U 1 for the endpoints at time t = 0 and t = 1 .

Original: [Equation]
Revised: [Equation]

Original: We generated 2048 and 512 samples from two endpoint distributions as training and validation data, respectively.
Revised: We generated 2048 and 512 samples from two endpoint distributions as training and validation data, respectively.

Original: S ETTINGS FOR NLSB
Revised: S ETTINGS FOR NLSB

Original: We set the Lagrangian for the random dynamical system with the mass m = 1 ,
Revised: We set the Lagrangian for the random dynamical system with the mass m = 1 ,

Original: i.e . L ( t, x , u ) =
Revised: i.e . L ( t, x , u ) =

Original: 12 || u
Revised: 12 || u

Original: ||
Revised: ||

Original: − U ( x ) and compared several practical potential functions U shown below.
Revised: − U ( x ) and compared several practical potential functions U shown below.

Original: Box-shaped obstacle 
Revised: Box-shaped obstacle U ( x, y )

Original: [Equation] − [Equation]
Revised: = (cid:26) − 100 − 0 .

Original: 5 ≤ x [Equation]
Revised: 5 ≤ x, y ≤ 0 .

Original: otherwise 
Revised: otherwise .

Original: [Equation] Slit-shaped obstacle 
Revised: Slit-shaped obstacle U ( x, y )

Original: [Equation]
Revised: = (cid:26) − 100 − 0 .

Original: [Equation]
Revised: 1 ≤ x ≤ 0 .

Original: [Equation] 25 ≤ y [Equation]
Revised: 1 ∧ ( y ≤ − 0 . 25 ∨ 0 .

Original:  otherwise [Equation]
Revised: 25 ≤ y ) otherwise .

Original: Hill potential
Revised: Well potential

Original: [Equation]
Revised: [Equation]

Original: Well potential
Revised: Hill potential

Original: [Equation]
Revised: [Equation]

Original: We defined the box and slit-shaped obstacle potential functions as differentiable by using the sigmoid functions.
Revised: We defined the box and slit-shaped obstacle potential functions as differentiable by using the sigmoid functions.

Original: We trained NLSB with a batch size of 512 and used the early stopping method, which monitors the validation loss value.
Revised: We trained NLSB with a batch size of 512 and used the early stopping method, which monitors the validation loss value.

Original: We used Adam optimizer to optimize all learnable parameters with a learning rate of 0 .
Revised: We used Adam optimizer to optimize all learnable parameters with a learning rate of 0 .

Original: 001 and the decay rate of β 1 = 0 .
Revised: 001 and the decay rate of β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: All weight coefficients of regularization terms were searched uniformly in { 0 . 01 , 0 . 001 } .
Revised: All weight coefficients of regularization terms were searched uniformly in { 0 . 01 , 0 . 001 } .

Original: All settings for the potential model Φ θ , the diffusion function model g ϕ , and the SDE solver were the same as in Appendix D.1.
Revised: All settings for the potential model Φ θ , the diffusion function model g ϕ , and the SDE solver were the same as in Appendix C.1.

Original: E A DDITIONAL E XPERIMENTAL R ESULTS E. S YNTHETIC P OPULATION D YNAMICS : T IME -D EPENDENT O RNSTEIN –U HLENBECK P ROCESS
Revised: D A DDITIONAL E XPERIMENTAL R ESULTS AND A NALYSIS D. S YNTHETIC P OPULATION D YNAMICS : T IME -D EPENDENT O RNSTEIN –U HLENBECK P ROCESS

Original: The figure 6 shows the visualization of both the original trajectory and the averaged trajectory by the
Revised: Figure 6 shows the visualization of both the original trajectory and the averaged trajectory by the

Original: SDE-based methods (Neural SDE, NLSB, and IPML) and the only original trajectory by the ODEbased methods (TrajectoryNet and OT-Flow).
Revised: SDE-based methods (Neural SDE, NLSB, and IPML) and the only original trajectory by the ODEbased methods (TrajectoryNet and OT-Flow).

Original: All trajectories were generated by all-step prediction from the initial samples at the time t = 0 .
Revised: All trajectories were generated by all-step prediction from the initial samples at the time t = 0 .

Original: The five colored point clouds in the background are the ground-truth data given at each time point.
Revised: The five colored point clouds in the background are the ground-truth data given at each time point.

Original: The pink area and the light blue line are the one-sigma empirical confidence intervals and their boundaries for each trajectory, respectively.
Revised: The pink area and the light blue line are the one-sigma empirical confidence intervals and their boundaries for each trajectory, respectively.

Revised: Figure 3 shows that NLSB and IPML outperform neural SDE and is comparable to other ODEbased methods in estimating populations with small variance.

Revised: In contrast, the SDE-based methods outperform ODE-based methods when estimating populations with a large variance.

Revised: That indicates that NLSB and IPML can estimate population-level dynamics even when the population variance is large or small.

Original: [Figure]
Revised: Furthermore, NLSB and IPML have a smaller CDD value than neural SDE. Figure 6b shows that the average behavior of samples E [ X ( t ) | X (0)] estimated by neural SDE is different from that of the ground-truth SDE (see Fig. 6a), especially in the interval [0 , 1] . In contrast, the predictions by NLSB and IPML in Figs.

Revised: 6c to 6e are much closer to the ground-truth, suggesting that IPML and

Revised: NLSB almost achieve the upper limit of performances.

Revised: These results show that the prior knowledge of the potential-free system helps to estimate the sample-level dynamics.

Revised: Note that the LSB problem solved by NLSB with the Lagrangian L = 12 ∥ u ∥ 2 and the SB problem solved by IPML are almost mathematically equivalent (see Appendices A.3 to A.5).

Revised: The result that

Revised: NLSB shows comparable performance to IPML, even though NLSB is trained differently from IPML, indicates that NLSB is a unified framework and can deal with SB problems as a special case.

Revised: [Figure]

Original: E. S INGLE -C ELL P OPULATION D YNAMICS
Revised: D. S INGLE -C ELL P OPULATION D YNAMICS

Original: The time evolution of the distribution for a single sample is visualized as a heat map in Figure 7.
Revised: The time evolution of the distribution for a single sample is visualized as a heat map in Fig.

Original: The sample population in Figure 8 and the trajectories are visualized in Figure 9.
Revised: The sample population in Fig. 8 and the trajectories are visualized in Fig.

Original: All figures are visualizations in the space of the first and second principal components.
Revised: All figures are visualizations in the space of the first and second principal components.

Original: The x-axis represents the first principal component and the y-axis the second principal component.
Revised: The x-axis represents the first principal component and the y-axis the second principal component.

Original: Figure 7 shows that the SDE-based methods, including NLSB, can handle the uncertainty of the trajectories in contrast to ODE-based methods.
Revised: Figure 7 shows that the SDE-based methods, including NLSB, can handle the uncertainty of the trajectories in contrast to ODE-based methods.

Original: Figure 8 shows that NLSB outperforms ODE-based methods in predicting the transitions with a high diffusion of samples from t 1 to t 2 and from t 3 to t 4 , indicating that the explicit modeling of diffusion is effective.
Revised: Figure 8 shows that NLSB outperforms ODE-based methods in predicting the transitions with a high diffusion of samples from t 1 to t 2 and from t 3 to t 4 , indicating that the explicit modeling of diffusion is effective.

Original: Figure 9 shows that the drift estimated by NLSB (E) is linear, the trajectories by NLSB (D) pass on the data manifold, and the trajectories by
Revised: Figure 9 shows that the drift estimated by NLSB (E) is linear, the trajectories by NLSB (D) pass on the data manifold, and the trajectories by

Original: NLSB (E+D+V) appear to reflect all other regularization effects.
Revised: NLSB (E+D+V) appear to reflect all other regularization effects.

Original: A GIF animation of the NLSB simulation in PCA space is also included in the supplemental materials.
Revised: A GIF animation of the NLSB simulation in PCA space is also included in the supplemental materials.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: E. S YNTHETIC P OPULATION D YNAMICS : T RAJECTORIES R EFLECTING THE P OTENTIAL F UNCTION
Revised: D. S YNTHETIC P OPULATION D YNAMICS : T RAJECTORIES R EFLECTING THE P OTENTIAL F UNCTION

Original: This section shows practical examples of NLSB with the Lagrangian for the random dynamical system using artificial synthetic data.
Revised: This section shows practical examples of NLSB with the Lagrangian for the random dynamical system using artificial synthetic data.

Original: The detailed experimental setup, including the data and potential functions, is described in Section D.3.
Revised: The detailed experimental setup, including the data and potential functions, is described in Appendix C.3.

Original: The visualization of trajectories is shown in Figure 10.
Revised: The visualization of trajectories is shown in Fig.

Original: Figure 10 shows that NLSB can estimate various trajectories that reflect information about obstacles or regions where samples cannot or are likely to exist represented by the potential function.
Revised: Figure 10 shows that NLSB can estimate various trajectories that reflect information about obstacles or regions where samples cannot or are likely to exist represented by the potential function.

Original: [Figure]
Revised: [Figure]

Original: E. P OPULATION D YNAMICS S IMULATION BY R EVERSE - TIME SDE
Revised: D. P OPULATION D YNAMICS S IMULATION BY R EVERSE - TIME SDE

Original: In this section, we describe a method for estimating the trajectories backwards through time and show experimental results using artificial synthetic data.
Revised: In this section, we describe a method for estimating the trajectories backwards through time and show experimental results using artificial synthetic data.

Original: According to the result from (Anderson, 1982), the reverse-time SDE of the original SDE (Eq.
Revised: According to the result from (Anderson, 1982), the reverse-time SDE of the original SDE (Eq.

Original: (3)) is given by
Revised: (3)) is given by

Original: [Equation]
Revised: [Equation]

Original: where ˜ W t is a Wiener process that flows backwards in time.
Revised: where ˜ W t is a Wiener process that flows backwards in time.

Original: However, the computation of the modification term u of the drift function requires an estimation of the score function and high computational cost.
Revised: However, the computation of the modification term u of the drift function requires an estimation of the score function and high computational cost.

Original: Therefore, we newly parameterize the modification term u ξ by NN with parameters ξ .
Revised: Therefore, we newly parameterize the modification term u ξ by NN with parameters ξ .

Original: Then, we train only the term u ξ using Sinkhorn divergence loss in Eq.
Revised: Then, we train only the term u ξ using Sinkhorn divergence loss in Eq.

Original: (21) by numerically simulating the reverse-time SDE.
Revised: (57) by numerically simulating the reverse-time SDE.

Original: For simplicity of implementation, we used the same model architecture for modification term u ξ as for the drift f θ .
Revised: For simplicity of implementation, we used the same model architecture for modification term u ξ as for the drift f θ .

Original: [Equation]
Revised: [Equation]

Original: Figure 11 shows the visualization of trajectory estimation by reverse-time SDE using artificial data.
Revised: Figure 11 shows the visualization of trajectory estimation by reverse-time SDE using artificial data.

Original: All trajectories from NLSB and TrajectoryNet were generated by all-step prediction from the initial samples at the time t = 4 to t = 0 .
Revised: All trajectories from NLSB and TrajectoryNet were generated by all-step prediction from the initial samples at the time t = 4 to t = 0 .

Original: [Figure]
Revised: [Figure]