Learning Population-Level Dynamics through a Latent Fokker-Planck Model and Discrepancy Transport Maps

Abstract

Many scientific and engineering systems are observed as time-indexed probability distributions whose governing dynamics are unknown and whose individual trajectories are unavailable. These settings challenge conventional system-identification approaches that rely on trajectory correspondence or prescribed evolution equations. This work presents a population-level inference framework that recovers latent stochastic dynamics directly from snapshot probability distributions by decomposing the observed evolution into an intrinsic latent stochastic process and a discrepancy transport map that captures geometric deformation between the latent and observed probability spaces. The latent dynamics are modeled using an Ornstein--Uhlenbeck process, providing a closed-form solution to the associated Fokker--Planck equation, while the discrepancy transport map is parameterized through the Knothe--Rosenblatt rearrangement with monotone neural networks. To mitigate the non-uniqueness inherent in the latent--transport decomposition, the transport map is regularized using a deformation energy motivated by hyperelasticity, promoting smooth, physically interpretable deformations while reducing unnecessary complexity. The latent stochastic model and discrepancy transport map are learned jointly through a unified optimization problem defined over probability distributions. Numerical examples involving nonlinear and multimodal distributional dynamics demonstrate that the proposed framework accurately reconstructs complex probability evolution while preserving a compact and analytically tractable latent representation. The proposed formulation provides a general framework for population-level dynamical inference and establishes a foundation for extending latent stochastic models and transport-based learning to more general and higher-dimensional systems.