Generative modeling with low-rank Wasserstein polynomial chaos expansions

A new Wasserstein multi-element polynomial chaos expansion (WPCE) is proposed, which is inspired by recent advances in computational optimal transport for estimating Wasserstein distances. The developed method combines unsupervised learning with the explicit functional representation of a random vector $Y$. Its training only relies on a finite set of samples from an unknown distribution, which is used to minimize a regularized empirical Wasserstein metric known as debiased Sinkhorn divergence. An interesting application that motivates the approach comes from the numerical upscaling of non-periodic random fields defined on a micro-scale. The WPCE can encode higher order stochastic information about the effective material behavior in contrast to the constant characterization with stochastic homogenization. A striking feature of the new method is the generalization of common diffeomorphic transport maps to the case of discontinuous and non-injective model classes $\mathcal{M}$ with possibly different input and output dimension. It computes a (functional) relation $Y=\mathcal{M}(X)$ in distribution with input random variables $X$ and target $Y$. The exponential growth of the PCE is alleviated by a new stacked tensor train (STT) format. By the choice of the model class $\mathcal{M}$ and the smooth loss function, higher-order optimization schemes and in particular Riemannian descent methods become possible. The proposed approach is illustrated numerically with a high-dimensional upscaling problem, which considers a microscopic random non-periodic composite material. It results in a computationally tractable effective macroscopic random field in adapted stochastic coordinates. By using a relaxation to a discontinuous model class, multimodal distributions also become tractable.