Title: Understanding the Expressivity and Trainability of Fourier Neural Operator: A Mean-Field Perspective
Venue: NeurIPS 2024
Accepted: yes
Keywords: Fourier Neural Operator

PDE

Mean-Field Theory
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 \emph{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.

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.

1 Introduction

The recent surge in interest in solving partial differential equations (PDEs) has led to the use of neural network (NN)-based surrogate models. One promising line of work is the neural operator (NO), which learns the solution operator of PDEs, thereby bypassing the need for mesh dependency. Among the variants of NO, the Fourier neural operator (FNO) (Li et al., 2020c) has gained popularity because of its advantageous cost/accuracy trade-off. The FNO can capture long-distance spatial interactions using the Fourier transform, whereas convolutional neural networks (CNNs) (Wen et al., 2019; Jiang et al., 2021b) and message-passing graph neural networks (GNNs) (Li et al., 2020a,b) are limited to operating solely on local variables. From a computational cost perspective, the Fourier transform is performed in quasi-linear time by the fast Fourier transform (FFT), making it significantly faster than the Transformer (Li et al., 2023).

Despite the widespread use of FNO as an architecture, there is a lack of comprehensive theoretical analysis on its expressivity and trainability. The universal approximation property (Kovachki et al., 2021), recognized as the basic expressivity of the FNO, is well-known; however, the exponential expressivity depending on the weight distribution, which are known for the densely connected network (DCN) (Schoenholz et al., 2016), a.k.a. fully connected network, and CNN (Xiao et al., 2018), remains unexplored. Regarding the trainability of FNO, the training instability in deep FNO has been experimentally reported by Tran et al. (2022), but the causes and conditions of the difficulty have not been clarified either theoretically or experimentally.

Issei Sato

Masahiro Fujisawa RIKEN AIP

Ordered phase

Chaotic phase

heterogeneous

Exploding

gradients

Vanishing

gradients

homogeneous

the edge of chaos

Backpropagation

2

2

*

H(1)

H(L)

H(0)

92

V

*

H(1)

H(L)

We analyze the exponential expressivity (how far can two different input vectors be pulled apart) and trainability (how much gradient explosion on average) of the random FNO from the perspective of whether the network is ordered or chaotic . This viewpoint is grounded in mean-field theory, an analytical framework for NN established by Poole et al. (2016); Schoenholz et al. (2016); Yang & Schoenholz (2017); Xiao et al. (2018). A network is considered ordered when it brings all representations of two different spatial positions closer together, and chaotic when it drives them apart during forward propagation. Furthermore, a network can only be stably trained when initialized close to the edge of chaos , which is the transition point between the ordered phase and the chaotic phase. In fact, He initialization (He et al., 2015) is an example of a commonly used edge of chaos initialization for the DCN with ReLU activation (Burkholz & Dubatovka, 2019).

In this study, we establish a mean-field theory to analyze the expressivity and trainability of the FNO. Our investigation reveals the expressivity of random FNO at initialization by examining the transition point between ordered and chaos phases. The phase transition exhibits FNO-specific characteristics induced by mode truncation, as well as similarities with the characteristics of DCN and CNN. We also find a link between expressivity and trainability: the ordered and chaotic phases correspond to regions of vanishing and exploding gradient, respectively. This discovery offers a practical initialization prerequisite for the stable training of the FNO.

2 Background

2.1 Fourier Neural Operators

The FNO (Li et al., 2020c) is one of the well-established methods for solving PDEs across many scientific problems (Yang et al., 2021; Wen et al., 2022b; Hwang et al., 2022; Jiang et al., 2021a; Pathak et al., 2022). An M -dimensional FNO with the number of hidden features D and the spatial size N for learning the operators between scalar-valued functions is defined as follows.

M

where X ( ℓ ) ∈ R ︷ ︸︸ ︷ N ×···× N × D is the ℓ -th hidden representation, M is the number of spatial dimensions, and ϕ is activation. The hidden representations X (0) and X ( L ) are the output of the lifting operator and the input of the projection operator, respectively. The architecture of these operators does not affect our analysis as long as the network stays shallow, as implemented in (Li et al., 2020c). The ℓ -th densely connected (DC) module D ( ℓ ) is an affine point-wise map in the physical space and the ℓ -th

Fourier convolution module K ( ℓ ) is a parameterized kernel integral operator using Fourier transform. The bias term is considered to be included in either or both modules K ( ℓ ) and D ( ℓ ) .

Several FNO variants have been developed to address specific challenges, such as geo-FNO (Li et al., 2022) for irregular regions and group equivariant FNO (G-FNO) (Helwig et al., 2023), which maintains equivariance to rotation, reflection, and translation. U-NO (Rahman et al., 2022) and U-FNO (Wen et al., 2022a) integrate FNO with U-Net for multiscale modeling. Additionally, WNO (Tripura & Chakraborty, 2022) utilizes wavelet bases, while CFNO (Brandstetter et al., 2023) enhances the use of geometric relations between different fields and field components through Clifford algebras. Adaptive FNO (Zhao et al., 2022; Guibas et al., 2021) 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 (Li et al., 2024), data augmentation (Brandstetter et al., 2022), 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 (Zhao et al., 2022) and training instability (Tran et al., 2022) have been reported. Tran et al. (2022) observed that training did not converge even at 24 layers. They successfully addressed the stability and accuracy degradation issues associated with an increase in the number of layers by implementing skip connections behind activation and introducing various training techniques. However, it is still unknown that the theoretical basis for why the original architecture of the FNO has problems with training instability and accuracy degradation.

2.2 Mean-field Theory for Neural Networks

The mean-field theory has been used to provide a mathematical framework for understanding the expressivity and trainability of neural networks (Poole et al., 2016; Schoenholz et al., 2016; Yang &Schoenholz, 2017; Hayou et al., 2018; Xiao et al., 2018). A series of papers (Poole et al., 2016; Schoenholz et al., 2016) delved into the average behavior of infinite-width random deep DCN, with weights and biases initialized by a zero-mean Gaussian distribution. The formulation is given below.

where x ( ℓ ) ∈ R D is the ℓ -th hidden representation, W ( ℓ ) ∈ R D × D , b ( ℓ ) ∈ R D are the ℓ -th learnable parameters, and the width is assumed to be sufficiently large D ≫ 1 .

Poole et al. (2016) and Schoenholz et al. (2016) explored the exponential expressivity of random DCN determined by two phases depending on the initial variance parameters σ 2 and σ 2 b , as shown in Fig. 5a. Poole et al. (2016) first examined the forward propagation of a random DCN with Tanh activation. They demonstrated that the covariance Σ ( ℓ ) of the ℓ -th pre-activation representations h ( ℓ ) and ˜ h ( ℓ ) corresponding to two different inputs x (0) and ˜ x (0) are obtained by

where the expectation is taken over the pre-activations [ h ( ℓ -1) d , ˜ h ( ℓ -1) d ] ∼ N ( 0 , Σ ( ℓ -1) ) . The covariance converges exponentially to a fixed point Σ ∗ determined by parameters σ 2 and σ 2 b .

A network is considered ordered when it brings two distinct representations closer together, which implies a state of small expressivity. Conversely, a network is chaotic when it drives them apart during forward propagation, implying a state of large expressivity. Networks with either excessively small or large expressivity can disrupt the structure of the input: the difference between two distinct inputs quickly becomes indistinguishable in networks with small expressivity, while similarities between inputs are no longer recognized in networks with large expressivity. The network is ordered if the initial variance of the weights is small. For larger values, and beyond a certain threshold, the phase shifts, and the network behaves chaotically. This phase shift point is termed the edge of chaos .

Subsequently, Schoenholz et al. (2016) discovered the connection between expressivity and trainability of DCN through analysis of the backpropagation behavior. In an ordered network, the expected

value of the gradient norm becomes exponentially small during backpropagation, while it becomes exponentially large in a chaotic network. This implies that the gradient vanishes/explodes in ordered or chaotic networks, respectively. These findings suggest that deep DCN can be stably trained only near the edge of chaos. Schoenholz et al. (2016) also provided an estimate of the maximum depth at which a network can be trained when initialized away from the edge of chaos. These insights are not limited to DCN and similar results have been observed for residual networks (Yang & Schoenholz, 2017) and CNN (Xiao et al., 2018).

3 A mean-field theory for FNO

In this section, we establish a mean-field theory of FNO. We demonstrate the exponential expressivity of random FNO by examining the ordered-chaos phase transition during the forward propagation. Furthermore, we identify the connection between expressivity and trainability by concentrating on backward propagation behaviors. Our analysis is an advanced version of the approach developed by Poole et al. (2016); Schoenholz et al. (2016); Yang & Schoenholz (2017); Xiao et al. (2018). In Section 3.1, we outline the problem setup. In Section 3.2, we analyze the forward and backward propagation behavior of random FNO at initialization. In Section 3.3, we discuss the practical prerequisites for initialization to stabilize the training of FNO, leveraging the similarities between FNO and DCN. The proofs for all the lemmas and theorems are provided in Appendices A and B.

3.1 Problem setting

Here, we consider a simplified one-dimensional (1D) FNO. Note that our theory is extensively applicable to the original FNO, as discussed in Section 3.3. The simplified 1D FNO, with a depth of L , is defined by the number of hidden features D , a spatial size N = 2 m (where m is an integer), the number of Fourier modes K ≤ N 2 + 1 , two learnable weights Θ ( ℓ,k ) ∈ R D × D and Ξ ( ℓ,k ) ∈ R D × D , and a bias b ( ℓ ) ∈ R D . Denote ϕ : R → R by the non-decreasing activation function. Let X ( ℓ ) ∈ R N × D and H ( ℓ ) ∈ R N × D be the post and pre-activation representations defined by

where δ a,b is the Kronecker-delta, c k = 2 -δ k, 0 -δ k,N/ 2 is a constant, 1 N is all-ones column vector with the size N , H ( ℓ,k ) is the conjugate of H ( ℓ,k ) corresponding to the ( N -k ) -th frequency components, † is the transpose conjugate, F ∈ C N × N is the Discrete Fourier Transform (DFT) matrix defined by F k,n = 1 N exp( -2 πk N n ) , and D ( k ) is a diagonal matrix with a 1 at position D ( k ) k,k .

There are two differences from the original FNO proposed by Li et al. (2020a): (1) the DC module is dropped for the simplicity, and (2) H ( ℓ,k ) is multiplied by √ 2 with respect to k = 0 , N 2 for appropriate normalization. We assume that the weights of FNO are initialized by i.i.d. samples from Gaussian distribution, i.e . Θ ( ℓ,k ) i,j i.i.d. ∼ N (0 , σ 2 2 D ) , Ξ ( ℓ,k ) i,j i.i.d. ∼ N (0 , σ 2 2 D ) , b ( ℓ ) i i.i.d. ∼ N (0 , σ 2 b ) . For k = 0 , N 2 , the parameter Ξ ( ℓ,k ) is set to zero exceptionally. For all d ∈ [ D ] = { 0 , . . . , D -1 } , the pre-activations H ( ℓ ) : ,d ∈ R N are i.i.d. random variables. When D ≫ 1 , by the central limit theorem, the variables H ( ℓ ) : ,d follow Gaussian distribution with mean 0 and covariance matrix Σ ( ℓ ) α,α ′ := E Θ 1: ℓ , Ξ 1: ℓ [ H ( ℓ ) α,d H ( ℓ ) α ′ ,d ] , where the expectation is taken over all random variables [ Θ 1: ℓ , Ξ 1: ℓ ] := { Θ ( ℓ ′ ,k ′ ) , Ξ ( ℓ ′ ,k ′ ) } ℓ ′ ∈ [ ℓ ] ,k ′ ∈ [ K ] . Our theory can be easily extended to 2D and 3D FNOs.

3.2 Expressivity and trainability of FNO

Firstly, the forward propagation of a single input signal with spatial features is described as follows.

Lemma 3.1 (Iterated map) . For all d ∈ [ D ] , the covariance Σ ( ℓ ) := E Θ 1: ℓ , Ξ 1: ℓ [ H ( ℓ ) : ,d H ( ℓ ) : ,d ⊤ ] is obtained recursively by the iterated map C defined by

where the expectation is taken over the pre-activations H : ,d ∼ N (0 , Σ ( ℓ -1) ) , θ ( k ) α,α ′ := 2 πk N ( α -α ′ ) represents the scaled positional difference.

The indices α and α ′ correspond to different spatial locations as with the mean-field theory for CNN (Xiao et al., 2018). Note that [ F ϕ ( H : ,d )] k is the k -th Fourier modes of the post-activation representation. When applying DCN (Poole et al., 2016; Schoenholz et al., 2016) or CNN (Xiao et al., 2018) to the spatial signal, the iterated map depends only on local spatial locations, while in the case of FNO, the iterated map depends on all spatial locations because of the global Fourier convolution. In addition, only periodic spatial correlations with shift-invariant are propagated, and high-frequency components exceeding mode K are eliminated.

Next, we explore the fixed point Σ ∗ of the iterated map C satisfying Σ ∗ = C ( Σ ∗ ) . By linearizing the dynamics of signal propagation around this fixed point and analyzing the stability and rate of convergence to the fixed point, we can determine the depth to which each component of the input can propagate. Schoenholz et al. (2016) showed that the iterated map of DCN defined in Eq. (2) has a fixed point of the form:

where q ∗ , c ∗ are the fixed points of variance and correlation, and I N is the identity matrix. Meanwhile, Xiao et al. (2018) showed that any fixed point for the iterated map of the DCN is also a fixed point for that of CNN. We show that random FNO has the same fixed points of the form of Eq. (5) with c ∗ = 1 in the following lemma.

Lemma 3.2 (Exsistance of fixed points) . When a random DCN defined by Eq. (2) has the fixed point ( q ∗ , c ∗ = 1) for the initial parameters ( σ 2 , σ 2 b ) , then a random simplified FNO defined by Eq. (3) has a fixed point Σ ∗ of the form

Lemma 3.2 indicates that the fixed point for the iterated map Σ ∗ of the DCN serves as a fixed point for the iterated map of the simplified FNO (as well as CNN). To analyze the stability and convergence rate, we linearly approximate the C-map around the fixed point Σ ∗ , i.e., C ( Σ ) ≈ Σ ∗ + J Σ ∗ ( Σ -Σ ∗ ) , where J Σ ∗ is the Jacobian linear map of the iterated map defined in Eq. (15). We then derive the eigenvalues and eigenvectors for the Jacobian linear map J Σ ∗ as follows.

Definition 3.3.

where the expectation is taken over the pre-activations H : ,d ∼ N (0 , Σ ∗ ) , and ϕ ′ , ϕ ′′ are the firstand second-order derivatives of the activation ϕ .

The bases ψ , ψ (1) , . . . , ψ ( K -1) ∈ R N × N using the above quantities are defined below.

From Lemma A.4, K -1 matrices in { ψ ( k ) } k ∈ [ K ] \{ 0 } are eigenbases with the eigenvalue χ c ∗ of the Jacobian linear map. From Lemma A.5, the matrix ψ is the eigenbases with the eigenvalue χ of the Jacobian linear map. Since the rank of the Jacobian linear map is at most K (Lemma A.3), the deviation from the fixed point Σ ( ℓ ) -Σ ∗ is spanned by K-dimensional eigenspace span ( { ψ ( k ) } k ∈ [ K ] \{ 0 } ∪ { ψ } ) . Then, the fixed point stability and the convergence rate are shown in the following theorem.

Theorem 3.4 (Exponential expressivity) . Let E ( ℓ ) := Σ ( ℓ ) -Σ ∗ be the deviation from the fixed point at the ℓ -th layer. Suppose that the deviation at the first layer is decomposed as E (0) = ϵ ψ + ∑ K -1 k =1 ϵ k ψ ( k ) + e . The scalars ϵ, ϵ k represent the scale of the perturbation for each eigencomponent of the linearly approximated map E ( ℓ ) ↦→ E ( ℓ +1) . The component e ∈ R N × N belongs to the orthogonal complements of the space span ( { ψ , ψ (1) , . . . , ψ ( K -1) } ) .

Then, the deviation at the ℓ -th layer is obtained by

In particular, when the Fourier mode is K = N 2 +1 , Eqs. (9) and (10) reduce to the following.

Theorem 3.4 shows the expressivity of the FNO, which is characterized by the ordered-chaos phases and varies exponentially with respect to the number of layers. Theorem 3.4 indicates that the asymptotic behavior of the zero-frequency deviation is mostly determined by χ and the periodic deviation is determined by χ c ∗ . If χ < 1 and χ c ∗ < 1 , the fixed point is stable as the deviation from the fixed point converges exponentially to zero. When the fixed point remains stable at c ∗ = 1 , a random network exists in an ordered phase, where all spatial representations are correlated in an asymptotic manner. Conversely, when the fixed point with c ∗ = 1 becomes unstable, the network transitions into a chaotic phase, exhibiting behavior dependent on the activation function ϕ . The boundary between these two phases is referred to as the edge of chaos .

The convergence rates χ q ∗ and χ c ∗ are the same as the convergence rates of the variance and correlation to the fixed point for DCN (Schoenholz et al., 2016) and CNN (Xiao et al., 2018). However, only periodic spatial correlations are propagated in the FNO, resulting in a different eigenspace of the map E ( ℓ ) ↦→ E ( ℓ +1) compared to the DCN and CNN. In DCN, the deviation belongs to a vector space with dimension N ( N -1) 2 in DCN, whereas in FNO, the dimension is K , or at most N 2 +1 . CNN possess diagonal eigenspaces associated with eigenvalues χ q ∗ and non-diagonal eigenspaces associated with eigenvalues χ c ∗ . In contrast, FNOs without mode truncation exhibit a similarity, possessing eigenspaces χ q ∗ for zero-frequency and eigenspaces χ c ∗ for k-frequencies with diagonal components removed. Furthermore, mode truncation increases the convergence rate of zero-frequency deviation from χ q ∗ to χ and affects all eigenbases as well. For further discussions on the similarities between CNN and FNO, please refer to Appendix C. A visualization of the covariance of the FNO with Tanh and ReLU activations is shown in Appendix F.

Finally, we demonstrate the connection between expressivity and trainability. By examining the covariance of the gradient in each layer during backpropagation, we investigate the conditions under which training is stable without gradient vanishing or exploding.

Theorem 3.5 (Trainability) . Let ˜ Σ ( ℓ ) ∈ R N × N be the gradient covariance with respect to some loss L , e.g . mean squared error, at the ℓ -th layer. Suppose that the gradient covariance at the L -th layer is decomposed as ˜ Σ ( L ) α,α ′ = ∑ K -1 k =0 ˜ ϵ k cos ( θ ( k ) α,α ′ ) + ˜ e , where ˜ ϵ k is the coefficient of each basis and ˜ e

(a) Simplified FNO with Tanh

(b) Simplified FNO with ReLU

(c) Original FNO with ReLU

10'

I0/1

1010

1 wuou yuspes8 parenbg

102

105

0.5

1.5

1.0

quared gradient norm

105

103

0.5

1.5

1.0

L

ou

106

102

0.5

1.5

1.0

108

2.0

2.5

1013

2.0

2.5

gradient

10~14

2.0

2.5

10-11

3.0

10-21

3.0

10~22

3.0

1014

3.5

4.0

1029

3.5

4.0

paren

1030

3.5

4.0

1017

10

37

10

10~20

0

10

20

30

0

50

60

0

10

20

40

50

60

0

10

20

30

40

50

09

Layer

Layer

Layer

belongs to the orthogonal complements of span( { cos ( θ ( k ) α,α ′ ) } K -1 k =0 ) . Then, the gradient covariance at the ℓ -th layer is obtained by

Theorem 3.5 shows that gradient vanishing occurs when χ c ∗ < 1 (ordered phase) and gradient explosion occurs when χ c ∗ > 1 (chaos phase). Thus, stable training of the FNO can be achieved close to the edge of chaos by setting the initial parameter σ 2 to satisfy χ c ∗ ≈ 1 . We specifically present the initial parameter choices that achieve the edge of chaos for several FNOs in Section 3.3.

When c ∗ = 1 , there is no change in the dynamics during backpropagation due to mode truncation. When using the full mode K = N 2 +1 , the condition χ c ∗ = 1 always achieves the edge of chaos , which is consistent with the results for the DCN and CNN. Despite the architectural and iterative map differences among FNO, DCN, and CNN, Theorem 3.4 and Theorem 3.5 demonstrate the similarities in the random behavior of FNO, DCN, and CNN. This allows existing results based on mean-field theory to be applied to the FNO.

3.3 Initialization requirements for stable training

For stable training, Theorem 3.5 suggests the necessity of initializing FNO near the edge of chaos , i.e ., initializing FNO so that χ c ∗ ≈ 1 . In this section, we present the initial parameter choices that achieve the edge of chaos for several FNOs, each with slightly different architectures such as activation functions. Furthermore, the behavior of the gradient norm Tr( ˜ Σ ( ℓ ) ) /D as a function of layer ℓ are visualized in Fig. 2 for several variants of random FNO with different initialization parameters σ 2 . We used FNO with a width of D = 32 and a number of layers L = 64 , and for simplicity, we computed the absolute value of the output as the loss with respect to the input sampled from the standard normal distribution.

Simplified FNO with Tanh activation .

The behavior of χ c ∗ for the parameters σ 2 and σ 2 b of the Tanh network has been extensively studied by Poole et al. (2016); Schoenholz et al. (2016). The phase diagram drawn by Pennington et al. (2017) is shown in Fig. 5a. By using parameters ( σ 2 , σ 2 b ) around the two phase boundaries of ordered and chaotic that achieve χ c ∗ = 1 , the training of the simplified FNO with Tanh activation can be stabilized. Figure 2a depicts the behavior of the gradient backpropagation in the simplified FNO with Tanh activation and the bias parameter being σ 2 b = 0 . 1 . Figure 2a shows that when σ 2 ⪅ 2 , the gradient diminishes exponentially; otherwise, it explodes exponentially.

Simplified FNO with ReLU activation .

The iterated map C of the DCN with ReLU activation is given by Cho & Saul (2009) as follows.

The edge of chaos initialization for the DCN with ReLU activation is known as He initialization (He et al., 2015), which sets the initial variance parameter as σ 2 = 2 and initial bias as b ( ℓ ) i = 0 for Eq. (2). From the similarity of the DCN and the FNO, we can derive the FNO version of the He initialization that achieves χ c ∗ = 1 by setting σ 2 = 2 , b ( ℓ ) i = 0 for Eq. (3). The He initialization scheme for the simplified FNO with the activation ϕ = ReLU is derived as follows.

Figure 2b demonstrates that the choice of σ 2 = 2 preserves the magnitude of the gradient norm during backpropagation of deep simplified FNO with ReLU activation.

Original FNO .

In the original architecture of the FNO proposed by Li et al. (2020c), the DC module is used together with the Fourier convolution module as shown in Eq. (1). We initialize the weights of both layers consistently as follows. For all ℓ ∈ [ L ] , k ∈ [ K ] , and i, j ∈ [ D ] ,

From the similarity of the initial network behavior of the FNO and the DCN, the fixed point with c ∗ = 1 of the simplified FNO is also a fixed point of the original FNO. In the neighborhood of the fixed point Σ ∗ , the eigencomponents spanned by { ψ ( k ) } K k =1 will decay or increase at the rate of χ c ∗ . Following the derivation of Theorem 3.5, the eigenvalues of the cos function eigencomponents of the gradient covariance are also χ c ∗ . These results show that the edge of chaos initialization scheme can be used for the original FNO with each activation function. Figure 2c shows that the original FNO with ReLU activation and the parameter fixed as b ( ℓ ) i = 0 exhibits similar backpropagation behavior as the simplified FNO, i.e ., σ 2 ≈ 2 is an appropriate choice.

4 Experiments

In this section, we experimentally demonstrate that deep FNO training requires appropriate initialization settings on a variety of datasets, consistent with the theory discussed in Section 3.

4.1 Datasets

We evaluated three models on commonly used PDEs: the advection equation, Burgers' equation, Darcy Flow equation, and incompressible Navier-Stokes (NS) equation. All datasets were generated by numerical simulations used in (Takamoto et al., 2022; Li et al., 2020c) and are publicly available. A summary of the dataset is provided in Table 1, and more details are given in Appendix D.

Advection equation and Burgers' equation . The linear advection equation for the function u ( x, t ) is given by

where u 0 is the initial condition and β = 2 . 0 is an advection speed. The non-linear Burgers' equation for the function u ( x, t ) is given by

where u 0 is the initial condition and ν = 4 . 0 is the diffusion coefficient.

Darcy Flow equation . The Darcy Flow equation for the function u ( x ) with a Dirichlet boundary is given by

where a is the diffusion coefficient and f ( x ) = 1 is the forcing function.

where ω ( x, t ) is the vorticity, ω 0 is the initial vorticity, u ( x, t ) is the velocity field for any r > 0 , ν ∈ R + is the viscosity, and f is the external forcing function defined by f ( x ) = 0 . 1 (sin(2 π ( x 1 + x 2 ) + cos(2 π ( x 1 + x 2 )) . We experimented with the viscosities ν = 1 e -3 , 1 e -4 , 1 e -5 . For the data with ν = 1 e -4 , the time resolution was also downsampled by half.

4.2 Experimental Settings

In the experiments on the 1D advection and Burgers' equation, we compared the results of simplified FNOs defined in Eq. (3) with Tanh and ReLU activations for varying number of layer L and initial parameters σ 2 . In the experiments on the 2D Darcy Flow and NS equations, we compared the results of the original FNO (Li et al., 2020c) as shown in Eq. (1) for varying number of layer L and initial parameters σ 2 . All models have a width of D = 32 and the Fourier mode of K = 12 . When using Tanh as the activation, we fixed the initial bias parameter σ 2 b = 0 . 1 , and when using ReLU activation, we fixed the initial bias b ( ℓ ) i = 0 . We used the AdamW optimizer and cosine annealing scheduler. Training was stopped early at the epoch of minimal normalized mean squared error (nMSE) on the validation data. Details of the experimental setup are given in Appendix D. In the experiments on the NS equation, we trained the original FNO with the autoregressive scheme using a teacher-forcing technique, input normalization, and regularization that adds small-Gaussian noise to the input (Tran et al., 2022). During the evaluation phase, only the initial 10 steps are provided as input, and the rollout results of all subsequent steps are evaluated. For all tasks, the mean squared error (MSE) is used as the training loss and nMSE as the validation and testing (Li et al., 2020a; Tran et al., 2022).

4.3 Results

Heatmaps of the training loss measured by MSE at the last epoch and the test performance measured by nMSE on six different datasets, for different depths L and initial parameters σ 2 , are shown in Figs. 3 and 4, respectively. Despite the differences in architectures and datasets, Figure 3 shows the same trend supporting the theory in all experiments. As the number of layers L increases, the range of acceptable initial σ 2 value settings becomes narrower, and initialization near the edge of chaos ( σ 2 ≈ 2 ) is required for stable training of deep FNO. Detailed analyses on training loss and test performance are presented in Appendix E.1 and Appendix E.2, respectively.

5 Conclusion

In this study, we developed a mean-field theory for FNO. We showed the expressivity and trainability of the FNO, which is characterized by the ordered-chaos phases. Furthermore, we observed both

100

100

100

001

30

30

30

I0I

30

1OI

25

10-1

25

10-1

25

102

25

10-2

10-3

20

15

10

103

102

20

15

10

102

104

103

20

15

10

c-01

104

105

106

20

15

10

10-4

106

OI

g~OI

5

104

105

5

10-7

109

1

2

3

4

1

2

3

4

1

2

3

4

1

2

3

30

105

+01

30

10

102

30

01

30

102

25

25

100

25

10

25

10

20

106

20

0I

20

20

10-7

102

100

100

15

108

15

103

15

15

10

109

10

104

105

10

10-1

10

10I

1010

5

106

5

102

5

102

2

3

4

1

2

3

4

1

2

4

1

2

3

4

9?

Figure 3: Training loss of FNOs at last epoch for four distinct PDEs. (a, b): the advection equation, (c, d): the Burgers' equation, (e): Darcy Flow, (f-h): the NS equation. The heatmaps represents the training loss values for varying depth L ∈ { 4 , 8 , 16 , 32 } and initial weight parameter σ 2 ∈ { 0 . 1 , 0 . 5 , 1 . 0 , 2 . 0 , 3 . 0 , 4 . 0 } , with lighter colors signifying lower training loss. The presented results are the mean training loss at the last epoch over three different seeds.

30

101

30

101

30

101

30

01

25

25

25

100

25

100

100

100

20

20

20

20

10-1

10-1

15

101

15

10-1

15

15

10

10

10

102

10

102

5

102

102

5

103

103

1

3

2

3

1

2

3

1

3

30

100

30

109

30

0.9

1.0

30

0.9

1.0

25

25

10-1

25

- 0.8

0.7

25

0.8

0.7

20

15

10-1

20

15

102

20

15

0.6

0.5

0.4

20

15

0.5

0.6

0.4

0.3

10

10

10

0.3

10

0.2

0.2

0.1

102

103

5

0.1

5

0.0

1

2

3

4

N

3

4

2

3

4

1

2

3

4

Figure 4: nMSE of FNOs on test datasets for four distinct PDEs. (a, b): the advection equation (c, d): the Burgers' equation, (e): Darcy Flow, (f-h): the NS equation. The heatmaps for each nMSE correspond to the results of each heatmap of training loss in Fig. 3. The lighter colors, the better the test performance. The presented results are the mean nMSE calculated over three different seeds.

unique FNO-specific behaviors caused by mode truncation, as well as common behaviors akin to those of DCN. With our analysis as a basis, we identified the necessity of initializing FNO near the edge of chaos for stable training of the FNO. Experimental results supported our theoretical results.

A limitation of our analysis is that it is limited to the network at initialization and does not address the stability of the entire optimization process. While we do not provide sufficient conditions for stable training, we do offer one necessary condition for achieving stable training. Future work may consider a mean-field analysis of the FNO when using skip-connection (Tran et al., 2022), Dropout and batch normalization, as well as initialization methods that ensure the input-output Jacobian of the FNO satisfies dynamical isometry (Pennington et al., 2017, 2018).

References

Brandstetter, J., Welling, M., and Worrall, D. E. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning , pp. 2241-2256. PMLR, 2022.

Brandstetter, J., van den Berg, R., Welling, M., and Gupta, J. K. Clifford neural layers for pde modeling. In The Eleventh International Conference on Learning Representations , 2023.

Burkholz, R. and Dubatovka, A. Initialization of relus for dynamical isometry. Advances in Neural Information Processing Systems , 32, 2019.

Cho, Y. and Saul, L. Kernel methods for deep learning. Advances in neural information processing systems , 22, 2009.

Guibas, J., Mardani, M., Li, Z., Tao, A., Anandkumar, A., and Catanzaro, B. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13587 , 2021.

Hayou, S., Doucet, A., and Rousseau, J. On the selection of initialization and activation function for deep neural networks. arXiv preprint arXiv:1805.08266 , 2018.

He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision , pp. 1026-1034, 2015.

Helwig, J., Zhang, X., Fu, C., Kurtin, J., Wojtowytsch, S., and Ji, S. Group equivariant fourier neural operators for partial differential equations. arXiv preprint arXiv:2306.05697 , 2023.

Hwang, R., Lee, J. Y., Shin, J. Y ., and Hwang, H. J. Solving pde-constrained control problems using operator learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, No. 4, pp. 4504-4512, 2022.

Jiang, P., Meinert, N., Jordão, H., Weisser, C., Holgate, S., Lavin, A., Lutjens, B., Newman, D., Wainright, H., Walker, C., et al. Digital twin earth-coasts: Developing a fast and physics-informed surrogate model for coastal floods via neural operators. In Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021) , 2021a.

Jiang, Z., Tahmasebi, P., and Mao, Z. Deep residual u-net convolution neural networks with autoregressive strategy for fluid flow predictions in large-scale geosystems. Advances in Water Resources , 150:103878, 2021b.

Kovachki, N., Lanthaler, S., and Mishra, S. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Research , 22(290):1-76, 2021.

Li, H., Xu, Z., Taylor, G., Studer, C., and Goldstein, T. Visualizing the loss landscape of neural nets. Advances in neural information processing systems , 31, 2018.

Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03485 , 2020a.

Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Stuart, A., Bhattacharya, K., and Anandkumar, A. Multipole graph neural operator for parametric partial differential equations. Advances in Neural Information Processing Systems , 33:6755-6766, 2020b.

Li, Z., Kovachki, N. B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., Anandkumar, A., et al. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations , 2020c.

Li, Z., Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209 , 2022.

Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research , 2023.

Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, A. Physics-informed neural operator for learning partial differential equations. ACM/JMS Journal of Data Science , 1(3):1-27, 2024.

Pathak, J., Subramanian, S., Harrington, P., Raja, S., Chattopadhyay, A., Mardani, M., Kurth, T., Hall, D., Li, Z., Azizzadenesheli, K., et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214 , 2022.

Pennington, J., Schoenholz, S., and Ganguli, S. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. Advances in neural information processing systems , 30, 2017.

Pennington, J., Schoenholz, S., and Ganguli, S. The emergence of spectral universality in deep networks. In International Conference on Artificial Intelligence and Statistics , pp. 1924-1932. PMLR, 2018.

Poli, M., Massaroli, S., Berto, F., Park, J., Dao, T., Ré, C., and Ermon, S. Transform once: Efficient operator learning in frequency domain. Advances in Neural Information Processing Systems , 35: 7947-7959, 2022.

Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., and Ganguli, S. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems , 29, 2016.

Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-no: U-shaped neural operators. arXiv preprint arXiv:2204.11127 , 2022.

Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. Deep information propagation. In International Conference on Learning Representations , 2016.

Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger, D., and Niepert, M. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems , 35:1596-1611, 2022.

Tran, A., Mathews, A., Xie, L., and Ong, C. S. Factorized fourier neural operators. In The Eleventh International Conference on Learning Representations , 2022.

Tripura, T. and Chakraborty, S. Wavelet neural operator: a neural operator for parametric partial differential equations. arXiv preprint arXiv:2205.02191 , 2022.

Wen, G., Tang, M., and Benson, S. M. Multiphase flow prediction with deep neural networks. arXiv preprint arXiv:1910.09657 , 2019.

Wen, G., Li, Z., Azizzadenesheli, K., Anandkumar, A., and Benson, S. M. U-fno-an enhanced fourier neural operator-based deep-learning model for multiphase flow. Advances in Water Resources , 163:104180, 2022a.

Wen, G., Li, Z., Long, Q., Azizzadenesheli, K., Anandkumar, A., and Benson, S. M. Accelerating carbon capture and storage modeling using fourier neural operators. arXiv preprint arXiv:2210.17051 , 3, 2022b.

Xiao, L., Bahri, Y., Sohl-Dickstein, J., Schoenholz, S., and Pennington, J. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning , pp. 5393-5402. PMLR, 2018.

Yang, G. and Schoenholz, S. Mean field residual networks: On the edge of chaos. Advances in neural information processing systems , 30, 2017.

Yang, Y., Gao, A. F., Castellanos, J. C., Ross, Z. E., Azizzadenesheli, K., and Clayton, R. W. Seismic wave propagation and inversion with neural operators. The Seismic Record , 1(3):126-134, 2021.

Zhao, J., George, R. J., Zhang, Y., Li, Z., and Anandkumar, A. Incremental fourier neural operator. arXiv preprint arXiv:2211.15188 , 2022.

0.25

0.25

Ordered

1.5

0.20

0.20

c*=1

0.15

Vanishing gradients

1.0

26

0.15

Ordered

Chaotic

2.6

b

bounded, c* = 1

unbounded, c* = 1

b

0.5

0.10

Vanishing gradients

Exploding gradients

0.10

Chaotic

0.05

c*<1

0.0

0.05

Exploding gradients

0.00

0.5

1.0

1.5

2.0

2.5

3.0

0.00

0.5

1.0

1.5

2.0

2.5

3.0

92

9

(a) Tanh activation (Pennington et al., 2017)

(b) ReLU activation

A Proof of Theorem 3.4

Theorem 3.4 (Exponential expressivity) . Let E ( ℓ ) := Σ ( ℓ ) -Σ ∗ be the deviation from the fixed point at the ℓ -th layer. Suppose that the deviation at the first layer is decomposed as E (0) = ϵ ψ + ∑ K -1 k =1 ϵ k ψ ( k ) + e . The scalars ϵ, ϵ k represent the scale of the perturbation for each eigencomponent of the linearly approximated map E ( ℓ ) ↦→ E ( ℓ +1) . The component e ∈ R N × N belongs to the orthogonal complements of the space span ( { ψ , ψ (1) , . . . , ψ ( K -1) } ) .

Then, the deviation at the ℓ -th layer is obtained by

In particular, when the Fourier mode is K = N 2 +1 , Eqs. (9) and (10) reduce to the following.

Proof. The theorem is obtained by the eigenvalue analysis on the the first-order Taylor approximation of the iterated map C at the fixed point Σ ∗ . The Jacobian matrix J ( Σ ∗ ) ∈ R N 2 × N 2 of the iterated map at the fixed point and the Jacobian linear map J Σ ∗ ( · ): R N × N → R N × N are defined as follows.

where vec performs the vectorization, i.e . transforming an N × N matrix into a vector of size N 2 , and mat performs the inverse operation of vec .

From Lemma A.3, K -1 matrices in { vec( ψ ( k ) ) } k ∈ [ K ] \{ 0 } are eigenbases with the eigenvalue χ c ∗ of the Jacobian linear map. From Lemma A.4, the matrix ψ is the eigenbases with the eigenvalue χ of the Jacobian linear map. Since the sets of K matrices in { ψ ( k ) } k ∈ [ K ] \{ 0 } ∪ { ψ } are linearly

independent (yet non-orthogonal) in R N × N , the subspace span ( { ψ ( k ) } k ∈ [ K ] \{ 0 } ∪ { ψ } ) is the K -dimensional eigenspace of the Jacobian linear map. From Lemma A.2, the rank of the Jacobian matrix J ( Σ ∗ ) at the fixed point Σ ∗ is at most K , thereby the rank of the Jacobian linear map in Eq. (15) is at most K . Therefore, we have

Lemma 3.1 (Iterated map) . For all d ∈ [ D ] , the covariance Σ ( ℓ ) := E Θ 1: ℓ , Ξ 1: ℓ [ H ( ℓ ) : ,d H ( ℓ ) : ,d ⊤ ] is obtained recursively by the iterated map C defined by

where the expectation is taken over the pre-activations H : ,d ∼ N (0 , Σ ( ℓ -1) ) , θ ( k ) α,α ′ := 2 πk N ( α -α ′ ) represents the scaled positional difference.

Proof. For simplicity, we introduce Y α,k,β,i as follows.

̸

Since the weights are sampled independently, for different k = k ′ ,

From Eq. (3), we have

̸

First, we calculate the term k = 0 , N 2 , where c k = 2 .

where ˆ X ( ℓ -1) k,i is the k -th Fourier mode of the representaion X ( ℓ -1) : ,i .

Since the Fourier modes | ˆ X ( ℓ -1) k,i | 2 are i.i.d. for each hidden dimension i ∈ [ D ] , we have

To obtain a more tractable expression in the subsequent proofs of theorems, we express the iterated map without using Fourier modes as follows.

Lemma 3.2 (Exsistance of fixed points) . When a random DCN defined by Eq. (2) has the fixed point ( q ∗ , c ∗ = 1) for the initial parameters ( σ 2 , σ 2 b ) , then a random simplified FNO defined by Eq. (3) has a fixed point Σ ∗ of the form

Proof. Using the properties of cosine functions, the following holds.

Then, the following holds for all α, α ′ ∈ [ N ] .

When the DCN is applied pointwise to each spatial representation H β, : , β ∈ [ N ] , the iterated map of the random DCN (Poole et al., 2016) is given by C DCN ( Σ ) := σ 2 E H : ,d ∼N (0 , Σ ) [ ϕ ( H β,d ) ϕ ( H β ′ ,d )] + σ 2 b . Since the covariance Σ ∗ is a fixed point with respect to the iterative map of the DCN, i.e . C DCN ( Σ ∗ ) = Σ ∗ , the following holds.

Thus, we confirm that the covariance Σ ∗ satisfies the definition of a fixed point in the map C , i.e . Σ ∗ = C ( Σ ∗ ) . This means that the fixed point for the iterated map Σ ∗ of the DCN also serves as a fixed point for the iterated map of the simplified FNO.

̸

LemmaA.1. Let Σ ∗ be the fixed point of the form in Eq. (5) . Suppose that the symmetric perturbation E ∈ R N × N , where E β,β = E β ′ ,β ′ and E β,β ′ are non-zero for some β, β ′ ∈ [ N ] , β = β , and all other elements are zero. Then, we have

where χ q ∗ , χ c ∗ , and χ κ are the constants defined by

Proof. Equation (19) is obviously shown by the result of Section 2 in (Poole et al., 2016) and Section 3 in (Schoenholz et al., 2016). We prove Eq. (20) with reference to the results of (Schoenholz et al., 2016).

where ∫ D d z = 1 √ 2 π ∫ d ze -1 2 z 2 is the measure for a standard Gaussian distribution, u 1 = √ qz 1 , u 2 = √ q ( c β,β ′ z 1 + √ 1 -c 2 β,β ′ z 2 ) , q = q ∗ + E β,β and c β,β ′ = c ∗ + E β,β ′ q .

We consider the case where c ∗ < 1 and c ∗ = 1 separately. Later, we will show that the two results agree with each other. First, we consider the case where c ∗ < 1 . Using Taylor expansion, we can approximate ϕ ( u 1 ) and ϕ ( u 2 ) as follows. For the simplicity, we assume O ( | E β,β | ) = O ( | E β,β ′ | ) .

where u ∗ 1 = √ q ∗ z 1 and u ∗ 2 = √ q ∗ ( c ∗ z 1 + √ 1 -( c ∗ ) 2 z 2 ) . Thus, we have

The results of the second term are obtained from the transformation of equations 36 to 39 in Appendix 7.2 of (Schoenholz et al., 2016).

Utilizing the identity, ∫ D zzf ( z ) = ∫ D zf ′ ( z ) , we obtain the third term as follows.

The last term is calculated as follows.

Summing the last two terms, we obtain the term E β,β χ κ .

Next, we consider the case where c ∗ = 1 . As with the discussion of (Schoenholz et al., 2016), the perturbed correlation is defined by c β,β ′ = c ∗ -E β,β ′ q where E β,β ′ > 0 and then ϕ ( u 2 ) is expanded as follows.

Thus, we have

Using the fact that u ∗ 2 = u ∗ 1 and u ∗ 2 is independent of z 2 ,

Substituting these facts into Eq. (21), we obtain

The above result agrees with that obtained by substituting c ∗ = 1 for the result obtained when the case c ∗ < 1 .

Lemma A.2. The Jacobian matrix J ( Σ ∗ ) of the iterated map C defined in Eq. (14) is obtained as follows.

Furthermore, the rank of the Jacobian matrix J ( Σ ∗ ) ∈ R N 2 × N 2 is at most K .

Proof. Let some semi-positive definite matrix E ∈ R N × N be a deviation from the fixed point Σ ∗ . From Lemma 3.1 and Lemma A.1, we have

Note that Eq. (22) is obtained by neglecting higher order terms in Eqs. (19) and (20).

By the definition of the fixed point,

By substituting the fact into Eq. (22), we obtain

The last equation can be rewritten using the matrix calculation as follows.

Since J ( Σ ∗ ) is the first-order coefficient to the deviaction E , J ( Σ ∗ ) is exactly the Jacobian matrix of the iterated map C .

Furthermore, the Jacobian matrix can be decomposed to two matricies A ∈ R N 2 × K and B ∈ R K × N 2 as follows.

Therefore, the rank of the Jacobian matrix is at most K .

Lemma A.3. Let Σ ∗ be the fixed point of the form in Eq. (5) and E ( k ) be the perturbation expressed as

where ϵ

Then, we have, for all k ∈ [ K ] \{ 0 }

k denotes the scale of the perturbation, assumed to be sufficiently small. ,

Proof. From Eq. (23), we have

From the definition of the fixed point, we obtain

By combining the above two results, we have

̸

Since cos( θ ( k ) β,β ) = 1 for all k ∈ [ K ] , we have ψ ( k ) β,β = 0 and E ( k ) β,β = 0 . Thus, only the term with E ( k ) β,β ′ remains.

̸

Given the orthogonality of the cosine functions, we have

̸

Utilizing Eq. (26) leads to the following two facts:

̸

̸

̸

̸

By substituting these facts into Eq. (24), we obtain

This completes the proof.

Lemma A.4. Let Σ ∗ be the fixed point of the form in Eq. (5) and E be the perturbation expressed as

where ϵ denotes the scale of the perturbation, assumed to be sufficiently small, 1 N 1 ⊤ N is all-one matrix with size N × N .

Then, we have

From Eq. (23), we have

̸

Our goal is to derive Eq. (27) from Eq. (28). The following results are useful for the computation of each term of Eq. (28).

̸

̸

̸

Using the above results, the first and second terms of Eq. (28) are calculated as follows.

̸

The third term of Eq. (28) is obtained as follows.

̸

̸

By substituting these facts into Eq. (28), we have

Finally, we obtain Eq. (27).

B Proof of Theorem 3.5

Theorem 3.5 (Trainability) . Let ˜ Σ ( ℓ ) ∈ R N × N be the gradient covariance with respect to some loss L , e.g . mean squared error, at the ℓ -th layer. Suppose that the gradient covariance at the L -th layer is decomposed as ˜ Σ ( L ) α,α ′ = ∑ K -1 k =0 ˜ ϵ k cos ( θ ( k ) α,α ′ ) + ˜ e , where ˜ ϵ k is the coefficient of each basis and ˜ e belongs to the orthogonal complements of span( { cos ( θ ( k ) α,α ′ ) } K -1 k =0 ) . Then, the gradient covariance

̸

at the ℓ -th layer is obtained by

Proof. Recall the definition of the gradient covariance. We first demonstrate the iterated map of the gradient covariance, starting from this definition.

The Jacobian ∂H ( ℓ +1) β,i ∂H ( ℓ ) α,j is calculated as follows.

The covariance of Jacobian ∂H ( ℓ +1) β,i ∂H ( ℓ ) α,j is as follows.

Then, the covariance of the gradient is given by the following recurrence relation for all α, α ′ ∈ [ N ] ,

As with (Schoenholz et al., 2016), we approximate Σ ( ℓ +1) ≈ Σ ∗ since the number of layer ℓ is assumed to be sufficiently large. Then, the linear iterated map of the gradient covariance Σ ( ℓ +1) ↦→ Σ ( ℓ ) is given as follows.

The rank of the linear iterated map Σ ( ℓ +1) ↦→ Σ ( ℓ ) is less than K since the matrix representation of the linear map can be decomposed into two matrices ˜ A ∈ R N 2 × K and ˜ B ∈ R K × N 2 as follows.

Next, we show that the subspace span ( { cos ( θ ( k ) α,α ′ ) } K -1 k =0 ) ⊂ R N × N is the K -dimensional eigenspace with eigenvalue χ c ∗ of the linear iterated map Σ ( ℓ +1) ↦→ Σ ( ℓ ) .

By substituting ˜ Σ ( ℓ +1) β,β ′ = ∑ K -1 k =0 χ L -( ℓ +1) c ∗ ˜ ϵ k cos ( θ ( k ) β,β ′ ) into Eq. (29), we obtain

From the orthogonality of the cosine and sine function, we obtain

Hence, we have

This completes the proof.

C Discussions on the similarity of DCN and CNN

While CNNs perform local convolutions in the spatial domain, FNOs execute convolutions in the frequency domain, thereby achieving global convolutions in the spatial domain. Our theory for the FNO and the mean-field theory for CNNs (Xiao et al., 2018) share a common focus on the correlation dynamics Σ (0) , Σ (1) , . . . , Σ ( L ) of the spatial representations H (0) , H (1) , . . . , H ( L ) ∈ R N × D . The iterated maps for CNNs and FNOs are obtained as follows.

where 2 r +1 is the number of filter width and ker = { β ∈ Z || β | ≤ r } is 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 ( α, α ) for any α ∈ [ N ] is equivalent to a CNN with filter size N performing global

convolution.

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.

Regarding the fixed point, Xiao et al. (2018) demonstrated that any fixed point for the iterated map of the DCN is also a fixed point for that of the CNN. Consequently, Lemma A.1 indicates that the fixed points for CNN and FNO are consistent.

The behaviour of the iterated map around fixed points reflects the nature of each architecture. CNN possess diagonal eigenspaces associated with eigenvalues χ q ∗ and non-diagonal eigenspaces associated with eigenvalues χ c ∗ . FNOs without mode truncation exhibit a similarity, possessing eigenspaces χ q ∗ for zero-frequency and eigenspaces χ c ∗ for k-frequencies with diagonal components removed.

Finally, the iterated map during the backpropagation for CNN and FNO are given by

where v β is the variance weight parameter dependent of the filter position β , i.e., w i,j ( β ) ∼ N (0 , σ 2 v β /D ) , ∑ v β = 1 .

Using the approximation Σ ( ℓ +1) ≈ Σ ∗ as with (Schoenholz et al., 2016), the backpropagation of the diagonal components ( α, α ) of the FNO without mode truncation is equivalent to that of the global CNN (see Eq. (2.16) in (Xiao et al., 2018)) with v β = 1 N , as shown below:

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 ˜ Σ .

D Details of Experimental Setup

In this section, we summarize the detailed setup of the all experiments, including the experiments in Section 4.

D.1 Datasets

D.1.1 Advection equation

We used the advection equation data published by Takamoto et al. (2022). The advection equation for the function u ( x, t ) ∈ L 2 ((0 , 1) × (0 , 2]; R ) is given by

where u 0 ∈ L 2 ((0 , 1); R ) is the initial condition and β ∈ R is an advection speed set to 2 . 0 . The exact solution is given as u ( x, t ) = u 0 ( x -βt ) for any initial condition u 0 .

Only periodic boundary conditions were used in this dataset. The initial conditions are the superposition of sinusoidal wave given by

where k i = 2 π ∑ N j =1 n i,j /L x are wave numbers whose n i,j are integer numbers randomly chosen in [1 , k max ] , L x = 1 is the calculation domain size, N = 2 is the number of wave to be added, and k max = 8 is the maximum wave number. The amplitude A i is uniformly chosen in [0 , 1] , and the phase ϕ i is the randomly chosen in (0 , 2 π ) . The 2nd-order temporal and spatial upwind finite difference scheme was used for generating the data. Settings are described in Appendix D of (Takamoto et al., 2022).

D.1.2 Burgers' equation

We used the Burgers' equation data published by Takamoto et al. (2022). The Burgers' equation for the function u ( x, t ) ∈ L 2 ((0 , 1) × (0 , 2]; R ) is given by

where u 0 ∈ L 2 ((0 , 1); R ) is the initial condition and ν is the diffusion coefficient set to 4 . 0 . The periodic boundary conditions and Equation (31) are used as the initial conditions. The 2nd-order temporal and spatial upwind finite difference scheme is used for generating the data. Settings are described in Appendix D of (Takamoto et al., 2022).

D.1.3 Darcy Flow equation

We used the data of 2D Darcy Flow equation on a regular grid published by Li et al. (2020a). The Darcy Flow equation for the function u ∈ H 1 0 ((0 , 1) 2 ; R + ) with a Dirichlet boundary is given by

where a ∈ L ∞ ((0 , 1) 2 ; R + ) is the diffusion coefficient and f ∈ L 2 ((0 , 1) 2 ; R ) is the forcing function. The coefficients a was generated by measure µ = ψ ♯ N (0 , ( -∆+9 I ) -2 ) using the Laplacian with zero Neumann boundary and the binary point-wise mapping ψ ( x ) = 12 ( x ≥ 0) , 3 ( x < 0) . The forcing function is fixed f ( x ) = 1 . Our aim is to predict the operator mapping the diffusion coefficient to the solution a → u . The solution function u was generated by using the second-order finite difference scheme on a 421 × 421 grid. Settings are described in Appendix A.3.2 of (Li et al., 2020c).

D.1.4 Incompressible Navier-Stokes equation

We used the 2D NS equation on the unit torus defined by

where ω ( x, t ) ∈ C ([0 , T ]; H r ((0 , 1) 2 ; R 2 )) is the vorticity, ω 0 ∈ H r ((0 , 1) 2 ; R 2 ) is the initial vorticity, u ( x, t ) ∈ C ([0 , T ]; H r ((0 , 1) 2 ; R 2 )) is the velocity field for any r > 0 , ν ∈ R + is the viscosity, and f ∈ L 2 ((0 , 1) 2 ; R ) is the external forcing function defined by f ( x ) = 0 . 1 (sin(2 π ( x 1 + x 2 ) + cos(2 π ( x 1 + x 2 )) . The initial vorticity ω 0 was generated by ω 0 ∼ µ where µ = N (0 , 7 3 2 ( -∆+49 I ) -2 . 5 ) with periodic boundary conditions. The viscosity was set to 1 e -3 , 1 e -4 , or 1 e -5 . Our aim is to predict the operator that maps a solution u up to time 10 to a solution up to some later time T > 10 . The data was generated by the pseudo-spectral Crak-Nicholson second-order method on 64 × 64 grid. For the data with ν = 1 e -4 , the time resolution was also downsampled by half. Settings are described in Section 5.3 of (Li et al., 2020c).

D.2 Training settings

Our detailed training settings of the experiments in Section 4 are provided in Table 2. Our experimental environment consists of an Intel Xeon Plantinum 8360Y (36-core) CPU and a single NVIDIA A100 GPU. Most of our code for experiments are based on the code of PDEBench ( https://github. com/pdebench/PDEBench ) (Takamoto et al., 2022). 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 √ 2 and to initialize the weights by Gaussian distribution, as described in Section 3.1.

E Detailed Experimental Analysis

E.1 Analysis of Training Loss

Figure 6a shows the training loss for each epoch of the 32-layer FNOs with parameters σ 2 ∈ { 0 . 1 , 0 . 5 , 1 . 0 , 2 . 0 , 3 . 0 , 4 . 0 } on the NS dataset with ν = 1 e -3 . When the initial parameter σ 2 is too small, the training loss is not well reduced due to gradient vanishing. On the other hand, when the initial parameter σ 2 is too large, the initial training loss blows up due to gradient exploding. The proposed edge of chaos initialization smoothly reduces the training loss in the initial epoch and enables stable training.

E.2 Analysis of Test Performance

The nMSE of the FNOs on test datasets for six distinct PDEs is presented in Tables 3 and 4. Results are shown only for the FNOs with initial parameters where training was successful in many cases. For the NS equation with viscosity values of ν = 1 e -3 , 1 e -4 , where sufficient data is available, Table 4 shows that best performance is achieved with 8 or 16 layers. This suggests that while shallow FNOs are currently prevalent, deep FNOs could be advantageous in certain tasks, underscoring the significance of our analysis of the bias in deep FNOs. Conversely, for other equations, the 4-layer FNO performs best and deeper FNOs result in a drop in performance even with the edge of chaos initialization. We will discuss this test performance deterioration in detail.

The over-fitting phenomenon is observed in the Darcy Flow and NS equation datasets with ν = 1 e -5 , where only limited training data is available. The training loss for each epoch on the NS equation is depicted in Fig. 6b. As demonstrated in Fig. 6b, the 16 and 32-layer FNOs yield a lower training loss than the 4-layer FNO, but exhibit poorer performance on the test dataset as shown in Table 4. These results suggest over-fitting to the training data, necessitating either abundant training data or appropriate regularization.

Conversely, the under-fitting phenomenon is apparent in the 1D advection and Burgers' equation datasets in Table 3. The training loss of the FNO with ReLU activation for each epoch on the Burgers' equation is presented in Fig. 6c. Figure 6c indicates that the larger the number of layers, the higher the training loss in the final epoch, and the worse the test performance. This under-fitting to the training data could be attributed to the escalating complexity of the loss landscape as the layer count increases, a known issue for DCN and CNN (Li et al., 2018). This may be due to the emergence of local minima corresponding to operators that generate too complex functions, preventing the

(a) NS eq. (1e-3)

(b) NS eq. (1e-5)

(c) Burgers' eq.

107

I°0

0.5

10

4

102

4

8

105

1.0

2.0

16

32

103

16

32

MSE

103

101

3.0

4.0

MSE

100

MSE

104

105

10-1

106

10-1

10-7

103

108

0

100

200

300

400

500

0

100

200

300

400

500

0

50

100

150

200

oodo

Table 3: Test performance measured by nMSE of 1D simplified FNO on 1D PDEs

attainment of parameters that achieve global minima. This issue could be mitigated by introducing more suitable regularization, an appropriate optimizer, or a skip connection (Tran et al., 2022).

Our theory and experiments suggest that the training of deep FNOs has suffered from problems including gradient vanishing and exploding due to improper initialization, over-fitting caused by insufficient training data, and under-fitting caused by loss landscapes with strong non-convexity. While our edge of chaos initialization prevents the gradient vanishing and exploding, techniques to solve over-fitting and under-fitting problems are still needed in practice.

F Visualization of Forward Propagation

We visualized the behavior of the simplified FNO's covariance matrix Σ ( ℓ ) with varying initialization parameters σ 2 ∈ { 0 . 1 , 1 . 0 , 2 . 0 , 4 . 0 } and activation functions. The FNO, with a width of D = 1024 , was used and the input was sampled from the standard normal distribution with a spatial size of N = 32 . The results of the FNO with Tanh activation, both with and without mode truncation, are shown in Figs. 7 and 8 and Figs. 9 and 10 respectively. Similarly, the results of the FNO with ReLU activation, both with and without mode truncation, are displayed in Figs. 11 and 12 and Figs. 13 and 14 respectively. In the ordered phase, all figures illustrate convergence to the fixed point Σ ∗ where c ∗ = 1 , with the rate of convergence increasing as the parameter σ 2 decreases. In the chaotic phase, the activation function dictates the covariance behavior. Without mode truncation, the covariance behavior of the FNO is identical to those of the DCN; otherwise non-uniform, FNO-specific periodic covariance is exhibited.

1.00

( =0

= 1

=4

( = 8

( = 16

= 32

o² = 0.1

0.75

3

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0.50

= 0

( = 1

( = 4

( = 8

= 16

(=32

0

o² = 1.0

0.25

spatial position

20

30

spatial position

20

30

spatial position

20

30

spatial position

20

30

spatial position

20

30

spatial position

20

0.00

=0

(=1

= 4

=8

= 16

=32

² = 2.0

0.25

30

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0.50

(= 0

(=1

= 4

=8

(= 16

= 32

0

g² = 4.0

0.75

30

0

spatial position

20

30

0

spatial position

30

0

spatial position

20

30

0

spatial position

20

30

spatial position

20

30

0

spatial position

20

1.00

2.00

= 16

=32

g² = 0.1

1.75

30

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

0

spatial position

20

1.50

=0

[=}

= 4

=8

= 16

= 32

o² = 1.0

1.25

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

1.00

= 0

=1

= 4

= 8

=16

=32

0

10

10

0.75

o² = 2.0

20

30

30

30

30

30

30

spatial position

20

spatial position

0

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0.50

= 0

=1

= 4

= 8

= 16

= 32

10

10

10

o² = 4.0

0.25

20

30

30

30

30

30

spatial position

20

spatial position

20

spatil psition

20

spatial position

20

0

spatial position

20

0

spatial position

20

0.00

1.00

= 0

= 16

= 32

o² = 0.1

0.75

30

spatial position

20

30

spatial position

20

0

spatial position

20

30

spatial position

20

30

0

spatial position

20

30

spatial position

20

0.50

= 0

( = 1

= 4

( = 8

= 16

=32

0

0

0.25

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0.00

=0

(=1

( = 4

(=8

( = 16

(=32

0

++**+

o² = 2.0

0.25

30

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0)

spatial position

20

spatial position

20

0.50

(= 0

(=1

( = 4

=8

= 16

= 32

0

0

g² = 4.0

0.75

30

30

30

30

30

30

0

spatial position

20

0

spatial position

0

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

1.00

2.00

=0

I=]

= 8

= 16

= 32

g² = 0.1

1.75

会

30

30

30

30

30

30

spatial psition

20

spatial psition

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

1.50

=0

f = 1

= 4

=8

= 16

=32

0

o² = 1.0

1.25

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

0

spatial position

20

1.00

= 0

[ =}

( = 4

( = 8

= 16

= 32

0

o² = 2.0

0.75

30

30

30

30

30

spatial poition

20

spatial position

20

spatial position

20

spatial psition

20

spatial psition

20

spatial position

20

0.50

= 0

= 1

( = 4

=8

= 16

= 32

0

0

g² = 4.0

0.25

30

30

30

30

30

30

spatial position

20

spatial position

20

0

spatil psition

20

spatial position

20

0

spatial position

20

0

spatial position

20

0.00

1.00

= 0

= 8

= 16

= 32

o² = 0.1

0.75

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0.50

= 0

( = 1

( = 4

( = 8

= 16

=32

0

0.25

spatial position

20

30

spatial position

20

spatial position

20

30

spatial position

20

30

spatial position

20

30

spatial position

20

0.00

=0

(=1

= 4

( = 8

( = 16

(=32

o² = 2.0

0.25

30

30

30

30

30

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0.50

(= 0

(=1

( = 4

( = 8

( = 16

= 32

0

o² = 4.0

0.75

30

30

30

30

30

0

spatial position

20

0

spatial position

0

spatial position

20

spatial position

20

0

spatial position

20

0

spatial position

20

1.00

(a) Initial parameters

σ

2

∈ {

0

.

1

,

1

.

0

,

2

.

0

}

.

(b) Initial parameter

σ

2

= 4

.

0

0 = 3

=8

=16

= 32

U

0

0

0

0

10

uosod peeds

10

10

o² = 0.1

2

30

30

30

30

30

30

0

20

0

20

0

20

0

20

0

20

0

20

spatial position

spatial position

spatial position

spatial position

spatial position

spatial position

= 0

( = 1

( = 4

=8

= 16

= 32

0

0

10

10

o² = 1.0

20

30

30

30

30

30

30

0

spatial position

20

0

spatial position

20

U

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

=0

( = 1

= 4

=8

=16

=32

U

0

0

0

0

10

10

10

o² = 2.0

20

20

20

20

20

30

30

30

30

30

30

0

20

0

20

0

20

0

20

0

20

0

20

spatial position

spatial position

spatial position

spatial position

spatial position

spatial position

Iglo

1010

101o

Ig10

U

0 = 3

1.0

0

D

108

0

108

0

f = 16

108

0

=32

10*

uoisod peeds

10

0.8

10

uopsod pepeds

10

106

10

106

uosd peexds

10

106

10

106

4.0

20

0.6

0.4

20

104

102

peeds

20

104

102

10

102

20

104

102

30

0

0Z

0.2

30

0

20

30

0

20

100

30

0

20

109

30

0

0Z

100

0

20

100

spatial position

spatial position

spatial position

102

spatial position

10-2

spatial position

102

spatial position

102

1.00

= 0

= 8

= 16

= 32

o² = 0.1

0.75

30

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0

spatial position

20

spatial position

20

0.50

= 0

= 1

( = 4

( = 8

( = 16

(= 32

0

0.25

30

30

30

30

spatial positio

20

spatial position

0

spatial position

20

spatial position

20

spatial position

20

spatial position

20

0.00

=0

(=1

= 4

=8

= 16

=32

0.25

o² = 2.0

30

30

30

30

30

spatial position

20

spatial position

20

spatial position

20

spatial position

20

spatial position

20

spatial psition

20

0.50

(=0

= 1

= 4

=8

( = 16

=32

0

0

o² = 4.0

0.75

30

0

spatial position

20

30

spatial position

20

0

spatial position

20

spatial position

20

30

spatial position

20

0

spatial position

20

1.00

10

(a) Initial parameters

σ

2

∈ {

0

.

1

,

1

.

0

,

2

.

0

}

.

(b) Initial parameter

σ

2

= 4

.

0

=0

=4

=8

=16

=32

U

0

U

0

0

0

osod peedss

10

10

10

o² = 0.1

2

2

20

2

20

30

30

30

30

30

30

0

20

0

20

0

20

0

20

0

20

0

20

spatial position

spatial position

spatial position

spatial position

spatial position

spatial position

=0

(= 1

= 4

=8

= 16

=32

0

3

10

ogisod pegeds

sos

10

o² = 1.0

20

20

2

20

30

30

30

30

30

30

0

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

= 0

( = 1

( = 4

( = 8

( = 16

=32

0

0

0

0

0

10

10

10

10

o² = 2.0

20

20

20

20

20

2

30

30

30

30

30

30

0

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

0

spatial position

20

spatial position

20

0 =3

1.0

0

1.0

1.2

0

100

105

0

f = 8

10

Iglo

0

f = 16

1gl0

108

0

f = 32

10*

spatial position

10

0.8

nosod peeds

10

0.8

10

106

gosod puds

10

106

10

106

uosod peeds

10

106

g² = 4.0

20

0.6

0.4

20

0.6

0.4

20

102

10*

20

104

102

pegeds

20

20

104

102

30

0

20

0.2

30

0

20

0.2

30

0

20

109

30

0

20

109

30

0

20

109

30

0

20

109

spatial position

spatial position

spatial position

102

spatial position

102

spatial position

102

spatial position

102

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The contribution and scope of the paper are described in the abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We mention the limitation of this paper in Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The full set of assumptions and a complete proof are provided in Section 3 and Appendices A and B.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All the information needed to reproduce the main experimental results are provided in Section 4 and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We use existing public datasets and cite them appropriately.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All training and test details are provided in Section 4 and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: In the experimental results depicted in Figs. 3 and 4, we determined that the variation in the reported values has minimal impact. In fact, the inclusion of error bars seems to obstruct intuitive understanding through visualization.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All sufficient information on the computer resources are provided in Section 4 and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We follow the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Our paper is foundational research and not tied to particular applications.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have stated the license and copyright of the data used.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The content of our public code is well described.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. **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.\n2. **Practical implications:** The theoretical findings lead to concrete recommendations for initializing FNOs to ensure stable training, particularly for deep FNOs.\n3. **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.\n\nWEAKNESSES:\n1. The adaptation of mean-field theory to FNOs, while a notable endeavor, may not sufficiently break new ground in the field.\n2. The significance of the theoretical aspects that FNOs share similarities with CNNs should be further discussed.\nOverall, 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.\n4. 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.\n\n[1] Helwig, J., Zhang, X., Fu, C., Kurtin, J., Wojtowytsch, S., and Ji, S. Group equivariant Fourier neural operators for partial differential equations.\n\nQUESTIONS:\n**Major Questions:**\n\n1. 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?\n\n2. In line 122, it states:\n   > \"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.\n\n   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?\n\n3. 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?\n\n4. 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?\n\n**Minor/Optional Questions and Suggestions:**\n\n1. 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).\n\n2. Given your insights into FNO expressivity and trainability, do you have any thoughts on potential architectural improvements or variations that could enhance FNO performance?\n\n3. 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.\n\nLIMITATIONS:\nYes, the authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper leverages mean-field theory, developed for neural networks, to study the expressivity and trainability of random Fourier neural operators (FNOs). It provides practical insights for the appropriate initialization of FNOs.\n\nSTRENGTHS:\nThe authors address a significant question in the neural operators literature: why does the performance of Fourier neural operators degrade as the number of layers increases? The authors take a first step toward answering this question by adopting an original theoretical perspective.\n\nWEAKNESSES:\n- The main body of the paper is not self-contained. Additionally, the technical section can be difficult to follow due to the need for back-and-forth referencing between multiple sections and references. Overall, the paper could benefit greatly from improved clarity. To enhance readability, I recommend reorganizing some parts of the text.\n  - For example, the covariance matrix $\\Sigma_{\\alpha,\\alpha'}^{(l)}$ should be introduced in section 3.2 where it is first examined.\n  - At first glance, one might question if the iterated map in equation (6) is a definition. However, its origin becomes clear upon reviewing the supplementary material.\n  - Lemma A.1 should be included in the main text, as it is a central result.\n  - Some parts, such as the detailed description of the datasets, can be moved to the appendix to make space.\n- I recommend providing further motivation for all concepts and referencing previous works where appropriate. For instance, the importance of studying the fixed point of the iterated map should be emphasized. Most importantly, please refer to the supplementary material for proofs when available.\n- 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. I am not convinced that the findings offer practical insights into the stable training of FNO, despite the preliminary numerical experiments provided. Additionally, the term _trainability_, as used by the authors, only pertains to the vanishing or exploding gradients of random FNO.\n- 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.\n\nQUESTIONS:\n- Does the fixed point of the iterated map always exist?\n- The representation of Fig.1(a) is unclear. Does it depict the gradient norm averaged over the training epochs? I assume it shows the values only at initialization, averaged over multiple seeds. Additionally, is the standard deviation relatively small? Clarification on these points would be helpful.\n\nLIMITATIONS:\nI understand that error bars are omitted from the figures for the sake of clarity. However, including standard deviations in the tables would enhance the presentation of the results. Additionally, it is concerning that the results are averaged over only three seeds, which is insufficient for robust statistical analysis.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a comprehensive theoretical analysis of the Fourier Neural Operator (FNO) using mean-field theory. The main contributions to me at least are:\n\n1. Establishing a mean-field theory for FNO, analyzing its behavior from an edge of chaos perspective.\n2. Examining the ordered-chaos phase transition of random FNO based on weight distribution.\n3. Identifying a connection between expressivity and trainability, showing that ordered and chaotic phases correspond to regions of vanishing and exploding gradients respectively.\n4. Providing practical initialization requirements for stable FNO training.\n\nSTRENGTHS:\n1. Theoretical Depth: The paper provides a rigorous mathematical analysis of FNOs, which has been lacking in the literature. The authors adapt and extend existing mean-field theory techniques to the specific structure of FNOs.\n2. Novel Insights: The paper reveals FNO-specific characteristics in the phase transition, induced by mode truncation, while also showing similarities with densely connected networks. This bridges the gap between FNOs and more traditional neural network architectures. This was very very insighting to me!\n3. Practical Implications: The theoretical analysis leads to practical guidelines for initializing FNOs, addressing a known issue of training instability in deep FNOs. This demonstrates the real-world value of the theoretical work.\n4. Comprehensive Experiments: The authors validate their theoretical findings through extensive experiments on various PDE-solving tasks, showing the practical relevance of their analysis.\n5. Clarity: Despite the technical nature of the content, the paper is well-structured and clearly written, making it accessible to a broader audience in the machine learning community.\n\nWEAKNESSES:\n1. There are no specific weaknesses that I see however, while the experiments are comprehensive, 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.\n2. Insufficient analysis of the impact of different Fourier modes: The paper doesn't provide a detailed analysis of how the number and distribution of Fourier modes affect the expressivity and trainability. This is crucial as the choice of modes is a key hyperparameter in FNOs.\n3. Insufficient exploration of error propagation: For time-dependent PDEs, there's limited analysis of how errors propagate through time steps and how this relates to the expressivity and trainability of FNOs.\n\nQUESTIONS:\n1. 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\n2. 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?\n\nLIMITATIONS:\nThe authors have described their limitations!", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: SUMMARY:
This work explores the expressivity and trainability of Fourier Neural Operators (FNOs) by establishing a mean-field theory and analyzing their behavior from an edge-of-chaos perspective.

STRENGTHS:
The study examines the ordered-chaos phase transition of the network based on weight distribution, highlighting unique characteristics induced by mode truncation and similarities to densely connected networks. It also identifies a connection between expressivity and trainability, noting that ordered and chaotic phases correspond to regions of vanishing and exploding gradients, respectively. This insight provides a practical prerequisite for the stable training of FNOs. The experimental results support the theoretical findings.

WEAKNESSES:
1. The summary of related work can be further improved, particularly by including recent neural operator architectures from recent conferences or journals.
2. While Figure 1.(c) confirms that the original setup for initialization is reasonable, our experience suggests that the performance of FNOs in both training and test accuracy is highly dependent on the initialization variances for different hyperparameters and tasks. The derivation and results in this paper do not provide sufficient insights into this issue.
3. The code is not available.

QUESTIONS:
1. Can the derivation in this paper for FNOs be generalized to more general neural operators?
2. Is it possible to identify (approximately) how the initialization depends on different datasets?

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12