Confronting Uncertainty in Mechanism-AI Coupled Dynamical Systems via Optimal Transport and Information Geometry: Part I. Distribution Learning

Pengfei Song, Yanni Xiao, Shuangshuang Yin

Mechanism-AI coupled dynamical systems are increasingly used to study complex biological systems. However, uncertainty associated with noisy, sparse, and partially observed data limits their reliability. This paper presents a systematic mathematical formulation for treating these uncertainties from the perspective of optimal transport and information geometry. We formulate the uncertainty quantification problem as the minimization of an energy functional on a probability density space equipped with a Riemannian metric, with the loss evaluated through a differential equation, and give a unified description of the gradient flows of the energy functional under specific metrics, such as the Wasserstein and Fisher-Rao metrics. Furthermore, we present corresponding numerical methods for solving these gradient flows using abstract adjoint sensitivity analysis. The formulation connects a range of Bayesian inference methods for dynamical systems; numerical experiments compare posterior approximation and predictive uncertainty under a common computational budget. Moreover, it also provides a geometric perspective on dynamical modeling and the study of heterogeneity evolution.