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

Due to the character limit, we have addressed some questions in the global rebuttal. Please refer to our global rebuttal. We appreciate your attention and review.

Q1. The discovery in their work, ...

The following part of the code is different from our He initialization, as it only initializes using the ‘nn.init.kaiming_uniform_’ function. The weights are initialized by a uniform distribution with different scales and the bias is not fixed to zero.

def reset_parameters(self):
        if self.Hermitian:
            for v in self.W.values():
                nn.init.kaiming_uniform_(v, a=math.sqrt(5))
        else:
            nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
        if self.B is not None:
            nn.init.kaiming_uniform_(self.B, a=math.sqrt(5))

---

Q2. If we treat FNO simply as a CNN ...

Thank you very much for your insightful and valuable comments. We added this discussion of the similarities between CNN and FNO to the appendix. The following is our answer to this question and also the part of appendix.

The iterated map for CNNs (Xiao et al., 2018) is obtained as follows:

$$\Sigma^{(\ell)}\_{\alpha, \alpha'} = \frac{\sigma^2}{2r+1} \sum\_{\beta \in \mathrm{ker}} \mathbb{E}\_{\boldsymbol{H}\_{:. d} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{\Sigma}^{(\ell-1)})}\left[ \phi(H\_{\alpha+\beta, d})\phi(H\_{\alpha'+\beta, d}) \right] + \sigma\_b^2$$

where $2r + 1$ is the number of filter width and $\operatorname{ker} = \{ \beta \in \mathbb{Z} \mid |\beta| \leq r\}$ is the set of indices referring to the elements of the filter. When we consider the FNO without mode truncation ($K=N/2 + 1$), the propagation of the diagonal components $\Sigma\_{\alpha, \alpha}$ for any $\alpha \in [N]$ is equivalent to the propagation in a CNN with filter size $N$ performing global convolution.

$$\begin{aligned} \Sigma^{(\ell)}\_{\alpha, \alpha} &= \sigma^2 \sum\_{k=0}^{\frac{N}{2}} c\_k \mathbb{E}\_{\boldsymbol{H}\_{:. d} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{\Sigma}^{(\ell-1)})}\left[\left|\left[\mathcal{F}\phi\left(\boldsymbol{H}\_{:, d}\right)\right]\_{k}\right|^2 \right] \underbrace{\cos\left( \theta^{(k)}\_{\alpha, \alpha}\right)}\_{=1} + \sigma^2\_b \\\\ &= \frac{\sigma^2}{N} \sum\_{\beta=0}^{N-1} \mathbb{E}\_{\boldsymbol{H}\_{:. d} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{\Sigma}^{(\ell-1)})}\left[ \left|\phi(H\_{\beta, d})\right|^2\right] + \sigma^2\_b \\\\ &= \frac{\sigma^2}{N} \sum\_{\beta=0}^{N-1} \mathbb{E}\_{\boldsymbol{H}\_{:. d} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{\Sigma}^{(\ell-1)})}\left[ \phi(H\_{\alpha+\beta, d})\phi(H\_{\alpha+\beta, d})\right] + \sigma^2\_b. \end{aligned}$$

The first equality follows from Perseval's equality and the second from the periodic boundary condition. In contrast, the propagation of the off-diagonal components differs between CNN and FNO, and the iterated map is different even in the presence of mode truncation. Similarly, the iterated map during the backpropagation for FNO can be recognized as the global CNN (see Eq. (2.16) in (Xiao et al., 2018)) with $v\_{\beta} = \frac{1}{N}$.

$$\tilde{\Sigma}\_{\alpha, \alpha}^{(\ell)} \approx \frac{1}{N^2} \sum\_{\beta, \beta'=0}^{N-1} \tilde{\Sigma}\_{\beta, \beta'}^{(\ell+1)} \chi\_{c^*} \underbrace{\left(\sum\_{k=0}^{N/2} c\_k \cos\left( \theta^{(k)}\_{\beta-\beta', \alpha-\alpha} \right) \right)}_{=N \delta\_{\beta, \beta'}} = \chi\_{c^*} \sum\_{\beta=0}^{N-1} \frac{1}{N} \tilde{\Sigma}\_{\beta, \beta}^{(\ell+1)}.$$

This equivalence suggests that the edge of chaos initialization (e.g., He initialization) is also valid for the FNO since the problems of gradient vanishing and explosion are determined by the diagonal components of $\tilde{\boldsymbol{\Sigma}}$.

(Xiao et al., 2018) Xiao, Lechao, et al."Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks."

---

Q3. In line 122, it states: ... & The weights in spectral convolution ...

For ease of handling in proofs, we define the complex weights $R^{(\ell)} \in \mathbb{C}^{D \times D \times K}$ of the $\ell$-th layer of the FNO by splitting them into their real and complex parts. We further divide these weights along the mode dimension using the index k, as $R^{(\ell)} = \bigg|\bigg|\_{k=1}^K \Theta^{(\ell, k)} + \sqrt{-1} \Xi^{(\ell, k)}$ where $\bigg|\bigg|$ denotes the concatenation operation along the mode axis k. For all $\ell \in [L]$ and $k \in [K]$, every element of the weights $\Theta^{(\ell, k)}$ and $\Xi^{(\ell, k)}$ is independently initialized from $\mathcal{N}(0, \sigma^2/2D)$.

---

Q4: The FNO is implemented with 𝐿=64 layers, ... Is such a large number of layers practical in real-world applications?

Networks with a large number of layers are practical for real-world applications. For instance, as demonstrated in Appendix D.2, deeper FNOs with 16 or 32 layers achieved better performance than the conventional 4-layer FNO in learning complex operators. The study by Tran et al. (2023) on F-FNO also explores the training of relatively deep FNOs. In the field of PDE solving, it is possible to generate a large amount of data through numerical computation. Therefore, we believe that in the future, architectures with larger model capacities will be employed for more complex operator approximations and the construction of large pre-trained models as in the fields of CV or NLP. Our theory provides valuable insights for such multi-layer FNOs and architectures that partially use FNO layers.

---

Q5. Some abbreviations ...

Thank you for pointing this out. We have revised this point.

---

Q6. Given your insights into FNO expressivity and trainability, do you have any thoughts on potential architectural improvements or variations that could enhance FNO performance?

It seems difficult to obtain suggestions for new architectures. We may be able to give theoretical guarantees for newly proposed architectures.

Due to questions similar to those from other reviewers, we have addressed these in the global rebuttal. Please refer to our global rebuttal. We appreciate your attention and review.

Q1. The covariance matrix should be introduced in section 3.2.

The covariance $\Sigma_{\alpha, \alpha’}^{(\ell)}$ is defined in Section 3.1, lines 135-136. For further clarity, definitions have also been added to the statements in Lemma 3.1 (Section 3.2).

---

Q2. Lemma A.1 should be included in the main text. Some parts, such as detailed description of datasets can be moved to the appendix.

I appreciate your valuable suggestions. We moved Lemma A.1 to the line 155 in the main and some details about the dataset to the appendix.

---

Q3. The proposed study focuses solely on the expressivity of random FNO at initialization, a rather limited scenario. This limitation should be more explicitly emphasized, particularly in the title.

Thank you for your helpful suggestions. We are considering the following titles. Which proposed title adequately reflects the scenario limitation?

• Understanding the expressivity and trainability of Fourier Neural Operator at initialization
• Understanding the impact of initialization on expressivity and trainability in Fourier Neural Operator
• Understanding the expressivity and trainability of Mean Field Fourier Neural Operator

---

Q4. I am not convinced that the findings offer practical insights into the stable training of FNO, despite the preliminary numerical experiments provided.

Our theory provides a universal insight, independent of the dataset, loss function, optimizer, and scheduler. While it does not offer sufficient conditions for stable training, it does provide a necessary condition for achieving stability in training. We believe that our results are of practical importance as they prevent unnecessary exploration of initialization. We have revisited the following.

line 44: This discovery provides the necessary conditions (not sufficient conditions) for stable training, which is practically significant as it eliminates the unnecessary initialization searches.

line 310: While we do not provide sufficient conditions for stable training, we do offer one necessary condition for achieving stable training.

---

Q5. The term trainability, as used by the authors, only pertains to the vanishing or exploding gradients of random FNO.

The term trainability is used in the following related references with the same meaning as ours:

• Xiao, Lechao, et al. "Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks."
• Yang, Ge, and Samuel Schoenholz. "Mean field residual networks: On the edge of chaos."
• Yang, Greg, et al. "A Mean Field Theory of Batch Normalization."
• Chen, Minmin, et al. "Dynamical isometry and a mean field theory of RNNs: Gating enables signal propagation in recurrent neural networks."

In these references, trainability refers to a necessary condition for successful training that is determined by the architecture, independent of the dataset, loss function, optimizer, or scheduler, primarily focusing on discussions related to gradient vanishing/exploding. However, we acknowledge your point that, in general, the term trainability could be interpreted to encompass the overall optimization process and the success of training. Therefore, we have revised the following sentences in the introduction for clarity. Thank you for your objective and valuable feedback.

Line 29: We analyze the exponential expressivity (how far representations of different positions can be pulled apart) and trainability (how much gradient vanishing/explosion on average) of the random FNO from the perspective of whether the network is ordered or chaotic.

---

Q6. One might question the impact of the optimizer, the initial learning rate, and its scheduler on the results. An ablation study addressing these factors would have been beneficial.

Thank you for your valuable suggestions. We are currently conducting additional experiments. If the results are ready in time for the discussion period, we will report them in the thread. However, we would like to reiterate that our theory does not tackle the analysis of the optimization process and our results represent necessary conditions for stable training, independent of these factors. It is important to emphasize that the failure of learning due to inappropriate choices of optimizer, learning rate, or scheduler is not inconsistent with our findings.

---

Q7. Does the fixed point of the iterated map always exist?

Lemma A.1 guarantees the existence of a fixed point for the FNO when the DCN has a fixed point with $(q^*, c^* = 1)$. The existence of the fixed point for the DCN depends on the activation function, so we need to check it on a case-by-case basis. For example, despite both being bounded activation functions, Tanh has a fixed point, whereas the sigmoid function does not. Among unbounded activation functions, ReLU has a fixed point, but softplus (a smoothed version of ReLU) does not. The fixed point analysis for bounded activation functions is conducted in [1], and the general case is discussed in [2].

[1] Poole, Ben, et al. "Exponential expressivity in deep neural networks through transient chaos."

[2] Roberts, Daniel A., Sho Yaida, and Boris Hanin. "The principles of deep learning theory. "

---

Q8. The representation of Fig.1(a) is unclear. ...

Fig. 1(a) experimentally verified Theorem 3.4. Instead of averaging over training epochs, It averages over multiple seeds at initialization (not over training epochs), with a relatively small standard deviation. We revised the first sentence of the title of Fig. 1 as follows.

fig1- title: Average gradient norm $\operatorname{Tr}(\tilde{\Sigma}^{(\ell)})$ at initialization over multiple seeds during the backpropagation of several FNOs plotted as a function of layer $\ell$

Q1. The authors don't fully explore very deep FNOs or a wide range of hyperparameters (in particular sweeping over the number of modes). Additional experiments on larger-scale problems or more complex PDEs could further strengthen the empirical validation.

Thank you for your valuable suggestions. We are currently conducting additional experiments. If the results are ready in time for the discussion period, we will report them in the thread.

---

Q2. I know this paper tackled FNO's, but I am curious if this could easily be adapted to other variants of FNOs like U-FNO, F-FNO etc

Our theory may be applicable to several variants of the FNO. Examples include the following:

• geo-FNO: For geo-FNO, the influence of deformation is limited to the beginning and the end of the architecture, with the majority of the architecture consisting of the original FNO. Thus, our theory can be directly applied.
• G-FNO: For G-FNO, similar results can be obtained if the initial distribution of the transformed weight $\hat{L}_{s_g} R$ by elements $s_g$ from the stabilizer $S_g$ satisfies the assumptions of our theory.
• WNO: Since the wavelet transform is a linear operator, the same analysis can potentially be applied to the WNO.

Conversely, applying our theory directly to the architectures of U-FNO and F-FNO is challenging due to the presence of skip-connections. In architectures with skip-connections, the degeneration of input information in the ordered phase does not occur, and chaotic behavior is expected to be prominent. Therefore, it is necessary to analyze the behavior of the input in the chaotic phase in detail, rather than focusing on the discussion of the fixed point. It may be possible to extend our theory by incorporating mean-field theory for DCNs with skip-connections (Yang & Schoenholz, 2017) alongside our theoretical framework. Once the discussion on skip-connections is resolved, similar arguments can be made as follows:

• U-FNO: For U-FNO, it is considered possible to discuss by integrating the results related to CNN, DCN, and simplified FNO, as discussed in Section 3.3 on the original FNO. Note that skip-connections are included in U-Net.
• F-FNO: The factorized Fourier convolution module (Eq. 8) can be discussed similarly to the conventional FNO. However, since the output of the inverse FFT is summed in the dimensional direction, it seems appropriate to consider the weight scale as 1/D.

(Yang & Schoenholz, 2017) Yang, Ge, and Samuel Schoenholz. "Mean field residual networks: On the edge of chaos." Advances in neural information processing systems 30 (2017).

---

Q3. The paper doesn't provide a detailed analysis of how different PDE characteristics (e.g., stiffness, non-linearity, multi-scale behavior) might affect the expressivity and trainability of FNOs. Could you explain more on this point?

In terms of expressivity, it depends on the architecture of the FNO and not on the characteristics of the dataset (PDE). Regarding trainability, our theory provides a universal insight, independent of the dataset, loss function, optimizer, and scheduler. While it does not offer sufficient conditions for stable training, it does provide a necessary condition for achieving stability in training. As discussed in Appendix D, the characteristics of the dataset may, for example, lead to strong non-convexity problems in the loss landscape, but this is not explained by our theory. Our findings suggest that the acceptable initial variance of the weights is already limited by the number of layers, so these problems need to be addressed by optimizers, schedulers, devising loss functions, introducing regularization, etc.

Q1. The summary of related work can be improved, particularly by including recent neural operator architectures

We have revised Section 2.1 as follows. (Only newly added references are cited in numbered form.)

line 48: The FNO (Li et al., 2021) is one of the well-established methods for solving PDEs across many scientific problems (Hwang et al., 2022; Jiang et al., 2021; Pathak et al., 2022; [1, 2]).

lines 58 - 71: Several FNO variants have been developed to address specific challenges, such as geo-FNO (Li et al., 2022a) for irregular regions and group equivariant FNO (G-FNO) (Helwig et al., 2023), which maintains equivariance to rotation, reflection, and translation. U-NO[3] and U-FNO[4] integrate FNO with U-Net for multiscale modeling. Additionally, WNO[5] utilizes wavelet bases, while CFNO[6] enhances the use of geometric relations between different fields and field components through Clifford algebras. Adaptive FNO[7, 8] and F-FNO (Tran et al., 2022) have improved computational and memory efficiency through incremental learning and architectural modifications. Other approaches for improving performance include methods with increasing physical inductive bias[9], data augmentation[10], and a variance-preserving weight initialization scheme (Poli et al., 2022). While numerous new models and learning methods have been proposed, relatively little research has been conducted to understand the intrinsic nature of these methods. Issues such as spectral bias[7] and training instability (Tran et al., 2022) have been reported.

[1] Yang, Y., et al. Seismic wave propagation and inversion with neural operators.

[2] Wen, G., et al. Accelerating carbon capture and storage modeling using fourier neural operators.

[3] Rahman, Md Ashiqur, et al. U-no: U-shaped neural operators

[4] Wen, Gege, et al. U-FNO—An enhanced Fourier neural operator-based deep-learning model for multiphase flow

[5] Tripura, Tapas, et al. Wavelet neural operator: a neural operator for parametric partial differential equations

[6] Brandstetter, Johannes, et al., Clifford Neural Layers for PDE Modeling

[7] Zhao, Jiawei, et al. Incremental spectral learning in Fourier neural operator.

[8] Guibas, John, et al. Adaptive fourier neural operators: Efficient token mixers for transformers.

[9] Li, Zongyi, et al. Physics-informed neural operator for learning partial differential equations.

[10] Brandstetter, Johannes, et al. Lie point symmetry data augmentation for neural PDE solvers.

---

Q2. The code is not available.

The paper does not propose any new models or learning methods and most of the code used in the experiments is based on the code of PDEBench. The only modifications to the model are to multiply the outputs (variable ‘out_ft’ in code of class FNO1d and FNO2d) corresponding to mode k=0, N/2 by $\sqrt{2}$ and to initialize the weights by Gaussian distribution, as described in Section 3.1. We added the same information as above to the detailed experimental setup description in the paper.

---

Q3. Can the derivation in this paper for FNOs be generalized to more general neural operators?

Our theory may be applicable to several variants of the FNO. Examples include the following:

• geo-FNO: For geo-FNO, the influence of deformation is limited to the beginning and the end of the architecture, with the majority of the architecture consisting of the original FNO. Thus, our theory can be directly applied.
• G-FNO: For G-FNO, similar results can be obtained if the initial distribution of the transformed weight $\hat{L}_{s_g} R$ by elements $s_g$ from the stabilizer $S_g$ satisfies the assumptions of our theory.
• WNO: Since the wavelet transform is a linear operator, the same analysis can potentially be applied to the WNO.

Conversely, applying our theory directly to the architectures of U-FNO and F-FNO is challenging due to the presence of skip-connections. In architectures with skip-connections, the degeneration of input information in the ordered phase does not occur, and chaotic behavior is expected to be prominent. Therefore, it is necessary to analyze the behavior of the input in the chaotic phase in detail, rather than focusing on the discussion of the fixed point. It may be possible to extend our theory by incorporating mean-field theory for DCNs with skip-connections (Yang & Schoenholz, 2017) alongside our theoretical framework. Once the discussion on skip-connections is resolved, similar arguments can be made as follows:

• U-FNO: For U-FNO, it is considered possible to discuss by integrating the results related to CNN, DCN, and simplified FNO, as discussed in Section 3.3 on the original FNO. Note that skip-connections are included in U-Net.
• F-FNO: The factorized Fourier convolution module (Eq. 8) can be discussed similarly to the conventional FNO. However, since the output of the inverse FFT is summed in the dimensional direction, it seems appropriate to consider the weight scale as 1/D.

---

Q4. Is it possible to identify (approximately) how the initialization depends on different datasets?

Our theory and experiments provide insight that, regardless of the dataset, loss function, optimizer, or scheduler, initialization around $\sigma^2 = 2.0$ is one of the necessary conditions (not sufficient condition) for stable training, especially for deep FNOs.

At the same time, we also agree with your experience that the performance is dependent on hyperparameters and datasets (tasks). As discussed in Appendix D, underfitting caused by a loss landscape with strong non-convexity and overfitting caused by an inappropriate setting of model capacity relative to the amount of data are both hyperparameter- and data-dependent. Our findings suggest that the acceptable initial variance of the weights is already limited by the number of layers, so these problems need to be addressed by optimisers, schedulers, devising loss functions, introducing regularization, etc. Thank you for your valuable comments.