Large-scale parametric PDE approximations with model-parallel Fourier neural operators

TitleLarge-scale parametric PDE approximations with model-parallel Fourier neural operators
Publication TypePresentation
Year of Publication2023
AuthorsRichard Rex, Thomas J. Grady II, Rishi Khan, Ziyi Yin, Felix J. Herrmann
KeywordsCCS, deep learning, Fourier neural operators, HPC, ML4SEISMIC, SLIM

Solving PDEs to simulate two-phase flow is expensive since it involves the inversion of large ill-conditioned matrices. FNOs represent a special type of neural network capable of approximating solutions to two-phase flow equations. In order to speed up this computation, we develop a high-level software abstraction tool to exploit the linearly separable property of Fourier Transforms via Kronecker products. We perform a series of all-to-all operations where we partition the data to apply the operations in a distributed fashion. We apply these FNOs to predict the evolution of CO2-plumes in subsurface environments. Our model takes an input permeability model, and outputs time-varying CO2 saturations in a quick and cost-effective manner. Additionally, our research involves developing a distributed matrix-free abstraction library that can be used to represent any generic linear and nonlinear operator. This library is scalable and auto-differentiable thanks to the hand-written customized AD rules, allowing us to represent and train any network. We provide an evaluation on the accuracy of FNO simulations compared to traditional PDE simulations in solving various classes of PDEs.

Citation Keyrex2023ML4SEISMIClsp