Correlational Lagrangian Schrodinger Bridge: Learning Dynamics with Population-Level Regularization

Modeling population dynamics is a fundamental problem with broad scientific applications.Motivated by real-world applications including biosystems with diverse populations, we consider a class of population dynamics modeling with two technical challenges: (i) dynamics to learn for individual particles are heterogeneous and (ii) available data to learn from are not time-series (i.e, each individual's state trajectory over time) but cross-sectional (i.e, the whole population's aggregated states without individuals matched over time).To address the challenges, we introduce a novel computational framework dubbed correlational Lagrangian Schr\"odinger bridge (CLSB) that builds on optimal transport to "bridge" cross-sectional data distributions. In contrast to prior methods regularizing all individuals' transport "costs" and then applying them to the population homogeneously, CLSB directly regularizes population cost allowing for population heterogeneity and potentially improving model generalizability.Specifically our contributions include (1) a novel population perspective of the transport cost and a new class of population regularizers capturing the temporal variations in multivariate relations, with the tractable formulation derived,(2) three domain-informed instantiations of population regularizers on covariance, and (3) integration of population regularizers into data-driven generative models as constrained optimization and an approximate numerical solution, with further extension to conditional generative models.Empirically, we demonstrate the superiority of CLSB in single-cell sequencing data analyses (including cell differentiation and drug-conditioned cell responses) and opinion depolarization.Codes will be released upon acceptance.