PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Lagrangian Schr\"odinger Bridge: Diffusion Modeling for Population Dynamics

Takeshi Koshizuka, Issei Sato

arXiv:2204.04853v5cs.LGq-bio.PE

TL;DR

Coarse cross-sectional observations leave individual population trajectories untracked, and CNF interpolation is deterministic despite stochastic diffusion in biological samples. This paper formulates LSB and approximates it with a Lagrangian-regularized neural SDE. NLSB predicts population dynamics better and trains faster than ODE methods as data become higher-dimensional, while modeling stochastic sample behavior.

  • Problem

    Coarse cross-sectional observations leave individual trajectories untracked, while CNF interpolation is deterministic despite stochastic diffusion in populations.

    Source: §1 · pp. 1–2

  • Method

    The paper formulates LSB as stochastic least-action dynamics and approximates it with a regularized neural SDE whose Lagrangian incorporates prior knowledge.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    NLSB predicts population dynamics better and trains faster than existing ODE-based methods as data become higher-dimensional, and SDE-based methods outperform ODE-based methods on highly diffuse transitions.

    Source: §5.2 · p. 8

  • Takeaways & Limitations

    The Lagrangian lets the method flexibly incorporate prior knowledge about the target system into stochastic sample trajectories.

    Source: §6 · p. 9

    The formulation assumes stochastic individual motion yields population diffusion and encourages individual samples to follow least-action trajectories.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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\"odinger 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

Population dynamics are difficult to track from sparse, fixed-point observations, leaving individual trajectories unobserved. NLSB models stochastic sample motion and uses a Lagrangian to encourage least-action trajectories.

1 INTRODUCTION

  • Coarse fixed-point observations often preclude long-term tracking; scRNA-seq destroys measured cells, leaving unaligned cross-sectional samples at a few time points.Source: §1 · p. 1
  • OT and CNF methods interpolate sample trajectories, but CNFs are deterministic despite the stochastic diffusion of population samples.Source: §1 · pp. 1–2
  • NLSB approximates the Lagrangian Schrödinger bridge with a regularized neural SDE, using prior knowledge to model stochastic, least-action trajectories.Its potential-function architecture speeds computation of the gradient and regularization term.Source: §1 · p. 2

2 BACKGROUND

The background connects optimal-transport-regularized flows, neural SDEs, and stochastic optimal transport as foundations for modeling population dynamics.

2.1 FLOWS REGULARIZED BY OPTIMAL TRANSPORT

  • OT-based CNF regularizations encourage straighter trajectories; potential functions and HJB regularization have been used in OT-Flow and Potential Flow.These regularizations have resulted in faster CNFs and improved modeling of cellular dynamics.Source: §2.1 · p. 2

2.2 NEURAL SDE AND APPLICATIONS TO DIFFUSION MODELING

  • Neural SDEs model drift and diffusion with neural networks, while related approaches learn Schrödinger-bridge solutions using iterative proportional fitting and drift regression.Source: §2.2 · p. 3

2.3 STOCHASTIC OPTIMAL TRANSPORT

  • Stochastic optimal transport finds a stochastic process minimizing expected cost under fixed time-point marginals; its Lagrangian specifies transport cost over system space-time.A variational formulation also describes marginal flows satisfying the Fokker–Planck advection-diffusion equation; the Schrödinger bridge is a special case.Source: §2.3 · pp. 3–4

3 PROPOSED METHOD

The method formulates stochastic least-action dynamics as an LSB problem and approximates it with a regularized neural SDE, using Lagrangians to encode system-specific prior knowledge.

3.1 LAGRANGIAN SCHR ¨ODINGER BRIDGE

  • LSB estimates stochastic sample paths between endpoint distributions by minimizing the time-integrated Lagrangian action, with both drift and diffusion optimized.Source: §3.1 · p. 4

3.2 NEURAL LAGRANGIAN SCHR ¨ODINGER BRIDGE

  • NLSB approximates LSB with a regularized neural SDE, combining endpoint distribution discrepancy, action cost, and an HJB PDE regularization.The action cost acts as regularization, while the discrepancy is a data-fitting term; a potential function parameterizes the drift.Source: §3.2 · pp. 4–5

3.3 EXAMPLES OF LAGRANGIAN IN NEURAL LAGRANGIAN SCHR ¨ODINGER BRIDGE

  • Lagrangian choices encode different priors: minimal travel and noise, cellular-manifold and velocity information, or random mechanics with potential energy and mass.For cellular systems, density-based penalties constrain trajectories to the data manifold and velocity penalties encode reference cell velocities.Source: §3.3 · p. 5

4 IMPLEMENTATION OF NEURAL LAGRANGIAN SCHR ¨ODINGER BRIDGE

The method trains neural SDEs on unaligned cross-sectional snapshots by matching predicted and observed distributions, while estimating regularization terms from simulated paths. Its potential-function architecture speeds computation of the HJB regularization term.

4.1 TRAINING FOR NEURAL LAGRANGIAN SCHR ¨ODINGER BRIDGE

  • Training matches predicted and observed distributions with L2-Sinkhorn divergence at every observation time except the initial point.The data comprise unaligned cross-sectional samples at K time points.Source: §4.1 · pp. 5–6
  • For each adjacent time interval, simulated SDE paths estimate the action and HJB terms while K − 1 approximate LSB problems are optimized jointly.The regularization weights are tuned separately for each interval.Source: §4.1 · p. 6

4.2 MODEL ARCHITECTURE SELECTION FOR SPEEDUP

  • The OT-Flow potential architecture enables explicit drift gradients and fast exact diagonal-Hessian computation, reducing the cost of the HJB regularization term.With diagonal diffusion, the relevant Hessian calculation reduces to its diagonal components; a two-layer network with final tanh is used for diffusion.Source: §4.2 · p. 6

5 EXPERIMENTS

Experiments evaluate NLSB on synthetic OU dynamics and embryoid-body single-cell data. The results report population-level comparisons and sample-level trajectory behavior, including interpolation robustness when an intermediate snapshot is omitted.

5 EXPERIMENTS

  • Evaluations cover synthetic one-dimensional SDE data and embryoid-body scRNA-seq, comparing NLSB with neural SDE, OT-Flow, TrajectoryNet, IPF, and SB-FBSDE.Source: §5 · p. 6, §5.2 · p. 8

5.1 SYNTHETIC POPULATION DYNAMICS: ORNSTEIN–UHLENBECK PROCESS

  • MDD evaluates marginal population distributions, while CDD evaluates initial-sample-conditioned evolution; both use Wasserstein-2 discrepancies across 12 equally spaced times.Source: §5.1 · pp. 7–8
  • For small-variance OU populations, NLSB and IPF outperform neural SDE and are comparable to ODE methods; at large variance, SDE methods outperform ODEs, and NLSB/IPF have smaller CDD than neural SDE.Source: Figure 3 · p. 7, §5.1 · p. 8
  • NLSB and IPF predictions track ground-truth sample behavior more closely than neural SDE, whose average trajectory differs especially over [0, 1].The authors attribute the improvement to prior knowledge of the potential-free system.Source: Figure 4 · p. 7, §5.1 · p. 8

5.2 SINGLE-CELL POPULATION DYNAMICS

  • On higher-dimensional single-cell data, NLSB predicts population dynamics better and trains faster than ODE methods; SDE methods outperform ODEs on highly diffuse transitions.NLSB also has lower variation in MDD scores than neural SDE.Source: §5.2 · p. 8
  • NLSB estimates sample-level dynamics more robustly than neural SDE under intermediate-snapshot removal, with energy regularization most stable and effective, though exceptions occur.Source: §5.2 · p. 8

6 DISCUSSION AND CONCLUSION

The conclusion presents NLSB as a stochastic, Lagrangian-informed framework for population dynamics; the appendices place it within stochastic optimal transport, optimal transport, and Schrödinger bridge theory.

6 DISCUSSION AND CONCLUSION

  • NLSB models diffusion with SDEs and uses a designed Lagrangian to incorporate system prior knowledge into stochastic sample trajectories.The conclusion also identifies inability to model cell birth and death and potentially reduced expressiveness from diagonal diffusion as limitations.Source: §6 · pp. 8–9

A THEORY OF STOCHASTIC OPTIMAL TRANSPORT

  • Stochastic optimal transport minimizes expected cost over stochastic processes with fixed marginals and includes the Schrödinger bridge as a special case.The appendix relates SOT to stochastic optimal control with a terminal constraint.Source: A THEORY OF STOCHASTIC OPTIMAL TRANSPORT · p. 13

A.1 DYNAMIC FORMULATION OF OPTIMAL TRANSPORT PROBLEM

  • The dynamic OT formulation represents transport between distributions as time-varying velocity-driven advection, with particle paths interpretable as shortest paths.Its Eulerian formulation describes global particle behavior, while its Lagrangian formulation describes individual particles.Source: §A.1 · pp. 13–14

A.2 SCHR ¨ODINGER BRIDGE

  • The classical Schrödinger bridge has a dynamic formulation with Fokker–Planck evolution, and the dynamic OT solution arises as its zero-noise limit.The appendix also notes the connection between the classical SB problem and dynamic OT with Fisher-information regularization.Source: §A.2 · p. 15

A.3 STOCHASTIC OPTIMAL TRANSPORT WITH TWO ENDPOINT MARGINALS

  • The appendix defines SOT with a Lagrangian cost and describes a marginal-flow variant whose minimizing drift is a conditional expectation of path drift.The endpoint-marginal formulation is grounded in the principle of least action (Mikami, 2021).Source: §A.3 · pp. 16–17

B SCHR ¨ODINGER BRIDGE AND GENERATIVE MODELING

B.1 LAGRANGIAN SCHR ¨ODINGER BRIDGE PROBLEM

The LSB formulation constrains stochastic transport with a Fokker–Planck equation and is solved practically through a relaxed objective. NLSB connects this formulation to OT-Flow through analogous optimality conditions.

B.1 LAGRANGIAN SCHR ¨ODINGER BRIDGE PROBLEM

  • The LSB problem is constrained by the Fokker–Planck equation, and its optimality conditions also apply when the terminal constraint is strictly satisfied.The paper states that the LSB problem has a solution based on existence for the more general SOT problem for marginal flows.Source: §B.1 · pp. 18, 20
  • The relaxed LSB objective replaces the terminal marginal constraint with a soft constraint and adds an HJB-PDE loss to the objective during NLSB training.The weak-form HJB-PDE loss is evaluated along simulated SDE paths for computational efficiency.Source: §B.1 · pp. 18–19, §B.2 · p. 21

B.2 DERIVATION OF NLSB AND ITS CONNECTION TO OT-FLOW

  • NLSB relaxes the LSB problem using a distribution discrepancy and neural SDEs, while OT-Flow relaxes the Brenier–Benamou OT formulation using KL divergence and neural ODEs.Their potential-function optimality conditions are analogues; the classical SB and OT formulations connect under the stated quadratic Lagrangian and diffusion setting.Source: §B.2 · p. 21

C MODEL ARCHITECTURE

The model architecture uses an OT-Flow potential-function model to support the drift and accelerate computation of the HJB-PDE loss.

C MODEL ARCHITECTURE

  • The potential network takes the concatenated state and time vector s = (x, t) and applies M iterative hidden-state updates.The activation is σ(x) = log(exp(x) + exp(−x)); the hidden representation has dimension m.Source: C MODEL ARCHITECTURE · p. 21
  • Computing the potential Hessian’s diagonal costs O(m^2d), creating a trade-off between DNN expressive power and computational cost.The appendix also states that OT-Flow uses a potential satisfying fθ = −∇xΦθ rather than directly modeling velocity.Source: C MODEL ARCHITECTURE · p. 22
  • The potential-function model computes the drift from the potential gradient and supports fast, exact computation of the potential Hessian’s diagonal.The paper applies this Hessian computation to speed up the NLSB HJB-PDE loss when the diffusion output is diagonal.Source: C MODEL ARCHITECTURE · p. 22

D BENEFITS OF GENERALIZATION TO LAGRANGIAN

The paper presents the Lagrangian as a way to incorporate prior knowledge and derive dynamics losses across coordinate systems, and describes reverse-time simulation and experimental setups.

D BENEFITS OF GENERALIZATION TO LAGRANGIAN

  • A designed Lagrangian provides a unified way to express the HJB and Hamilton equations across coordinate systems, allowing the loss function to be derived from the Lagrangian.The examples show that Lagrangian design can reflect varied prior knowledge in sample trajectories, and the loss is nearly invariant under linear coordinate transformations when using W2.Source: D BENEFITS OF GENERALIZATION TO… · pp. 22–24

D.1 POPULATION DYNAMICS SIMULATION BY REVERSE-TIME SDE

  • For reverse-time simulation, the paper parameterizes the drift modification term with a neural network and trains it using Sinkhorn divergence.The authors note that computing the modification directly requires estimating the score function and has high computational cost.Source: §D.1 · p. 24

E.1 IMPLEMENTATION

  • The implementation compares NLSB with neural SDE, OT-Flow, TrajectoryNet, IPF variants, and SB-FBSDE using specified shared model and optimization settings.The experiments used an Intel Xeon Platinum 8360Y CPU and a single NVIDIA A100 GPU.Source: §E.1 · p. 25

E.2 SYNTHETIC POPULATION DYNAMICS: TIME-DEPENDENT ORNSTEIN–UHLENBECK PROCESS

  • The synthetic validation uses a one-dimensional time-dependent Ornstein–Uhlenbeck SDE to compare predicted trajectories and uncertainty with ground truth.The experiment is intended to test NLSB on a simple linear SDE.Source: §E.2 · p. 26

E.2.1 DETAILS OF EXPERIMENTAL SETUP

The appendix reports the shared model settings and synthetic-data regularization setup used for the 1D OU experiments.

E.2.1 DETAILS OF EXPERIMENTAL SETUP

  • The potential model uses two ResNet layers, step size 1.0, matrix rank 10, and hidden-vector dimension 2.The same potential-model settings are specified for OT-Flow, neural SDE, NLSB, and IPF (NN).Source: POTENTIAL MODEL IN OT-FLOW, NEURAL… · p. 26
  • The experiments include 1D OU process data and predictions (Fig. 6).Table 4 lists the regularization weight coefficients for synthetic-data experiments.Source: Figure 6 · p. 27, Table 4 · p. 27

E.2.2 RESULTS

The OU experiments compare prediction across variance regimes and in reverse time; the following single-cell experiment evaluates NLSB on embryo scRNA-seq populations observed at five snapshots.

E.2.2 RESULTS

  • For large-variance OU transitions, SDE-based methods outperform ODE-based methods, while NLSB and IPF are comparable to ODE-based methods at small variance.NLSB and IPF estimate population dynamics at both variance levels and have smaller CDD values than neural SDE; their sample-level predictions are closer to ground truth than neural SDE in [0, 1].Source: §E.2.2 · pp. 27–28
  • NLSB with the potential-free Lagrangian performs comparably to IPF, whose SB problem is almost mathematically equivalent in this setting.The authors interpret this as evidence that NLSB can handle SB problems as a special case.Source: §E.2.2 · p. 28
  • The reverse-time method recovers population-level dynamics when predicting from samples at t = 4 back to t = 0.The authors leave suitable evaluation of reverse-time sample-level dynamics to future work.Source: §E.2.2 · p. 28

E.3 SINGLE-CELL POPULATION DYNAMICS

  • The single-cell evaluation applies NLSB to human-embryo cell populations and introduces a quantitative metric for this practical setting.The data track embryoid-body stem-cell differentiation over 27 days using five snapshots; cells are represented in a PCA feature space.Source: §E.3 · p. 29, §E.3.1 · p. 29

E.3.1 DETAILS OF DATASET

  • The scRNA-seq dataset comprises five snapshots of embryoid-body differentiation into multiple cell lineages over 27 days.The data were split into training, validation, and test sets separately for each time point.Source: §E.3.1 · p. 29

E.3.2 DETAILS OF EXPERIMENTAL SETUP

The setup defines density and velocity regularizers for the cellular experiment; results describe uncertainty handling and high-diffusion transitions, followed by synthetic trajectory tests of Lagrangian priors.

E.3.2 DETAILS OF EXPERIMENTAL SETUP

  • The cellular setup combines energy, density, and velocity regularization terms; the density term is a GMM-estimated data log-likelihood.The mixture parameters are estimated from adjacent time points, with component count selected by BIC; the velocity term uses TrajectoryNet's reference velocity.Source: NLSB · p. 29

E.3.3 RESULTS

  • The scRNA-seq results are visualized in the first two principal components, while Table 6 reports MDD using EMD with L1 cost.The sample-population visualization is Figure 10, and single-sample distribution evolution is shown in Figure 9.Source: Figure 9 · p. 30, Table 6 · p. 30, §E.3.3 · p. 31
  • NLSB handles trajectory uncertainty, and it outperforms ODE-based methods on high-diffusion transitions from t1 to t2 and t3 to t4.The trajectories also differ by regularization: NLSB (E) has linear drift, NLSB (D) passes along the data manifold, and NLSB (E+D+V) reflects the combined regularization effects.Source: §E.3.3 · p. 31

E.4 SYNTHETIC POPULATION DYNAMICS: TRAJECTORIES REFLECTING THE POTENTIAL FUNCTION

  • The synthetic experiments test whether Lagrangians designed from prior knowledge can produce trajectories reflecting potential functions.The section also emphasizes potential-function design and flexibility in drift-penalty design.Source: §E.4 · p. 31

E.4.1 DATASET

The appendix tests NLSB with designed Lagrangians on synthetic populations and opinion dynamics, including whether trajectory costs encode obstacles, directional penalties, or a depolarizing target.

E.4.1 DATASET

  • The synthetic obstacle experiments use two-dimensional uniform endpoint distributions and 2,048 training samples plus 512 validation samples.The box and slit potentials are implemented differentiably using sigmoid functions.Source: §E.4.1 · p. 31, §E.4.2 · p. 33

E.4.2 DETAILS OF EXPERIMENTAL SETUP

  • The synthetic setup tests four potential functions and also varies dimension-specific drift penalties through the diagonal matrix R.The two R settings are diag([10.0, 0.1]) and diag([0.1, 10.0]).Source: §E.4.2 · p. 33, WELL POTENTIAL · p. 33

E.4.3 RESULTS

  • NLSB trajectories reflect obstacle or region information encoded by a potential, and drift is larger along the axis with the smaller R penalty.The directional-drift comparison uses L(t, u, x) = 1/2 u^T R u with R set to diag([10.0, 0.1]) or diag([0.1, 10.0]).Source: §E.4.3 · p. 33, WELL POTENTIAL · p. 33

E.5 OPINION DYNAMICS

  • The opinion-dynamics example applies NLSB as optimal control using a prior drift that causes polarization in the party model.The model describes agents' opinions interacting in Euclidean opinion space.Source: §E.5 · p. 33, §E.5.1 · p. 33

E.5.1 DATASET

  • The opinion experiment uses two-dimensional Gaussian distributions N0 and N1 as the endpoint populations.The agreement function encodes whether two opinions agree on independently sampled information.Source: §E.5.1 · p. 34

E.5.2 DETAILS OF EXPERIMENTAL SETUP

  • For opinion dynamics, the potential represents the population's averaged interaction, and an entropy potential c log p(x, t) is proposed to control population-diversity changes.NLSB's optimal drift combines the learned potential gradient with the negative polarizing drift to target an ideal terminal opinion distribution.Source: §E.5.2 · p. 34

E.5.3 RESULTS

  • After NLSB is applied, the learned drift prevents polarization caused by f_polarize and converges to the ideal terminal distribution N1.The directional similarity plots represent cosine angles between paired opinions, with red and blue indicating disagreement and agreement.Source: §E.5.3 · p. 34

Related papers

Loading 2204.04853v5…