Highly-scalable, physics-informed GANs for learning solutions of stochastic PDEs

--Uncertainty quantification for forward and inverse problems is a central challenge across physical and biomedical disciplines. We address this challenge for the problem of modeling subsurface flow at the Hanford Site by combining stochastic computational models with observational data using physics-informed GAN models. The geographic extent, spatial heterogeneity, and multiple correlation length scales of the Hanford Site require training a computationally intensive GAN model to thousands of dimensions. We develop a hierarchical scheme for exploiting domain parallelism, map discriminators and generators to multiple GPUs, and employ efficient communication schemes to ensure training stability and convergence. We developed a highly optimized implementation of this scheme that scales to 27,500 NVIDIA V olta GPUs and 4584 nodes on the Summit supercomputer with a 93.1% scaling efficiency, achieving peak and sustained half-precision rates of 1228 PF/s and 1207 PF/s. Index T erms --Stochastic PDEs, GANs, Deep Learning I. O VERVIEW A. Parameter estimation and uncertainty quantification for subsurface flow models Mathematical models of subsurface flow and transport are inherently uncertain because of the lack of data about the distribution of geological units, the distribution of hydrological properties (e.g., hydraulic conductivity) within each unit, and initial and boundary conditions. Here, we focus on parameter-ization and uncertainty quantification (UQ) in the subsurface flow model at the Department of Energy's Hanford Site, one of the most contaminated sites in the western hemisphere. During the Hanford Site's 60-plus years history, there have been more than 1000 individual sources of contaminants distributed over 200 square miles mostly along Columbia River [1]. Accurate subsurface flow models with rigorous UQ are necessary for assessing risks of the contaminants reaching the Columbia river and numerous wells used by agriculture and as sources of drinking water, as well as for the design of efficient remediation strategies. B. UQ with Stochastic Partial Differential Equations Uncertain initial and boundary conditions and model parameters render the governing model equations stochastic. In this context, UQ becomes equivalent to solving stochastic PDEs (SPDEs). Forward solution of SPDEs requires that all model parameters as well as the initial/boundary conditions are prescribed either deterministically or stochastically, which is not possible unless experimental data are available to provide additional information for critical parameters, e.g. the field conductivity.