Correlational Lagrangian Schrödinger Bridge: Learning Dynamics with Population-Level Regularization

TL;DR

CLSB addresses learning dynamics from cross-sectional population data with heterogeneous agents by introducing population-level regularizers that track temporal changes in multivariate relations. It replaces traditional per-particle action with correlational Lagrangians, deriving tractable analytical forms (e.g., for $k=1$) and domain priors such as Covariance Kinetics and Covariance Potential, integrated into neural SDE-based generation via unconstrained optimization. The method demonstrates improved modeling of developmental trajectories and dose-response cellular perturbations, outperforming ODE/SDE baselines on Wasserstein-distance metrics and illustrating the benefits of population guidance for heterogeneity. By encoding domain knowledge at the population level, CLSB offers a generalizable framework for simulating population dynamics in biology and related fields, with extensions to conditional generation and broader domain priors.

Abstract

Accurate modeling of system dynamics holds intriguing potential in broad scientific fields including cytodynamics and fluid mechanics. This task often presents significant challenges when (i) observations are limited to cross-sectional samples (where individual trajectories are inaccessible for learning), and moreover, (ii) the behaviors of individual particles are heterogeneous (especially in biological systems due to biodiversity). To address them, we introduce a novel framework dubbed correlational Lagrangian Schrödinger bridge (CLSB), aiming to seek for the evolution "bridging" among cross-sectional observations, while regularized for the minimal population "cost". In contrast to prior methods relying on \textit{individual}-level regularizers for all particles \textit{homogeneously} (e.g. restraining individual motions), CLSB operates at the population level admitting the heterogeneity nature, resulting in a more generalizable modeling in practice. To this end, our contributions include (1) a new class of population regularizers capturing the temporal variations in multivariate relations, with the tractable formulation derived, (2) three domain-informed instantiations based on genetic co-expression stability, and (3) an integration of population regularizers into data-driven generative models as constrained optimization, and a numerical solution, with further extension to conditional generative models. Empirically, we demonstrate the superiority of CLSB in single-cell sequencing data analyses such as simulating cell development over time and predicting cellular responses to drugs of varied doses.

Figures (21)

• Figure 1: In-silico simulation of the expressions of gene ABCA3 (x-axis) and A1BG (y-axis) during the embryonic stem cell development with different regularizations. Notably, our proposed population-level regularization facilitates a more accurate estimation of distributions, with quantitative evidence detailed in Sec. \ref{['sec:experiments']}.
• Figure 2: Visualization of the simulated gene expressions and trajectories with different methods. The trajectories are plotted for the gene pairs with the highest correlation (ABCA3 and A1BG), and for the first two principal components.
• Figure 3: Visualization of temporal variations of the covariance of embryonic stem cell expression. The first row of figures presents direct plots of the covariance at time $t$, while the second row displays violin plots illustrating the differences between time $t$ and $t-1$.
• Figure 4: Visualization of local dynamics for genes A2M-AS1 and A1BG.
• Figure 5: Visualization of local dynamics for genes A2ML1 and A2M.