Understanding the Expressivity and Trainability of Fourier Neural Operator: A Mean-Field Perspective

Abstract: In this paper, we explores the expressivity and trainability of the Fourier Neural Operator (FNO). We establish a mean-field theory for the FNO, analyzing the behavior of the random FNO from an edge of chaos perspective. Our investigation into the expressivity of a random FNO involves examining the ordered-chaos phase transition of the network based on the weight distribution. This phase transition demonstrates characteristics unique to the FNO, induced by mode truncation, while also showcasing similarities to those of densely connected networks. Furthermore, we identify a connection between expressivity and trainability: the ordered and chaotic phases correspond to regions of vanishing and exploding gradients, respectively. This finding provides a practical prerequisite for the stable training of the FNO. Our experimental results corroborate our theoretical findings.

• The paper develops a mean-field theory that characterizes the initialization bias in Fourier Neural Operators due to mode truncation.
• The paper reveals that initializing FNOs at the edge of chaos enhances training stability and mitigates gradient issues.
• The paper validates its approach on multiple PDE benchmarks, demonstrating stable deep FNO training without traditional skip connections.

Analysis of Initialization Bias in Fourier Neural Operators

The paper, titled "Initialization Bias of Fourier Neural Operator: Revisiting the Edge of Chaos," discusses the influence of initialization bias in Fourier Neural Operators (FNOs), a novel framework used for solving partial differential equations (PDEs). This work builds a mean-field theory to explore how FNOs behave at initialization, considering the interplay between randomness and architecture. It underscores characteristics unique to FNOs due to mode truncation and demonstrates parallels with densely connected networks (DCNs). Through this theoretical lens, the authors propose a new initialization scheme at the edge of chaos that aims to mitigate negative initialization bias, thereby enhancing the stability of FNO training. The efficacy of this approach is demonstrated through experiments that eliminate the need for skip connections in deep FNO architectures.

Key Insights

1. Mean-field Theory for FNOs: The paper develops a mean-field theory that characterizes FNOs' forward and backward propagation behaviors, especially how mode truncation affects them. This theory permits an analysis of the FNO's initialization bias from the edge of chaos perspective.
2. Edge of Chaos and Initialization: The research revisits the edge of chaos theory, which indicates that neural networks function optimally when initialized at the boundary between order and chaos. The authors' analysis shows that FNOs require a specific initialization strategy to achieve stable training, analogous to He initialization in DCNs with ReLU activation.
3. Impact of Mode Truncation: FNOs utilize the Fourier transform to manage spatial dependencies, which introduce unique behaviors arising from mode truncation. These effects, along with similarities to DCNs, guide the proposed initialization scheme to reach the edge of chaos, facilitating better generalization and training stability.
4. Experimental Validation: Multiple datasets and PDEs, such as Burgers', Darcy Flow, and Navier-Stokes equations, validate the theoretical claims. The results confirm that suitable initialization helps stabilize the training of deep FNOs, preventing gradient vanishing or exploding issues and enabling large depth models without the necessity of skip connections.

Practical and Theoretical Implications

This study offers a profound theoretical contribution by extending the mean-field analysis to FNOs, elucidating the importance of initialization in complex networks. From a practical standpoint, the proposed initialization scheme can significantly improve the performance and training stability of FNOs in scientific computing and neural PDE solvers. These insights can guide future work on optimizing network architectures and initialization methods for complex and high-dimensional data modeling.

Future Directions

Given the rapid advancements in scientific machine learning and neural PDE solvers, several directions for future research emerge:

• Extended Architectures: Investigate the applicability of the edge of chaos initialization to other neural operator variants or in combination with advanced architectures like transformers and graph neural networks.
• Optimization Dynamics: Analyze how training dynamics evolve beyond initialization to convergence, especially when employing regularization techniques like dropout or batch normalization.
• Dynamical Isometry: Explore alternate initialization paradigms that achieve dynamical isometry, thereby promoting efficient information propagation even with very deep architectures.
• Robustness and Generalization: Evaluate how initialization impacts robustness to noise and generalization to unseen data distributions across different PDE domains.

Overall, this work provides a substantial theoretical foundation and practical strategies for improving deep models solving PDEs, contributing to both the neural operator framework and broader deep learning paradigms.