Intrinsic Models in Wasserstein Space with Applications to Molecular Dynamics

We study the problems of efficient modeling and representation learning for probability distributions in Wasserstein space. We consider a general barycentric coding model in which data are represented as Wasserstein-2 (W2) barycenters of a set of fixed reference measures. Leveraging the geometry of W2-space, we develop a tractable optimization program to learn the barycentric coordinates and provide a consistent statistical procedure for learning these coordinates when the measures are accessed only by i.i.d. samples. Our consistency results and algorithms exploit entropic regularization of the optimal transport problem, and the statistical convergence of entropic optimal transport maps will be discussed. We also consider the problem of learning reference measures given observed data. Our regularized approach to dictionary learning in W2-space addresses core problems of ill-posedness and in practice learns interpretable dictionary elements and coefficients useful for downstream tasks. Applications of optimal transport to compression of molecular dynamics simulations will be considered.