Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs
Jean Kossaifi Nikola Kovachki Kamyar Azizzadenesheli Anima Anandkumar
Memory complexity and data scarcity have so far prohibited learning solution
operators of partial differential equations (PDEs) at high resolutions. We
address these limitations by introducing a new data efficient and highly
parallelizable operator learning approach with reduced memory requirement and
better generalization, called multi-grid tensorized neural operator (MG-TFNO).
MG-TFNO scales to large resolutions by leveraging local and global structures
of full-scale, real-world phenomena, through a decomposition of both the input
domain and the operator's parameter space. Our contributions are threefold: i)
we enable parallelization over input samples with a novel multi-grid-based
domain decomposition, ii) we represent the parameters of the model in a
high-order latent subspace of the Fourier domain, through a global tensor
factorization, resulting in an extreme reduction in the number of parameters
and improved generalization, and iii) we propose architectural improvements to
the backbone FNO. Our approach can be used in any operator learning setting. We
demonstrate superior performance on the turbulent Navier-Stokes equations where
we achieve less than half the error with over 150x compression. The
tensorization combined with the domain decomposition, yields over 150x
reduction in the number of parameters and 7x reduction in the domain size
without losses in accuracy, while slightly enabling parallelism.