Learning diffeomorphisms of real spaces to infer transport equations of probability densities via discrepancy maps

  • Han, Saem (Google Inc.)
  • Tan, Jingye (University of Southern California)
  • Garikipati, Krishna (University of Southern California)

Please login to view abstract download link

In many applications, it is of interest to infer transport equations governing a probability density in high dimensions. Such situations arise also in the physical context of the dynamics of populations ranging from biological cells, through human and other agents. Challenges arise from the  sparsity of data in the high-dimensional space over which the density is defined, and the consequent difficulty of learning the parameters of the attendant Fokker-Planck equations. Here we approach this problem via a two-step approach, in which we separate the learning of the transport map of the density from the Fokker-Planck parameters. We use Bayesian neural networks with variational inference for this step. Because this circumvents the estimation of derivatives in discretization-based methods, this approach is robust in the presence of sparsity and noise. We further introduce the notion of a diffeomorphism of the high-dimensional space, which can manifest in  discrepancies between the  observed dynamics and the evolution predicted by Fokker-Planck equations in this  space. Learning the diffeomorphism simultaneously with the transport map and Fokker-Planck equation can resolve the complexity of many high-dimensional transport problems.