Enhancing Solutions for Complex PDEs: Introducing Translational Equivariant Attention in Fourier Neural Operators

W1: Our work shares a similar philosophy with some recent works, with translation equivariant attention serving as the core of our approach (detailed in W2). Compared to the methods you mentioned, such as LSM and HANO, our method has two key advantages:

1. Preserving FNO Architecture: We retain FNO's foundational architecture while addressing its limitations using translation equivariant attention. This ensures our approach remains consistent with FNO's principles. In contrast, methods like LSM and HANO are built on entirely new frameworks. While these methods may achieve good performance, they do not directly solve FNO's inherent issues.
2. Novel Equivariant Attention Mechanism: Our equivariant attention introduces a new mechanism that significantly enhances FNO’s performance. For example, our approach improves results on Trigonometric 256 and 512 by 17% and 26%, respectively, compared to the second-best method. Substantial improvements are also observed across other tasks, such as reverse solving (see W2 for details).

W2: Thank you for bringing up this discussion. First, we regret that we did not come across the [Helwig et al., 2023] paper you mentioned. To provide clarity, let’s use G-FNO as an example. Upon revisiting the methods section of G-FNO, we searched for mentions of “attention” but found none.

The key distinction lies in the following: G-FNO employs group convolution to achieve its structure through FFT. In contrast, we transform group convolution into attentive group convolution, incorporating translation-equivariant attention via group equivariance and FFT. This design fundamentally differentiates our translational attention from G-FNO, as it is specifically built from attentive group convolution to address FNO’s challenges.

W3: Apologies for not articulating our starting point clearly earlier. Let me provide a more detailed explanation.

We first observed that FNO exhibits a bias toward low-frequency components. This is evident in Figure 3(i), where the center of the figure represents high-frequency components. As shown, FNO demonstrates relatively high errors in capturing these high-frequency details.

Recognizing this limitation, we sought to address it through an attention mechanism. However, due to FNO's architecture being deeply rooted in Green's function, directly adding attention would not align with its design. Therefore, we targeted the specific part of FNO that approximates the kernel function with a convolution. By introducing the Attentive Equivariant Convolution (refer to Equations (3) and (4)), we successfully built translation-equivariant attention that resolves this issue.

This innovation demonstrates how translation equivariant attention can enhance FNO while staying consistent with its original principles.

W1: Thanks for your advice; we will cite this paper in our revised version and discuss it. The problems we focus on indeed are the same, but the main solutions are not quite different. We mainly solve the translation equivariant attention proposed from the perspective of the kernel function in Green's function. In this paper, relevant calculations are carried out directly from the local properties of convolution. By the way, I think our work and related work may belong to the same period, but we made some modifications because of the imperfections of the previous version.

Q1: We apologize for the issue with the FNO results. We have re-run the latest FNO code using the official implementation from the NeuralOperator GitHub repository. The updated results are 0.88 and 6.60 for $1\times10^{-3}$ and $1\times10^{-4}$, respectively.

Q2: Apologies for the earlier imprecise explanation. In our paper, all experiments were conducted three times, and the results were averaged. We have also added the standard deviations of our model's results in Table 2 for better clarity:

Q3: Thank you for your insightful question. Our title emphasizes translation equivariance as a key concern. This begins with the design of the FNO, which originates from Green's function to solve PDEs. A critical step in FNO is approximating the kernel of Green’s function using a convolution function. One of the essential properties of convolution is translation equivariance.

When we observed FNO's limitations in capturing high-frequency information, we aimed to address this through an attention mechanism. However, because FNO's architecture is fundamentally derived from Green’s function, directly integrating attention would be incompatible. Thus, we focused on the part of FNO that approximates the kernel function with a convolution. By introducing the Attentive Equivariant Convolution (Equations (3) and (4)), we demonstrated that translation-equivariant attention could address this issue. Consequently, our primary focus is on the translation group.

Q4: We believe there might be some misunderstanding regarding our convolution operator. A more detailed explanation is provided in Q5. To clarify, we use fixed-size $3\times3$ local convolution kernels with padding and a stride of 1.

Q5: $\text{Conv}(\boldsymbol{v}^k)$ represents the convolution operation applied to $\boldsymbol{v}^k$.

Q6: It seems there may have been some confusion. We do not perform convolution in the Fourier domain see Figure 2. As detailed in Equation (9), the convolution operation is performed in the spatial domain. The variable $k$ denotes the $k$-th hierarchical layer, where each layer utilizes Equivariant Attention, a Fourier Layer, and an MLP.

Q7: We are sorry we may not have read these two papers before. However, we hope to take UNO as an example. In our understanding, U-net-shaped networks may pay more attention to multi-scale feature capture through up-sampling and down-sampling. Of course, this method can solve part of the problems, but we found through experiments that Individual calculations at each scale allow for better modeling of the potential space at each scale, resulting in better results.

Q8: We regret the writing errors in our paper. All the tables in our work use N-MSE as the evaluation metric, consistent with the relative L2 loss employed in prior work with FNO. To ensure fairness and prevent potential misunderstanding, we have adhered to the same writing conventions as previous papers.

Q9: Thank you for raising this interesting question. It is important to note that even without equivariant attention, our structure differs from others. The primary issue is that removing equivariant attention reduces the number of model parameters, as the Fourier modes of $R$ in our Fourier layer diminish progressively (as detailed in the appendix). Consequently, the model’s performance deteriorates without equivariant attention.

To understand the role of attention, please refer to Table 4. In the rows for $\text{Att}{\tiny{\rightarrow}}\text{d-Conv}$ and $\text{Att}{\tiny{\rightarrow}}\text{MLP}$, you'll notice that the parameter counts remain roughly unchanged. For instance, in the Trigonometric dataset, a comparison of the two ablation experiments demonstrates that attention transforms the performance of TE-FNO into a state-of-the-art (SOTA) level.

Q10: Yes, we acknowledge an important limitation. We will add a note indicating that TE-FNO is better suited for solving PDEs with identical input and output resolutions. This limitation is generally shared by U-shaped network architectures.

As the rebuttal period ends soon, would you provide more justification into why you need to have invariant attentions? Additionally, as I understand, if you have local fixed sized convolutions, it will break the discretization convergence of FNO. Can you elaborate more on this?

Additionally, FNO also updated their architectures for Fourier_2d (i.e. for Darcy-Smooth in your Table 1). I would suggest to update that result as well besides that of the NS equation.

Thank you for the comments. Overall, I think the justification for why a translational invariant attention mechanism is needed is weak.

"This means shifting the computation of attention from the frequency domain to the spatial domain"

Basically, the complexity of the kernel will be $O(n^2)$ despite it being a convolution kernel. I don't see any intuitive advantage of this over either 1) using the full integral of attention (more expressive power) or 2) using the Fourier convolution theorem (better complexity).

Thank you for pointing this out. This result is the one we used in LSM [2]. We have updated the results for the Fourier_2D architecture as applied to Darcy-Smooth in Table 1. The updated value is 0.83.

Thank you for this update. If you haven't done so, I would recommend updating all other results. (Due to the short time frame of the rebuttal, I know this is not a reasonable request, but to ensure fairness, I think it's necessary to update them in future iterations. I would reevaluate this work based on the promise that the authors will do their best to ensure fairness in the future.)

Overall, I am willing to increase my score to a $5$, reflecting the clarifications and updates provided by the authors. However, due to the aforementioned issues, including the relatively weak motivation and some minor practical concerns surrounding translational invariant attention kernels, I have to restrain myself from assigning a higher score.

W1: Thank you for your question. In most cases, we directly utilize the open-source implementations of the relevant methods as baselines. We do not impose constraints requiring identical parameter counts or training durations across methods.

W2: Yes, there is an important limitation we acknowledge. We will include a note indicating that TE-FNO is more suitable for solving PDEs where the input and output resolutions are identical. This limitation is commonly observed in U-shaped network architectures. Additionally, the main limitation of translational equivariant attention lies in its reliance on convolution-based local attention, which makes it less effective in capturing global features compared to certain other approaches.

W3: We have already discussed the two mentioned methods in our related work section (refer to the related work and more related work in the appendix). We will discuss this further in the revised version of our manuscript. Our method differs significantly from these two approaches:

• GT primarily utilizes Fourier transformations to reduce the computational complexity of self-attention.
• AFNO performs self-attention-like computations in the frequency domain after dividing the input into patches.

In contrast, our translational equivariant attention mechanism avoids traditional self-attention in order to maintain translation equivariance. Therefore, translational equivariant attention is better categorized as a visual attention mechanism (see more related work in the appendix for further details). This distinction highlights that our approach has a fundamentally different starting point and structure compared to the methods that rely on self-attention.

Q1: Apologies for the lack of clarity in our writing. The effect of "add hier" refers to adding an additional hierarchical layer. In our current model, we use a structure with four hierarchical layers. In the ablation experiment, we modified the model to include a fifth hierarchical layer. The setup of this fifth layer is similar to the previous layers.

Q2: First, it is important to note that AFNO requires input patching before computation, which introduces additional complexities. In comparison, our translational equivariant attention is more theoretically grounded. It does not rely on directly integrating self-attention or adaptive token mixing, which can feel more like combining existing structures. We conducted ablation studies on the performance of self-attention and adaptive token mixing, which demonstrate the theoretical and practical advantages of our approach.

W: Thank you for your questions. Our work can also be applied to large-scale problems, and in our recent work we have applied methods to related problems such as spatial-temporal prediction.

Q: Thank you for your recognition of our work. We provide detailed parameter counts for our models below and other baselines.