Summary
In this work, the authors investigate the expressive power and trainability of the Fourier Neural Operator (FNO). A mean-field theory is adapted for the FNO to examine the behavior of randomly initialized FNOs from the perspective of the 'edge of chaos'. The investigation focuses on understanding the expressive behavior of randomly initialized FNOs by examining the transition between ordered and chaotic phases. This phase transition demonstrates specific characteristics unique to the FNO, influenced by mode truncation, and shares similarities observed in densely connected networks and CNNs. Additionally, the authors establish a correlation between expressive power and trainability: the ordered and chaotic phases correspond to regions where gradients vanish or explode, respectively. The authors identify the necessity of initializing FNO near the edge of chaos for stable training theoretically and experimentally.
Strengths
1. **Theoretical novelty:** The paper adapts mean-field theory specifically for Fourier Neural Operators (FNOs), extending existing theories from other neural network architectures such as densely connected networks and CNNs.
2. **Practical implications:** The theoretical findings lead to concrete recommendations for initializing FNOs to ensure stable training, particularly for deep FNOs.
3. **Comprehensive experiments:** The authors validate their theoretical results across multiple PDE-solving tasks, including advection, Burgers, Darcy Flow, and the Navier-Stokes equations. All these datasets are publicly available.
Weaknesses
1. The adaptation of mean-field theory to FNOs, while a notable endeavor, may not sufficiently break new ground in the field.
2. The significance of the theoretical aspects that FNOs share similarities with CNNs should be further discussed.
Overall, this work addresses how the network should be initialized for stable training of FNOs; however, its contribution might be somewhat limited. A stability analysis beyond the initialization phase, such as on the architecture, would be interesting to see.
4. There are barely any experimental results shown in the main text, and the discussion of how their theoretical analysis is useful in practical applications is lacking. The discovery in their work, the He distribution for FNOs, has already been applied in prior works such as [1] (see their GitHub implementation, `reset_parameters()` method under models); hwoever, it is always good to have a rigorous justification, which is shown in this work.
[1] Helwig, J., Zhang, X., Fu, C., Kurtin, J., Wojtowytsch, S., and Ji, S. Group equivariant Fourier neural operators for partial differential equations.
Questions
**Major Questions:**
1. If we treat FNO simply as a CNN, where global convolution occurs in the frequency domain, with Parseval theorems, it also implies, though less rigorously, that we can initialize the (complex) weights based on the He Distribution. This also hints on the similarities between FNOs and CNNs. Could the authors comment on this perspective?
2. In line 122, it states:
> "two learnable weights $\Theta^{(\ell, k)} \in \mathbb{R}^{D \times D}$ and $\Xi^{(\ell, k)} \in \mathbb{R}^{D \times D}$". If I understand correctly, these are complex weights for spectral convolution.
Firstly, why is $\Xi^{(\ell, k)} \in \mathbb{R}^{D \times D}$? Is this a typo where a $K$ is missing? Secondly, why are there two sets of learnable weights? Isn't there just one for the Fourier transform of the convolution kernel? Do you treat the real and complex parts separately, resulting in one set of parameters for the real and another for the complex parts?
3. The weights in spectral convolution are complex. By initializing the weights according to $\mathcal{N}\left(0, \frac{\sigma^2}{4 D}\right)$, do you treat the real and complex parts separately?
4. The FNO is implemented with $L = 64$ layers, which is particularly deep for FNOs involving $64$ forward and backward FFTs. Is such a large number of layers practical in real-world applications? If not, how are your findings useful for shallower FNOs?
**Minor/Optional Questions and Suggestions:**
1. Some abbreviations are introduced before their full names, such as DCN in line 25, whereas its full name is introduced in line 37 (densely connected network, DCN).
2. Given your insights into FNO expressivity and trainability, do you have any thoughts on potential architectural improvements or variations that could enhance FNO performance?
3. It would be beneficial to provide a clearer explanation and discussion of the theoretical results and their implications for readers without a background in mean-field theory.
Limitations
Yes, the authors adequately addressed the limitations.