Neural Lagrangian Schr\"{o}dinger Bridge: Diffusion Modeling for Population Dynamics
Authors: Takeshi Koshizuka
Issei Sato
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Population Dynamics
Trajectory Inference
Neural SDEs
Stochastic Optimal Transport
Schrödinger Bridge

Abstract: 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. 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. 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. While the sample behavior in CNFs is deterministic, the actual sample in biological systems moves in an essentially random yet directional manner. 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. To satisfy these requirements of the sample trajectories, we formulate the Lagrangian Schrödinger bridge (LSB) problem and propose to solve it approximately by modeling the advection-diffusion process with regularized neural SDE. We also develop a model architecture that enables faster computation of the loss function. 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.

Published as a conference paper at ICLR 2023
NEURAL LAGRANGIAN SCHRÖDINGER BRIDGE: DIF- FUSION MODELING FOR POPULATION DYNAMICS
Takeshi Koshizuka, Issei Sato
The University of Tokyo (koshizua-takkehhiitak
autokyg..cc.utttkyo........
ABSTRACT
Population dynamics is the study of temporal and spatial variation in the size populations organisms and is a major part of population ecology. One 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 measurement constraints. Recently, modeling population dynamics by using continuous normalizing flows (CNFs) and dynamic optimal transport has been proposed to infer the sample trajectories from fixed-point observed population. While the sample behavior CNFss deterministic, the actual sample biological systems moves in essentially random yet directional manner. Moreover, when sample moves from point A to point B in dynamical systems, trajectory typically follows the principle of least action in which the corresponding action has the smallest possible value. To satisfy these requirements of the sample trajectories, we formulate the Lagrangian Schrödinger bridge (LSB) problem and propose solve it approximately by modeling the ddlvectionddiffusion process with regularized neural SDE. We also develop a model architecture that enables faster computation of the loss function. Experimental results show that the proposed method can efficiently approximate the population-level dynamics even for nigh-dimensional data and that using the prior knowledge introduced by the Lagrangian enables us to estimate the sample-level dynamics with stochastic behavior.
INTRODUCTION
The population dynamics of time-evolving individuals appears in various scientific fields, such as cell population in biology (Schiebinger et al. 2019 Yang & Uhler 2018 2018), air meteorology (Fisher et et 2009), and healthcare statistics (Manton et al. 2008) in medicine. However, tracking individuals over long period is often difficult due to experimental costs. Furthermore, it sometimes be impossible track the time evolution. 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. Instead, we only obtain individual samples from cross- sectional populations without alignment across time steps few distinct time points. Under these constraints on data measurements, our goal is to better understand the time evolution of samples the populations.
Figure 1: Example of trajectories by NLSB.
Existing methods attempt to estimate population-level dynamics following the Wasserstein gradient flow using a recurrent neural network (RNN) (Hashimoto et al., (2066) or the Bindta--indeeleheer--oo (JKO) flow (Bunne et al.) 2021). Recent studies have attempted to interpolate the trajectories of individual samples between cross-sectionss populations at multiple time points by using optimal transport (OT) (Schiebinger et al. 2019 Yang & Uhler 2018), CNF (Tong et al. 2020). Using CNF generates continuous-tinm non-lines sample trajectories from multiple time points. addition,
Tong et al. 2020) proposed regularization for CNF that encourages straight trajectory on the basis of the OT theory. 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

---

Published as a conference paper at ICLR 2023
completely deterministic manner. However, samples in population are known to move stochastically and diffuse in nature, e.g., biological system (Kolomgorov et al., 1937).
To handle the stochastic and complex behavior of individual samples, we propose to model the udvection-diffusion processes by using SDEs to describe the time evolution of the sample. Further- more, the basis of the principle least action, we estimate the sample trajectories that minimize action, defined by the time integral of the Lagrangian determined from the prior knowledge. We formulate this problem the Lagrangian Schrödinger bridge (LSB) problem, which is a special case of the stochastic optimal transport (SOT) problem, and propose an approximate solution method neural Lagrangian Schrödinger bridge (NLSB). In NLSB, we train regularized neural SDE (Li et al.. 2020; Tzen & Raginsky 2211a b) by minimizing the Wasserstein loss between the ground-trun and predicted population. 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 al. (2020). In addition, we parameterize a potential function instead of the drift function. Adopting the model architecture of the potential function from OT-Flow (Onken et al. 202T) will speed up the computation of the potential function's gradient and the regularization term. As result, capture the population-level dynamics well or better than conventional methods, and can more accurately predict the sample trajectories.
In short, our contributions are summarized follows.
1. We formulate the LSB problem to estimate the stochastic sample trajectory according to the principle of least action.
2. We propose NLSB to approximate the LSB problem practically by modeling the advection- diffusion process with regularized neural SDE on the basis of the prior knowledge introduced by the Lagrangian.
3. We adopt the model architecture of the potential function from OT-Flow OOken al. 2021 to speed up the computation of the regularization term to minimize HJB-PDE loss.
2
BACKGROUND
2.1 we introduce the method of combining CNF and the OT theory, the basis of our method. Then, we explain dynamics modeling techniques: the diffusion modeling using neural SDE (Section 2.2) and the SOT theory (Section 2.3
In Section
2.1 FLOWS REGULARIZED BY OPTIMAL TRANSPORT
CNFs (Chen et al. 2018) are method for learning the continuous transformation between two distributions p and y by modeling the ordinary differential equation (ODE):
dx(() (x, subject x(lo) ~ p. x(11) ~ q.
dt
where fo is the velocity model with learnable parameters 0.
Several regularizations CNF' leading to straight trajectories have been proposed on the basis of the OT theory. The likelihood maximization problem of regularized CNF is derived by replacing the terminal constraint of the Kullback-Leibler (KL. divergence. RNODE (Finlay et al., 2020), OT-Flow (Onken et al. 2021). and
formulation (Benamou & Brenier 2000) with
TrajectoryNet (Tong et al. 2020) introduced a regularization R. in Eq. (1). Potential Flow (Yang & Karniadakis] 2020) and FT-Flow modeled the potential function &b that satisfies f = - V instead
of modeling the velocity function f. They also proposed an additional OT-based regularization R, (HJJ) equation (Evans 1983) satisfied by the potential function shown in Eq. (2). These OT-based regularizations have resulted in faster CNFs (Finlay et al.
derived from the Hamilton-Iacoob-Bellnn
2020 Onken et al. 2021 and improved the modeling of cellular dynamics (Tong et al. 2020).
ph 1
R_
[x, 1) (x) dt,
Jro Jan 2
Rr, = pr: (x,,)
1
v oo 1211 dpe(x) dt.
(2)
to Jzn
2
2

---

Published as a conference paper at ICLR 2023
L2-OT problem
SB problem
SOT problem V(Jo.
Sec. 3.1
min (kI(Q))P) QEP s.t. 4i
LSB problem (Ours)
Wa(mas (tr)
Q
inf Jo
B| X 6e]
Transport problem Xo Pr Xy ~
dynamic
dynamic
a.t. X, SDE( Ba)
Fix
Dynamic OT
inf
01P
((Ho-Ma)
continuity equation
SDEIa.
v2r
Dynamic OT by CNF
SB problem by neural SDE
LSB problem by neural SDE
Solutions
[Tong al. 2020]
[Vargas et al. 2021] / [Bunne et al. 2022]
Sec. 3.2
Ours
Figure 2: Overview
where Pt is the law of the sample x at the time t.
2.2 NEURAL SDE AND APPLICATIONS TO DIFFUSION MODELING
SDE has been used to model weal-world random phenomena a wide range of areas, such as chemistry, biology, mechanics, economics and finance (Higham 2001). As an extension of neural ODEs, neural SDE has been proposed to model drift and diffusion functions with neural networks (NNs), follows.
dx. = fo (X,, dt go (X,, t) dW.
where {Xclre]to:tx is a continuous valuud stochastic process, fg Rd x [tu. tk-il + Rd is drift function, go Rd x [to. tx-1] + Rdxm is diffusion function and (Wi helltn an m-dimensional Wiener process. Score-based generative models (SGM) (Song et al. 20200. which use score matching to learn the reverse diffusion process of generating images from noise as SDE, have demonstrated the ability produce high-quality data. Other recent studies on diffusion modeling using SDE (De Bortoli et al., [02 Wang et al. 2021] Vargas et al., 2021 Chen et al.) 2021a Zhang & Chen 2022) have proposed learning SDE solutions to the SB problem. Vargas et al. (2021) and Bortoli et al. (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.
2.3 STOCHASTIC OPTIMAL TRANSPORT
Mikami (0088 generalized the OT problem and defined the SOT problem as a random mechanics problem determined by the principle of least action. The SOT problem with the endpoint marginals fixed to and μ1 is represented as
(Pa, pi) = inf
E
fx1,X,;xx(X,,1 ()) di X.o umo X1, 11 J.
Xear
where L(t, u) is continuous and convex in and of is the set of all -Rd-ullued, continuous semimartingales {oo on complete filtered probability space such that there exists a Borel measurable drift function fx (X,, f) for which satisfies several conditions (see Appendix A.3 The function L, called the Lagrangian, the transport cost defined on the space-time of the system and allows us describe phenomena consistently regardless of the choice of coordinate system.
Mikami (2008) also introduced another variational version of the SOT problem for a flow of marginal distributions which satisfies the Bokker-Planck (FP) equation, which is addection-difffusee equation describing the time evolution of the probability density function.
v(po, v1) ipp i
inf S.h. / L (t,xx dx)(x)dt
Prr Mkc (f. (c) satisfies the FP eq
(5)

---

Published as a conference paper at ICLR 2023
where Pr and Pr the law and density of the random variable X,, respectively.
When Lit, x, u) = 1 [u|l 1,2 the SOT problem is regarded as the special case of Schrödinger bridge
(SB) JJamison 1975 Chen et al. 2021b; Dai Pra 1991 Léonard 2013
Mikami 1990) problem.
See Appendix A.3 for details on the SOT theory.
PROPOSED METHOD
In this section, first introduce Sarraaggian-Schrödingg bridge (LSB) problem. We next present an approximate solution for the LSB problem: neural Lagrangian Schrödinger bridge (NLSB). Finally, we give specific examples of our method for specific use cases. Figure illustrates the position of the LSB problem in the transport problem and the relationship between the NLSB and existing methods.
3.1 LAGRANGIAN SCHRÖDINGER BRIDGE
We consider the problem of estimating dynamics at both the population and individual sample levels by using probability distributions at two known end points. We make two realistic assumptions about the target system model.
1. The stochastic behavior of individual samples yields the population diffusion phenomenon.
2. Individual samples are encouraged to move according to the principle of least action.
Samples in populations, such biological system (Kolomgorov et al. 1937). are known to move stochastically and diffuse. 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. We formulate this dynamics estimation problem on the target system as special case of the SOT problem and call it the Serrraniann-chhridiin The position of the LSB problem in the SOT problem is clarified in Appendices) A.3 and B.1 Definition 3.1 (LSB problem).
Bridge (LSB) problem.
J'M [ L(t, x, (ft,,,((x,)) dx:(x; f. dt,
minimize f.g
(6)
subject to
dX, = f(X,,() () dl + (x,,, dW.,
X Pro Po: X Pr P1,
where Pr is the law of the random variable X,, depending on the functions and g.
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 the two endpoints. The stochastic movement of samples and diffusion phenomena are explicitly modeled by using SDEs. The diffusion coefficient g is also optimized together with the drift to reveal the effect of noise in real environments.
3.2 NEURAL LAGRANGIAN SCHRÖDINGER BRIDGE
In the setting of our paper, it is difficult to solve the LSB problem because the exact endpoint constraints Pn and μ1 are unknown and only samples from them are available. Therefore, we propose a practical solution method for the LSB problem using neural SDEs (Eq. (3)) with regularization. We propose to approximate the LSB problem by learning neural SDEs with the gradients of the loss (See Appendix B.2 for theoretical justification and connections to existing works.):
minimize Dei,, Pt,) (0; 41) Rn, de 40, 6,0
2 [ L(t,x, f(t,x, (())))
(0:
(8)
/") J&
d
R₂ (0,4; lo, (1)
0,00(xx 1) +
Σ (i,((x ()) V200)as - Hx(x,)) dpi(x)dl,
1,f=
Hī (xx t): = H(t,,, -x t)) (- V Ex E fx(,tt)
(9)
4

---

Published as a conference paper at ICLR 2023
where D is the distribution discrepancy measure, (t, is the entry in the i-th row and j-th column of the diffusion coefficient matrix D 2808g. H is the Hamiltonian defined by (t, := [(t, x, fo), and &g is the potential function satisfying fa (t, (x))(xxt)).
(z, fo)
We briefly explain the action cost RS and the HJB regularization RA. Equation 7 is actually computed by Eq. {II. First. relax the constraint at time t,, i.e. R. + D( Pt, ). by using Wasserstein
distance (Bunne et al.
2021 Hashimoto et al. 2016). KL--divergence (Finlay et al., 2020 Onken et al. 2021 Tong et al. 2020). or a combination of these (Lavenant et 20.1(2021) for the discrepaney measure D. This objective function can be reinterpreted with D( Pt.) as the data-fitting term and
the action cost R. as the regularization. Second, we exploit SOT theory by incorporating further structure into the modeling. Similar to OT-Flow (Onken et al. 2021), we model the drift functions by using the potential function fa V Ht X, V (x t) and encourage the potential function & to satisfy the (stochastic) HJB equation (Yong & Zhou, 1999) by minimizing the PDE loss R2 The relation between the drift and the potential function is an analogue of Hamilton's 's equations of motion. The HJB equation represents Bellman's 's principle of optimality in continuous-time optimization.
3.3 EXAMPLES OF LAGRANGIAN IN NEURAL LAGRANGIAN SCHRÖDINGER BRIDGE
In this section, we provide the three Lagrangian examples of the NLSB and their use cases. The examples demonstrate that the Lagrangian design allows a variety of prior knowledge to be reflected in the sample trajectories. See Appendix E.5 for more examples.
Potential-free system. Without a specific external force on the individual sample, the principle of least action typically indicates that when sample moves from point A to point B, it tries minimize energy by reducing the travel distance. This means that the drift is encouraged to be straight and the stochastic movement be small, i.e., the Lagrangian is formulated by L(t, x. u) |||| and the drift function is given as fo = Vx
Cellular system. Tong al. (220)) proposed to introduce the prior knowledge of manifold geometry and local velocity arrows such as RNA-velocity for modeling cellular systems. In the NL.BB, these prior knowledge be handled consistently by designing the Lagrangian. First, we introduce density-based penalty to constrain the sample trajectories on the data manifold. 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. Next, we redefine the velocity regularization, which is formulated as cosine similarity maximization by Tong al. (2020), as squared error minimization and add the Lagrangian as well. Therefore, the Lagrangian for the cellular system is defined by
L(t,x, u) = 2 ||u(x,t)11 U(x t) + |llu(x, t) v(x, t)1??
(10)
Density
Energy
Velocity
where v(x, t) is the reference velocity of the cell at the position x and the time t. The drift function obtained by fo = (v(x, 1) (x t)).
Random dynamical system. Let (xx be the potential energy of the system, R be the mass matrix,
which is symmetric and L(t, x, u) 2 u Ru Uxx) be the Lagrangian in the random dynamical system. The drift function is given by f, - X t). Then, the individual samples are encouraged to follow stochastic analogue of the equations of motion of the Newtonian mechanics.
In practice, the potential function can be used to roughly incorporate information such as obstacles and regions where samples are not likely to exist. Practical examples are shown in Figs. and
IMPLEMENTATION OF NEURAL LAGRANGIAN SCHRÖDINGER BRIDGE
4.1 TRAINING FOR NEURAL LAGRANGIAN SCHRÖDINGER BRIDGE
In this section, we describe a practical learning method for neural SDEs by minimizing loss in Eq. on the training data, as described in Algorithm| I 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. } be a set of time points and denote the data set at time t, as X,, Next,
Let T = (to. { describe how calculate the loss in Eq. ( using the training data (Xe Jt. We compute
5

---

Published as a conference paper at ICLR 2023
the distribution discrepancy using the L2- Sinkhom divergence W. between the observed data and the predicted sample at all observation points in except the initial point at time to. The Sinkhorn divergence can efficiently approximate the Wasserstein distance and solve the entropic bias problem when using the Sinkhorn algorithm, ie. μ = v < (1. v) = 0. The prediction sample is obtained by numerically simulating neural SDE from the training data at one previous time point by standard SDE solver such as the Euler-Maruyama method. The second and third terms, R. and Rr. are also approximated on the sample paths obtained from the numerical simulations of SDE. Therefore, the computational cost of empirical R. and R, is not high and can be further accelerated by the model architecture described in the next section. To summarize, we simultaneously solve K approximated LSB problems by minimizing the subsequent loss for the model parameters 0, Φb
((0,0)=
E Welme, pi:?) de (tx-11 RR. (ts-1,ts) An (t)- (eee:(tttt10 (ta-1, tx), (11) VAT
0.0 where Pks Pta are the ground-truth and predicted probability measures at time t. expressed by data Rh
samples, respectively. (tx-1, and R,, (fa-1- (.11) f2) are empirical quantities computed on the simulated sample paths from fx-1 to & with data Vx((x-1) Xia_ as the initial value. The weight coefficients (e (x) and A4 (1)) are tuned for each interval [4x-1,4x] 4x] respectively.
4.2 MODEL ARCHITECTURE SELECTION FOR SPEEDUP
We adopt the model of the potential function b proposed by OT-Flow (Onken al. 2021 because it has two advantages in our framework. First, it can compute the gradient Vx 4 explicitly, which enables us to calculate the drift function easily. Second, the model designed for the fast and exact computation of the diagonal component of the potential function's ' Hessian. Moreover, when we assume that the diffusion model's 's output a diagonal matrix, L.e. go (x, t) E R and gi, = 0 (i of j), the function D & is also diagonal matrix, and the Eij in the second term of Rr shown in Eq. (9)
turns into the sum of the diagonal components only. Combining these two tricks enables speeding up
the computation of RR, which requires the expensive computation of the Hessian of the potential function Φ. 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, in neural SDE. We use two-layer fully connected NN for the diffusion function. We also adopt the device used by Kidger et al. (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. See Appendix C for details on the model of h
EXPERIMENTS
We evaluated our methods on two datasets. 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. We set the Lagrangian for the potential-fre system,i.e. L(t, x, u)= = |//l// Second, we used the evolution of single-cell populations obtained from a developing human embryo system. We used the Lagrangian for the cellular system and compared several combinations of the regularization terms. In Tables and "E" is the energy term, "D" is the density term, and "V" is the velocity term. The density term U (x, t) is the liceeiooddo function of the data estimated by GMM. We compared our methods against standard neural SDE, OT-Flow (Onken et al. 2021). TrajectoryNet (Tong et al.] 2020). IPF with GP (Vargas et al.) 2021 and NN DD Bortoli et al. (022]. and SB-FBSDE (Chen et al., 2021a). We trained the standard neural SDE using only the Sinkhorn divergence. The velocity model of TrajectoryNet includes the concatsquash layers used in Grathwohl et al. (2018). The base models of OT-Flow and TrajectoryNet were trained with the standard neural ODEs scheme. +OT represents a model trained with the
bT-based regularization defined by Eqs. 1 and (2) We used only R. for TrajectoryNet;; we used both R. and Rr for OT-Flow. We set the interval-dependdet coefficients Àe, An for the OT-based regularization as well as Eq. (T].. The drift model of IPF (GP) was changed to sparse GP from vanilla GP (Vargas et al. 2)21) save computation cost. The drift model of IPF (NN) and SB-FBSDE are the same networks as the NLSB for a fair comparison. The diffusion coefficients of IPF and SB-FBSDE were tuned hyperparameters. See Appendix F for more details hyperparameters.

---

Published as a conference paper at ICLR 2023
Neural SDE
IPF (GP)
Neural SDE
NLSB (Qurs)
NLSB (Qurs)
IPF (NN)
0.5
TrajectorNeet
0.30
OT-FLoW
TrajectoryNet
0.4
0.25
OT FLow
IPF (GP)
FMD-17
0.20
IPF (NN)
EMD.11
0.3
0.2
0.10
0.1
0.05
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
time
(a) MDD 12 time points synthetic OU process dhta.(h) CDD at 12 time points on synthetic OU process data.
Figure 3: Numerical evaluation on synthetic OU process data. All MDD and CDD values were computed between the ground-truth and the estimated samples within generated trajectories all-step ahead from initial samples x(to).
alkt xX01)
0
0
(a) ground-truth SDE
(b) Neural SDE
(c) NLSB (Ours)
(d) TrajectoryNet OT
Figure 4: ID OU process data and predictions. The five colored point clouds in the background the ground-truth data given at each time point. The pink area and the light blue line are the one-sigma empirical confidence intervals and their boundaries for each trajectory, respectively. All trajectories were generated by all-step prediction from the initial samples at the time & = 0.
5.1 SYNTHETIC POPULATION DYNAMICS:
PROCESS
Data. We used time-dependent one-dimensional Ienstein-Uhlenberd (OU) process defined by:
2ta
dx, (ut 0X,) dt
dW,
fk-1
where μ 0.4, -= 0.1, o 0.8. and tx-1 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} snapshots for training. generated 2048 and 512 samples for each time point as training and validation data, respectively.
Performance metrics. We evaluated the estimation performance of the dynamics in the time interval (to, 4x-1] by using two metrics on test data: marginal distribution discrepancy (MDD) and conditional distribution discrepancy (CDD).
smaller MDD between μ4 and Pt calculated with the Wasserstein-2 distance indicates better prediction of population-level dynamics at time f. We calculated MDD at 12 equally spaced time points and the square root of the earth mover's 's distance with L cost (EMD-L2) between 1000 samples generated from the ground-truth SDE and predicted by the model. When evaluating the based method, we ran 100 simulations from the same initial values, computing the MDD value each time and computing their mean and variance. A smaller CDD between (x(0)xx(ta) and Px(t)x((t))
7

---

Published as a conference paper at ICLR 2023
using the Wasserstein-2 - distance indicates better prediction of the time evolution of the initial sample x(to). In short, it is a metric for evaluating the time evolution of at the individual sample level. In the actual CDD calculation, we first prepared 1000 samples x(to) the initial time point. We then generated 100 trajectories from each initial sample by the trained model and the ground-truth SDE, and calculated the EMD-L22 for the samples from (x(1))x(to) and (x(t))x(ta) 12 equally spaced time points.
Results. The evaluation results are shown in Fig. 3 and the visualization of trajectories is shown in Fig. Figure | shows that NLSB and IPF outperform neural SDE and is comparable to other ODE-based methods in estimating populations with small variance. In contrast, the SDE-based methods outperform ODE- based methods when estimating populations with a large variance. That indicates that NL.B and IPF can estimate population-leva dynamies even when the population variance is large small. Furthermore, NLSB and IPF have smaller CDD value than neural SDE. Figure (4b shows that the average behavior of samples (t) estimated by neural SDE different from that of the ground-truth SDE (see Fig. 4ap. especially in the interval [0, 1J. In contrast, the predictions by NLSB and IPF in Figs. bc to 6e are much closer to the ground-truth. These results show that the prior knowledge of the potential-free system helps to estimate the sample-level dynamics. See Appendix EE2 for further results and analysis.
5.2 SINGLF-CELLL POPULATION DYNAMICS
Data. We evaluated on embryoid body ScRNA-seq data (Moon et al. 2019). 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. During this period, cells were collected at five different snapshots (lo: day 0 to 3, f.: day 6 to 9, t2: day 12 to 15, t3: day 18 to 21, lg: day 24 to 27). The collected cells were then measured by scRNAseq, filtered at the quality control stage, and mapped analysis (PCA). For details, see Appendix E.2 in (Tong et al. 2020). We split the dataset into train, validation (~~ 8.5%) and test data (~ 15%).
dimenssiiionl feature space using a principal component
Performance metrics. Unlike the experiment described in Section 5.1 there are no ground-truth trajectories in the real data. Thus, MDD can be calculated only at the time of observation, and CDD between the ground truth and predicted trajectories cannot be calculated. To evaluate the sample-level dynamics, the model was trained on the full data and the data without only one intermediate snapshot, respectively. We then calculated CDD between the predicted trajectories by those models. Let D-tx be the training data without a snapshot at time tg. 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. When evaluating the SDEe-based methods, we calculated the mean and standard deviation of 100 MDD scores. All performance metrics are calculated on the test data.
Results. Table 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. In particular, the SDE-based methods significantly outperform the ODE-basce ones in predicting the transitions where the samples from t to t2 and from t3 to t4 are highly diffuse, indicating that the explicit modeling of diffusion is effective. (see Fig. TO in Appendix HD. Overall, the standard deviation of the MDD values smaller for NLSB than for neural SDE, demonstrating less variation in the approximation accuracy of the marginal distribution. Table shows that NLSB estimates the sample-level dynamics robustly with and without population at the intermediate time point than neural SDE with some exceptions. Especially, energy regularization is the most stable and effective. This result suggests that the LSB- based regularization helps interpolate the sample- level dynamics. See Appendix E.3 for further results and analysis.
DISCUSSION AND CONCLUSION
In this work, we proposed a novel framework for estimating population dynamics that reflect prior knowledge about the target system. Unlike existing methods with OT (Schiebinger et al.) 2019 Yang & Uhler 2018), or CNF (Tong et al.) (020) for a biological system, we explicilly modeled the diffusion phenomena by using SDES for the samples with stochastic behavior. This allowed to
8

---

Published as a conference paper at ICLR 2023
Table I: Evaluation results for population-level dynamics on five-dimensional (5D) PCA space at time of observation for ScRNA-seq data. The MDD value at & is computed between the ground-truit -trut and the samples predicted from the previous ground-trute samples at - for each k 1, 2, 3 and 4.
MDD (EMD-L2) &
t
t2
t3
NLSB (E)
0.71 0.020
0.86 + 0.027
0.83 + 0.016
0.79 + 0.012
NLSB (D)
0.67 + 0.017
0.90 = 0.029
=.015
NLSB (V)
0.00
0.89 = 0.030
0.83 : 0.022
0.81 0.019
NLSB (E+D+V)
0.68 = 0.016
0.84 = 0.030
0.81 + 0.018
0.79 + 0.017
Neural SDE
0.69 + 0.020
0.91 +: 0.029
0.85 + 0.025
0.81 0.017
OT-Flow
0.83
1.10
1.07
1.05
OT-Flow + OT
0.85
1.05
1.09
1.00
TrajectoryNet
0.73
1.06
0.90
1.01
TrajectoryNet OT IPF (GP)
0.76
1.05
0.88
1.10
0.70 + 0.015
1.04 - 0.041
0.94 + 0.029
0.98 + 0.033
IPF (NN)
0.73 0.019
0.89 0.89
= 0.030
0.83 0.020
SB-FBSDE
0.56 = 0.010
0.80 - 0.017
1.00 + 0.019
1.00 + 0.010
(E+DeV)
TrajrtoryNett OT
IPF (GF)
Table Mean value of CDD on 5D PCA space evalu- ated at 7 equally spaced time points within time pe- around excluded time point ta for each k = 1,2 and 3. The CDD value at the time I 14-1, 4a+11 was computed between the two groups of predicted samples generated from the samples at 4--1 using the model trained on all data and the data D D-1+*
OT-FLow
5B-F65DE
4.5
riod [t
dlm20e
3.3
dmm2
dim20p
dim23
EMD-12
25
2.0
dm10
dim1
dim10
dmm10
Mean CDD 4
No, [2]
[4, [[]
((2, (4]
1.5
dim
dims
NLSB (E)
0.88
0.72
0.79
1.0
clm5
dmm
clm5
Gim5
dini
NLSB (D)
1.64
0.76
0.77
10
10"
10
NLSB (V)
1.15
0.82
0.86
training time [s] (log scalei
NLSB (E+D+V)
0.96
0.83
0.84
Neural SDE
1.36
0.85
0.87
Figure 5: Relationship between data dimen- sion, performance, and training time. The x-axis the learning time until sufficient convergence. The y-axis is the mean value of MDD over time points 4 to (4. represent- ing population-level performance.
IPF (GP)
0.97
1.03
1.06
IPF (NN)
0.90
0.95
0.97
SB-FBSDE
0.84
1.06
1.49
handle the uncertainty of the trajectory (Fig. 9 and to successfully capture the diffuse transitions
(Table and Fig. 10). In contrast, Vargas et al. 2021 and Bunne et al. (002)) estimated the sample
trajectories of biological systems using the SDE solution to the SB problem. Vargas et al. (2021) proposed GP-based IPF to solve the SB problem. [unne al. (222) proposed GSB-Flow, in
which two SB problems are solved sequentially. Compared with these methods, our method handled wider class of SDEs. and the diffusion function of the SDE was learned from the data using backpropagation. In addition, designing the Lagrangian enables us to flexibly incorporate prior knowledge about the target system into the model. Our Lagrangian- based regularization also treated the biological constraints proposed by Tong et al. (2020) and Maoutsa & Opper (2021) in unified manner. 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. Our method limited in that it cannot model reaction phenomena such cell birth and death, and restriction to the diagonal diffusion matrix (Section 4.2) may cause the model to less expressive. Future work includes developing methods that can handle reaction phenomena with advection and diffusion.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank the lab members for their discussions and inspiration for our research.
REFERENCES
Brian DO Anderson. Reverse-time diffusion equation models. Applications, 12(3):313-328 1982.
Stochastic Processes and their
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3)3375-393, 2000.
Charlotte Bunne, Laetitia optimal transport modeling of population dynamics. arXiv preprint arXiv:2106.06344 2021.
Andreas Krause, and Marco Cuturi. Jkonet: Proximal
Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. Recovering stochastic dynamics via gaussian schrödinger bridges. arXiv preprint (rXiv:2202.05722. 2022.
Kenneth Caluya and Abhishek Halder. Wasserstein proximal algorithms for the schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using Norward-bacweer sdes theory. In International Conference on Learning Representa- lions, 2021a.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, (99(2):671699.) 2016.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. SIAM Review, 63(2):249-31 2021b.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-329 1991.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
Lawrence C Evans. An introduction mathematical optimal control theory version 0.2. Lecture notes available at http://math. berkeley. edu/ evans/control. course. pdf. 1983.
Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Soc., 2010.
Chris Finlay, Jörn Hennik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pp. 3154-3164. PMLR, 2020.
Mike Fisher, Jorge Nocedal, Yannick Trémolet, and Stephen J Wright. Data assimilation in weather forecasting: a case study in pde-constrained optimization. Optimization and Engineering, 10(3): 409-426. 2009.
Robert Fortet. Résolution d'un systeme d' d'équations de m. schrödinger. J. Math. Pure Appl. IX, 1: 83-105, 1940.
Jason Gaitonde, Jon Kleinberg, and Éva Tardos. Polarization geometric opinion dynamics. In Proceedings of the 22nd ACM Conference Economics and Computation, pp. 499-519, 2021.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In International Confer- on Learning Representations, 2018.
10

---

Published as a conference paper at ICLR 2023
Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rmss. In International Conference on Machine Learning, pp. 2417-2426. PMLR, 2016.
Desmond J Higham. An algorithmic introduction to numerical simulation of stochastic differential equations. SIAM review, 43(3):525-546, 2001.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.). Advances Neural Information Processing Systems, 2022. URL https::/ net (afoormmmd=aahhhhooo
Benton Jamison. The markov processes of schrödinger. Zeitschrift für Wahnrshhemlicckeettsseeeei und Verwandte Gebiete, 32(4):323-32) 331,, 1975.
Patrick Kidger, James Foster, Xuechen Li, and Terry Lyons. Neural sdes infinite-dimensionl gans. In International Conference on Machine Learning, pp. 5453-5463. PMLR, 2021.
AN Kolomgorov, IG Petrovskii, and NS Piskunov. Study of diffusion equation that is related the growth of a quality matter, and its application to a biological problem, byul. Mosk. Gos. Univ. Ser. A Mat. Mekh, 1(1):26, 1937.
Solomon Kullback. Probability densities with given marginals. The Annals of Mathematical Statistics, 39(4): 1968.
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger. Towards a mathemati- cal theory of trajectory inference. stat, 1050:18, 2021.
Christian Léonard. From the schrödinger problem to the monge-kantorovich problem. Journal of Functional Analysis, 262(4):1879-1920 2012.
Christian Léonard. A survey of the schr\ odinger problem and some of its connections with optimal transport. arXiv preprint
2013.
Xuechen Li, Ting- Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Staristics, pp.
PMLR, 2020.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schr\ odinger bridge. arXiv preprint arXiv:2209. 09893, 2022.
Kenneth G Manton, XiLiang Gu, and Gene R Lowrimore. Cohort changes active life expectancy in the us elderly population: Experience from the 1982-2004 national long-term care survey. The Journals of Gerontology Series Psychological Sciences and Social Sciences, 63(5))8266 $281, 2008.
Dimitra Maoutsa and Manfred Opper. Deterministic particle flows for constraining stochastic nonlinear systems. arXiv preprint arXiv:21 2021.
Toshio Mikami. Variational processes from the weak forward equation. Communications in mathe- matical physics, 135(1):19-40, 1990.
Toshio Mikami. Monge's problem with quadratic cost by the zero-noise limit of h-path processes. Probability theory and related fields, 129(2):245-260) 2004.
Toshio Mikami. Optimal transportation problem as stochastic mechanics. Selected Papers on Probability and Statistics, Amer. Math. Soc. Transl. Ser, 2(227):75 94, 2008
Toshio Mikami. Stochastic optimal transport revisited. SN Partial Differential Equations and Applications, 2(1):1-26, 2021.
Toshio Mikami and Michèle Thieullen. Duality theorem for the stochastic optimal control problem. Stochastic processes and their applications, (16(12):1115-183) 2006.
11

---

Published as a conference paper at ICLR 2023
Kevin R Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew Him, Ronald R Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology. 37(12): 1482-1492, 2019.
Derek Onken, Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, 2021.
Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of Statistics, pp. 1160-1174, 1995.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-943, 2019.
Simon Schweighofer, David Garcia, and Frank Schweitzer. An agent-based model of multi- dimensional opinion dynamics and opinion alignment. Chaos: An Interdisciplinary Journal of Nonlinear Science, 30(9):093139, 2020.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: dynamic optimal transport network for modeling cellular dynamics. In International Conference on Machine Learning, pp. 9526-9536. PMLR, 2020.
Dario Trevisan. Well-posedness of multidimensional diffusion processes with weakly differentiable coefficients. Electronic Journal of Probability, 21:1-41, 2016.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arXiv preprint 190.:93905.09883 2019a.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pp. 5084-3114, PMLR, 2019b.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
Gefei Wang. Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference Machine Learning, pp. 10794-10804, PMLR, 2021.
Karren D Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2018.
Liu Yang and George Em Karniadakis. Potential flow generator with 12 optimal transport regularity for generative models. IEFF Transactions on Neural Networks and Learning Systems, 2020.
Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https:// openreview net forum id= uCb2ynRu7y
12

---

Published as a conference paper at ICLR 2023
Algorithm 1 Training of NLSB
Input:
Dataset (KiHer potential function 49, diffusion function gø, Lagrangian regularization coefficients Ae, An-
while 0, have not converged do
Set loss L = 0.
for k 0 to K do
Sample mini-batch (*)) from Lee Set Yr, = =x())
Numerically solve augmented SDE ds; = fo () dl (sc. dW, from la to 1x+1.
V. (o Lt. (tH - (V)(
(sc,
|0 1) + Eng Das V By t) + Hā t)
(ye.)e)t) t)
Yt
E (8:, ()
Rt, (R((s:)).
Stx
0
Sample mini-batch (k+1) from Ness and calculate the loss function E(tu; (x+1).
Elle, (++1 We (y ta x (k+1)) Ao(le, t ) (Ka, (x+1) Aulles le+1 Rr (4x, (x+1),
Yotts
Ro tx+2) 10 (tk, tx+1)
Stet1
Accumulate loss + C+ f(tx,tk+1). Update 0, é with gradient VoL and V VdC, respectively.
Output: to. g+
A THEORY OF STOCHASTIC OPTIMAL TRANSPORT
The SOT problem is the problem of finding a stochastic process that minimizes the expected cost under fixed marginal distributions at several time points. 'hee SOT problem is stochastic analog of the OT problem, especially related to the dynamic formulation of the OT problem introduced by Benamou & Brenier (2000). It is also considered as stochastic optimal control (SOC) problem with additional terminal constraint. Furthermore, the classical Schrödinger Bridge (SB) problem is a special case of the SOT problem and is related to Nelson's stochastic mechanics, reformulation of quantum mechanics using diffusion processes. An overview of the relationships among these problems is illustrated in Fig. -
We first explain the dynamic formulation of the OT problem in Appendix A. the SB problem in Appendix A.2 and then describe the SOT problem given fixed marginal distributions generalization of the SB problem in Appendix A.3 Next, we discuss the LSB problem. which is the main focus of our paper among the SOT problems Appendix B.I and finally, we discuss the relationship between NLSB and OT-Flow in Appendix B.2
A.I DYNAMIC FORMULATION OPTIMAL TRANSPORT PROBLEM
Benamou & Brenier (2000) redefined the OT problem with quadratic cost in continuum mechanics
framework. They introduced time parameter t E [0, 1] and considered the transport process from to the advection in the time interval [0. 1]. The OT problem was then reformulated minimization problem with respect to the time-varying velocity vector field v and the time evolution of the distribution p. advected by Uz.
13

---

Published as a conference paper at ICLR 2023
Theorem A.1 (Brenier-Benanl Benamou formulation (Eularian formalism); (Benamou & Brenier 2000)).
W2 (1, ( J [|v(xx (x) dxdt,
inf (v.px)
(12)
subject to dep. = div(pev), Po = P1 = q.
(13)
where p and q are the densities of probability measures y and v, respectively.
The first condition in Eq. (3) is known the continuity equation and represents the conservation of mass in time evolution. We can consider an equivalent formulation of Eq. (12) by introducing
Lagrangian coordinates (l.,x), Lagrangian coordinates X(1, x) represent the position at time of
the particle whose initial position is x, i.e. ((0, x.
Theorem A.2 (Brenier-Be Benamou formulation (Lagrangian formalism); (Benamou Brenier 2000)).
W2(1, inf I'' | ()) dxdf, v ax(t, = v(X(,,,), Po 4
(14)
subject to dt
The right-hand side of Eq. (14) can be viewed as problem of finding the shortest path ka.k.a. geodesic) for each particle in the sense Euclidean space between probability distributions specified at times & 0, 1. In continuum mechanics, Lagrangian formalism (Eq. (14) describes the motion of each individual particle, while Eulerian formalism (Eq. (12) focuses on the global property of particles.
Benamou & Brenier (2000) also showed the optimality conditions of the dynamic formulation. They introduced the Lagrangian multiplier of the constraint of Eq. (13) and obtained the saddle point conditions using the variational method.
Theorem A.3 (Optimality conditions for the dynamic formulation; (Benamou Brenier 2000
There exists space-time dependent potential function
functiion 4: IR [0, 1] IR which satisfies:
(X ),t) - Vx tt,,, (- 1 (x, v (x,
(15)
1 || (x, t)|? 0,
0,4(x,t)
(16)
2
ox = (x), iy(y) - (y),
(17)
where f and g" are Kantorovich potentials. Equation 15 is Hamilton's 's equation of motion with the Hamiltonian defined by (p. := sup, {p. v(x, t)) v[v t)]? and the momentum as
= V and equation 16 is Hamilton-Jacob- Bellmma (HJB) equation.
Equations 15 and 17 indicate that the particle, whose initial position is Xg X. moves straight ahead the constant velocity Vo ~o vx ff(((xx (xo) during t [0, 1]. Even without using the variational method, Hamilton's 's equation of motion (Eq. (15) and the HJB equation (Eq. (16) can be also derived from the Bellman's 's principle of optimality. The existence of potential function
satisfying the equation 15 is guaranteed from the Pontryagin Maximum Principle (Evans) 1983 2010
A.2 SCHRÖDINGER BRIDGE
Definition A.4 (SB problem; (Jamison 1975 Let C(OO, 1], Rd) be the space of vaduud continuous functions on time interval [0. IJ. Denote by P()) the probability measures space on the path pace Q. The SB problem is defined by
min DKI. (Q])), QEP(N)
subject to
Co ~ n4o Qi U1.
(18)
where Mo, 1/1 PP( are the probability measures at the time 0 and 1, respectively, and the relative entropy DKL. log ((P) dQ if & P. and DKL. otherwise.
14

---

Published as a conference paper at ICLR 2023
The case of no prior dynamics is considered classically, Le., the reference path measure P is the Brownian diffusion SDE dX; v2e dW,, where W, is the standard Wiener process. We refer to this setting as the classical SB problem, following the reference (Caluya & Halder [021). Chen et al. (202Ib) derived the optimality condition for problem Eq. 418 which is characterized by the forward and backward time- harmooic equations.
Theorem A.5. Let (x, t) and L (xx t) be the solutions the following PDEs:
-
eg
s.t. V. ((,0))(((-0) = Po. W0 (.
(19)
where Po and pr are the probability density of pi and pr, respectively. Then, the solution to the SB problem (Eg. (18) can be described by the following forward or backward SDE:
dX, 2 log (X,,1) di + 22
X pos
(20)
dx, 2 log V (X,,t) dt + /2c dw
X1 P1-
Dai Pra 1991) considered a dynamic formulation of the classic SB problem (Eq. it as a SOC problem with the additional terminal constraint.
18 by interpreting
Theorem A.6 (Dynamic formulation; (Dai Pra 1991)). Let f = V log V where T satisfies the SB optimality (Eq. (19). Then, f is the minimizer of the following optimization problem:
Vs,e 11) == inf E f
2 |ff (X,, t)l? dt
(21)
subject to
dX, tX dt + d2e, Xo po. X1 #1.
Léonard (2013) formulated the SB problem in Eq. (8) into a variational SOC problem equivalent to the above problem (Eq. (2I).
Theorem A.7 (Dynamic formulation (Eularian formalism); (Léonard 2013 Let fx t) =
Ve log E L (Xr, T) X, x). where satisfies the SB optimality (Eq. 19
Then, f is the
minimizer of the following optimization problem:
/ [1x t)]1 dp((x))(,
vs.c(po: 111 = inf
(22)
(f.pt)
subject to d.p: div(p:f) EApt, Po = po: P1 H1.
(23)
where Pt the probability density of the probability measure PL.
The first condition in Eq. 23 is Fokker-Planck (FP) equation. The two versions of the dynamic
formulation (Eq. 21 and Eq. (22)) are equivalent, i.e. Vs m1) = (po. (1). Furthermore, the dynamic solution of the OT problem (Eq. (I5)) is obtained as the zero-noise limit (Mikami 2004
Léonard 2012) of the classical SB problem.
Theorem A.8 (Zero-noise limit of the classical SB problem; MMikami) 2004)). Let (t) be the solution of the SB problem (Eq. (20)). Suppose that po: P1 P P(R have fniit second moments and the density function Po(x)= = ((dd)/ddr exists. Then, the following holds
lim (so, m1) We(Ho: p1)? ++0
and there exists a convex function satisfying
-xo (VV X
lim 600
sup
The map V x is the optimal transport map of the OT problem with quadratic cost.
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 weol-known. Finally, we introduce a simpler variant of the SB problem for which closed-form solution exists.
15

---

Published as a conference paper at ICLR 2023
Theorem A.9 (Gaussian SB problem; (Bunne et al. 2022 ). Let Q be the solution of the Gaussian SB problem defined by
min (QI(P) QEP(1))
subject
2(0) ~ No. Q(1)
where No = NNo, Eo) and N1 = [N Σ1) are Gaussian distributions. The reference path measure IP is described by the linear SDE as follows.
dX,- = (c(t) + f(t) dt + g(t) dW., Xo No.
where 0: [0,1] -
f: 00 + TRd g: [0, 1] + R are smooth functionn. We define the following
notation from Bunne et al. 2022
exp (f, c(s) ((o)))ds ds
Do = 42 D
+ +o+1) 3
C. = 2 -2) Do -021)
1
:=: k(t,T) T)
v k(T,T)
Fe: Ti rTTr
x(I,T)'
((1) = To a(s)ds, A: = S0 T. gg
J Ts g (s)
P: i (r:2i +F,Co.), Q: == -. (FrLi +r.C..)
S. P. -Q. + [c(())((t f) (1 - pr) - 9 (1)p.) Fip4 TIpT + ((t) (T((T)
1
E- , + t (Co. + CC,) + x(t,t) pt) Then, the solution Q is a Markov Gaussian process where the marginal X; N 27), and follows the following SDE:
dx S. 21 (X;
(Xi - (G) dt g(t) dW..
A.3 STOCHASTIC OPTIMAL TRANSPORT WITH TWO ENDPOINT MARGINALS
generalized the OT problem and defined the SOT problem as a random mechanics problem determined by the principle of least action. The SOT problem with the endpoint marginals fixed to and μ1 is given by the following.
Mikami 2008
Definition A.10 (SOT problem; (Mikami 2021)). Let L be continuous function, the Lagrangian and let u + L(t, x, u) be convex. The SOT problem with two endpoint marginals is defined by
S ((t.Xx,x(Xxtt) dt subject to Xo ~ Po: X, ~ Pi.
Vio. M1) = inf E xed
(24)
where is the set of all IRd valued, continuous semimartingales {XeJosesı on a complete filtered probability space such that there exists a Borel measurable drift function fx (X, t) for which satisfies the following conditions:
for all t. 2. X; = Xx 1o' fx (X, 8) ds + So s),, dW., 05t51
1. w + fx (t, w) is
3. E [ (1 s)1 lg( 2)) dt 00
Definition A.11 (SOT problem for marginal flows; (Mikami 2008
So [ L (t,x,f(xxt)) dp.(x) dt
vI)
inf (EA((P))SS))) subject Pa uo: pt pl:
(25)
(26)
A (1petosess) = f(x,1) Oep. - div(pef) Σ
8²
(D ()((()]]
(27)
dr,Or,
-11=
16

---

Published as a conference paper at ICLR 2023
where pu and p1 are the densities of probability measures Pu and f1, respectively.
The minimizer of the SOT problem for marginal flows (Eq. (25) is obtained by (* t) E [x, X, = x)]. We introduce assumptions from (Mikami) 2021).
(A-1) L E C²RRd" IR (0, ool). u + L(t, u) is strictly convex. L(t, x, u) L(t,y, and Lt, x, u) + x(t,x,u) are bounded on t E [0, 1] and x,y, E R SuPxER VxL(t, x, u))l is locally bounded. lim inf L(t, x, u)/ | co.
(A-2) (t,, u) is bounded uniformly nondegenerate on [0. R R
Theorem A.12 (Trevisan's Superposition Principle; (Trevisan 2016)) ). Assume that there exists
f Rd x [0 Rd and {pe)osesı C(R")) such that f satisfies the FP equation, i.e. f
A {pe)osesi))
Then, there exists a semimartingale (X,)osisa for which the following holds:
x, = Xo+ + [ f(X,, 8) ds + { g(X,, dW..
X s0
From Theorem A.12 can be easily shown that (1uo, = v(No: (11) and there exist minimizers X of the SOT problem (Eq. (24)) for which fx- (X', .) = (X;, .t).
Theorem A.13 (Duality theorem; (Mikami 2021)). Suppose that (A-1) holds. Then, for any po M1 P(R"),
ULe
Φ(x, 0; duo(x) - -J f(x) ).
V(uo, P1) = v(Ho,
sup rEOt (Rd)
(28)
where Cg (R.) is the set of all infinitely differentiable functions on IR, which have bounded continu- derivative and & (xx, t: f) is the viscosity solution the HJB equation:
0. 0x0(x,t;f)
Σ tx,
(x (x f) H (t,x, - -Vx0(x,t;; 0,
(xx 1:f) f(x).
The Hamiltonian H is defined by [t, x, z) == sup (zz, u) L(t, x. u) Theorem A.14 (Mikami & Thieullen 2006 Suppose that (A-I) and (A-2) hold and V (Wo: p1) is finite. Then, there exists a minimizer
E a of (10, 11) given by
1 ('X:;, ds S g(x,, dW.
x; =
X
For any maximizing sequence {#n]n21 of Eq. (28). there exists a subsequence {nx]a21 such that I" ()..,8)= lim V (s, V 8)) k0000
B SCHRÖDINGER BRIDGE AND GENERATIVE MODELING
The theory of the SB problem is mostly mature shown in Appendix A.2 but scalable numerical methods for estimating SB are still actively studied. In particular, there have been many recent studies (De Bortoli al. 2021 Wang al. 2021 Vargas al. 2021 Chen et al.) 202 la Bunne et al. 2022). which uses the SB as process for generating data. These studies other than (Wang et al. 2021) proposed methods learn SDE solutions of the SB problem between the prior distribution and the target data distribution. These methods combine the classical multi-stage optimization method called Iterative proportional fitting (IPF) (Fortet) 1940; Kullback] 1968, Ruschendort] (995) for solving the SB with machine learning methods for optimization of subproblems. In IPF, the tollowing subproblems are solved alternately and iteratively. The reference path measure set to the initial
measure ())
(i) argmin DLL [R] Q (1 "
subject to
R(1) μ1,
(29)
FEP())
argmin DKL. (Q. (R)). QEP(())
subject to
((0) ~ p40:
(30)
17

---

Published as a conference paper at ICLR 2023
where [t),, Q(t) are the probability measures the time t on the path measures RR and Q The
path measures (!)) R at the i-th step are simulated by the following forward-backward SDEs
in Eqs. and (3 respectively.
f (X, dt + 22c dW, dx,
Xo Po:
(31)
dx, = b ( (X,.t t) dt + 2 dW,
X1 p1
(32)
The convergence of IPF was proved Ruschendori 1995).
The sub-optimization problems (Eqs. (29) and (30)) approached differently for each method.
First, Vargas et al. (2021) and De Bortoli et al. (221)) proposed to solve them by mean-matching regression of the SDE drift function using Gaussian process (GP) and NN, respectively. They find the
drift function of the SDEs that minimizes the following losses for some sampled time t.
b argmin x--!!!")
(x., - X + p(--1) (X, At) (x.)
((() argmin x f((X,) - (x, -(x,-
(4) (X (X,) - 6/2 (X))
f
In contrast, Chen et al. divergence- based losses are a modified version the approximate of SGM for use in alternating optimization schemes.
a0 proposed to use the divergence-based losses shown in Eq. 33 The
training
[T" 2
v (+) argmax
||v (X,, ()1 g gdivx(v) + ( v di vdl.
X
1 2 (X, 0)| g + u
argmax R
where are learnable drift terms of the forward-backward SDEs that redefines Eqs. 31
and 32 with fixed prior drift Eprior. simulating the path measures Q! R(.
dX, = [Eprior (2e, t) + g(t) (9) (X.,t) () dt + / dW,, (X,,
Xo
(33)
dx, = Sprior Xx, t) g ((t))(((x..))) dt + dW, X1
g(t)
To estimate more complex dynamics, Bunne et al. (2022) proposed to solve a general SB problem (Eq. (88) in which the solution of the Gaussian SB problem used reference measure P. As shown in Theorem A.9 the Gaussian SB problem has a closed-form solution, and the general SB problem is solved using the alternating optimization with divergence-based losses proposed by Chen
et al. a2)
B. LAGRANGIAN SCIIRÖDINGER BRIDGE PROBLEM
We consider the LSB problem constrained by the FP equation corresponding SDE (Eq. over the Euclidean space Rdd Definition B.1 (LSB problem constrained by the FP equation).
((u,m)) = inf ('' Jga L(t, x, (ft,x,f(x,t t)) dx:(x) dt,
(34)
(f.p:)ES subject to Pa = uo: Pt = -iI
(35)
d 82 0.p. div(pif) + Σ 1-=1
(f,m (
S s:
dr,Ün,
The LSB problem in Eq. 34 is more general problem that does not fix D. in the SOT problem
for marginal flows (Eq. (25) Thus, solution to the LSB problem in Eq. (34 clearly exists on the
basis of the existence of solution to the SOT problem in Eq. (25).
We practically solve the following relaxed LSB problem, where the terminal constraint are replaced
by the soft constraint.
18

---

Published as a conference paper at ICLR 2023
Definition B.2 (Relaxed LSB problem constrained by the FP equation).
inf ('' J x(t, f(x, t)) dpx(x) dt f G(x) dpi(x), (F,Ps)ES
V(Ho, M1)
(36)
subject to Po μo:
(37)
G(x) = 8p1
D(i (PI (x))xx)x)
where G is the terminal cost introduced by relaxing the constraint P1 μal and 5 is the variational derivative with respect P1-
derive optimality conditions for the LSB problems in Eqs. and (36) using variational
method in the following theorem. The derivation procedure is similar to that for the variational formulation of the SB problem by (Chen et al. 2021b).
Theorem B.3 (Optimality conditions for the LSB problem). There exists a space-time dependent potential function 4 Rd [0, R. which satisfies:
(x, = VtH((t,,, vxx(xxt)) t)),
(38)
(x0(x,t)+
D (x, f; f;) (2)(x)))) H =0,
(39)
5==
Φ(x, 1) G(x),
where (f",Pt) is the minimizer of the relaxed LSB problem Eq. (36)
The equation 38 and 39arr are the optimality conditions for both the relaxed LSB problem in Eq. 36
and the LSB problem where G(x) in Eq. (34)
Proof. We derive the optimality conditions for the relaxed LSB problem shown in Eq. (36) by reformulating it as saddle point problem for (pt, m.) == (efe). Let L be the the Lagrangian with the ume-space-dependent Lagrange multiplier (x,,)).
C(p, m, ②) = f (.. / 1, m.
(x(x) dxdl + / G(x)pı (x) dx
- -fo f t(x,t)
0²
dept
div(m.) (a)
[D,, ())x
dx de.
Ox Or,Or,
(6)
+==1
The term (a) is transformed by performing partial integral over t:
S. [ Φ(x, (x0)p.()) dx dt
gM Φ(x, ())) dx f F(x, dx
1.' / 0,0(x,t)pL dx dt.
The term (b) is simplified follows.
[ 22 Φ(x, t) div(m:) dx = / div(Am:) - -x((x) TV Vx 1) dx
[ mi(x) Vx dx.
19

---

Published as a conference paper at ICLR 2023
The term (c) is transformed by performing partial integral over
0 Dx, t)px(x) (xx t) dx = -
d Σ
(Da,(x, ()))
04((x, t)
dx
11=1
i1=11
Rd
Ozj
Ovi
d
(x,( t) 0x, dx
Σ
(x,
(x)
JR
Σ (x,( (20(xx))) px(x) dx.
nd
Then, rewrite the LSB problem
inf sup L(p, m, $), p.m +
(40)
C(p, m. 4) 4)=
/ ked
me t,x, L
+ 0,0((x,))+
Σ Das (xx (x)
pc(x) dx dt
Pt
gg G(x)p1(XX dx + fo' J mi(x) Vx t) dx df
[ $(x, 1)pp(( (x) dx / T(x, (x) dx.
Bd
The saddle point (p".m", x*) of the problem in Eq. 40
satisfies the following conditions:
8²
0+Ll (p- + O.p. +
m (x,
[xi,(x, t)p?] = 0,
dr,Oz,
i===1
0, (p-"(p,,m-000 "m=)) =0
(x, + Σ (x,( [ H(t,x, =
i=
8m fl ,0)) =
VeL(t, x, x,) f.) = ~>x;0
Op. f1 -0- =0
-
where 2. Vnl(l, x, u) and the Hamiltonian is defined by H((,x,%)
supu ((2, - L(I, x, u)). The Lagrangian satisfies L(4, x, u) = sup, {(u, 2) H(1,x,2))
Therefore, the optimal drift function is given by m
fr (x t)
p V,H(t,,x,2))
(41)
The potential function & is the solution of the HJB equation:
d
(x0(x,0) +
(x,(xx))
[V ( Ol., H(I, x, 8' (x, (x, = 00
(42)
1リ=1
Equation 41 and 42 are also optimal conditions even for the LSB problem in Eq.
34 where the
terminal constraint strictly satisfied, i.e. G(x)
From Theorem A.12 we can show that there exists a semimartingale (XtJosesı for which the following holds:
x / (x,, ds /. 8X( dW,, X, P. (0
where (f*, is the minimizer of the LSB problems in Eq. 34 or Eq. (36) and g" is the diffusion function determined from the minimizer (f" pt). Therefore, the LSB problems constrained by the FP equation (Eqs. (34) and (36)) and the original LSB problem defined by Eq. (6) are equivalent.
Note that named Eq. 6 the LSB problem, the SB problem is well-known in the machine learning field as the problem of finding the most likely stochastic process between sampled time points.
20

---

Published as a conference paper at ICLR 2023
B.2 DERIVATION OF NLSB AND ITS CONNECTION TO OT-FLOW
The loss function of the NLSB shown in Eq. 7 is the objective function of the relaxed LSB prob- lem (Eq. (6)) with additional PDE loss applied satisfy the optimality condition, HJB equation.
Theorem B.3 justifies the minimization of HJB-PDI loss and parameterization that strictly sat-
isfies Eq. (38 during training when the terminal constraint is not satisfied, i.e. G(x) * O. For computational efficiency, we practically define the HJB-PDE loss R2 in a weak form as shown in Eq. and evaluate R2, on the path of the simulated SDE.
NLSB and OT-Flow (Onken et al. 2021 have strong theoretical connections. NLSB is the practical solution to the LSB problem (Eq. (6) relaxed by the distribution discrepancy measure ID using neural SDE, while OT-Flow is the solution to the Brenier-Benamou formulation of the OT problem (Eq. ( relaxed by the KL-diiergence using neural ODE. The optimality conditions for the LSB problem shown in Theorem B.3 are analogues of the optimality conditions for the Benamrr problem shown in Theorem A.3 In NLSB and OT-Flow. the potential function is modeled using NN and optimized by both action cost (Eqs. (1) and (8) and HJB-PDE loss (Eqs. ( and ()) based on the respective optimality conditions. Theoretically. the LSB problem with the Lagrangian L(t, x. u) = |u|| and the diffusion function = V2e is reduced to the classical SB problem (Eq. (2T) and the solution to the Benammu problem is recovered from the zero-noise limit of the solution to the classical SB problem (Theorem A.8)
C MODEL ARCHITECTURE
The model structure of the potential function proposed in OT-Flow (Onke et al. 2021 is shown below.
((s) w NN N (s; (K, bilosss)))
T (ATA)s A)s b 2-
N (s: (K,, b. UM.
u, u + ho K,u, + b,) (1 M), u. = (Kos + bu).
where (x Rd+1 is input vector, w, Kg md K, (1 iisis M) m= b, (0 s i S M) E mmxm A, b, and e 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) = log(exp(x) + exp(-x)).
The gradient of the potential function is described.
(s V.N (s: (K, b. + (ATA)s b, V.N (s; (K, b, JosisM) = KJ (ios
(Kos + bo))z1,
2: = 2:+1 + AK; diag(o' (K, + b,) (1 < M). Z = 1. We can write o' down as tanh since a defined as above. Finally, the diagonal components of the potential function's Hessian is shown below.
(2 Vs [K[ diag(o'(Kos (Kos + bu))z1)
M
+ Vau:-Vas
(KK; + bi)) V
KJ diag(o" (os + bu) 3 z1 Ko
+ V.u, K. diag(o" (K diag(o"(K,u,-1 b,) 2.44)
[2 [Kos(Kos Ko) 0z) (Ko Kn)];
h
(o" + b,) 2411) (KVV,u, O
21

---

Published as a conference paper at ICLR 2023
where @ is the element- product, V u, can be obtained by using the following update equation:
Vau + Vaul1 + diag(ho' (K; b
The diagonal component of the potential function Hessian can be computed at computation cost of (mm2d). The complexity dm)²d) indicates that there a trade-off between the expressive power of the DNN and the computational cost of the Hessian.
OT-Flow modeled the potential function 2g that satisfies fa - Vx instead directly modeling the velocity function f of the neural ODE. The continuous transformation from x(0) to x(1) is described by
((1) = x(0) J/' (o((()))) dx
x(0) + - vxx))( dt.
OT-FFlow is trained by likelihood maximization in the well-known CNF framework, The likelihood computation is performed as follows.
pi (x(()) = po(x(()) - S. Tr (Vx(ooox(tt) dt,
= px - / Tr (=
dt,
where pl the density of the target distribution and po is the density the prior distribution, which usually a Gaussian or Laplace distribution. The fast and exact computation of the diagonal component of the potential function's 's Hessian useful for the computation of Appendix C
In contrast, propose for the first time to this technique to speed up the computation of the HUB-PDE loss R, in the NLSB under the assumption that the diffusion model's output is a diagonal matrix. The computed loss R2 is given by
S." f (x, E (x, Ve Hx, dpe(x)dtt H3 1) (x,, Lo ())
(xx,)))
D BENEFITS OF GENERALIZATION TO LAGRANGIAN
We first present the advantages of generalization by the Lagrangian. Next, define the general form of the Lagrangian considered in this thesis and explain its generality theoretically. We also provide several examples of Lagrangians to demonstrate their applicability as models for wide range of real-world systems.
The benefit of generalization by the Lagrangian is that the HJB equation and Hamilton's equation can be described in unified manner using Lagrangian, independent of the coordinate system of the space on SDE. In other words, the loss function can always be easily derived once the Lagrangian is designed.
The general form of the Lagrangian considered in this paper is given by
C (u-m) - U(x,)) t),
(u
(43)
where E RRdx x is the positive definite since the Lagrangian convex with respect to u, U the potential function defined from the prior knowledge on the target system. The optimal drift function obtained by fx t) = -2(R+ + R - (vx2, (x, + c) + v.
In Lagrangian mechanics, deterministic Newtonian dynamical systems can be described by the principle of least action using the general Lagrangian (Eq. (43), independent of the coordinate system. In contrast, since describe an even wider range of systems than Newtonian dynamical systems. The general Lagrangian (Eq. (43)) also includes as the special case the Lagrangian cost used in the optimal control (OC)
dealing with random dynamical system using SDEs, can
22

---

Published as a conference paper at ICLR 2023
problems such as the linear- queaar-uuarraicc regulator (LQR). To show the generality of the Lagrangian of Eq. (43), we demonstrate that the Lagrangian in linearly transformed coordinates can also be written the same form.
We define SDE on the coordinate-transfomonn state variable x R from the observed variable y R...
ta "//e J.
minimize fg
f(t,x,,(xtt)) (xx dt,
to nd
{ dx, = f(X(,t) dt + g(x,,t) dW.,
(44)
subject to
y Px,
Yo PX, ~ Pta = po: Y1 PX1 Ph = Pl.
The linear coordinate transformation is given by y Px represented by the projection matrix P E Rexd In particular, when we use the Wasserstein-2 distance W/2 as the distribution discrepancy measure in Eq. (II), the loss function of NLSB is almost invariant for the linear coordinate transformations.
Regular matrix. When the projection matrix Rdxd regular, the general Lagrangian defined on the observed variable space Id is computed on the state variable space E IR as follows.
(2(11,Pa)= = inf
Ily: y22l? dr (2),
XRdR
= =inf JRexR xBd
Px1 - Px2ll ?| dx1
inf /
(x1-x2) TR(x - X2) (x1,,,x2);
Bdd
L(,,y, uy; Ry: Cy, Vy, my Uy)
2 (uy vy) -yy"Ry (uy -vy) cy
cy(uy (uy -my) Uy(y,t),
1
2 [ux vx) P.P Puux - -xx + c P(ux - mxx) - Uy (Px,t),
[ux vx) T Rx (ux - vx) + (ux - mxx - Uxx(x,t),
2
L(,, x, uxi Rxx Oxi Vxi mx: Ux),
where R P and R2 P daxd are guaranteed to be positive definite, P."" the push-forward operator of the linear map represented by P-1, Cx = PTC, E Rd and Ux (x, t) Uy (Py, t) is the potential function. Especially when P is an orthogonal transformation. i.e. - I. then R I and R Ry hold.
The general Lagrangian has sufficient representational capacity to describe Lagrangians in coordinate systems that be linearly transformed into each other in unified manner. Furthermore, similar results are obtained for the PCA projection, which is an irregular matrix.
PCA projection matrix. We consider inverse projection from the latent space the data space y as the projection matrix e Rxd (d & p). The columns of the PCA projection matrix are orthogonal, i.e. P P = I. The cost functions defined on the observed variable space y E RP is
23

---

Published as a conference paper at ICLR 2023
efficiently computed on the low-dimensional space R as follows.
S( R
W2( - =inf
I - y2ll d
inf J J
PPx Px2l1 d
[
inf
||[| - -x2l1 x2),
RddRRd
(2[PP] MI,
pt.) Pb,))
L(t,y, My: Ry, Cy, Vy, my: Uy)
2 7/, - vyy) Ty)Ry(u (uy vy) + +Cy c my) Uy(y,t),
2 [ux vx) xP RyP vx) + cy P(ux - mx) (Px,t))
1
2 [1x vx) Rxx (ux vx) + cx (ux mx) - Uxx(x,)),
Lt,, x, uxi Rxi Cxi Vx: myx Ux),
where Rx
RR dxd is guaranteed to be positive definite, P.7 is the push-forward operator
of the PCA projection represented by PT, = PTcyy E Rd, and (x(x,)) t) = Uy (Py,t) t) is the potential function.
In Section 3.3 and Appendix we provided the four Lagrangian examples of the NLSB and their use cases. The examples demonstrate that the Lagrangian design allows a variety prior knowledge to be reflected in the sample trajectories.
In this appendix, we have described the LSB problem on the linear coordinate transformed space. In order to improve the modeling of dynamics, it is recommended to investigate the potential benefits of incorporating nonlinearities the projection from x to y or imposing specific geometric structures on the space of as considered in (Huguet et al. 2022). These considerations are deemed promising avenues for the advancement of NLSB.
D.1 POPULATION DYNAMICS SIMULATION BY REVERSE-TIME SDE
We describe a method for estimating the trajectories backwards through time and show experimental results using artificial synthetic data. According to the result from (Anderson 1982). the reverse-time SDE of the original SDE (Eq.
3 is given by
dx,- = (fo X f) us (X,,t) dt + go (X,,t) dW,, Ug.0 (X,,)= div (g. (X,, t) (X,;t)T) (X, t)g (X,, t) Vx (X,),
where W, is Wiener process that flows backwards in time.
However, the computation of the modification term of the drift function requires an estimation of the score function and high computational cost. Therefore, we newly parameterize the modification term u. by NN with parameters Ç. Then, we train only the term ug using Sinkhorn divergence loss in Eq. (45) by numerically simulating the eeverse-time SDE, For simplicity of implementation, used the same model architecture for modification term us as for the drift fa.
i(E)
Σ
Welmes PL.)
(45)
The experimental results for the reverse-time SDE learning method shown in Appendix E.2.2
24

---

Published as a conference paper at ICLR 2023
Table 3: Comparison of implementation
velocity/drift
diffusion
ODE/SDE solver
TrajectoryNet
FFJORD Grathwohl al.
2018
dopri5
OT-Flow
~V
Euler
Neural SDE
(t, V H (t, -V
FCNN ge (x, t)
Euler-Maruyama
NLSB (Ours)
FCNN g t)
Euler-Maruyama
IPF (GP)
sparse GP x)) (x
Hyperparameter g(()
Euler-Maruyama
IPF (NN)
Hyperparameter g(()
Euler-Maruyama
SB-FBSDE
Hyperparameter g(t)
Maler-Maruyama
E EXPERIMENTAL DETAILS AND ADDITIONAL. RESULTS
E. IMPLEMENTATION
We compared our methods against standard neural SDE, TrajectoryNet (Tong et al., (020)) OT- Flow (Onken et al. (021), IPF with GP (Vargas et al. 2021) and NN (De Bortoli et al. 2021), SB-FBSDE (Chen et al. 2021a). TrajectoryNet and Ol-Flow are examples of existing ODE-based methods, neural SDE is an example of learning SDE using only data without prior information by the Lagrangian, IPF and SB BB-BBSDE are methods to find SDE solutions to the classical SB problem. While IPF and SB-FBSDE solve the SB problem, which is special case of the LSB problem, they are different algorithm from NLSB. The parameterization and numerical solvers for ODEs and SDEs are summarized in Table 3 For a fair comparison, we used the same potential model Pg described in Section 4.2 for OT-Flow, neural SDE, NLSB, IPF (NN), and SB-FBSDE. We the number of ResNet layers M = 2, the step size h = 1.0, the rank of matrix rank (ank(A) (A) = 10, and the dimension of the hidden vector z 2. training all models, we used Adam optimizer to optimize all learnable parameters with a learning rate of 0.001 and the decay rate of B1 0.9, 82 = 0.999. We searched all weight coefficients of regularization terms Ae, A in (0.0, 0.5) and selected those with the largest possible coefficients among those with sufficiently small EMD--.2 values on the validation data.
In our implementation, we modified code in the Nea OT-FLow IPF and SB-FBSDE repositories, which were released under the MIT license. Our experimental environment consists of an Intel Xeon Plantinum 8360Y (36-core) CPU and a single NVIDIA A A100 GPU. Our code is available at https: / qithub...buu..ggg toon/ttae-kosii www.ksiiikaaaaaaa zuka/nnsb TrajectoryNet and OT-Flow were trained using the torchdiffeq library and neural SDE and NLSB with the torchsde library The settings common to all experiments for each method are described below.
NEURAL SDE AND NLSB
We trained the drift and diffusion models of the standard neural SDE using only the Sinkhorn divergence W.. We used the Euler-Maruyama method with the constant step size 0.01 as an SDE solver. Backpropagation was performed without using the adjoint method.
TRAJECTORYNET
The velocity model of TrajectoryNet includes three concatsquash layers with hyperbolic tangent activations. A concatsquash layer es was defined in the released code of FFJORD (Grathwohl et al.)
2018 by:
cs(x, t) (W , + by (W,)) b.) + (Wet+ + bbt),
where a is the sigmoid function, and W.. W,, W.. ba b, b, are learnable parameters.
25

---

Published as a conference paper at ICLR 2023
The base models of TrajectoryNet was trained using only the negative log-likelihood loss in standard CNF scheme. +OT represents model trained with the OT-based regularization R. defined by Eq. (1p We set the interval-dependent coefficients Åe for the OT-based regularization as well as Eq. 11 For the ODE solver, the dopri5 solver with both absolute and relative tolerances set to 10 was used.
OT-FFLOW
The base models of OT-Flow also trained with CNF scheme. We used both R, and Rs with the interval-dependen coefficients Åe, X in OT-Flow + OT. We used the Euler method as the ODE solver with constant step size of 0.01 and both absolute and relative tolerances set to 10-3
IPF
The drift model of IPF (GP) was changed to sparse GP with the exponential kernel from vanilla GP VVargas al. 202T to save computation cost. We selected 100 inducing points using the K-means algorithm in IPF (GP). The SB problem was solved by using IPF algorithm with 15 iterations. We used the Muler-Maruyama method as the SDE solver. The diffusion coefficients IPF were tuned as hyperparameters.
SB-FBSDE
employed alternating training and solved the classical SB problem specifying Brownian motion the prior stochastic process and did not use collectors.
E.2 SYNTHETIC POPULATION DYNAMICS: PROCESS
We validated NLSB artificial synthetic data generated from one-dimensional SDEs, the time- dependent Cemstein-Uhlenben process used in Kidger et al.) 2021). In this experiment, the predicted trajectory and uncertainty can be compared with the ground truth and easily evaluated by visualization. The purpose this experiment is to confirm that the NLSB works well on simple linear SDE.
E.2.1 DETAILS OF EXPERIMENTAL SETUP
We trained all models with batch size of 512 for each time point. The tuned weight coefficients shown in Table
POTENTIAL MODEL IN OT-FLOW, NEURAL SDE, NLSB, AND IPF (NN)
For the potential model (x, t), set the number of ResNet layers M 2, the step size 1.0, the rank of matrix rank(A) (A) = 10, and the dimension of the hidden vector z to 2.
NEURAL SDE NLSB
used two-layer FCNN hidden dimension 16 for the diffusion function (x, t). The activations functions were LipSwish. In NLSB, we used the Lagrangian for the potential-free system L(I, x, u) ||
NRAJECTORY
We used the concatsquash layers of hidden dimension 16.
IPF
We set the diffusion coefficients as follows.
0.2
E [0.0, 1.0)
0.6
E
2.0.
g(t)
1.0
t
3.0
1.4
E [3.0,4.0]
Other experimental settings have not been changed from (Vargas al. 2021
26

---

Published as a conference paper at ICLR 2023
Table Weight coefficients for regularization terms in experiments on synthetic data
[to [1]
[41, [2]
[12, ta]
[ta, t4]
de, An for NLSB
0.3, 0.2
0.1, 0.01
0.01, 0.0001
0.01, 0.0001
te for Neat + OT
0.1
0.1
0.001
0.001
des dn for OT-Flow + OT
0.1, 0.01
0.1,0.01
0.001, 0.001
0.001, 0.001
Xto (X00
El XXX1
Xto ((0)
E00o (2001)
-2
ground-truth SDE
(b) Neural SDE
Xmn x0o)
EL XIX
Xto (X0]
| XXO))
-2
(c) NLSB (Ours)
(d) IPF (GP)
Xu1 X1o)
ELXUD XXOI
Att) (xo)
Att1
2
2
-2
(c) IPF (NN)
(f) TrajectoryNet OT
(g) OT-Floww OT
Figure 6: ID OU process data and predictions.
E.2.2 RESULTS
Figure shows the visualization of both the original trajectory and the averaged trajectory by the SDE--based methods (neural SDE, NLSB, and IPF) and the only original trajectory by the ODE-based methods (TrajectoryNet and OT-Flow). All trajectories initial samples at the time t = 0. The five colored point clouds in the background are the ground truund-truth data given each time point. The pink area and the light blue line the one-sigma empirical confidence intervals and their boundaries for each trajectory, respectively.
generated by all-step prediction from the
Figure indicates that NLSB and IPF outperform neural SDE and is comparable to other ODE- based methods in estimating populations with small variance. In contrast, the SDE-based methods outperform ODE-based methods when estimating populations with large variance. That indicates
27

---

Published as a conference paper at ICLR 2023
x(t) X
X(e) X(O)
X X(O)
8
4
2
2
(a) ground-truth forward SDE
(b) NLSB (Ours)
(c) Trajectory OT
Figure 7: ID OU process data and predictions backwards through time.
NLSB (Ours)
TrajectoryNet OT
TrajectoryNet
CT-FLow OT
OT-FLow
0.14
0.12
EMD 1.2
0.08
0.06
0.04
0.02
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
time
Figure 8: Numerical evaluation for the reverse- reverse-time simulation on synthetic OU process data. All MDD values were computed between the ground-trute and the estimated samples within generated trajectories all-step behind from initial samples x(44).
that NL.B and IPF can estimate population-level dynamies when the population variance is large or small. Furthermore, NLSB and IPF have smaller CDD value than neural SDE. Figure 6b shows that the average behavior of samples [X[ (t) (0)] estimated by neural SDE is different from that of the ground-truth SDE (see Fig. 6a). especially in the interval [0, 1]. In contrast, the predictions by NLSB and IPF Figs. 6c to are much closer the ground- -ruth. These results show that the prior knowledge of the potential-fre free system helps to estimate the sample-level dynamics.
Note that the LSB problem solved by NLSB with the Lagrangian solved by IPF are almost mathematically equivalent (see Appendices A.3 and B. The result that NL.BB shows comparable performance to IPF, even though NI.SB trained differently from IPF, indicates that NLSB is unified framework and can deal with SB problems special case.
] |u|2 and the SB problem
The quantitative evaluation results the reverse-time ODE/SDE using MDD are shown in Fig. S and the visualization of trajectories is shown in Fig. All trajectories from NLSB and TrajectoryNet were generated by all-step prediction from the initial samples at the time = 4 to t = 0. Experimental results show that the proposed method described in Appendix..IIsucceeeee. level dynamics the reverse-time SDE. The development of appropriate evaluation methods for sample-level dynamics of the reverse-time SDE is included in future work.
recovers the population-
28

---

Published as a conference paper at ICLR 2023
Table 5: Weight coefficients for regularization terms in experiments on scRNA-seq data
[o: #1]
[+1, t2]
[t2, ta]
[ta, t4]
Ae A for NLSB (E)
0.1, 0.01
0.01, 0.01
0.001, 0.0001
0.01, 0.001
Ae. An for NLSB (V)
0.01, 0.001
0.01, 0.001
0.01,
0.01, 0.001
Aes An for NLSB (D)
0.01, 0.001
0.01,
0.01,
0.01, 0.001
Xe An for NLSB (E+D+V)
0.01, 0.001
0.001,
0.001, 0.001
0.001, 0.001
Xe for TrajectoryNet OT
0.01
0.01
0.1
0.1
Xe Å for OT-Flow + OT
0.01, 0.01
0.01,0.01
0.001, 0.01
0.001,0.01
E.3 SINGLE-CELL POPULATION DYNAMICS
We evaluated our method on the time-evolution of single-cell populations obtained from a developing human embryo system. In this experiment, we presented a new quantitative evaluation metric in practical setting and validated the effectiveness of the NLSB on real data of single-cell population.
E.3.1 DETAILS OF DATASET
We evaluated our method on embryoid body scRNA-seq data (Moon al. 2019), which is also used in (Tong et al. 22200] Vargas et al. 2021; Bunne et al. [021 2022). 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. During this period, cells were collected at five different snapshots (4p: day 3, 41 day 9, 12: day to 15, 13: day to 27). The collected cells were then measured by scRNAseq. filtered at the quality control stage, and mapped to a low-dimensional feature space using a principal component analysis (PCA). For details, Appendix E.2 in (Tong et al. 2020). We reused the rre-processed data available in the released repository of TrajectoryNet and split the data into 200 samples (~ 8.5%) of validation data, 350 samples (~~ 15%) of test data, and the rest train data for each time point. The scRNA-seq data licensed under Creative Commons Attribution 4.0 International license.
21, 14: day 24
E.3.2 DETAILS OF EXPERIMENTAL. SETUP
We trained all models with batch size of 1000 for each time point and used the early stopping method, which monitors the EMD-L2 value on the validation data. The tuned weight coefficients are shown in Table 5 We used the same potential model o in OT-Flow, neural SDE, NLSB and IPF (NN) with the same hyperparameters described Appendix E.2.1 In the following, we describe the different settings from the experiment in Section 5.1
NLSB
We used the Lagrangian for the cellular system and compared several combinations the regulariza- tion terms. In Tables and "E" is the energy term, "D" is the density term, and "V" is the velocity term. The density term U (x, t) is the log-likelihood function of the data estimated by GMM. For the calculation of the density regularization term, the time-dependent density function U(x,)) t) was defined by
M. (x,t) clog O.), θ, (u 2m Jm=1
where e is hyperparameter to change the scale, and (1) (1) are mean and variance parameters of the mixed Gaussian distribution, respectively. When t E [++]) the parameters O were estimated with the data at la and f++1 by GMM and the number of mixture components M, was determined by the value of Bayesian information criterion (BIC). The hyperparameter e was searched among {0.1, 1.0, 10.0} and we
10.0 for NLSB (D) and 0.1 for NLSB (E+D+V), For the calculation of the velocity regularization term, we used the same reference velocity as those used in TrajectoryNet (Tong al. 2020).
29

---

Published as a conference paper at ICLR 2023
15
15
15
15
10
10
10
10
50
50
0
0
t: Day
Day
-5
-5
-5
Day 28
-10
Day 24
-10
-10
MLSB D
-15
55
45
50
55
50
PC1
PC1
PC1
PC1
15
15
10
10
10
5
5-
g
0
0
-5
Day
-5
5
Day
110
-10
-10
-10
-15
-15
55
50
50
PC1
PC1
PC1
PC1
Figure 9: Visualization of the time evolution for single sample on scRNA-seq data. The upper and lower images are predictions from the same initial sample to. with the upper predicted by NLSB and the lower row predicted by TrajectoryNet + OT. The color gradients depict the magnitude of the probability density. The probability density function is estimated by GMM with five mixture components.
Table 6: The MDD value (EMD-L1) for population-level dynamics five-dimensional (5D) PCA space at time observation for ScRNA-seq data.
MDD (EMD-LI) &
t1
t2
f3
t4
NLSB (E)
1.13 + 0.025
1.36 + 0.035
1.30 + 0.018
NLSB (D)
1.08 + 0.021
1.40 0.043
1.38 0.030
1.29 + 0.024
NLSB (V)
1.13 + 0.030
1.39 = 0.043
1.34 + 0.029
1.35 + 0.025
NLSB (E+D+V)
1.09 0.023
1.34 = 0.037
1.32 + 0.024
1.30
0.025
Neural SDE
1.11 + 0.028
1.41 + 0.041
1.38 + 0.033
1.34 0.025
OT-Flow
1.31
1.73
1.68
1.69
OT-Flow + OT
1.33
1.65
1.69
1.56
TrajectoryNet Net
1.15
1.60
1.42
1.58
TrajectoryNer Net OT IPF (GP)
1.20
1.60
1.41
1.72
1.14 0.024
1.59 - 0.052
1.49 0.037
1.57 0.042
IPF (NN)
1.16 0.027
1.42 0.037
1.37 + 0.030
1.37 + 0.027
SB-FBSDE
0.016 0.89 0.016
1.32 0.025
1.63 + 0.030
1.57 0.015
TRA
We used the velocity model with 64 hidden dimensions.
IPF
We set the diffusion coefficients to g(() 0.5 (1 E [0.0, 1.0)). 1.0 E [1.0, 4.0)). Other experi- mental settings have been changed from (Vargas al. 2021).
SB-FBSDE
We set the diffusion coefficients g()) 0.5 (1 E (0.0,
0.1 (E 4.0).
30

---

Published as a conference paper at ICLR 2023
20
20
20
20
15
15
15
15
10
10
10
5
-10
Day tcatel
tostal
10
toate) teate)
Day tcatal
Day tcatal
(date)
teatal
30
55
30
50
55
-10
45
50
55
-150
45
PC1
15
10
10
5
-10
-10
-10
-10
-10
10
20
-10
10
-10
10
20
-10
10
20
PC3
(a) Ground --utt data
E.3.3 RESULTS
The time evolution of the distribution for single sample is visualized as heat map in Fig. 9 The sample population in Fig. 10 and the trajectories are visualized in Fig. []] All figures are visualizations the space of the first and second principal components. The x-axis represents the first principal component and the y-axis the second principal component. The experimental results using MDD using EMD with L' cost is shown Table
Figure 9 shows that the SDE-based methods, including NLSB, can handle the uncertainty of the trajectories contrast to ODE-based methods. Figure TO shows that NL.B outperforms ODE-based methods in predicting the transitions with high diffusion of samples from i to 12 and from 13 to f4- indicating that the explicit modeling of diffusion effective. Figure shows that the drift estimated by NLSB (E) is linear, the trajectories by NLSB (D) pass the data manifold, and the trajectories by NLSB (E+D+V) appear to reflect all other regularization effects.
A GIF animation of the NLSB simulation in PCA space is also included in the supplemental materials.
E.4 SYNTHETIC POPULATION DYNAMICS: TRAJECTORIES REFLECTING THE POTENTIAL FUNCTION
This section presents the experimental results of applying the NLSB to numerical simulations of synthetic population dynamics. Through this experiment, we show that the NLSB can model wide range of phenomena by designing the Lagrangian based on prior knowledge. In particular, we emphasize the importance of designing potential functions and the flexibility of penalty design for drift functions.
E.4.1 DATASET
We used two-dimensional uniform distributions Up and Ui for the endpoints time 0 and 1.
(Xo, Yo) ~ Ub -1.25 < X, -1, -1 1, (X1, Y1) ~ 24 : 5 x1 < 1.25, -1 -1 11551
We generated 2048 and 512 samples from two endpoint distributions as training and validation data, respectively.
31

---

Published as a conference paper at ICLR 2023
20
20
20
20
15
15
15
10
-5
-5
100
Day 1o lcatei
10
by Day to tostal
-10
t: Day 12 to toabo
-10
Day 18 21 tcabe)
Day to (ored!
b: Day 12 to coree)
t: Day loned)
Day 24 27 (predi
-150 30
45
50
55
300
50
55
-15.0
45
50
55
-15
45
50
55
PC1
10
10
10
-5
-5
-5
-10
-10
-10
-10
PC3
(b) NLSB (E+D+V)
15
10
10
5
5
5
0
0
-5
-5
-5
-10
Day
-10
Day instal
-10
ty (eate)
-10
Day Day (pred)
Day (pred)
t: Day
t: toredl
10
55
100
50
55
-15
50
55
-15
45
50
PC1
15
15
10
10
0
-5
10
-10
10
20
PC3
(c) TrajectoryNett OT
Figure 10: The ground trrut data and step ahead sample prediction on scRNA-seg data.
32

---

Published as a conference paper at ICLR 2023
E.4.2 DETAILS OF EXPERIMENTAL SETUP
We trained NLSB with batch size of 512 and used the early stopping method, which monitors the validation loss value. We adopt the Lagrangian for the random dynamical system. All weight coefficients of regularization terms were searched {0.01, 0.001}.
We first conducted experiments applying the NLSB with the Lagrangian of the form L x, u) = |||u|12 U (x) defined by four different potential functions shown below. We implemented the box and slit-shaped obstacle potential functions as differentiable by using the sigmoid functions.
BOX-SHAPED OBSTACLE
-100
-0.5 <
U(r 9)
otherwise
SLIT-SHAPED OBSTACLE
100 5 50 S -0.25 V V.25 S v)
U(z,9)
otherwise
HILL POTENTIAL
U(r,y) -2.((x² + )²)
WELL POTENTIAL
Ulz y) -0 100xxp-((x= + (xxp(-(=++z2)
Next, we conducted experiment using the Lagrangian T(,, u, Ju Ru. The matrix R can be used to penalize the magnitude of the drift differently for each dimension. We set R to diag([10.0, 0.1]) and diag((0.1, 10.0)).
E.4.3 RESULTS
The visualization results are shown in Figs. 12 and 13 Figure 12 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. Figure 13a and show that the drift of the trajectories generated by the NLSB with the Lagrangian smaller penalties defined by Rt and vice versa.
Ru larger on the axis with
E.5
OPINION DYNAMICS
We demonstrated the application NLSB to optimal control on a party model of opinion dynam- ics (Schweighofer et al. 2020 Gaitonde et al. 2021 Liu et al. 2022).
E.5.1
DATASET
Opinion dynamics (Schweighofer et al. 2020 Gaitonde et al.) 2021) is the time evolution of each agent's opinions interacting with each other. MFGs theory provides mathematical analytical framework for the opinion dynamics of large agent populations, which are very difficult to handle computationally. The dynamics is modeled using SDEs defined the opinion representation space of each agent embedded in Euclidean space. In recent years, the phenomenon of strong polarization (Gaitonde al. 2021), which agents are divided into groups with opposite opinions, has attracted particular attention. We the drift pplariie, which causes polarization defined in the party model (Gaitonde et al. 2021), as prior information of the target system.
polariiee := Epolarize/ /| ||fpolatize || = y |
fpolarize (x, p: E) = Eynp ()x
sign ((x, E)) = sign ((y E)) otherwise
(: =
33

---

Published as a conference paper at ICLR 2023
where p is the probability measure of the population, & is random information from some distribution independent of p. a(x. g) is the agreement function, which represents whether the two opinions x and y agree on the information
We used two-dimensional Gaussian distributions Ny and N₁ for the endpoints at time and t 1.
0.5
0.0
(Xo, Yo) ~ No =N
0,
0.0
0.25
3.0
0.0
(X1, Vi) ~ M₁ N
0,
0.0
3.0
We generated 2048 and 512 samples from endpoint distributions training and validation data, respectively.
E.5.2 DETAILS OF EXPERIMENTAL SETUP
The Lagrangian for the opinion dynamics is defined by
[tt,x, u)
2 E polarize (x, t) u(x, t)ll Uxx,t) t)-
The potential function U (x, t) represents the averaged interaction that each agent receives from the population. The entropy function U(x. f; c) clog t) with constant coefficient e is candidate for a useful potential function and helps control changes in population diversity. The optimal drift function is given by fa x (x, t) F poolarize (x, t). By setting the ideal opinion distribution as the terminal condition, NLSB can be used as a method to find the optimal drift converging to the ideal opinion distribution.
We trained NL.BB with batch size of 512 and used the early stopping method, which monitors the validation loss value. weight coefficients of regularization terms were searched (0.01, 0.001). We conducted experiments using the NLSB with the Lagrangian for the opinion dynamics.
E.5.3 RESULTS
Visualization of polarized opinion dynamics driven by the drift f polarize and the time variation of directional similarity are shown in Figs. 14a and 15a The directional similarity (Schweighofer et al. 2020) is the distribution of cosine angles between paired opinions, with the red and blue color gradients representing the degree of disagreement and agreement, respectively. Figures [4b and 15b show the results after applying NL.BB. These results show that NLSB can learn a drift function that prevents the polarization caused by polarize, which also converges to the ideal terminal distribution Ni.
34

---

Published as a conference paper at ICLR 2023
X(t) ((0)
EIXIO [(O0
20
15
PC2
for Day to
fi Day 6 to
Est Day 18 to
Day
10
55
(a) NLSB (E+D+V)
Xte (X10)
| ((0))
Xtn (X00)
EDit (X(0))
0
to:
N:
-10
-10
30
45
50
55
40
45
50
55
40
45
50
55
40
45
50
55
PC1
PC1
PC1
PC1
(b) Neural SDE
(c) NLSB (E)
Xto (X0)
EDito X(D)
Xmn (00)
20
15
10
10
5
0-
0
4: Day ta 6. Day to
Re Day to c: Day to
b: Day
by
10
30
40
(d) NLSB (D)
(c) NLSB (V)
Stp
20
20
10
10
10
10
-5
-5
-10
-10
-10
-10
-100
-130
-130
-130
(f) TrajectoryNet
TrajectoryNet OT
(h) OT-Flow
(1) OT-Flow + OT
Figure 11: scRNA-seq data and predictions. The x- and y-axes denote the first and second principal components, respectively. The five colored point clouds in the background are the ground-truth data given at each time point. All five trajectories are generated by all-step prediction from the initial samples at to-
35

---

Published as a conference paper at ICLR 2023
X041
208
60021
1.0
0.5
0.5
0.0
0.0
-0.5
-0.5
-1.0
(a) Neural SDE
(b) Potential-free (U = 0)
X03
-15
-15
-15
30
30
30
45
0.0
0.0
860
-00
60
-60
00.5
-75
-0.5
75
-90
-90
-1.0
-90
-90
-105
-005
-1.5
-1.3
(c) Box-shaped obstacle
(d) Slit-shaped obstacle
X00
E(X0n)
X00
1.0
10
1.5
1.5
-30
45
-.5
4.5
0.0
600
-60
-0.5
-90
-7
-75
9.0
-0.5
7.5
-75
10.5
10.5
-1.5
-1
12.0
12.0
-1.5
10.55
10.5
(e) Well potential
(I) Hill potential
Figure 12: Visualization of trajectories reflecting the potential function. The color gradients depict the magnitude of the potential function.
2(k)
X163
0.
0.0
0.0
-0.5
-0.5
1.0
-10
(a) R diag([10.0.0.) 0.1))
(b) R diag([0.1,
Figure 13: Visualization of trajectories by NLSB using the Lagrangian L = fuu Ru. The blue and point clouds are the source and target distributions, respectively.
36

---

Published as a conference paper at ICLR 2023
100
50
-50
300 10
30
(a) Polarized dynamics
160
Ee0
1-027
bwger
50
-1
-10
-10
40
(b) Depolarized dynamics by NLSB
3.0
t=027
25
2.0
1.52
1.0
Drectional Siiiaity t-]
t- 0.5T
t-0.87
3.0
2.5
0.5
0g
0.2
0.4
0.6
1.0
0.0
0.6
0.8
3.0
0.2
0.4
0.6
08
1.0
arty
Similarity
Siilartty
(a) Directional similarity of polarized dynamics
t-0
t=027
t- 0.4T
3.0
2.55
152
0.5
0.0
Directional Smi anly
Directioral Simibrity
Directiooal Smilaity
3.0
203
15
10
0.
00
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
10
Directional Smiiarly
Directional Similarity
Direciional Simiarity
(b) Directional similarity of the depolarized dynamics by NLSB
37


Reviews:
[{"review_id": "tcVsd_6ljso", "reviewer": "Reviewer_8YwN", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes the Lagrangian Schrodinger bridge (LSB) problem, a framework to recover the population dynamics from temporal data. LSB is a stochastic optimal control problem, which enables the authors to leverage the structure and propose the training objective (7). This method is evaluated on synthetic and single-cell population datasets.", "summary_of_the_review": "This paper proposed an interesting framework for population dynamics. However, in its present state, it seems like gluing together several pieces (SB, stochastic optimal control, regularization from previous work, etc.) without any core contribution. I'm therefore inclined to rejection.\n\nI suggest the following improvements:\n\n1. Revise the presentation to highlight the contribution, focusing especially on the benefits of the neural SDE component in (6). \n\n2. Expand Section 3.3 to include more examples, even designing new ones, to reflect the generality of the proposed LSB. Otherwise, if the authors merely included one previous example (the cellular system), then the applicability of LSB is not really convincing to me.\n\n3. Since the theoretical contribution is fairly limited, the authors should perform more extensive experiments. In particular, for the single-cell population dynamics, the authors might consider including the dataset of Schiebinger et al. 2019. In addition, following the point above, experiments on more applications, instead of just synthetic and single-cell population dynamics, are recommended.", "strength_and_weaknesses": "Strength:\nThe proposed framework is quite flexible. Moreover, it is simple, and it seems to work in experiments. The application of sing-cell population dynamics is interesting.\n\nWeaknesses:\nThe major issue of the paper for me is that the contributions of the submission are not clear: \n- The objective (7) is not really new. The idea of relaxing the endpoint constraint by a softer one ($\\mathcal{W}_2$ in (7)) was exploited by the prior work of Lavenant et al. 2021 in a similar context, although without the neural SDE; the authors seem have missed this reference. The idea of incorporating the Hamiltonian term (9) was also present in existing literature (Yang & Karniadakis 2020). As for the Lagrangian term $\\mathcal{R}_e$, the authors simply reuse the regularization of Tong et al. 2020 without designing new ones.\n\n- There is also little contribution from the algorithmic side, since the authors essentially combined various known methods to treat each term separately in (7). \n\n- I'm not sure why the authors mentioned various SB related training but did not compare to them in the experiments. For instance, what is IPML (NN) in Table 1? Is it from Vargas et al. 2021 or De Bortoli et al. 2021? Is there any reason of including these methods, instead of, say, Chen et al. 2021a?\n\n- The $\\mathcal{W}_2$ in the proposed objective (7) suddenly turns into the entropy-regularized $\\mathcal{W}_\\epsilon$ in (11). Why not propose $\\mathcal{W}_\\epsilon$ to begin with?\n\n-----\nReference\nLavenant et al. 2021, Towards a mathematical theory of trajectory inference.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "sDIKq76Ikc", "reviewer": "Reviewer_6WTu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors formulated the Lagrangian Schrodinger bridge problem and proposed to solve it approximately by the advection-diffusion process with regularized neural SDE. The expensive trace computation operation was also alleviated by adopting a model architecture motivated by OT-Flow. A few experiments were conducted on population dynamics and showed the efficiency in modeling stochastic behavior.", "summary_of_the_review": "An interesting extension of SB with speed-ups on divergence estimation.", "strength_and_weaknesses": "**Pros:** \n\n1. An interesting extension of Schr\\\"{o}dinger bridge by including additional regularizers, which ends up with an interesting Lagrangian Schr\\\"{o}dinger bridge that generalized the quadratic transport cost function.\n\n2. avoids the expensive computation of the divergence suffered by the Hutchinson estimator; the adoption of architecture in OT-Flow is interesting.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "pgOesiX0GBz", "reviewer": "Reviewer_hMLS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper aims to learn population dynamics using neural stochastic differential equations.  The paper formulates the Lagrangian Schrodinger bridge (LSB) of optimal transport. Experiments are conducted on synthetic data sets and single-cell embryoid body scRNA-seq data.\n", "summary_of_the_review": "The paper presents an interesting approach to learn neural SDEs for the problems of population dynamic modelling. However, the paper is not well-motivated and the technical contributions seem incremental.", "strength_and_weaknesses": "Strength:\n- To my knowledge, the approach using the LSB is interesting and novel.\n \n\nWeaknesses:\n- The paper claims that one of the main contributions is the model architecture of potential functions. It seems that the paper adopts such an architecture from OT-Flow.\n- The paper does not highlight or motivate the reason for using Lagrangian Schrodinger bridge.\n- The empirical results, unfortunately, only give marginal improvements compared to models (i.e. IPML).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BXKmrZ0HZYA", "reviewer": "Reviewer_WFgo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work tackles the problem of learning dynamics from population observations over time. Existing works have used Continuous normalizing flows (CNFs) or schrodinger bridges (SBs) to tackle this problem. This work (NLSB) combines adhoc regularizations in Tong et al. 2020 and more recent SDE formulations to learn a principled and regularized SDE to model cell populations and incorporate additional regularizations. The method is tested on a synthetic OU SDE and a single-cell RNA-seq dataset measuring embryoid body development over time. ", "summary_of_the_review": "This paper introduces a new architecture and training procedure for neural Lagrangian Schrodinger bridges. This provides a more principled approach to cellular systems than Tong et al., and includes additional prior information not present in the general Schrodinger bridge formulations. This is an incremental but solid contribution to the literature on modelling population dynamics with additional constraints. ", "strength_and_weaknesses": "Strengths:\n\n- Clearly explained motivation formalizing previously adhoc regularizations in Tong et al. 2020\n- Adds ideas from cellular modeling to SDE and Schrodinger bridge models for improved modeling of a cellular system.\n- Good tradeoff of time vs. accuracy in figure 5 across dimensionalities.\n- Interesting architecture adaptation for fast gradient computation.\n\nWeaknesses:\n\n- Limited evaluation (2 settings) with relatively minor improvements in performance.\n- The network is not as easy to invert as other models which use ODE dynamics (Tong et al. 2020) or train both a forward and backwards model. I see that a reverse time simulation was performed in E.4, but it is unclear to me how well this works in practice (quantitatively). I have had numerical difficulties training this backwards network in the past.\n\nComments / questions:\n\nThere is something strange in these comparisons… what are the differences with NeuralSDE?\n\nFigure 10, are the captions correct? It looks to me that (e) and (f) are switched.\n\n“The weight coefficients $\\lambda_e$ and $\\lambda_h$ are tuned for each interval $[t_{k-1}, t_k]$ respectively”. Can the authors elaborate here? This seems concerning as there may be potential for data leakage.\n\nIt might be interesting to visualize the potential landscape (as is done in Bunne et al. 2022). Biologists are quite interested in what is termed “Waddington’s landscape” for developing systems. \n\nWould it be possible to include an additional evaluation metric (besides the 2-Wasserstein distance)? The 2-Wasserstein distance mostly cares about outliers, that the distribution is “mostly right” which might miss some of the other differences between the predicted and true distributions. Common choices I believe are 1-Wasserstein distance (Tong et al. 2020, Vargas et al. 2021) and Maximum mean discrepancy (Huguet et al. 2022). Since you directly optimize for the 2-Wasserstein distance this may be an overoptimistic metric. If this is not possible it would at least be good to note the limitations of the current evaluation.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%