Multi-Grid Tensorized Fourier Neural Operator for High- Resolution PDEs
Abstract: Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDE) 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 150× compression. The tensorization combined with the domain decomposition, yields over 150× reduction in the number of parameters and 7× reduction in the domain size without losses in accuracy, while enabling parallelism.
Submission Length: Long submission (more than 12 pages of main content)
Previous TMLR Submission Url: https://openreview.net/forum?id=oFqHIkw8sd
Changes Since Last Submission: Dear area editor and reviewers,
Following the advice from the editors, we are resubmitting our manuscript as our response was missed as it was not posted in the expected format. All reviewers have appreciated the novelty and improved performance of our approach.
To summarize, we proposed a novel approach that combines a (multi-grid) decomposition of the input domain and a low-rank tensor decomposition of the parameters of the model. This enables: i) improved performance (over one order of magnitude) ii) reduced overfitting iii) improves learning in low-data regime iv) parallelization of the learning over very large instances thus enabling large-scale training. v) a large reduction in the number of parameters, over 500x compression ratio
We have updated our manuscript to incorporate all the comments from th reviewers and we ran additional experiments.
Specifically, we added samples generated by our model in various scenarios (super-resolution, super-evaluation) and at different resolutions. In particular, we generate test samples at resolution 1024x1024 and compare super-evaluation, which leverages the neural operator properties, with down-sampling and re-upsampling (section 4.6 in the updated manuscript).
To demonstrate that our approach can be applied to any neural operator and that it helps with overfitting and performance, we ran a new set of experiments with U-NO and applied our proposed approach to this U-shaped neural operator (table 2 in the updated manuscript). We show that it helps with performance there too.
We benchmarked our proposed TFNO approach and compared it with a regular FNO in terms of memory usage and runtime (section 4.5 in the updated manuscript). We show that our approach can increase memory usage and lead to computational speedup depending on the choice of rank and parameters. We found that in some cases we can get some memory savings, and up to 1.7x speedup over the baseline FNO and up to 3.7x over the FFNO.
We added a figure (Figure 3 in the updated manuscript) to illustrate the indexing of the Fourier modes as requested. We also updated the architecture diagrams.
We also include here a detailed response to the reviewers since this was missed in the previous review.
Lastly, we also want to point out that our proposed method has already been successfully used in several other academic publications, including some published in ICLR and NeurIPS [1,2,3,4,5].
[1] Geometry-Informed Neural Operator for Large-Scale 3D PDEs, NeurIPS 2023
[2] GUARANTEED APPROXIMATION BOUNDS FOR MIXED-PRECISION NEURAL OPERATORS, ICLR 2024
[3] Scattering with Neural Operators, Sebastian Mizera, preprint
[4] An Operator Learning Framework for Spatiotemporal Super-resolution of Scientific Simulations, Duruisseaux, V., & Chakraborty, A.
[5] Neural Operators for Accelerating Scientific Simulations and Design, Anandkumar et al
Assigned Action Editor: ~Lorenzo_Orecchia1
Submission Number: 2510
Loading