NEURAL LAGRANGIAN SCHR ÖDINGER BRIDGE: DIF-FUSION MODELING FOR POPULATION DYNAMICS

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.

1. 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) , air in meteorology (Fisher et al., 2009) , and healthcare statistics (Manton et al., 2008) in medicine. However, tracking individuals over a long period is often difficult due to experimental costs. Furthermore, it can sometimes be impossible to 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 crosssectional populations without alignment across time steps at a few distinct time points. Under these constraints on data measurements, our goal is to better understand the time evolution of samples in the populations. Existing methods attempt to estimate population-level dynamics following the Wasserstein gradient flow using a recurrent neural network (RNN) (Hashimoto et al., 2016) or the Jordan-Kinderlehrer-Otto (JKO) flow (Bunne et al., 2021) . Recent studies have attempted to interpolate the trajectories of individual samples between cross-sectional populations at multiple time points by using optimal transport (OT) (Schiebinger et al., 2019; Yang & Uhler, 2018 ), or CNF (Tong et al., 2020) . Using a CNF generates continuous-time non-linear sample trajectories from multiple time points. In addition, Tong et al. ( 2020) proposed a regularization for CNF that encourages a 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 a 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 advection-diffusion processes by using SDEs to describe the time evolution of the sample. Furthermore, on the basis of the principle of least action, we estimate the sample trajectories that minimize action, defined by the time integral of the Lagrangian determined from the prior knowledge. We formulate this problem as 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, 2019a; b) by minimizing the Wasserstein loss between the ground-truth and the 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 et 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., 2021) will speed up the computation of the potential function's gradient and the regularization term. As a result, we capture the population-level dynamics as well as or better than conventional methods, and can more accurately predict the sample trajectories. In short, our contributions are summarized as 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 advectiondiffusion 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(Onken et al., 2021) to speed up the computation of the regularization term to minimize HJB-PDE loss.

2. BACKGROUND

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

2.1. FLOWS REGULARIZED BY OPTIMAL TRANSPORT

CNFs (Chen et al., 2018) are a method for learning the continuous transformation between two distributions p and q by modeling the ordinary differential equation (ODE): dx(t) dt = f θ (x, t), subject to x(t 0 ) ∼ p, x(t 1 ) ∼ q, where f θ is the velocity model with learnable parameters θ. 2). These OT-based regularizations have resulted in faster CNFs (Finlay et al., 2020; Onken et al., 2021) and improved the modeling of cellular dynamics (Tong et al., 2020) .

Several regularizations of

Re = t1 t0 R d 1 2 ∥f θ (x, t)∥ 2 dρ t (x) dt, Rh = t1 t0 R d ∂ t Φ θ (x, t) - 1 2 ||∇ x Φ θ (x, t)|| 2 dρ t (x) dt,



Figure 1: Example of trajectories by NLSB.

CNFs 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 Brenier-Benamou formulation (Benamou & Brenier, 2000) with Kullback-Leibler (KL) divergence. RNODE (Finlay et al., 2020), OT-Flow (Onken et al., 2021), and TrajectoryNet (Tong et al., 2020) introduced a regularization Re in Eq. (1). Potential Flow (Yang & Karniadakis, 2020) and OT-Flow modeled the potential function Φ that satisfies f = -∇Φ instead of modeling the velocity function f . They also proposed an additional OT-based regularization Rh derived from the Hamilton-Jacobi-Bellman (HJB) equation (Evans, 1983) satisfied by the potential function shown in Eq. (

