Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs

TMLR Paper1499 Authors

Abstract:

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.

Submission Length: Long submission (more than 12 pages of main content)

Changes Since Last Submission:

We should like to thank the reviewers for all their helpful comments which have helped us improve the paper. A detailed response to each reviewer comment can be found here:

https://drive.google.com/file/d/1Jh1cviwB_MedOPpe3KR2-oIxc702cQjQ/view?usp=sharing

We have updated the submitted manuscript. We hope to have answered all questions and addressed all concerns appropriately.

Assigned Action Editor: ~Mingsheng_Long2

Submission Number: 1499