Incremental Spatial and Spectral Learning of Neural Operators for Solving Large-Scale PDEs

Robert Joseph George,Jiawei Zhao,Jean Kossaifi,Zongyi Li,Anima Anandkumar
2024-03-05
Abstract:Fourier Neural Operators (FNO) offer a principled approach to solving challenging partial differential equations (PDE) such as turbulent flows. At the core of FNO is a spectral layer that leverages a discretization-convergent representation in the Fourier domain, and learns weights over a fixed set of frequencies. However, training FNO presents two significant challenges, particularly in large-scale, high-resolution applications: (i) Computing Fourier transform on high-resolution inputs is computationally intensive but necessary since fine-scale details are needed for solving many PDEs, such as fluid flows, (ii) selecting the relevant set of frequencies in the spectral layers is challenging, and too many modes can lead to overfitting, while too few can lead to underfitting. To address these issues, we introduce the Incremental Fourier Neural Operator (iFNO), which progressively increases both the number of frequency modes used by the model as well as the resolution of the training data. We empirically show that iFNO reduces total training time while maintaining or improving generalization performance across various datasets. Our method demonstrates a 10% lower testing error, using 20% fewer frequency modes compared to the existing Fourier Neural Operator, while also achieving a 30% faster training.
Machine Learning
What problem does this paper attempt to address?
The paper aims to address the following issues: 1. **Computational efficiency of Fourier Neural Operator (FNO) on high-resolution inputs**: FNO is very time-consuming and computationally complex when performing Fourier transforms on high-resolution inputs, which is crucial for solving many partial differential equations (such as fluid flow). 2. **Selection of frequency modes in the spectral layer**: Choosing the appropriate number of frequency modes K is very challenging; too many or too few modes can lead to overfitting or underfitting. To address these issues, the authors propose the Incremental Fourier Neural Operator (iFNO), which gradually increases the number of frequency modes K used and the resolution of the training data during the training process. In this way, iFNO can reduce the total training time while maintaining or even improving generalization performance. Experimental results show that compared to existing Fourier Neural Operators, iFNO reduces test error by 10%, decreases the number of frequency modes used by 20%, and increases training speed by 30%.