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

Yuning You, Ruida Zhou, Yang Shen

Keywords: generative models, diffusion models

Abstract:

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.

Submission Number: 63