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 ρ 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 Appendix A.3).
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: 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 , these two SOT problems are regarded as the special case of Schr¨odinger bridge (SB)
Revised: When L ( t, x , u ) = 12 || u || 2 , these two SOT problems are regarded as the special case of Schr¨odinger bridge (SB)

Original: (Jamison, 1975; Chen et al., 2021b; Dai Pra, 1991;
Revised: (Jamison, 1975; Chen et al., 2021b; Dai Pra, 1991;

Original: L´eonard, 2013; Mikami, 1990) problem.
Revised: L´eonard, 2013; Mikami, 1990) problem.

Original: See Appendix A for details on the SOT theory.
Revised: 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 distributions at two known end points.
Revised: We consider the problem of estimating dynamics at both the population and individual sample levels by using probability distributions at two known end points.

Original: We make two realistic assumptions about the target system model.
Revised: We make two realistic 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: Samples in populations, such as a biological system (Kolomgorov et al., 1937), are known to move stochastically and diffuse.
Revised: Samples in populations, such as a biological system (Kolomgorov et al., 1937), are known to move stochastically and diffuse.

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

Original: We formulate this dynamics estimation problem on the target system as a special case of the SOT problem and call it the Lagrangian-Schr¨odinger Bridge (LSB) problem .
Revised: We formulate this dynamics estimation problem on the target system as a special case of the SOT problem and call it the Lagrangian-Schr¨odinger Bridge (LSB) problem .

Original: The position of the LSB problem in the SOT problem is clarified in Appendices A.3 and A.4.
Revised: The position of the LSB problem in the SOT problem is clarified in Appendices A.3 and A.4.

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

Original: [Equation]
Revised: [Equation]

Original: t(cid:90)
Revised: t(cid:90)

Original: R d L ( t, x , f θ ( x , t ))
Revised: R d L ( t, x , f θ ( x , t ))

Original: d ρ t ( x ; θ, ϕ ) d t, subject to d X t = f θ ( X t , t ) d t + g ϕ
Revised: d ρ t ( x ; θ, ϕ ) d t, subject to d X t = f θ ( X t , t ) d t + g ϕ

Original: ( X t , t ) d W t ,
Revised: ( X t , t ) d W t ,

Original: X 0 ∼ ρ t 0 = µ 0 , X 1 ∼ ρ t 1 = µ 1 , (6) where ρ t is the law of the random variable X t , depending on the functions f θ and g ϕ .
Revised: X 0 ∼ ρ t 0 = µ 0 , X 1 ∼ ρ t 1 = µ 1 , (6) where ρ t is the law of the random variable X t , depending on the functions f θ and g ϕ .

Original: The LSB problem is the problem of exploring the sample paths that minimize action defined by the time integral of the Lagrangian, given the distributions at the two endpoints.
Revised: The LSB problem is the problem of exploring the sample paths that minimize action defined by the time integral of the Lagrangian, given the distributions at the two endpoints.

Original: The stochastic movement of samples and diffusion phenomena are explicitly modeled by using SDEs.
Revised: The stochastic movement of samples and diffusion phenomena are explicitly modeled by using SDEs.

Original: The diffusion coefficient g is also optimized together with the drift f to reveal the effect of noise in real environments.
Revised: The diffusion coefficient g is also optimized together with the drift f to reveal the effect of noise in real environments.

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

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: 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 ϕ
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: = 12 g ϕ g ⊤ ϕ , H is the Hamiltonian defined by H ( t, x , z ) :
Revised: = 12 g ϕ g ⊤ ϕ , H is the Hamiltonian defined by H ( t, x , z ) :

Original: = ⟨ z , f θ ⟩ − L ( t, x , f θ ) , and Φ θ is the potential function satisfying f θ = ∇ z H ( t, x , −∇ x Φ θ ( x , t )) .
Revised: = ⟨ 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, we relax the constraint at time
Revised: First, we relax the constraint at time

Original: t 1 , i.e . R
Revised: t 1 , i.e . R

Original: e
Revised: e

Original: + D ( µ 1 , ρ t 1 ) , by using Wasserstein distance (Bunne et al., 2021;
Revised: + D ( µ 1 , ρ t 1 ) , by using Wasserstein distance (Bunne et al., 2021;

Original: Hashimoto et al., 2016), KL-divergence (Finlay et al., 2020; Onken et al., 2021;
Revised: Hashimoto et al., 2016), KL-divergence (Finlay et al., 2020; Onken et al., 2021;

Original: Tong et al., 2020), or a combination of these (Lavenant et al., 2021) for the discrepancy measure D .
Revised: 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 D ( µ 1 , ρ t 1 ) as the data-fitting term 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 θ = ∇ z H ( t, x , −∇ x
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: Φ θ ( x , t ))
Revised: Φ θ ( x , t ))

Original: and encourage the potential function Φ θ to satisfy the (stochastic)
Revised: 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: See Appendix D.5 for more examples.
Revised: See Appendix D.5 for more examples.

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 θ = 12 ( 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 )
Revised: Let m be the mass of the sample, U ( x ) be the potential energy of the system, and L ( t, x , u )

Original: = m 2 ∥ u ( x , t ) ∥ 2 − 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 Figs.
Revised: Practical examples are shown in Figs.

Original: 1 and 10.
Revised: 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 in Eq.
Revised: In this section, we describe a practical learning method for neural SDEs by minimizing loss in Eq.

Original: (7) on the training data, as described in Algorithm 1.
Revised: (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 in Eq.
Revised: Next, we describe how we calculate the loss in Eq.

Original: (7) using the training data {X t i } t i ∈ T .
Revised: (7) using the training data {X t i } t i ∈ T .

Original: 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 .
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
Revised: The prediction sample is obtained by numerically simulating neural SDE from the training data at one previous time point by a standard

Original: SDE solver such as the Euler-Maruyama method.
Revised: 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 −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: A RCHITECTURE 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 for the first time to use it as a technique to speed up the computation of R h in neural SDE.
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 NN 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: See Appendix B for details on the model of Φ .
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,
Revised: We compared our methods against standard neural SDE,

Original: OT-Flow (Onken et al., 2021), TrajectoryNet (Tong et al., 2020) and IPF with GP (Vargas et al., 2021) and NN (De Bortoli et al., 2021).
Revised: OT-Flow (Onken et al., 2021), TrajectoryNet (Tong et al., 2020) and IPF with GP (Vargas et al., 2021) and NN (De Bortoli et al., 2021).

Original: We trained the drift and diffusion models of the standard neural SDE using only the Sinkhorn divergence.
Revised: We trained the drift and diffusion models of the standard neural SDE using only the Sinkhorn divergence.

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
Revised: +OT represents a model trained with the

Original: OT-based regularization defined by Eqs.
Revised: 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 OT-Flow.
Revised: We used only ˜ R e for TrajectoryNet; we used both ˜ R e and ˜ R h for 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 model of IPF (GP) was changed to sparse GP from vanilla
Revised: The drift model of IPF (GP) was changed to sparse GP from vanilla

Original: GP (Vargas et al., 2021) to save computation cost.
Revised: GP (Vargas et al., 2021) to save computation cost.

Original: The drift model of IPF (NN) is the same network as the NLSB for a fair comparison.
Revised: The drift model of IPF (NN) is the same network as the NLSB for a fair comparison.

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

Original: See Appendix C 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:
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 distance 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 L 2 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 Wassersteindistance 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 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
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 Fig.
Revised: The evaluation results are shown in Fig.

Original: 3, and the visualization of trajectories is shown in Fig.
Revised: 3, and the visualization of trajectories is shown in Fig.

Original: Figure 3 shows that NLSB and IPF outperform neural SDE and is comparable to other
Revised: Figure 3 shows that NLSB and IPF outperform neural SDE and is comparable to other

Original: ODE-based methods in estimating populations with small variance.
Revised: 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 and IPF can estimate population-level dynamics even when the population variance is large or small.
Revised: That indicates that NLSB and IPF can estimate population-level dynamics even when the population variance is large or small.

Original: Furthermore, NLSB and IPF have a smaller CDD value than neural SDE.
Revised: Furthermore, NLSB and IPF 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 Fig. 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 predictions by NLSB and IPF in Figs.
Revised: In contrast, the predictions by NLSB and IPF in Figs.

Original: 4c, 6d and 6e are much closer to the ground-truth.
Revised: 4c, 6d and 6e are much closer to the ground-truth.

Original: These results show that the prior knowledge of the potential-free system 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.

Original: See Appendix D.1 for further results and analysis.
Revised: 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 : day to
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

Original: , t 4 : day 24 to
Revised: , 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
Revised: The collected cells were then measured by scRNAseq, filtered at the quality control stage and mapped to a low

Original: -dimensional feature space using a principal component analysis (PCA).
Revised: -dimensional feature space using a principal component analysis (PCA).

Original: For details, see Appendix E.2 in (Tong et al.,
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 t 4 are 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 Fig.in Appendix D).
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: 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.
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: See Appendix D.2 for further results and analysis.
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 (Table 1 and Fig.
Revised: This allowed us to handle the uncertainty of the trajectory (Fig. 7) and to successfully capture the diffuse transitions (Table 1 and Fig.

Original: In contrast, Vargas et al.
Revised: In contrast, Vargas et al.

Original: (2021) and Bunne et al.
Revised: (2021) and Bunne et al.

Original: (2022) estimated the sample
Revised: (2022) estimated the sample

Original: [Table]
Revised: [Table]

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

Original: The CDD value at the time t ∈
Revised: The CDD value at the time t ∈

Original: [ 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 .
Revised: [ 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: IPF (GP) 0 .0 .0 .
Revised: IPF (GP) 0 .0 .0 .

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

Original: 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 GP-based IPF to solve the SB problem.
Revised: (2021) proposed GP-based IPF 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.

Original: A T HEORY OF S TOCHASTIC O PTIMAL T RANSPORT
Revised: A T HEORY OF S TOCHASTIC O PTIMAL T RANSPORT

Original: 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 the problem of finding a stochastic process that minimizes the expected cost under fixed marginal distributions at several time points.

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

Original: Benamou & Brenier (2000).
Revised: Benamou & Brenier (2000).

Original: It is also considered as a stochastic optimal control (SOC) problem with additional terminal constraint.
Revised: It is also considered as a stochastic optimal control (SOC) problem with additional terminal constraint.

Original: 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.
Revised: 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.

Original: An overview of the relationships among these problems is illustrated in Fig.
Revised: An overview of the relationships among these problems is illustrated in Fig.

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

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

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

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

Original: They introduced a time parameter t ∈ [0 , 1] and considered the transport process from µ to ν as the advection in the time interval [0 , 1] .
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: 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 .
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 .

Original: Theorem 2 (Brenier-Benamou formulation (Eularian formalism); (Benamou & Brenier, 2000)) .
Revised: Theorem 2 (Brenier-Benamou formulation (Eularian formalism); (Benamou & Brenier, 2000)) .

Original: [Equation]
Revised: [Equation]

Original: subject to ∂
Revised: subject to ∂

Original: [Equation]
Revised: [Equation]

Original: where p and q are the densities of probability measures µ and ν , respectively.
Revised: where p and q are the densities of probability measures µ and ν , respectively.

Original: The first condition in Eq.
Revised: The first condition in Eq.

Original: (13) is known as the continuity equation and represents the conservation of mass in time evolution.
Revised: (13) is known as the continuity equation and represents the conservation of mass in time evolution.

Original: We can consider an equivalent formulation of Eq.
Revised: We can consider an equivalent formulation of Eq.

Original: (12) by introducing
Revised: (12) by introducing

Original: Lagrangian coordinates X ( t, x ) .
Revised: Lagrangian coordinates X ( t, x ) .

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

Original: Theorem 3 (Brenier-Benamou formulation (Lagrangian formalism); (Benamou & Brenier, 2000)) .
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]

Original: The right-hand side of Eq.
Revised: The right-hand side of Eq.

Original: (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: (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 .

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

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: 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: Equations 16 and 18 indicate that the particle, whose initial position is x 0 ∈ X , moves straight ahead at the constant velocity v 0 =

Original: −∇ x Φ 0 ( x 0 )
Revised: −∇ x Φ 0 ( x 0 )

Original: =
Revised: =

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

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

Original: and the HJB equation (Eq.
Revised: and the HJB equation (Eq.

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

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

Original: A. S CHR ¨ ODINGER B RIDGE
Revised: A. S CHR ¨ ODINGER B RIDGE

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

Original: Let Ω
Revised: Let Ω

Original: =
Revised: =

Original: C
Revised: C

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

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

Original: The SB problem is defined by
Revised: The SB problem is defined by

Original: [Equation]
Revised: [Equation]

Original: 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: 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)

Original: d Q d P (cid:1)
Revised: d Q d P (cid:1)

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

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

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

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

Original: Chen et al.
Revised: Chen et al.

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

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

Original: Theorem 6.
Revised: Theorem 6.

Original: Let Ψ( x , t ) and ˆΨ( x , t ) be the solutions to the following PDEs:
Revised: Let Ψ( x , t ) and ˆΨ( x , t ) be the solutions to the following PDEs:

Original: [Equation]
Revised: [Equation]

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

Original: Then, the solution to the SB problem (Eq.
Revised: Then, the solution to the SB problem (Eq.

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: SB optimality (Eq.
Revised: SB optimality (Eq.

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

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Let u ∗ t ( x ) =
Revised: Let u ∗ t ( x ) =

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

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

Original: [Equation]
Revised: [Equation]

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

Original: The first condition in Eq.
Revised: The first condition in Eq.

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

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

Original: (22) and Eq.
Revised: (22) and Eq.

Original: (24)) are equivalent, i.e . V S,ϵ ( µ 0 , µ 1 )
Revised: (24)) are equivalent, i.e . V S,ϵ ( µ 0 , µ 1 )

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

Original: Furthermore, the dynamic solution of the OT problem (Eq.
Revised: Furthermore, the dynamic solution of the OT problem (Eq.

Original: (16)) is obtained as the zero-noise limit (Mikami, 2004;
Revised: (16)) is obtained as the zero-noise limit (Mikami, 2004;

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

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

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

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

Original: = µ 0 (d x ) /
Revised: = µ 0 (d x ) /

Original: d x exists.
Revised: d x exists.

Original: Then, the following holds
Revised: Then, the following holds

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: and there exists a convex function φ satisfying
Revised: and there exists a convex function φ satisfying

Original: [Equation]
Revised: [Equation]

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

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

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

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 given by the following.
Revised: The SOT problem with the endpoint marginals fixed to µ 0 and µ 1 is given by the following.

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

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

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

Original: [Equation]
Revised: [Equation]

Original: 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: 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:

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: (A-1)
Revised: (A-1)

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

Original: u (cid:55)→ L ( t, x , u ) is strictly convex.
Revised: u (cid:55)→ L ( t, x , u ) is strictly convex.

Original: L ( t, x , u ) / (1 + L ( t, y , u ) and
Revised: L ( t, x , u ) / (1 + L ( t, y , u ) and

Original: |∇
Revised: |∇

Original: x
Revised: x

Original: L
Revised: L

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

Original: |
Revised: |

Original: + L ( t,
Revised: + L ( t,

Original: x
Revised: x

Original: , u ))
Revised: , u ))

Original: are bounded
Revised: are bounded

Original: on
Revised: on

Original: t
Revised: t

Original: ∈
Revised: ∈

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

Original: sup
Revised: sup

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

Original: locally bounded
Revised: locally bounded

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

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

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

Original: 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: 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,

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: Suppose that (A-1) holds.
Revised: Suppose that (A-1) holds.

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

Original: P ( R d ) ,
Revised: P ( R d ) ,

Original: [Equation]
Revised: [Equation]

Original: 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: 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:

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: Then, there exists a minimizer X ∗ ∈
Revised: Then, there exists a minimizer X ∗ ∈

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

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: The the LSB problem in Eq.
Revised: The the LSB problem in Eq.

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

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

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

Original: We solve the following relaxed
Revised: We solve the following relaxed

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

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

Original: [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 = µ

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

Original: We derive the optimality conditions for the LSB problems in Eqs.
Revised: We derive the optimality conditions for the LSB problems in Eqs.

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: LSB problem in Eq.
Revised: LSB problem in Eq.

Original: Proof.
Revised: Proof.

Original: We derive the optimality conditions for the relaxed LSB problem as shown in Eq.
Revised: We derive the optimality conditions for the relaxed LSB problem as shown in Eq.

Original: (42) by reformulating it as a saddle point problem for ( p t , m t ) :
Revised: (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 :
Revised: The term (a) is transformed by performing a partial integral over t :

Original: [Equation]
Revised: [Equation]

Original: The term (b) is simplified as follows.
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: (47) satisfies the following conditions:
Revised: (47) satisfies the following conditions:

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: Therefore
Revised: Therefore

Original: [Equation]
Revised: [Equation]

Original: the optimal drift function is given by
Revised: the optimal drift function is given by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: The potential function Φ is the solution of the HJB equation:
Revised: The potential function Φ is the solution of the HJB equation:

Original: [Equation]
Revised: [Equation]

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

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

Original: From Theorem 12, we can show that there exists a semimartingale { X t } 0 ≤ t ≤ 1 for which the following holds:
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: where ( f ∗ , ρ ∗ t ) is the minimizer of the LSB problems in Eqs.
Revised: where ( f ∗ , ρ ∗ t ) is the minimizer of the LSB problems in Eqs.

Original: (39) and (42))
Revised: (39) and (42))

Original: and g ∗ is the diffusion function determined from the minimizer ( f ∗ , ρ ∗ t ) .
Revised: and g ∗ is the diffusion function determined from the minimizer ( f ∗ , ρ ∗ t ) .

Original: Therefore, the LSB problems constrained by the
Revised: Therefore, the LSB problems constrained by the

Original: FP equation (Eqs.
Revised: FP equation (Eqs.

Original: (39) and (42))
Revised: (39) and (42))

Original: and the original LSB problem defined by Eq.
Revised: and the original LSB problem defined by Eq.

Original: (6) are equivalent.
Revised: (6) are equivalent.

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

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

Original: (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: (7) is the objective function of the relaxed LSB problem (Eq. (42)) with additional PDE loss applied to satisfy the optimality condition, HJB equation.

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: Eqs.
Revised: Eqs.

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

Original: Theoretically, the LSB problem with the Lagrangian
Revised: Theoretically, the LSB problem with the Lagrangian

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

Original: 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).
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: B M ODEL A RCHITECTURE
Revised: B M ODEL A RCHITECTURE

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 (1 ≤ i ≤ M ) ∈
Revised: R m × ( d +1) , K i (1 ≤ i ≤ M ) ∈

Original: 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 ) =
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, ∇ 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: 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.
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: OT-Flow modeled the potential function Φ θ that satisfies v θ = −∇ x
Revised: OT-Flow modeled the potential function Φ θ that satisfies v θ = −∇ x

Original: Φ θ instead of directly modeling the velocity function v of the neural ODE.
Revised: Φ θ instead of directly modeling the velocity function v of the neural ODE.

Original: The continuous transformation from x (0) to x (1) is described by
Revised: The continuous transformation from x (0) to x (1) is described by

Original: [Equation]
Revised: [Equation]

Original: OT-Flow is trained by likelihood maximization in the well-known CNF framework.
Revised: OT-Flow is trained by likelihood maximization in the well-known CNF framework.

Original: The likelihood computation is performed as follows.
Revised: The likelihood computation is performed as follows.

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

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

Original: The fast and exact computation of the diagonal component of the potential function’s Hessian is useful for the computation of Eq.
Revised: The fast and exact computation of the diagonal component of the potential function’s Hessian is useful for the computation of Eq.

Original: In contrast, we propose for the first time to use this technique to speed up the computation of the
Revised: In contrast, we propose for the first time to use this technique to speed up the computation of the

Original: HJB-PDE loss R h in the NLSB under the assumption that the diffusion model’s output is a diagonal matrix.
Revised: HJB-PDE loss R h in the NLSB under the assumption that the diffusion model’s output is a diagonal matrix.

Original: The computed loss R h is given by
Revised: The computed loss R h is given by

Original: [Equation]
Revised: [Equation]

Original: C 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 IPF 3 repositories, which were released under the MIT license.
Revised: In our implementation, we modified code in the TrajectoryNet 1 , OT-FLow 2 and IPF 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: C. 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.
Revised: We describe the details of the training setup for the experiments on synthetic data conducted in Section 5.1.

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 ) , 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 .
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: 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
Revised: We set the weight coefficients λ e , λ h shown in Table

Original: [Table]
Revised: [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
Revised: A concatsquash layer cs was defined in the released code of

Original: FFJORD (Grathwohl et al., 2018) as:
Revised: FFJORD (Grathwohl et al., 2018) as:

Original: [Equation]
Revised: [Equation]

Original: where σ is the sigmoid function
Revised: where σ is the sigmoid function

Original: [Table]
Revised: [Table]

Original: and W x
Revised: and W x

Original: [Table]
Revised: [Table]

Original: W t
Revised: W t

Original: [Table]
Revised: [Table]

Original: W b
Revised: W b

Original: [Table]
Revised: [Table]

Original: b x
Revised: b x

Original: [Table]
Revised: [Table]

Original: b t
Revised: b t

Original: [Table]
Revised: [Table]

Original: b b are all learnable parameters.
Revised: b b are all learnable parameters.

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: 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
Revised: In the training for TrajectoryNet with OT-based regularization

Original: [Table]
Revised: [Table]

Original: we set the weight coefficients shown in Table 3.
Revised: 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 IPF
Revised: S ETTINGS FOR IPF

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 IPF (GP) is sparse GP with the exponential kernel.
Revised: The drift model of IPF (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 IPF (NN) is the same network as for NLSB, including hyperparameters.
Revised: The drift model of IPF (NN) is the same network as for NLSB, including hyperparameters.

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

Original: We used the Euler-Maruyama method as the SDE solver and set the diffusion coefficients, 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: 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: C. 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 C.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 C.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 IPF
Revised: S ETTINGS FOR IPF

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 C.1.
Revised: All other settings were the same as in Appendix C.1.

Original: C. 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 Appendix D.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 U ( x, y )
Revised: Box-shaped obstacle U ( x, y )

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

Original: 5 ≤ x, y ≤ 0 .
Revised: 5 ≤ x, y ≤ 0 .

Original: otherwise .
Revised: otherwise .

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

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

Original: 1 ≤ x ≤ 0 .
Revised: 1 ≤ x ≤ 0 .

Original: 1 ∧ ( y ≤ − 0 .
Revised: 1 ∧ ( y ≤ − 0 .

Original: 25 ∨ 0 .
Revised: 25 ∨ 0 .

Original: 25 ≤ y ) otherwise .
Revised: 25 ≤ y ) otherwise .

Original: Well
Revised: Well

Original: potential
Revised: potential

Original: [Equation]
Revised: [Equation]

Original: 4. Hill potential
Revised: 4. 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 C.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: 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
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: 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 IPF) and the only original trajectory by the ODE-based methods (TrajectoryNet and OT-Flow).
Revised: SDE-based methods (Neural SDE, NLSB, and IPF) and the only original trajectory by the ODE-based 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.

Original: Figure 3 shows that NLSB and IPF outperform neural SDE and is comparable to other ODEbased methods in estimating populations with small variance.
Revised: Figure 3 shows that NLSB and IPF outperform neural SDE and is comparable to other ODEbased 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 and IPF can estimate population-level dynamics even when the population variance is large or small.
Revised: That indicates that NLSB and IPF can estimate population-level dynamics even when the population variance is large or small.

Original: Furthermore, NLSB and IPF have a smaller CDD value than neural SDE.
Revised: Furthermore, NLSB and IPF have a smaller CDD value than neural SDE.

Original: 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.
Revised: 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.

Original: 6a), especially in the interval [0 , 1] .
Revised: 6a), especially in the interval [0 , 1] .

Original: In contrast, the predictions by NLSB and IPF in Figs.
Revised: In contrast, the predictions by NLSB and IPF in Figs.

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

Original: NLSB almost achieve the upper limit of performances.
Revised: NLSB almost achieve the upper limit of performances.

Original: These results show that the prior knowledge of the potential-free system 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.

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

Original: The result that
Revised: The result that

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

Original: D. 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 Fig.
Revised: The time evolution of the distribution for a single sample is visualized as a heat map in Fig.

Original: The sample population in Fig.
Revised: The sample population in Fig.

Original: 8 and the trajectories are visualized in Fig.
Revised: 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: The experimental results using MDD using EMD with L 1 cost is shown in Table 5.
Revised: The experimental results using MDD using EMD with L 1 cost is shown in Table 5.

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] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: D. 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 Appendix C.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 Fig.
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: D. 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: (57) 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] [Equation]
Revised: [Equation] [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: Experimental results using the MDD between reverse-time simulations and ground-truth are shown in Fig.
Revised: Experimental results using the MDD between reverse-time simulations and ground-truth are shown in Fig.

Original: Figure 12 shows that the proposed method in Appendix D.4 successfully recovers the population-level dynamics of the reverse-time SDE.
Revised: Figure 12 shows that the proposed method in Appendix D.4 successfully recovers the population-level dynamics of the reverse-time SDE.

Original: The development of appropriate evaluation methods for sample-level dynamics of the reverse-time SDE is included in future work.
Revised: The development of appropriate evaluation methods for sample-level dynamics of the reverse-time SDE is included in future work.

Original: [Figure]
Revised: [Figure]

Original: D. T HE QUADRATIC - FORM L AGRANGIAN
Revised: D. T HE QUADRATIC - FORM L AGRANGIAN

Original: In this section, we introduce the quadratic-form Lagrangian inspired by the Lagrangian cost commonly used in the optimal control (OC) problems such as the linear-quadratic regulator (LQR).
Revised: In this section, we introduce the quadratic-form Lagrangian inspired by the Lagrangian cost commonly used in the optimal control (OC) problems such as the linear-quadratic regulator (LQR).

Original: To demonstrate the applicability of the quadratic-form Lagrangian, we present a new scenario and experimental results in Appendix D.5.
Revised: To demonstrate the applicability of the quadratic-form Lagrangian, we present a new scenario and experimental results in Appendix D.5.

Original: Furthermore, we show the invariance of the NLSB  to various coordinate transformations.
Revised: Furthermore, we show the invariance of the NLSB using the quadratic-form Lagrangian to the linear coordinate transformations.

Original: We first consider the LSB problem with explicit coordinate transformation by the matrix P .
Revised: We first consider the LSB problem with explicit coordinate transformation by the matrix P .

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where y ∈ R p is the observed variable, x ∈ R
Revised: where y ∈ R p is the observed variable, x ∈ R

Original: d is the state variable, and P ∈ R p ×
Revised: d is the state variable, and P ∈ R p ×

Original: d is the projection matrix, a.k.a.
Revised: d is the projection matrix, a.k.a.

Original: the output matrix in the state-space model.
Revised: the output matrix in the state-space model.

Original: We propose the quadratic-form Lagrangian defined by
Revised: We propose the quadratic-form Lagrangian defined by

Original: [Equation]
Revised: [Equation]

Original: (59) where R ∈ R d × d is the positive definite, U Q is the quadratic function in x .
Revised: (59) where R ∈ R d × d is the positive definite, U Q is the quadratic function in x .

Original: The drift function is obtained by f θ ( x , t ) =
Revised: The drift function is obtained by f θ ( x , t ) =

Original: − ( R + R ⊤ ) − 1 ( ∇ x Φ θ ( x , t )+ c ) .
Revised: − ( R + R ⊤ ) − 1 ( ∇ x Φ θ ( x , t )+ c ) .

Original: The Lagrangians for the potential-free and the cellular system described in Section 3.3 are the special cases of the quadraticform Lagrangian.
Revised: The Lagrangians for the potential-free and the cellular system described in Section 3.3 are the special cases of the quadraticform Lagrangian.

Original: GMM used in Section 5.2, the well and hill potentials used in Appendix D.3 are examples of the quadratic potential U Q .
Revised: GMM used in Section 5.2, the well and hill potentials used in Appendix D.3 are examples of the quadratic potential U Q .

Original: The matrix R can be used to penalize the magnitude of the drift differently for each dimension.
Revised: The matrix R can be used to penalize the magnitude of the drift differently for each dimension.

Original: Figure 13a and 13b show that the drift of the trajectories generated by the NLSB with the Lagrangian L = 12 u ⊤ Ru is larger on the axis with smaller penalties defined by R and vice versa.
Revised: Figure 13a and 13b show that the drift of the trajectories generated by the NLSB with the Lagrangian L = 12 u ⊤ Ru is larger on the axis with smaller penalties defined by R and vice versa.

Original: Next, we present that the loss function of NLSB  is essentially invariant for the various coordinate transformations represented by the projection matrix P .
Revised: Next, we present that the loss function of NLSB using the quadratic-form Lagrangian is essentially invariant for the various coordinate transformations represented by the projection matrix P .

Original: In particular, we use the Wassersteindistance W 2 as the distribution discrepancy measure D .
Revised: In particular, we use the Wasserstein-2 distance W 2 as the distribution discrepancy measure D , as shown in Eq.

Original: [Figure]
Revised: [Figure]

Original: Regular matrix .
Revised: Regular matrix .

Original: When the projection matrix P ∈ R
Revised: When the projection matrix P ∈ R

Original: d
Revised: d

Original: ×
Revised: ×

Original: d is regular, the cost functions defined on the observed variable space
Revised: d is regular, the cost functions defined on the observed variable space

Original: y ∈ R d is computed on the state variable space x ∈ R d as follows.
Revised: y ∈ R d is computed on the state variable space x ∈ R d as follows.

Original: [Equation]
Revised: [Equation]

Original: where R = P ⊤ P and R x =
Revised: where R = P ⊤ P and R x =

Original: P ⊤ R y P ∈ R
Revised: P ⊤ R y P ∈ R

Original: d × d are guaranteed to be positive definite, P − 1 ♯ is the push-forward operator of the linear map represented by P − 1 ,
Revised: d × d are guaranteed to be positive definite [Equation] P − 1 ♯ is the push-forward operator of the linear map represented by P − 1 ,

Original: c x = P ⊤ c y ∈ R d
Revised: c x = P ⊤ c y ∈ R d

Original: [Equation]
Revised: [Equation]

Original: and U x Q is also the quadratic function in x .
Revised: and U x Q is also the quadratic function in x .

Revised: Especially when P is an orthogonal transformation, i.e . P ⊤ P

Revised: =

Revised: I , then

Original: PCA projection matrix .
Revised: R = I and R x = R y hold. The quadratic-form Lagrangian has sufficient representational capacity to describe Lagrangians in coordinate systems that can be linearly transformed into each other in a unified manner.

Revised: Then, the loss function of the NLSB can be easily derived without depending on coordinates of the state variable x , except for the Lagrangian design.

Revised: Furthermore, similar results are obtained for the PCA projection, which is an irregular matrix.

Revised: PCA projection matrix .

Original: We consider an inverse projection from the latent space x to the data space y as the projection matrix P ∈ R p ×
Revised: We consider an inverse projection from the latent space x to the data space y as the projection matrix P ∈ R p ×

Original: d ( d ≪ p ) .
Revised: d ( d ≪ p ) .

Original: The columns of the PCA projection matirx are orthogonal, i.e . P ⊤ P
Revised: The columns of the PCA projection matrix are orthogonal, i.e . P ⊤ P

Original: = I .
Revised: = I .

Original: The cost functions defined on the observed variable space y ∈ R p is efficiently computed on the low-dimensional space x ∈ R d as follows.
Revised: The cost functions defined on the observed variable space y ∈ R p is efficiently computed on the low-dimensional space x ∈ R d as follows.

Original: [Equation]
Revised: [Equation]

Original: where R x =
Revised: where R x =

Original: P ⊤ R y P ∈ R d × d is guaranteed to be positive definite, P − 1 ♯ is the push-forward operator of the linear map represented by P ⊤ , c x = P ⊤ c y ∈ R d , and U x Q is also the quadratic function in x .
Revised: P ⊤ R y P ∈ R d × d is guaranteed to be positive definite, P ⊤ ♯ is the push-forward operator of the PCA projection represented by P ⊤ , c x = P ⊤ c y ∈ R d , and U x Q is also the quadratic function in x .

Original: For the Lagrangian of the cellular system used in Section 5.2, R y
Revised: These results with R y

Original: =
Revised: =

Original: R x
Revised: R x

Original: = I holds from R y = I .
Revised: = I are the justifications for learning NLSB using the Lagrangian for the cellular system on the PCA space in Section 5.2.